| 242 | } // namespace |
| 243 | |
| 244 | StatusOr<py::capsule> BufferToDLPackManagedTensor(PyLocalBuffer* buffer) { |
| 245 | auto pack = absl::make_unique<DLPackTensor>(); |
| 246 | pack->buffer = buffer->DeviceBuffer(); |
| 247 | if (!pack->buffer) { |
| 248 | return InvalidArgument( |
| 249 | "Cannot convert deleted/invalid buffer to DLPack tensor."); |
| 250 | } |
| 251 | pack->tensor.manager_ctx = pack.get(); |
| 252 | pack->tensor.deleter = DLPackTensorDeleter; |
| 253 | DLTensor& dt = pack->tensor.dl_tensor; |
| 254 | if (buffer->on_device_shape().IsTuple()) { |
| 255 | return Unimplemented( |
| 256 | "unsafe_buffer_pointer is not implemented for tuple " |
| 257 | "buffers."); |
| 258 | } |
| 259 | TF_RET_CHECK(pack->buffer->device_memory().size() == 1); |
| 260 | dt.data = pack->buffer->device_memory().front().opaque(); |
| 261 | TF_ASSIGN_OR_RETURN(dt.ctx, DLContextForDevice(*buffer->device())); |
| 262 | dt.ctx.device_id = buffer->device()->local_device_state()->device_ordinal(); |
| 263 | dt.ndim = buffer->on_host_shape().dimensions_size(); |
| 264 | TF_ASSIGN_OR_RETURN(dt.dtype, PrimitiveTypeToDLDataType( |
| 265 | buffer->on_host_shape().element_type())); |
| 266 | |
| 267 | pack->shape = std::vector<int64>(buffer->on_host_shape().dimensions().begin(), |
| 268 | buffer->on_host_shape().dimensions().end()); |
| 269 | pack->strides = StridesForShape(buffer->on_host_shape()); |
| 270 | dt.shape = reinterpret_cast<std::int64_t*>(pack->shape.data()); |
| 271 | dt.strides = reinterpret_cast<std::int64_t*>(pack->strides.data()); |
| 272 | dt.byte_offset = 0; |
| 273 | |
| 274 | py::capsule capsule(&pack.release()->tensor, kDlTensorCapsuleName, |
| 275 | [](PyObject* obj) { |
| 276 | DLManagedTensor* dlmt = static_cast<DLManagedTensor*>( |
| 277 | PyCapsule_GetPointer(obj, kDlTensorCapsuleName)); |
| 278 | if (dlmt) { |
| 279 | DLPackTensorDeleter(dlmt); |
| 280 | } else { |
| 281 | // The tensor has been deleted. Clear any error from |
| 282 | // PyCapsule_GetPointer. |
| 283 | PyErr_Clear(); |
| 284 | } |
| 285 | }); |
| 286 | |
| 287 | TF_RETURN_IF_ERROR(buffer->BlockHostUntilReady()); |
| 288 | return capsule; |
| 289 | } |
| 290 | |
| 291 | StatusOr<std::unique_ptr<PyLocalBuffer>> DLPackManagedTensorToBuffer( |
| 292 | const pybind11::capsule& tensor, std::shared_ptr<PyLocalClient> client) { |
nothing calls this directly
no test coverage detected