[mlir] Get array from the dense elements attribute with buffer protocol.

[zhanghb97]

Summary

• Add mlirDenseElementsAttrGetRawData C API.
• Add def_buffer binding to PyDenseElementsAttribute.
• Implement the protocol to access the buffer.

[zhanghb97]

mlir/lib/Bindings/Python/IRModules.cpp

1695

I'm not sure about how to free the data pointer properly. Should we add a private member data and a destructor to class PyDenseElementsAttribute, or use a smart pointer to manage the dynamically allocated array? I think I need suggestions here.

[stellaraccident]

mlir/lib/Bindings/Python/IRModules.cpp

1695

I am pretty sure that we don't want to be creating this element-by-element copy and that we should instead have an API to get the const void* of the backing buffer of the attribute (basically the inverse of the mlirDenseElementsAttr*Get functions that take a pointer to a linear buffer). The C++ method to use is getRawData(). For known int/float data types, this is a pretty critical optimization, and then we can opt to use an element-by-element copy as a fallback if needed.

If you were to do that (just accessing the raw data), then you wouldn't have this allocated array to hang on to.

The other way to do this is to not have the attribute itself implement the buffer protocol but instead have a property array that returns an array for access. This is the approach I take in IREE (https://github.com/google/iree/blob/f9c11f20fc97154a7079a3809025241f5f467183/bindings/python/pyiree/rt/host_types.cc#L179). If doing that, then you get to set a base object (this attribute object) when constructing the array, thereby keeping the backing data alive.

[zhanghb97]

• Add mlirDenseElementsAttrGetRawData C API.
• Use the raw data to avoid the allocated array in the buffer protocol.

mlir/lib/Bindings/Python/IRModules.cpp

1695

At the very beginning, I saw the raw data way in the pybind11 doc, so I implement this first. Now I use the C API to avoid the allocated array, PTAL.

Also, I'm wondering that which method (raw data & property array) is better for our situation.

[ftynse]

mlir/include/mlir-c/StandardAttributes.h

297

There is already mlirAttributeGetType, I don't think we want to duplicate it on each subclass of MlirAttribute.

mlir/lib/Bindings/Python/IRModules.cpp

1628

Style nit: I wouldn't check for == 1 for "booleans" in C, the usual convention that 0 is false and non-zero (which may or may not be one) is true. And`mlirTypeIsAF32(elementType)` just reads better.

1636

Nit: braces must be symmetric, if one "block" uses them, all blocks should. Also, the code above isn't really trivial and would mandate braces anyway.

1692

Can't we extract rank from shapedType?

size looks unused in the function.

1698

Nit: why ssize_t? We use intptr_t for sizes in C bindings, let's be consistent.

1703

Why int64_t here? There are 3 different types in 5 lines, this just adds complexity.

[zhanghb97]

• Rebase.
• Remove the mlirElementsAttrGetType C API.
• Modify checking conditions and unify the use of braces of the if and if else in accessBuffer function.
• Unify the type into intptr_t in bufferInfo function.

[zhanghb97]

mlir/include/mlir-c/StandardAttributes.h

297

Remove the mlirElementsAttrGetType and use mlirAttributeGetType in the accessBuffer function.

mlir/lib/Bindings/Python/IRModules.cpp

1628

Remove the ==1 for all the checking.

1636

Add braces to all the blocks.

1692

Remove the size and rank parameters, extract rank from shapeType in the function.

1703

Unify the above types into intptr_t.

[zhanghb97]

Rebase.

[mehdi_amini]

mlir/include/mlir-c/StandardAttributes.h

409

Can you add C-unit-tests for these?

[zhanghb97]

Add unit tests for mlirDenseElementsAttrGetRawData C API.

[zhanghb97]

According to Stella's suggestion, there are two methods to get the array from dense elements attribute:

• Buffer protocol
• Array property

This patch shows the buffer protocol method, and I also implement the array property method in D91326.
Now, we can compare those two methods and choose one. Personally, I prefer the buffer protocol method, because I think the use of np.array(attr) is more pythonic than attr.array for an array.

[stellaraccident]

Thanks - I prefer this version, but have a material request wrt readonly access.

mlir/lib/Bindings/Python/IRModules.cpp

1697

This is a problem: by default, the buffer will be exposed as writable, which will corrupt the uniqued data in the context if ever done.

A readonly flag was added to pybind ~one year ago (11/23/2019), which means that we no longer can just build on an ancient pybind as from regular 2018 era LTS distros. Not the end of the world, but we should document it. I believe this puts the minimum version at v2.5.0. Since we haven't launched yet, we might as well advertise v2.6.0, since as far as I can tell, mainstream distros don't include either and going as recent as possible gives us headroom. Requiring v.2.6.0 will also let me resolve a hack elsewhere.

Can you:

• In the docs/Bindings/Python.md file, add a note to the pre-requisite for pybind11 noting that v2.6.0 is required
• Pass true to the readonly flag in the py::buffer_info constructor (it follows strides).

1699

Use llvm::SmallVector<intptr_t, 4> here and for strides.

1715

strides, /*readonly=*/true);

[zhanghb97]

• Set the buffer to read-only.
• Add note to Python binding doc to note that pybind11 v2.6.0 is required.
• Use llvm::SmallVector instead of std::vector.

[stellaraccident]

Thanks - a few nits. Maybe put a message on discord when you land indicating the new version requirement (and tag ftynse since he is working on this)

mlir/docs/Bindings/Python.md

10 ↗

(On Diff #305950)

Maybe just: "minimum version required: :2.6.0"

mlir/lib/Bindings/Python/IRModules.cpp

1696

Please use static_cast when casting to our from void* (it is stricter than reinterpret and explicitly works for void casts)

1697

Since const_cast is "impolite", add a comment: buffer is configured for read-only access below

[zhanghb97]

• Rebase.
• Use static_cast and add a comment.

mlir/include/mlir-c/StandardAttributes.h

 mlirElementsAttrIsValidIndex(MlirAttribute attr, intptr_t rank, uint64_t *idxs);
 
 /** Gets the total number of elements in the given elements attribute. In order
  * to iterate over the attribute, obtain its type, which must be a statically
  * shaped type and use its sizes to build a multi-dimensional index. */
 MLIR_CAPI_EXPORTED int64_t mlirElementsAttrGetNumElements(MlirAttribute attr);
 
 //===----------------------------------------------------------------------===//
 // Dense elements attribute.
 //===----------------------------------------------------------------------===//
 
 // TODO: decide on the interface and add support for complex elements.
 /* TODO: add support for APFloat and APInt to LLVM IR C API, then expose the
  * relevant functions here. */
 
 /// Checks whether the given attribute is a dense elements attribute.
 MLIR_CAPI_EXPORTED int mlirAttributeIsADenseElements(MlirAttribute attr);

 mlirDenseElementsAttrGetUInt64Value(MlirAttribute attr, intptr_t pos);
 MLIR_CAPI_EXPORTED float mlirDenseElementsAttrGetFloatValue(MlirAttribute attr,
                                                             intptr_t pos);
 MLIR_CAPI_EXPORTED double
 mlirDenseElementsAttrGetDoubleValue(MlirAttribute attr, intptr_t pos);
 MLIR_CAPI_EXPORTED MlirStringRef
 mlirDenseElementsAttrGetStringValue(MlirAttribute attr, intptr_t pos);
 
+/** Returns the raw data of the given dense elements attribute. */
+MLIR_CAPI_EXPORTED const void *
+mlirDenseElementsAttrGetRawData(MlirAttribute attr);
+
 //===----------------------------------------------------------------------===//
 // Opaque elements attribute.
 //===----------------------------------------------------------------------===//
 
 // TODO: expose Dialect to the bindings and implement accessors here.
 
 /// Checks whether the given attribute is an opaque elements attribute.
 MLIR_CAPI_EXPORTED int mlirAttributeIsAOpaqueElements(MlirAttribute attr);

mlir/lib/Bindings/Python/IRModules.cpp

if (!DerivedTy::isaFunction(orig.attr)) {
       throw SetPyError(PyExc_ValueError,
                        llvm::Twine("Cannot cast attribute to ") +
                            DerivedTy::pyClassName + " (from " + origRepr + ")");
     }
     return orig.attr;
   }
 
   static void bind(py::module &m) {
-    auto cls = ClassTy(m, DerivedTy::pyClassName);
+    auto cls = ClassTy(m, DerivedTy::pyClassName, py::buffer_protocol());
     cls.def(py::init<PyAttribute &>(), py::keep_alive<0, 1>());
     DerivedTy::bindDerived(cls);
   }
 
   /// Implemented by derived classes to add methods to the Python subclass.
   static void bindDerived(ClassTy &m) {}
 };
 
static PyDenseElementsAttribute getSplat(PyType shapedType,
 
     MlirAttribute elements =
         mlirDenseElementsAttrSplatGet(shapedType.type, elementAttr.attr);
     return PyDenseElementsAttribute(contextWrapper->getRef(), elements);
   }
 
   intptr_t dunderLen() { return mlirElementsAttrGetNumElements(attr); }
 
+  py::buffer_info accessBuffer() {
+    MlirType shapedType = mlirAttributeGetType(this->attr);
+    MlirType elementType = mlirShapedTypeGetElementType(shapedType);
+
+    if (mlirTypeIsAF32(elementType)) {
+      // f32
+      return bufferInfo(shapedType, mlirDenseElementsAttrGetFloatValue);
+    } else if (mlirTypeIsAF64(elementType)) {
+      // f64
+      return bufferInfo(shapedType, mlirDenseElementsAttrGetDoubleValue);
+    } else if (mlirTypeIsAInteger(elementType) &&
+               mlirIntegerTypeGetWidth(elementType) == 32) {
+      if (mlirIntegerTypeIsSignless(elementType) ||
+          mlirIntegerTypeIsSigned(elementType)) {
+        // i32
+        return bufferInfo(shapedType, mlirDenseElementsAttrGetInt32Value);
+      } else if (mlirIntegerTypeIsUnsigned(elementType)) {
+        // unsigned i32
+        return bufferInfo(shapedType, mlirDenseElementsAttrGetUInt32Value);
+      }
+    } else if (mlirTypeIsAInteger(elementType) &&
+               mlirIntegerTypeGetWidth(elementType) == 64) {
+      if (mlirIntegerTypeIsSignless(elementType) ||
+          mlirIntegerTypeIsSigned(elementType)) {
+        // i64
+        return bufferInfo(shapedType, mlirDenseElementsAttrGetInt64Value);
+      } else if (mlirIntegerTypeIsUnsigned(elementType)) {
+        // unsigned i64
+        return bufferInfo(shapedType, mlirDenseElementsAttrGetUInt64Value);
+      }
+    }
+
+    std::string message = "unimplemented array format.";
+    throw SetPyError(PyExc_ValueError, message);
+  }
+
   static void bindDerived(ClassTy &c) {
     c.def("__len__", &PyDenseElementsAttribute::dunderLen)
         .def_static("get", PyDenseElementsAttribute::getFromBuffer,
                     py::arg("array"), py::arg("signless") = true,
                     py::arg("context") = py::none(),
                     "Gets from a buffer or ndarray")
         .def_static("get_splat", PyDenseElementsAttribute::getSplat,
                     py::arg("shaped_type"), py::arg("element_attr"),
                     "Gets a DenseElementsAttr where all values are the same")
         .def_property_readonly("is_splat",
                                [](PyDenseElementsAttribute &self) -> bool {
                                  return mlirDenseElementsAttrIsSplat(self.attr);
-                               });
+                               })
+        .def_buffer(&PyDenseElementsAttribute::accessBuffer);
   }
 
 private:
   template <typename ElementTy>
   static MlirAttribute
   bulkLoad(MlirContext context,
            MlirAttribute (*ctor)(MlirType, intptr_t, ElementTy *),
            MlirType mlirElementType, py::buffer_info &arrayInfo) {
     SmallVector<int64_t, 4> shape(arrayInfo.shape.begin(),
                                   arrayInfo.shape.begin() + arrayInfo.ndim);
     auto shapedType =
         mlirRankedTensorTypeGet(shape.size(), shape.data(), mlirElementType);
     intptr_t numElements = arrayInfo.size;
     const ElementTy *contents = static_cast<const ElementTy *>(arrayInfo.ptr);
     return ctor(shapedType, numElements, contents);
   }
+
+  template <typename Type>
+  py::buffer_info bufferInfo(MlirType shapedType,
+                             Type (*value)(MlirAttribute, intptr_t)) {
+    intptr_t rank = mlirShapedTypeGetRank(shapedType);
+    // Prepare the data for the buffer_info.
+    Type *data = reinterpret_cast<Type *>(
+        const_cast<void *>(mlirDenseElementsAttrGetRawData(this->attr)));
+    // Prepare the shape for the buffer_info.
+    std::vector<intptr_t> shape;
+    for (intptr_t i = 0; i < rank; ++i)
+      shape.push_back(mlirShapedTypeGetDimSize(shapedType, i));
+    // Prepare the strides for the buffer_info.
+    std::vector<intptr_t> strides;
+    intptr_t strideFactor = 1;
+    for (intptr_t i = 1; i < rank; ++i) {
+      strideFactor = 1;
+      for (intptr_t j = i; j < rank; ++j) {
+        strideFactor *= mlirShapedTypeGetDimSize(shapedType, j);
+      }
+      strides.push_back(sizeof(Type) * strideFactor);
+    }
+    strides.push_back(sizeof(Type));
+    return py::buffer_info(data, sizeof(Type),
+                           py::format_descriptor<Type>::format(), rank, shape,
+                           strides);
+  }
 };
 
 /// Refinement of the PyDenseElementsAttribute for attributes containing integer
 /// (and boolean) values. Supports element access.
 class PyDenseIntElementsAttribute
     : public PyConcreteAttribute<PyDenseIntElementsAttribute,
                                  PyDenseElementsAttribute> {
 public:

mlir/lib/CAPI/IR/StandardAttributes.cpp

 MlirStringRef mlirDenseElementsAttrGetStringValue(MlirAttribute attr,
                                                   intptr_t pos) {
   return wrap(
       *(unwrap(attr).cast<DenseElementsAttr>().getValues<StringRef>().begin() +
         pos));
 }
 
 //===----------------------------------------------------------------------===//
+// Raw data accessors.
+
+const void *mlirDenseElementsAttrGetRawData(MlirAttribute attr) {
+  return static_cast<const void *>(
+      unwrap(attr).cast<DenseElementsAttr>().getRawData().data());
+}
+
+//===----------------------------------------------------------------------===//
 // Opaque elements attribute.
 //===----------------------------------------------------------------------===//
 
 int mlirAttributeIsAOpaqueElements(MlirAttribute attr) {
   return unwrap(attr).isa<OpaqueElementsAttr>();
 }
 
 //===----------------------------------------------------------------------===//

mlir/test/Bindings/Python/ir_array_attributes.py

 def testGetDenseElementsF32():
   with Context():
     array = np.array([[1.1, 2.2, 3.3], [4.4, 5.5, 6.6]], dtype=np.float32)
     attr = DenseElementsAttr.get(array)
     # CHECK: dense<{{\[}}[1.100000e+00, 2.200000e+00, 3.300000e+00], [4.400000e+00, 5.500000e+00, 6.600000e+00]]> : tensor<2x3xf32>
     print(attr)
     # CHECK: is_splat: False
     print("is_splat:", attr.is_splat)
+    # CHECK: {{\[}}[1.1 2.2 3.3]
+    # CHECK: {{\[}}4.4 5.5 6.6]]
+    print(np.array(attr))
 
 run(testGetDenseElementsF32)
 
 
 # CHECK-LABEL: TEST: testGetDenseElementsF64
 def testGetDenseElementsF64():
   with Context():
     array = np.array([[1.1, 2.2, 3.3], [4.4, 5.5, 6.6]], dtype=np.float64)
     attr = DenseElementsAttr.get(array)
     # CHECK: dense<{{\[}}[1.100000e+00, 2.200000e+00, 3.300000e+00], [4.400000e+00, 5.500000e+00, 6.600000e+00]]> : tensor<2x3xf64>
     print(attr)
+    # CHECK: {{\[}}[1.1 2.2 3.3]
+    # CHECK: {{\[}}4.4 5.5 6.6]]
+    print(np.array(attr))
 
 run(testGetDenseElementsF64)
 
 
 ### 32 bit integer arrays
 # CHECK-LABEL: TEST: testGetDenseElementsI32Signless
 def testGetDenseElementsI32Signless():
   with Context():
     array = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.int32)
     attr = DenseElementsAttr.get(array)
     # CHECK: dense<{{\[}}[1, 2, 3], [4, 5, 6]]> : tensor<2x3xi32>
     print(attr)
+    # CHECK: {{\[}}[1 2 3]
+    # CHECK: {{\[}}4 5 6]]
+    print(np.array(attr))
 
 run(testGetDenseElementsI32Signless)
 
 
 # CHECK-LABEL: TEST: testGetDenseElementsUI32Signless
 def testGetDenseElementsUI32Signless():
   with Context():
     array = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.uint32)
     attr = DenseElementsAttr.get(array)
     # CHECK: dense<{{\[}}[1, 2, 3], [4, 5, 6]]> : tensor<2x3xi32>
     print(attr)
+    # CHECK: {{\[}}[1 2 3]
+    # CHECK: {{\[}}4 5 6]]
+    print(np.array(attr))
 
 run(testGetDenseElementsUI32Signless)
 
 # CHECK-LABEL: TEST: testGetDenseElementsI32
 def testGetDenseElementsI32():
   with Context():
     array = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.int32)
     attr = DenseElementsAttr.get(array, signless=False)
     # CHECK: dense<{{\[}}[1, 2, 3], [4, 5, 6]]> : tensor<2x3xsi32>
     print(attr)
+    # CHECK: {{\[}}[1 2 3]
+    # CHECK: {{\[}}4 5 6]]
+    print(np.array(attr))
 
 run(testGetDenseElementsI32)
 
 
 # CHECK-LABEL: TEST: testGetDenseElementsUI32
 def testGetDenseElementsUI32():
   with Context():
     array = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.uint32)
     attr = DenseElementsAttr.get(array, signless=False)
     # CHECK: dense<{{\[}}[1, 2, 3], [4, 5, 6]]> : tensor<2x3xui32>
     print(attr)
+    # CHECK: {{\[}}[1 2 3]
+    # CHECK: {{\[}}4 5 6]]
+    print(np.array(attr))
 
 run(testGetDenseElementsUI32)
 
 
 ## 64bit integer arrays
 # CHECK-LABEL: TEST: testGetDenseElementsI64Signless
 def testGetDenseElementsI64Signless():
   with Context():
     array = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.int64)
     attr = DenseElementsAttr.get(array)
     # CHECK: dense<{{\[}}[1, 2, 3], [4, 5, 6]]> : tensor<2x3xi64>
     print(attr)
+    # CHECK: {{\[}}[1 2 3]
+    # CHECK: {{\[}}4 5 6]]
+    print(np.array(attr))
 
 run(testGetDenseElementsI64Signless)
 
 
 # CHECK-LABEL: TEST: testGetDenseElementsUI64Signless
 def testGetDenseElementsUI64Signless():
   with Context():
     array = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.uint64)
     attr = DenseElementsAttr.get(array)
     # CHECK: dense<{{\[}}[1, 2, 3], [4, 5, 6]]> : tensor<2x3xi64>
     print(attr)
+    # CHECK: {{\[}}[1 2 3]
+    # CHECK: {{\[}}4 5 6]]
+    print(np.array(attr))
 
 run(testGetDenseElementsUI64Signless)
 
 # CHECK-LABEL: TEST: testGetDenseElementsI64
 def testGetDenseElementsI64():
   with Context():
     array = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.int64)
     attr = DenseElementsAttr.get(array, signless=False)
     # CHECK: dense<{{\[}}[1, 2, 3], [4, 5, 6]]> : tensor<2x3xsi64>
     print(attr)
+    # CHECK: {{\[}}[1 2 3]
+    # CHECK: {{\[}}4 5 6]]
+    print(np.array(attr))
 
 run(testGetDenseElementsI64)
 
 
 # CHECK-LABEL: TEST: testGetDenseElementsUI64
 def testGetDenseElementsUI64():
   with Context():
     array = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.uint64)
     attr = DenseElementsAttr.get(array, signless=False)
     # CHECK: dense<{{\[}}[1, 2, 3], [4, 5, 6]]> : tensor<2x3xui64>
     print(attr)
+    # CHECK: {{\[}}[1 2 3]
+    # CHECK: {{\[}}4 5 6]]
+    print(np.array(attr))
 
 run(testGetDenseElementsUI64)