| 177 | }; |
| 178 | |
| 179 | int PyLocalBufferGetBuffer(PyObject* exporter, Py_buffer* view, int flags) { |
| 180 | auto& buffer = |
| 181 | py::reinterpret_borrow<py::object>(exporter).cast<PyLocalBuffer&>(); |
| 182 | Status status = [&]() { |
| 183 | // Py_buffer objects are POD C structures, so we don't need to hold the GIL. |
| 184 | // Additionally we call BlockHostUntilReady() below, which may block. |
| 185 | py::gil_scoped_release gil_release; |
| 186 | |
| 187 | if (buffer.device()->platform_name() != "cpu") { |
| 188 | return InvalidArgument( |
| 189 | "Python buffer protocol is only defined for CPU buffers."); |
| 190 | } |
| 191 | if (!buffer.on_device_shape().IsArray()) { |
| 192 | return InvalidArgument( |
| 193 | "Python buffer protocol is only defined for array buffers."); |
| 194 | } |
| 195 | // If we allowed exports of formatted BF16 buffers, consumers would get |
| 196 | // confused about the type because there is no way to describe BF16 to |
| 197 | // Python. |
| 198 | if (buffer.on_host_shape().element_type() == BF16 && |
| 199 | ((flags & PyBUF_FORMAT) == PyBUF_FORMAT)) { |
| 200 | return InvalidArgument( |
| 201 | "bfloat16 buffer format not supported by Python buffer protocol."); |
| 202 | } |
| 203 | if ((flags & PyBUF_WRITEABLE) == PyBUF_WRITEABLE) { |
| 204 | return InvalidArgument("XLA buffers are read-only."); |
| 205 | } |
| 206 | std::shared_ptr<SharedDeviceBuffer> device_buffer = buffer.DeviceBuffer(); |
| 207 | if (!device_buffer) { |
| 208 | return InvalidArgument("Deleted buffer used in buffer protocol."); |
| 209 | } |
| 210 | const Shape& shape = buffer.on_host_shape(); |
| 211 | if (((flags & PyBUF_C_CONTIGUOUS) == PyBUF_C_CONTIGUOUS || |
| 212 | (flags & PyBUF_STRIDES) == PyBUF_ND) && |
| 213 | !LayoutUtil::IsMonotonicWithDim0Major(shape.layout())) { |
| 214 | return InvalidArgument("Buffer is not in C-contiguous layout."); |
| 215 | } else if ((flags & PyBUF_F_CONTIGUOUS) == PyBUF_F_CONTIGUOUS && |
| 216 | !LayoutUtil::IsMonotonicWithDim0Minor(shape.layout())) { |
| 217 | return InvalidArgument("Buffer is not in F-contiguous layout."); |
| 218 | } else if ((flags & PyBUF_ANY_CONTIGUOUS) == PyBUF_ANY_CONTIGUOUS && |
| 219 | !LayoutUtil::IsMonotonicWithDim0Major(shape.layout()) && |
| 220 | !LayoutUtil::IsMonotonicWithDim0Minor(shape.layout())) { |
| 221 | return InvalidArgument("Buffer is not in contiguous layout."); |
| 222 | } |
| 223 | std::memset(view, 0, sizeof(Py_buffer)); |
| 224 | CHECK_EQ(device_buffer->device_memory().size(), 1); |
| 225 | view->buf = |
| 226 | const_cast<void*>(device_buffer->device_memory().front().opaque()); |
| 227 | auto extra = absl::make_unique<ExtraBufferInfo>(); |
| 228 | extra->device_buffer = std::move(device_buffer); |
| 229 | view->itemsize = ShapeUtil::ByteSizeOfPrimitiveType(shape.element_type()); |
| 230 | view->len = ShapeUtil::ByteSizeOf(shape); |
| 231 | view->readonly = 1; |
| 232 | if ((flags & PyBUF_FORMAT) == PyBUF_FORMAT) { |
| 233 | TF_ASSIGN_OR_RETURN(extra->format, FormatDescriptorForPrimitiveType( |
| 234 | shape.element_type())); |
| 235 | view->format = const_cast<char*>(extra->format.c_str()); |
| 236 | } |
nothing calls this directly
no test coverage detected