ndarray.py 10.8 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16
# 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.
17
# pylint: disable=invalid-name, unused-import
18
"""Runtime NDArray api"""
19
from __future__ import absolute_import
20 21

import sys
22 23
import ctypes
import numpy as np
eqy committed
24
from .base import _LIB, check_call, c_array, string_types, _FFI_MODE, c_str
25 26
from .runtime_ctypes import TVMType, TVMContext, TVMArray, TVMArrayHandle
from .runtime_ctypes import TypeCode, tvm_shape_index_t
27

28

29
IMPORT_EXCEPT = RuntimeError if _FFI_MODE == "cython" else ImportError
30

31 32 33 34 35
try:
    # pylint: disable=wrong-import-position
    if _FFI_MODE == "ctypes":
        raise ImportError()
    if sys.version_info >= (3, 0):
36
        from ._cy3.core import _set_class_ndarray, _make_array, _from_dlpack
37
        from ._cy3.core import NDArrayBase as _NDArrayBase
38
        from ._cy3.core import _reg_extension
39
    else:
40
        from ._cy2.core import _set_class_ndarray, _make_array, _from_dlpack
41
        from ._cy2.core import NDArrayBase as _NDArrayBase
42
        from ._cy2.core import _reg_extension
43 44
except IMPORT_EXCEPT:
    # pylint: disable=wrong-import-position
45
    from ._ctypes.ndarray import _set_class_ndarray, _make_array, _from_dlpack
46
    from ._ctypes.ndarray import NDArrayBase as _NDArrayBase
47
    from ._ctypes.ndarray import _reg_extension
48

49

50 51
def context(dev_type, dev_id=0):
    """Construct a TVM context with given device type and id.
52 53 54

    Parameters
    ----------
55 56
    dev_type: int or str
        The device type mask or name of the device.
57 58 59 60

    dev_id : int, optional
        The integer device id

61 62 63 64
    Returns
    -------
    ctx: TVMContext
        The corresponding context.
65

66 67 68 69
    Examples
    --------
    Context can be used to create reflection of context by
    string representation of the device type.
70

71
    .. code-block:: python
72

73 74 75
      assert tvm.context("cpu", 1) == tvm.cpu(1)
      assert tvm.context("gpu", 0) == tvm.gpu(0)
      assert tvm.context("cuda", 0) == tvm.gpu(0)
76
    """
77
    if isinstance(dev_type, string_types):
78 79 80 81 82 83 84
        if '-device=micro_dev' in dev_type:
            dev_type = 'micro_dev'
        else:
            dev_type = dev_type.split()[0]
            if dev_type not in TVMContext.STR2MASK:
                raise ValueError("Unknown device type %s" % dev_type)
            dev_type = TVMContext.STR2MASK[dev_type]
85
    return TVMContext(dev_type, dev_id)
86

eqy committed
87

88 89 90 91 92 93
def numpyasarray(np_data):
    """Return a TVMArray representation of a numpy array.
    """
    data = np_data
    assert data.flags['C_CONTIGUOUS']
    arr = TVMArray()
94
    shape = c_array(tvm_shape_index_t, data.shape)
95 96 97
    arr.data = data.ctypes.data_as(ctypes.c_void_p)
    arr.shape = shape
    arr.strides = None
98
    arr.dtype = TVMType(np.dtype(data.dtype).name)
99 100
    arr.ndim = data.ndim
    # CPU device
101
    arr.ctx = context(1, 0)
102 103
    return arr, shape

104 105

def empty(shape, dtype="float32", ctx=context(1, 0)):
106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123
    """Create an empty array given shape and device

    Parameters
    ----------
    shape : tuple of int
        The shape of the array

    dtype : type or str
        The data type of the array.

    ctx : TVMContext
        The context of the array

    Returns
    -------
    arr : tvm.nd.NDArray
        The array tvm supported.
    """
124 125
    shape = c_array(tvm_shape_index_t, shape)
    ndim = ctypes.c_int(len(shape))
126
    handle = TVMArrayHandle()
127
    dtype = TVMType(dtype)
128
    check_call(_LIB.TVMArrayAlloc(
129 130 131 132 133 134 135
        shape, ndim,
        ctypes.c_int(dtype.type_code),
        ctypes.c_int(dtype.bits),
        ctypes.c_int(dtype.lanes),
        ctx.device_type,
        ctx.device_id,
        ctypes.byref(handle)))
136
    return _make_array(handle, False, False)
137

eqy committed
138 139 140 141 142 143 144 145 146 147

def from_dlpack(dltensor):
    """Produce an array from a DLPack tensor without memory copy.
    Retreives the underlying DLPack tensor's pointer to create an array from the
    data. Removes the original DLPack tensor's destructor as now the array is
    responsible for destruction.

    Parameters
    ----------
    dltensor : DLPack tensor
148
        Input DLManagedTensor, can only be consumed once.
eqy committed
149 150 151 152 153 154

    Returns
    -------
    arr: tvm.nd.NDArray
        The array view of the tensor data.
    """
155
    return _from_dlpack(dltensor)
eqy committed
156 157


158
class NDArrayBase(_NDArrayBase):
159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175
    """A simple Device/CPU Array object in runtime."""

    @property
    def dtype(self):
        """Type of this array"""
        return str(self.handle.contents.dtype)

    @property
    def ctx(self):
        """context of this array"""
        return self.handle.contents.ctx

    @property
    def context(self):
        """context of this array"""
        return self.ctx

176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201
    def __hash__(self):
        return ctypes.cast(self.handle, ctypes.c_void_p).value

    def __eq__(self, other):
        return self.same_as(other)

    def __ne__(self, other):
        return not self.__eq__(other)

    def same_as(self, other):
        """Check object identity equality

        Parameters
        ----------
        other : object
            The other object to compare to

        Returns
        -------
        same : bool
            Whether other is same as self.
        """
        if not isinstance(other, NDArrayBase):
            return False
        return self.__hash__() == other.__hash__()

202 203 204 205 206 207 208 209 210 211
    def __setitem__(self, in_slice, value):
        """Set ndarray value"""
        if (not isinstance(in_slice, slice) or
                in_slice.start is not None
                or in_slice.stop is not None):
            raise ValueError('Array only support set from numpy array')
        if isinstance(value, NDArrayBase):
            if value.handle is not self.handle:
                value.copyto(self)
        elif isinstance(value, (np.ndarray, np.generic)):
212
            self.copyfrom(value)
213 214 215
        else:
            raise TypeError('type %s not supported' % str(type(value)))

216
    def copyfrom(self, source_array):
217 218 219 220 221 222
        """Peform an synchronize copy from the array.

        Parameters
        ----------
        source_array : array_like
            The data source we should like to copy from.
223 224 225 226 227

        Returns
        -------
        arr : NDArray
            Reference to self.
228
        """
229 230 231 232
        if isinstance(source_array, NDArrayBase):
            source_array.copyto(self)
            return self

233 234 235 236 237 238
        if not isinstance(source_array, np.ndarray):
            try:
                source_array = np.array(source_array, dtype=self.dtype)
            except:
                raise TypeError('array must be an array_like data,' +
                                'type %s is not supported' % str(type(source_array)))
239

240 241 242 243 244 245
        t = TVMType(self.dtype)
        shape, dtype = self.shape, self.dtype
        if t.lanes > 1:
            shape = shape + (t.lanes,)
            t.lanes = 1
            dtype = str(t)
246

247 248 249
        if source_array.shape != shape:
            raise ValueError("array shape do not match the shape of NDArray {0} vs {1}".format(
                source_array.shape, shape))
250
        source_array = np.ascontiguousarray(source_array, dtype=dtype)
251 252
        assert source_array.flags['C_CONTIGUOUS']
        data = source_array.ctypes.data_as(ctypes.c_void_p)
253
        nbytes = ctypes.c_size_t(source_array.size * source_array.dtype.itemsize)
254
        check_call(_LIB.TVMArrayCopyFromBytes(self.handle, data, nbytes))
255
        return self
256

257 258 259 260 261 262 263 264
    def __repr__(self):
        res = "<tvm.NDArray shape={0}, {1}>\n".format(self.shape, self.context)
        res += self.asnumpy().__repr__()
        return res

    def __str__(self):
        return str(self.asnumpy())

265 266 267 268 269 270 271 272
    def asnumpy(self):
        """Convert this array to numpy array

        Returns
        -------
        np_arr : numpy.ndarray
            The corresponding numpy array.
        """
273 274 275 276 277 278 279
        t = TVMType(self.dtype)
        shape, dtype = self.shape, self.dtype
        if t.lanes > 1:
            shape = shape + (t.lanes,)
            t.lanes = 1
            dtype = str(t)
        np_arr = np.empty(shape, dtype=dtype)
280 281
        assert np_arr.flags['C_CONTIGUOUS']
        data = np_arr.ctypes.data_as(ctypes.c_void_p)
282
        nbytes = ctypes.c_size_t(np_arr.size * np_arr.dtype.itemsize)
283
        check_call(_LIB.TVMArrayCopyToBytes(self.handle, data, nbytes))
284 285 286 287 288 289 290
        return np_arr

    def copyto(self, target):
        """Copy array to target

        Parameters
        ----------
291
        target : NDArray
292 293 294
            The target array to be copied, must have same shape as this array.
        """
        if isinstance(target, NDArrayBase):
295 296 297 298 299
            return self._copyto(target)
        elif isinstance(target, TVMContext):
            res = empty(self.shape, self.dtype, target)
            return self._copyto(res)
        raise ValueError("Unsupported target type %s" % str(type(target)))
300

eqy committed
301

302 303
def register_extension(cls, fcreate=None):
    """Register a extension class to TVM.
304 305 306 307 308 309 310

    After the class is registered, the class will be able
    to directly pass as Function argument generated by TVM.

    Parameters
    ----------
    cls : class
311
        The class object to be registered as extension.
312

313 314 315
    fcreate : function, optional
        The creation function to create a class object given handle value.

316 317
    Note
    ----
318 319 320 321 322
    The registered class is requires one property: _tvm_handle.

    If the registered class is a subclass of NDArray,
    it is required to have a class attribute _array_type_code.
    Otherwise, it is required to have a class attribute _tvm_tcode.
323 324

    - ```_tvm_handle``` returns integer represents the address of the handle.
325 326
    - ```_tvm_tcode``` or ```_array_type_code``` gives integer represents type
      code of the class.
327 328 329 330 331 332 333 334 335 336 337 338 339

    Returns
    -------
    cls : class
        The class being registered.

    Example
    -------
    The following code registers user defined class
    MyTensor to be DLTensor compatible.

    .. code-block:: python

340
       @tvm.register_extension
341
       class MyTensor(object):
342 343
           _tvm_tcode = tvm.TypeCode.ARRAY_HANDLE

344 345 346 347
           def __init__(self):
               self.handle = _LIB.NewDLTensor()

           @property
348
           def _tvm_handle(self):
349 350
               return self.handle.value
    """
351 352 353 354
    assert hasattr(cls, "_tvm_tcode")
    if fcreate and cls._tvm_tcode < TypeCode.EXT_BEGIN:
        raise ValueError("Cannot register create when extension tcode is same as buildin")
    _reg_extension(cls, fcreate)
355
    return cls