Make it possible to create DenseElementsAttrs with arbitrary shaped types in Python bindings

[apaszke]

Summary

Right now the bindings assume that all DenseElementsAttrs correspond to tensor values,
making it impossible to create vector-typed constants. I didn't want to change the API
significantly, so I opted for reusing the current signature of .get. Its type argument
now accepts both element types (in which case shape and signless can be specified too),
or a shaped type, which specifies the full type of the created attr (shape cannot be specified
in that case).

[jpienaar]

Could you add a test to show both usages/avoid regressions?

[apaszke]

Added tests

mlir/lib/Bindings/Python/IRAttributes.cpp

if (explicitType) {
       } else if (arrayInfo.itemsize == 2) {
         // i16
         bulkLoadElementType = signless
                                   ? mlirIntegerTypeGet(context, 16)
                                   : mlirIntegerTypeUnsignedGet(context, 16);
       }
     }
     if (bulkLoadElementType) {
-      auto shapedType = mlirRankedTensorTypeGet(
+      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.");

mlir/test/python/dialects/builtin.py

 # RUN: %PYTHON %s | FileCheck %s
 
 from mlir.ir import *
 import mlir.dialects.builtin as builtin
 import mlir.dialects.func as func
+import numpy as np
 
 
 def run(f):
   print("\nTEST:", f.__name__)
   f()
   return f
 
 
def testFuncArgumentAccess():
   # CHECK: f32 {custom_dialect.res1 = 4.200000e+01 : f32},
   # CHECK: f32 {custom_dialect.res2 = 2.560000e+02 : f64})
   # CHECK: return %[[ARG0]], %[[ARG1]] : f32, f32
   #
   # CHECK: func @other_func(
   # CHECK: %{{.*}}: f32 {custom_dialect.foo = "qux"},
   # CHECK: %{{.*}}: f32)
   print(module)
+
+
+# CHECK-LABEL: testDenseElementsAttr
+@run
+def testDenseElementsAttr():
+  with Context(), Location.unknown():
+    values = np.arange(4, dtype=np.int32)
+    i32 = IntegerType.get_signless(32)
+    print(DenseElementsAttr.get(values, type=i32))
+    # CHECK{LITERAL}: dense<[0, 1, 2, 3]> : tensor<4xi32>
+    print(DenseElementsAttr.get(values, type=i32, shape=(2, 2)))
+    # CHECK{LITERAL}: dense<[[0, 1], [2, 3]]> : tensor<2x2xi32>
+    print(DenseElementsAttr.get(values, type=VectorType.get((2, 2), i32)))
+    # CHECK{LITERAL}: dense<[[0, 1], [2, 3]]> : vector<2x2xi32>