[MLIR:Python] Make DenseElementsAttr.get() only request a buffer format if no explicit type was provided.

[phawkins]

Summary

Not every NumPy type (e.g., the ml_dtypes.bfloat16 NumPy extension type) has a type in the Python buffer protocol, so exporting such a buffer with PyBUF_FORMAT may fail.

However, we don't care about the self-reported type of a buffer if the user provides an explicit type. In the case that an explicit type is provided, don't request the format from the buffer protocol, which allows arrays whose element types are unknown to the buffer protocol to be passed.

[jpienaar]

Thanks!

[phawkins]

Hold off merging for a minute.

mlir/lib/Bindings/Python/IRAttributes.cpp

701–703

Hold off a second before merging, I realized there's a bug here.

[phawkins]

Fixed bug where we didn't correctly error on all failing cases.

[phawkins]

This should be ready to merge now, I think.

mlir/lib/Bindings/Python/IRAttributes.cpp

 //===- IRAttributes.cpp - Exports builtin and standard attributes ---------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include <optional>
+#include <string_view>
 #include <utility>
 
 #include "IRModule.h"
 
 #include "PybindUtils.h"
 
+#include "llvm/ADT/ScopeExit.h"
+
 #include "mlir-c/BuiltinAttributes.h"
 #include "mlir-c/BuiltinTypes.h"
 #include "mlir/Bindings/Python/PybindAdaptors.h"
 
 namespace py = pybind11;
 using namespace mlir;
 using namespace mlir::python;
 
public:
   using PyConcreteAttribute::PyConcreteAttribute;
 
   static PyDenseElementsAttribute
   getFromBuffer(py::buffer array, bool signless,
                 std::optional<PyType> explicitType,
                 std::optional<std::vector<int64_t>> explicitShape,
                 DefaultingPyMlirContext contextWrapper) {
     // Request a contiguous view. In exotic cases, this will cause a copy.
-    int flags = PyBUF_C_CONTIGUOUS | PyBUF_FORMAT;
-    Py_buffer *view = new Py_buffer();
-    if (PyObject_GetBuffer(array.ptr(), view, flags) != 0) {
-      delete view;
+    int flags = PyBUF_ND;
+    if (!explicitType) {
+      flags |= PyBUF_FORMAT;
+    }
+    Py_buffer view;
+    if (PyObject_GetBuffer(array.ptr(), &view, flags) != 0) {
       throw py::error_already_set();
     }
-    py::buffer_info arrayInfo(view);
+    auto freeBuffer = llvm::make_scope_exit([&]() { PyBuffer_Release(&view); });
     SmallVector<int64_t> shape;
     if (explicitShape) {
       shape.append(explicitShape->begin(), explicitShape->end());
     } else {
-      shape.append(arrayInfo.shape.begin(),
-                   arrayInfo.shape.begin() + arrayInfo.ndim);
+      shape.append(view.shape, view.shape + view.ndim);
     }
 
     MlirAttribute encodingAttr = mlirAttributeGetNull();
     MlirContext context = contextWrapper->get();
 
     // Detect format codes that are suitable for bulk loading. This includes
     // all byte aligned integer and floating point types up to 8 bytes.
     // Notably, this excludes, bool (which needs to be bit-packed) and
     // other exotics which do not have a direct representation in the buffer
     // protocol (i.e. complex, etc).
     std::optional<MlirType> bulkLoadElementType;
     if (explicitType) {
       bulkLoadElementType = *explicitType;
-    } else if (arrayInfo.format == "f") {
+    } else {
+      std::string_view format(view.format);
+      if (format == "f") {
         // f32
-      assert(arrayInfo.itemsize == 4 && "mismatched array itemsize");
+        assert(view.itemsize == 4 && "mismatched array itemsize");
         bulkLoadElementType = mlirF32TypeGet(context);
-    } else if (arrayInfo.format == "d") {
+      } else if (format == "d") {
         // f64
-      assert(arrayInfo.itemsize == 8 && "mismatched array itemsize");
+        assert(view.itemsize == 8 && "mismatched array itemsize");
         bulkLoadElementType = mlirF64TypeGet(context);
-    } else if (arrayInfo.format == "e") {
+      } else if (format == "e") {
         // f16
-      assert(arrayInfo.itemsize == 2 && "mismatched array itemsize");
+        assert(view.itemsize == 2 && "mismatched array itemsize");
         bulkLoadElementType = mlirF16TypeGet(context);
-    } else if (isSignedIntegerFormat(arrayInfo.format)) {
-      if (arrayInfo.itemsize == 4) {
+      } else if (isSignedIntegerFormat(format)) {
+        if (view.itemsize == 4) {
           // i32
-        bulkLoadElementType = signless ? mlirIntegerTypeGet(context, 32)
-                                       : mlirIntegerTypeSignedGet(context, 32);
-      } else if (arrayInfo.itemsize == 8) {
+          bulkLoadElementType = signless
+                                    ? mlirIntegerTypeGet(context, 32)
+        } else if (view.itemsize == 8) {
           // i64
-        bulkLoadElementType = signless ? mlirIntegerTypeGet(context, 64)
-                                       : mlirIntegerTypeSignedGet(context, 64);
-      } else if (arrayInfo.itemsize == 1) {
+          bulkLoadElementType = signless
+                                    ? mlirIntegerTypeGet(context, 64)
+        } else if (view.itemsize == 1) {
           // i8
           bulkLoadElementType = signless ? mlirIntegerTypeGet(context, 8)
                                          : mlirIntegerTypeSignedGet(context, 8);
-      } else if (arrayInfo.itemsize == 2) {
+        } else if (view.itemsize == 2) {
           // i16
-        bulkLoadElementType = signless ? mlirIntegerTypeGet(context, 16)
-                                       : mlirIntegerTypeSignedGet(context, 16);
+          bulkLoadElementType = signless
+                                    ? mlirIntegerTypeGet(context, 16)
         }
-    } else if (isUnsignedIntegerFormat(arrayInfo.format)) {
-      if (arrayInfo.itemsize == 4) {
+      } else if (isUnsignedIntegerFormat(format)) {
+        if (view.itemsize == 4) {
           // unsigned i32
           bulkLoadElementType = signless
                                     ? mlirIntegerTypeGet(context, 32)
                                     : mlirIntegerTypeUnsignedGet(context, 32);
-      } else if (arrayInfo.itemsize == 8) {
+        } else if (view.itemsize == 8) {
           // unsigned i64
           bulkLoadElementType = signless
                                     ? mlirIntegerTypeGet(context, 64)
                                     : mlirIntegerTypeUnsignedGet(context, 64);
-      } else if (arrayInfo.itemsize == 1) {
+        } else if (view.itemsize == 1) {
           // i8
-        bulkLoadElementType = signless ? mlirIntegerTypeGet(context, 8)
-                                       : mlirIntegerTypeUnsignedGet(context, 8);
-      } else if (arrayInfo.itemsize == 2) {
+          bulkLoadElementType = signless
+                                    ? mlirIntegerTypeGet(context, 8)
+        } else if (view.itemsize == 2) {
           // i16
           bulkLoadElementType = signless
                                     ? mlirIntegerTypeGet(context, 16)
                                     : mlirIntegerTypeUnsignedGet(context, 16);
         }
       }
-    if (bulkLoadElementType) {
-      MlirType shapedType;
-      if (mlirTypeIsAShaped(*bulkLoadElementType)) {
-        if (explicitShape) {
-          throw std::invalid_argument("Shape can only be specified explicitly "
-                                      "when the type is not a shaped type.");
-        }
-        shapedType = *bulkLoadElementType;
-      } else {
-        shapedType = mlirRankedTensorTypeGet(
-            shape.size(), shape.data(), *bulkLoadElementType, encodingAttr);
-      }
-      size_t rawBufferSize = arrayInfo.size * arrayInfo.itemsize;
-      MlirAttribute attr = mlirDenseElementsAttrRawBufferGet(
-          shapedType, rawBufferSize, arrayInfo.ptr);
-      if (mlirAttributeIsNull(attr)) {
-        throw std::invalid_argument(
-            "DenseElementsAttr could not be constructed from the given buffer. "
-            "This may mean that the Python buffer layout does not match that "
-            "MLIR expected layout and is a bug.");
-      }
-      return PyDenseElementsAttribute(contextWrapper->getRef(), attr);
-    }
+      if (!bulkLoadElementType) {
+        throw std::invalid_argument(
+            std::string("unimplemented array format conversion from format: ") +
+            std::string(format));
+      }
+    }
+
+      shapedType = mlirRankedTensorTypeGet(shape.size(), shape.data(),
+                                           *bulkLoadElementType, encodingAttr);
+    size_t rawBufferSize = view.len;
+    MlirAttribute attr =
+        mlirDenseElementsAttrRawBufferGet(shapedType, rawBufferSize, view.buf);
 
-    throw std::invalid_argument(
-        std::string("unimplemented array format conversion from format: ") +
-        arrayInfo.format);
-  }
-
   static PyDenseElementsAttribute getSplat(const PyType &shapedType,
                                            PyAttribute &elementAttr) {
     auto contextWrapper =
         PyMlirContext::forContext(mlirTypeGetContext(shapedType));
     if (!mlirAttributeIsAInteger(elementAttr) &&
         !mlirAttributeIsAFloat(elementAttr)) {
       std::string message = "Illegal element type for DenseElementsAttr: ";
       message.append(py::repr(py::cast(elementAttr)));
c.def("__len__", &PyDenseElementsAttribute::dunderLen)
                  throw py::value_error(
                      "get_splat_value called on a non-splat attribute");
                return mlirDenseElementsAttrGetSplatValue(self);
              })
         .def_buffer(&PyDenseElementsAttribute::accessBuffer);
   }
 
 private:
-  static bool isUnsignedIntegerFormat(const std::string &format) {
+  static bool isUnsignedIntegerFormat(std::string_view format) {
     if (format.empty())
       return false;
     char code = format[0];
     return code == 'I' || code == 'B' || code == 'H' || code == 'L' ||
            code == 'Q';
   }
 
-  static bool isSignedIntegerFormat(const std::string &format) {
+  static bool isSignedIntegerFormat(std::string_view format) {
     if (format.empty())
       return false;
     char code = format[0];
     return code == 'i' || code == 'b' || code == 'h' || code == 'l' ||
            code == 'q';
   }
 
   template <typename Type>

mlir/test/python/ir/array_attributes.py

def testGetDenseElementsUnsupported():
     with Context():
         array = np.array([["hello", "goodbye"]])
         try:
             attr = DenseElementsAttr.get(array)
         except ValueError as e:
             # CHECK: unimplemented array format conversion from format:
             print(e)
 
+# CHECK-LABEL: TEST: testGetDenseElementsUnSupportedTypeOkIfExplicitTypeProvided
+@run
+def testGetDenseElementsUnSupportedTypeOkIfExplicitTypeProvided():
+    with Context():
+        array = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.int64)
+        # datetime64 specifically isn't important: it's just a 64-bit type that
+        # doesn't have a format under the Python buffer protocol. A more
+        # realistic example would be a NumPy extension type like the bfloat16
+        # type from the ml_dtypes package, which isn't a dependency of this
+        # test.
+        attr = DenseElementsAttr.get(array.view(np.datetime64),
+                                     type=IntegerType.get_signless(64))
+        # CHECK: dense<{{\[}}[1, 2, 3], [4, 5, 6]]> : tensor<2x3xi64>
+        print(attr)
+        # CHECK: {{\[}}[1 2 3]
+        # CHECK: {{\[}}4 5 6]]
+        print(np.array(attr))
+
 
 ################################################################################
 # Splats.
 ################################################################################
 
 # CHECK-LABEL: TEST: testGetDenseElementsSplatInt
 @run
 def testGetDenseElementsSplatInt():