Implementation of the CUDA array interface for sharing GPU buffers with other Python libraries.
| 274 | // Implementation of the CUDA array interface for sharing GPU buffers with other |
| 275 | // Python libraries. |
| 276 | StatusOr<py::dict> PyLocalBufferCudaArrayInterface( |
| 277 | const PyLocalBuffer& buffer) { |
| 278 | if (buffer.device()->local_device_state()->executor()->platform_kind() != |
| 279 | se::PlatformKind::kCuda) { |
| 280 | return InvalidArgument( |
| 281 | "__cuda_array_interface__ is only defined for NVidia GPU buffers."); |
| 282 | } |
| 283 | if (!buffer.on_device_shape().IsArray()) { |
| 284 | return InvalidArgument( |
| 285 | "__cuda_array_interface__ is only defined for array buffers."); |
| 286 | } |
| 287 | if (buffer.on_host_shape().element_type() == BF16) { |
| 288 | return InvalidArgument( |
| 289 | "__cuda_array_interface__ is not supported for bfloat16 buffers."); |
| 290 | } |
| 291 | TF_RET_CHECK( |
| 292 | LayoutUtil::IsMonotonicWithDim0Major(buffer.on_host_shape().layout())); |
| 293 | TF_ASSIGN_OR_RETURN(ShapedBuffer shaped_buffer, buffer.AsShapedBuffer()); |
| 294 | |
| 295 | py::dict result; |
| 296 | result["shape"] = IntSpanToTuple(shaped_buffer.on_host_shape().dimensions()); |
| 297 | TF_ASSIGN_OR_RETURN(py::str typestr, |
| 298 | TypeDescriptorForPrimitiveType( |
| 299 | shaped_buffer.on_host_shape().element_type())); |
| 300 | result["typestr"] = std::move(typestr); |
| 301 | py::tuple data(2); |
| 302 | data[0] = py::int_( |
| 303 | absl::bit_cast<std::uintptr_t>(shaped_buffer.root_buffer().opaque())); |
| 304 | data[1] = py::bool_(true); // read-only |
| 305 | result["data"] = std::move(data); |
| 306 | result["version"] = py::int_(2); |
| 307 | return result; |
| 308 | } |
| 309 | |
| 310 | void BuildOpsSubmodule(py::module* m) { |
| 311 | // ops submodule, containing free functions that add operators to an |
nothing calls this directly
no test coverage detected