Spaces:
Runtime error
Runtime error
| # Licensed to the Apache Software Foundation (ASF) under one | |
| # or more contributor license agreements. See the NOTICE file | |
| # distributed with this work for additional information | |
| # regarding copyright ownership. The ASF licenses this file | |
| # to you under the Apache License, Version 2.0 (the | |
| # "License"); you may not use this file except in compliance | |
| # with the License. You may obtain a copy of the License at | |
| # | |
| # http://www.apache.org/licenses/LICENSE-2.0 | |
| # | |
| # Unless required by applicable law or agreed to in writing, | |
| # software distributed under the License is distributed on an | |
| # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY | |
| # KIND, either express or implied. See the License for the | |
| # specific language governing permissions and limitations | |
| # under the License. | |
| """TVM Runtime NDArray API. | |
| tvm.ndarray provides a minimum runtime array API to test | |
| the correctness of the program. | |
| """ | |
| # pylint: disable=invalid-name,unused-import | |
| from __future__ import absolute_import as _abs | |
| import numpy as _np | |
| from ._ffi.ndarray import TVMContext, TVMType, NDArrayBase | |
| from ._ffi.ndarray import context, empty, from_dlpack | |
| from ._ffi.ndarray import _set_class_ndarray | |
| from ._ffi.ndarray import register_extension, free_extension_handle | |
| class NDArray(NDArrayBase): | |
| """Lightweight NDArray class of TVM runtime. | |
| Strictly this is only an Array Container (a buffer object) | |
| No arthimetic operations are defined. | |
| All operations are performed by TVM functions. | |
| The goal is not to re-build yet another array library. | |
| Instead, this is a minimal data structure to demonstrate | |
| how can we use TVM in existing project which might have their own array containers. | |
| """ | |
| def cpu(dev_id=0): | |
| """Construct a CPU device | |
| Parameters | |
| ---------- | |
| dev_id : int, optional | |
| The integer device id | |
| Returns | |
| ------- | |
| ctx : TVMContext | |
| The created context | |
| """ | |
| return TVMContext(1, dev_id) | |
| def gpu(dev_id=0): | |
| """Construct a CPU device | |
| Parameters | |
| ---------- | |
| dev_id : int, optional | |
| The integer device id | |
| Returns | |
| ------- | |
| ctx : TVMContext | |
| The created context | |
| """ | |
| return TVMContext(2, dev_id) | |
| def rocm(dev_id=0): | |
| """Construct a ROCM device | |
| Parameters | |
| ---------- | |
| dev_id : int, optional | |
| The integer device id | |
| Returns | |
| ------- | |
| ctx : TVMContext | |
| The created context | |
| """ | |
| return TVMContext(10, dev_id) | |
| def opencl(dev_id=0): | |
| """Construct a OpenCL device | |
| Parameters | |
| ---------- | |
| dev_id : int, optional | |
| The integer device id | |
| Returns | |
| ------- | |
| ctx : TVMContext | |
| The created context | |
| """ | |
| return TVMContext(4, dev_id) | |
| def metal(dev_id=0): | |
| """Construct a metal device | |
| Parameters | |
| ---------- | |
| dev_id : int, optional | |
| The integer device id | |
| Returns | |
| ------- | |
| ctx : TVMContext | |
| The created context | |
| """ | |
| return TVMContext(8, dev_id) | |
| def vpi(dev_id=0): | |
| """Construct a VPI simulated device | |
| Parameters | |
| ---------- | |
| dev_id : int, optional | |
| The integer device id | |
| Returns | |
| ------- | |
| ctx : TVMContext | |
| The created context | |
| """ | |
| return TVMContext(9, dev_id) | |
| def vulkan(dev_id=0): | |
| """Construct a Vulkan device | |
| Parameters | |
| ---------- | |
| dev_id : int, optional | |
| The integer device id | |
| Returns | |
| ------- | |
| ctx : TVMContext | |
| The created context | |
| """ | |
| return TVMContext(7, dev_id) | |
| def opengl(dev_id=0): | |
| """Construct a OpenGL device | |
| Parameters | |
| ---------- | |
| dev_id : int, optional | |
| The integer device id | |
| Returns | |
| ------- | |
| ctx : TVMContext | |
| The created context | |
| """ | |
| return TVMContext(11, dev_id) | |
| def ext_dev(dev_id=0): | |
| """Construct a extension device | |
| Parameters | |
| ---------- | |
| dev_id : int, optional | |
| The integer device id | |
| Returns | |
| ------- | |
| ctx : TVMContext | |
| The created context | |
| Note | |
| ---- | |
| This API is reserved for quick testing of new | |
| device by plugin device API as ext_dev. | |
| """ | |
| return TVMContext(12, dev_id) | |
| def micro_dev(dev_id=0): | |
| """Construct a micro device | |
| Parameters | |
| ---------- | |
| dev_id : int, optional | |
| The integer device id | |
| Returns | |
| ------- | |
| ctx : TVMContext | |
| The created context | |
| """ | |
| return TVMContext(13, dev_id) | |
| cl = opencl | |
| mtl = metal | |
| def array(arr, ctx=cpu(0)): | |
| """Create an array from source arr. | |
| Parameters | |
| ---------- | |
| arr : numpy.ndarray | |
| The array to be copied from | |
| ctx : TVMContext, optional | |
| The device context to create the array | |
| Returns | |
| ------- | |
| ret : NDArray | |
| The created array | |
| """ | |
| if not isinstance(arr, (_np.ndarray, NDArray)): | |
| arr = _np.array(arr) | |
| return empty(arr.shape, arr.dtype, ctx).copyfrom(arr) | |
| _set_class_ndarray(NDArray) | |