[mlir][python] Support buffer protocol for splat dense attributes

[rkayaith]

Summary

These can be made to work by setting the buffer strides to 0.

[rkayaith]

rebase

[stellaraccident]

Nice - thanks!

mlir/lib/Bindings/Python/IRAttributes.cpp

static PyDenseElementsAttribute getSplat(const PyType &shapedType,
     MlirAttribute elements =
         mlirDenseElementsAttrSplatGet(shapedType, elementAttr);
     return PyDenseElementsAttribute(contextWrapper->getRef(), elements);
   }
 
   intptr_t dunderLen() { return mlirElementsAttrGetNumElements(*this); }
 
   py::buffer_info accessBuffer() {
-    if (mlirDenseElementsAttrIsSplat(*this)) {
-      // TODO: Currently crashes the program.
-      // Reported as https://github.com/pybind/pybind11/issues/3336
-      throw std::invalid_argument(
-          "unsupported data type for conversion to Python buffer");
-    }
-
     MlirType shapedType = mlirAttributeGetType(*this);
     MlirType elementType = mlirShapedTypeGetElementType(shapedType);
     std::string format;
 
     if (mlirTypeIsAF32(elementType)) {
       // f32
       return bufferInfo<float>(shapedType);
     }
py::buffer_info bufferInfo(MlirType shapedType,
     Type *data = static_cast<Type *>(
         const_cast<void *>(mlirDenseElementsAttrGetRawData(*this)));
     // Prepare the shape for the buffer_info.
     SmallVector<intptr_t, 4> shape;
     for (intptr_t i = 0; i < rank; ++i)
       shape.push_back(mlirShapedTypeGetDimSize(shapedType, i));
     // Prepare the strides for the buffer_info.
     SmallVector<intptr_t, 4> strides;
-    intptr_t strideFactor = 1;
+    if (mlirDenseElementsAttrIsSplat(*this)) {
+      // Splats are special, only the single value is stored.
+      strides.assign(rank, 0);
+    } else {
       for (intptr_t i = 1; i < rank; ++i) {
-      strideFactor = 1;
-      for (intptr_t j = i; j < rank; ++j) {
+        intptr_t 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));
+    }
     std::string format;
     if (explicitFormat) {
       format = explicitFormat;
     } else {
       format = py::format_descriptor<Type>::format();
     }
     return py::buffer_info(data, sizeof(Type), format, rank, shape, strides,
                            /*readonly=*/true);

mlir/test/python/ir/array_attributes.py

 def testRepeatedValuesSplat():
   with Context():
     array = np.array([[1.0, 1.0, 1.0], [1.0, 1.0, 1.0]], dtype=np.float32)
     attr = DenseElementsAttr.get(array)
     # CHECK: dense<1.000000e+00> : tensor<2x3xf32>
     print(attr)
     # CHECK: is_splat: True
     print("is_splat:", attr.is_splat)
-    # TODO: Re-enable this once a solution is found to raising an exception
-    # from buffer protocol.
-    # Reported as https://github.com/pybind/pybind11/issues/3336
-    # print(np.array(attr))
+    # CHECK{LITERAL}: [[1. 1. 1.]
+    # CHECK{LITERAL}:  [1. 1. 1.]]
+    print(np.array(attr))
 
 
 # CHECK-LABEL: TEST: testNonSplat
 @run
 def testNonSplat():
   with Context():
     array = np.array([2.0, 1.0, 1.0], dtype=np.float32)
     attr = DenseElementsAttr.get(array)