Spaces:
Sleeping
Sleeping
File size: 4,273 Bytes
66c9c8a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 | import ctypes
_c_str_dltensor = b"dltensor"
class DLDeviceType(ctypes.c_int):
"""The enum that encodes the type of the device where
DLTensor memory is allocated.
"""
kDLCPU = 1
kDLCUDA = 2
kDLCUDAHost = 3
kDLOpenCL = 4
kDLVulkan = 7
kDLMetal = 8
kDLVPI = 9
kDLROCM = 10
kDLROCMHost = 11
kDLCUDAManaged = 13
kDLOneAPI = 14
def __str__(self):
return {
self.kDLCPU: "CPU",
self.kDLCUDA: "CUDA",
self.kDLCUDAHost: "CUDAHost",
self.kDLOpenCL: "OpenCL",
self.kDLVulkan: "Vulkan",
self.kDLMetal: "Metal",
self.kDLVPI: "VPI",
self.kDLROCM: "ROCM",
self.kDLROCMHost: "ROMCHost",
self.kDLCUDAManaged: "CUDAManaged",
self.kDLOneAPI: "oneAPI",
}[self.value]
class DLDevice(ctypes.Structure):
"""Represents the device where DLTensor memory is allocated.
The device is represented by the pair of fields:
device_type: DLDeviceType
device_id: c_int
"""
_fields_ = [
("device_type", DLDeviceType),
("device_id", ctypes.c_int),
]
class DLDataTypeCode(ctypes.c_uint8):
"""An integer that encodes the category of DLTensor elements' data type."""
kDLInt = 0
kDLUInt = 1
kDLFloat = 2
kDLOpaquePointer = 3
kDLBfloat = 4
kDLComplex = 5
def __str__(self):
return {
self.kDLInt: "int",
self.kDLUInt: "uint",
self.kDLFloat: "float",
self.kDLBfloat: "bfloat",
self.kDLComplex: "complex",
self.kDLOpaquePointer: "void_p",
}[self.value]
class DLDataType(ctypes.Structure):
"""Descriptor of data type for elements of DLTensor.
The data type is described by a triple, `DLDataType.type_code`,
`DLDataType.bits`, and `DLDataType.lanes`.
The element is understood as packed `lanes` repetitions of
elements from `type_code` data-category of width `bits`.
"""
_fields_ = [
("type_code", DLDataTypeCode),
("bits", ctypes.c_uint8),
("lanes", ctypes.c_uint16),
]
TYPE_MAP = {
"bool": (DLDataTypeCode.kDLUInt, 1, 1),
"int8": (DLDataTypeCode.kDLInt, 8, 1),
"int16": (DLDataTypeCode.kDLInt, 16, 1),
"int32": (DLDataTypeCode.kDLInt, 32, 1),
"int64": (DLDataTypeCode.kDLInt, 64, 1),
"uint8": (DLDataTypeCode.kDLUInt, 8, 1),
"uint16": (DLDataTypeCode.kDLUInt, 16, 1),
"uint32": (DLDataTypeCode.kDLUInt, 32, 1),
"uint64": (DLDataTypeCode.kDLUInt, 64, 1),
"float16": (DLDataTypeCode.kDLFloat, 16, 1),
"float32": (DLDataTypeCode.kDLFloat, 32, 1),
"float64": (DLDataTypeCode.kDLFloat, 64, 1),
"complex64": (DLDataTypeCode.kDLComplex, 64, 1),
"complex128": (DLDataTypeCode.kDLComplex, 128, 1),
}
class DLTensor(ctypes.Structure):
"""Structure describing strided layout of DLTensor.
Fields are:
data: void pointer
device: DLDevice
ndim: number of indices needed to reference an
element of the tensor
dtype: data type descriptor
shape: tuple with lengths of the corresponding
tensor dimensions
strides: tuple of numbers of array elements to
step in each dimension when traversing
the tensor
byte_offset: data + byte_offset gives the address of
tensor element with index (0,) * ndim
"""
_fields_ = [
("data", ctypes.c_void_p),
("device", DLDevice),
("ndim", ctypes.c_int),
("dtype", DLDataType),
("shape", ctypes.POINTER(ctypes.c_int64)),
("strides", ctypes.POINTER(ctypes.c_int64)),
("byte_offset", ctypes.c_uint64),
]
class DLManagedTensor(ctypes.Structure):
"""Structure storing the pointer to the tensor descriptor,
deleter callable for the tensor descriptor, and pointer to
some additional data. These are stored in fields `dl_tensor`,
`deleter`, and `manager_ctx`."""
_fields_ = [
("dl_tensor", DLTensor),
("manager_ctx", ctypes.c_void_p),
("deleter", ctypes.CFUNCTYPE(None, ctypes.c_void_p)),
]
|