| 193 | } |
| 194 | |
| 195 | DLManagedTensor* mgb::imperative::to_dlpack(const ValueRef src) { |
| 196 | DeviceTensorND dv = src.dev_tensor()->as_nd(true); |
| 197 | DLMTensor* TensorHandler(new DLMTensor); |
| 198 | size_t ndim = dv.shape().ndim; |
| 199 | TensorHandler->value = dv; |
| 200 | TensorHandler->tensor.manager_ctx = TensorHandler; |
| 201 | TensorHandler->tensor.deleter = &deleter; |
| 202 | TensorHandler->tensor.dl_tensor.data = TensorHandler->value.raw_ptr(); |
| 203 | TensorHandler->tensor.dl_tensor.device = get_dl_device(dv); |
| 204 | TensorHandler->tensor.dl_tensor.ndim = ndim; |
| 205 | TensorHandler->tensor.dl_tensor.dtype = get_dl_datatype(dv); |
| 206 | |
| 207 | auto src_shape = TensorHandler->value.layout().shape; |
| 208 | auto src_stride = TensorHandler->value.layout().stride; |
| 209 | for (size_t i = 0; i < ndim; i++) { |
| 210 | if (src_shape[i] > std::numeric_limits<int64_t>::max()) { |
| 211 | mgb_throw( |
| 212 | MegBrainError, "unsupported input shape: %s", |
| 213 | TensorHandler->value.layout().to_string().c_str()); |
| 214 | } |
| 215 | TensorHandler->shape[i] = static_cast<int64_t>(src_shape[i]); |
| 216 | TensorHandler->stride[i] = static_cast<int64_t>(src_stride[i]); |
| 217 | } |
| 218 | TensorHandler->tensor.dl_tensor.shape = TensorHandler->shape; |
| 219 | TensorHandler->tensor.dl_tensor.strides = TensorHandler->stride; |
| 220 | TensorHandler->tensor.dl_tensor.byte_offset = 0; |
| 221 | return &(TensorHandler->tensor); |
| 222 | } |
| 223 | |
| 224 | TensorShape ptr2shape(const int64_t* ptr, size_t ndim) { |
| 225 | TensorShape shape; |