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

[zhanghb97]

Summary

• Add mlirDenseElementsAttrGetRawData C API.
• Add array property binding to PyDenseElementsAttribute.

[zhanghb97]

This patch implements array property for dense elements attribute to get array. There is also a buffer protocol method in D91021.
Let's compare and choose one.

mlir/include/mlir-c/StandardAttributes.h

 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

 #include "Globals.h"
 #include "PybindUtils.h"
 
 #include "mlir-c/Bindings/Python/Interop.h"
 #include "mlir-c/Registration.h"
 #include "mlir-c/StandardAttributes.h"
 #include "mlir-c/StandardTypes.h"
 #include "llvm/ADT/SmallVector.h"
+#include <pybind11/numpy.h>
 #include <pybind11/stl.h>
 
 namespace py = pybind11;
 using namespace mlir;
 using namespace mlir::python;
 
 using llvm::SmallVector;
 
static PyDenseElementsAttribute getSplat(PyType shapedType,
 
     MlirAttribute elements =
         mlirDenseElementsAttrSplatGet(shapedType.type, elementAttr.attr);
     return PyDenseElementsAttribute(contextWrapper->getRef(), elements);
   }
 
   intptr_t dunderLen() { return mlirElementsAttrGetNumElements(attr); }
 
+  py::array accessArray() {
+    MlirType shapedType = mlirAttributeGetType(this->attr);
+    MlirType elementType = mlirShapedTypeGetElementType(shapedType);
+
+    if (mlirTypeIsAF32(elementType)) {
+      // f32
+      return arrayInfo(py::dtype("f"), shapedType,
+                       mlirDenseElementsAttrGetFloatValue);
+    } else if (mlirTypeIsAF64(elementType)) {
+      // f64
+      return arrayInfo(py::dtype("d"), shapedType,
+                       mlirDenseElementsAttrGetDoubleValue);
+    } else if (mlirTypeIsAInteger(elementType) &&
+               mlirIntegerTypeGetWidth(elementType) == 32) {
+      if (mlirIntegerTypeIsSignless(elementType) ||
+          mlirIntegerTypeIsSigned(elementType)) {
+        // i32
+        return arrayInfo(py::dtype("i"), shapedType,
+                         mlirDenseElementsAttrGetInt32Value);
+      } else if (mlirIntegerTypeIsUnsigned(elementType)) {
+        // unsigned i32
+        return arrayInfo(py::dtype("I"), shapedType,
+                         mlirDenseElementsAttrGetUInt32Value);
+      }
+    } else if (mlirTypeIsAInteger(elementType) &&
+               mlirIntegerTypeGetWidth(elementType) == 64) {
+      if (mlirIntegerTypeIsSignless(elementType) ||
+          mlirIntegerTypeIsSigned(elementType)) {
+        // i64
+        return arrayInfo(py::dtype("l"), shapedType,
+                         mlirDenseElementsAttrGetInt64Value);
+      } else if (mlirIntegerTypeIsUnsigned(elementType)) {
+        // unsigned i64
+        return arrayInfo(py::dtype("L"), 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_property_readonly("array", &PyDenseElementsAttribute::accessArray);
   }
 
 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::array arrayInfo(const py::dtype dtype, 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));
+    // Set the base object to keep the backing data alive.
+    auto base = py::cast(this, py::return_value_policy::take_ownership);
+    return py::array(dtype, shape, strides, data, std::move(base));
+  }
 };
 
 /// 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(attr.array)
 
 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(attr.array)
 
 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(attr.array)
 
 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(attr.array)
 
 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(attr.array)
 
 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(attr.array)
 
 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(attr.array)
 
 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(attr.array)
 
 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(attr.array)
 
 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(attr.array)
 
 run(testGetDenseElementsUI64)
 

mlir/test/CAPI/ir.c

if (mlirDenseElementsAttrGetBoolSplatValue(splatBool) != 1 ||
       mlirDenseElementsAttrGetInt32SplatValue(splatInt32) != 1 ||
       mlirDenseElementsAttrGetUInt64SplatValue(splatUInt64) != 1 ||
       mlirDenseElementsAttrGetInt64SplatValue(splatInt64) != 1 ||
       fabsf(mlirDenseElementsAttrGetFloatSplatValue(splatFloat) - 1.0f) >
           1E-6f ||
       fabs(mlirDenseElementsAttrGetDoubleSplatValue(splatDouble) - 1.0) > 1E-6)
     return 17;
 
+  uint32_t *uint32RawData =
+      (uint32_t *)mlirDenseElementsAttrGetRawData(uint32Elements);
+  int32_t *int32RawData =
+      (int32_t *)mlirDenseElementsAttrGetRawData(int32Elements);
+  uint64_t *uint64RawData =
+      (uint64_t *)mlirDenseElementsAttrGetRawData(uint64Elements);
+  int64_t *int64RawData =
+      (int64_t *)mlirDenseElementsAttrGetRawData(int64Elements);
+  float *floatRawData =
+      (float *)mlirDenseElementsAttrGetRawData(floatElements);
+  double *doubleRawData =
+      (double *)mlirDenseElementsAttrGetRawData(doubleElements);
+  if (uint32RawData[0] != 0u || uint32RawData[1] != 1u ||
+      int32RawData[0] != 0 || int32RawData[1] != 1 ||
+      uint64RawData[0] != 0u || uint64RawData[1] != 1u ||
+      int64RawData[0] != 0 || int64RawData[1] != 1 ||
+      floatRawData[0] != 0.0f || floatRawData[1] != 1.0f ||
+      doubleRawData[0] != 0.0 || doubleRawData[1] != 1.0)
+    return 18;
+
   mlirAttributeDump(splatBool);
   mlirAttributeDump(splatUInt32);
   mlirAttributeDump(splatInt32);
   mlirAttributeDump(splatUInt64);
   mlirAttributeDump(splatInt64);
   mlirAttributeDump(splatFloat);
   mlirAttributeDump(splatDouble);
   // CHECK: dense<true> : tensor<1x2xi1>