[mlir][Python] Finish adding RankedTensorType support for encoding.

[stellaraccident]

Rebase

[aartbik]

mlir/test/python/dialects/sparse_tensor/dialect.py

78

typo 2x: Tensor

86

do we want to check the contents too, or is that too detailed for these kinds of tests?

[stellaraccident]

Comments

[stellaraccident]

PTAL

mlir/test/python/dialects/sparse_tensor/dialect.py

86

I feel like that is covered elsewhere and the equality check should be enough here. Mainly, as the primary example of an encoding, this test is doing double duty as a test of the RankedTensorType machinery for manipulating encodings.

Mainly making sure there is a test that uses encoding= and .encoding.

mlir/include/mlir-c/BuiltinTypes.h

MLIR_CAPI_EXPORTED MlirType mlirRankedTensorTypeGet(intptr_t rank,
                                                     MlirAttribute encoding);
 
 /// Same as "mlirRankedTensorTypeGet" but returns a nullptr wrapping MlirType on
 /// illegal arguments, emitting appropriate diagnostics.
 MLIR_CAPI_EXPORTED MlirType mlirRankedTensorTypeGetChecked(
     MlirLocation loc, intptr_t rank, const int64_t *shape, MlirType elementType,
     MlirAttribute encoding);
 
+/// Gets the 'encoding' attribute from the ranked tensor type, returning a null
+/// attribute if none.
+MLIR_CAPI_EXPORTED MlirAttribute mlirRankedTensorTypeGetEncoding(MlirType type);
+
 /// Creates an unranked tensor type with the given element type in the same
 /// context as the element type. The type is owned by the context.
 MLIR_CAPI_EXPORTED MlirType mlirUnrankedTensorTypeGet(MlirType elementType);
 
 /// Same as "mlirUnrankedTensorTypeGet" but returns a nullptr wrapping MlirType
 /// on illegal arguments, emitting appropriate diagnostics.
 MLIR_CAPI_EXPORTED MlirType
 mlirUnrankedTensorTypeGetChecked(MlirLocation loc, MlirType elementType);

mlir/lib/Bindings/Python/IRTypes.cpp

public:
   static constexpr IsAFunctionTy isaFunction = mlirTypeIsARankedTensor;
   static constexpr const char *pyClassName = "RankedTensorType";
   using PyConcreteType::PyConcreteType;
 
   static void bindDerived(ClassTy &c) {
     c.def_static(
         "get",
         [](std::vector<int64_t> shape, PyType &elementType,
+           llvm::Optional<PyAttribute> &encodingAttr,
            DefaultingPyLocation loc) {
-          MlirAttribute encodingAttr = mlirAttributeGetNull();
           MlirType t = mlirRankedTensorTypeGetChecked(
-              loc, shape.size(), shape.data(), elementType, encodingAttr);
+              loc, shape.size(), shape.data(), elementType,
+              encodingAttr ? encodingAttr->get() : mlirAttributeGetNull());
           // TODO: Rework error reporting once diagnostic engine is exposed
           // in C API.
           if (mlirTypeIsNull(t)) {
             throw SetPyError(
                 PyExc_ValueError,
                 Twine("invalid '") +
                     py::repr(py::cast(elementType)).cast<std::string>() +
                     "' and expected floating point, integer, vector or "
                     "complex "
                     "type.");
           }
           return PyRankedTensorType(elementType.getContext(), t);
         },
-        py::arg("shape"), py::arg("element_type"), py::arg("loc") = py::none(),
+        py::arg("shape"), py::arg("element_type"),
+        py::arg("encoding") = py::none(), py::arg("loc") = py::none(),
         "Create a ranked tensor type");
+    c.def_property_readonly(
+        "encoding",
+        [](PyRankedTensorType &self) -> llvm::Optional<PyAttribute> {
+          MlirAttribute encoding = mlirRankedTensorTypeGetEncoding(self.get());
+          if (mlirAttributeIsNull(encoding))
+            return llvm::None;
+          return PyAttribute(self.getContext(), encoding);
+        });
   }
 };
 
 /// Unranked Tensor Type subclass - UnrankedTensorType.
 class PyUnrankedTensorType
     : public PyConcreteType<PyUnrankedTensorType, PyShapedType> {
 public:
   static constexpr IsAFunctionTy isaFunction = mlirTypeIsAUnrankedTensor;

mlir/lib/CAPI/IR/BuiltinTypes.cpp

MlirType mlirRankedTensorTypeGetChecked(MlirLocation loc, intptr_t rank,
                                         const int64_t *shape,
                                         MlirType elementType,
                                         MlirAttribute encoding) {
   return wrap(RankedTensorType::getChecked(
       unwrap(loc), llvm::makeArrayRef(shape, static_cast<size_t>(rank)),
       unwrap(elementType), unwrap(encoding)));
 }
 
+MlirAttribute mlirRankedTensorTypeGetEncoding(MlirType type) {
+  return wrap(unwrap(type).cast<RankedTensorType>().getEncoding());
+}
+
 MlirType mlirUnrankedTensorTypeGet(MlirType elementType) {
   return wrap(UnrankedTensorType::get(unwrap(elementType)));
 }
 
 MlirType mlirUnrankedTensorTypeGetChecked(MlirLocation loc,
                                           MlirType elementType) {
   return wrap(UnrankedTensorType::getChecked(unwrap(loc), unwrap(elementType)));
 }

mlir/test/CAPI/ir.c

static int printBuiltinTypes(MlirContext ctx) {
   mlirTypeDump(vector);
   fprintf(stderr, "\n");
   // CHECK: vector<2x3xf32>
 
   // Ranked tensor type.
   MlirType rankedTensor = mlirRankedTensorTypeGet(
       sizeof(shape) / sizeof(int64_t), shape, f32, mlirAttributeGetNull());
   if (!mlirTypeIsATensor(rankedTensor) ||
-      !mlirTypeIsARankedTensor(rankedTensor))
+      !mlirTypeIsARankedTensor(rankedTensor) ||
+      !mlirAttributeIsNull(mlirRankedTensorTypeGetEncoding(rankedTensor)))
     return 16;
   mlirTypeDump(rankedTensor);
   fprintf(stderr, "\n");
   // CHECK: tensor<2x3xf32>
 
   // Unranked tensor type.
   MlirType unrankedTensor = mlirUnrankedTensorTypeGet(f32);
   if (!mlirTypeIsATensor(unrankedTensor) ||

mlir/test/python/dialects/sparse_tensor/dialect.py

with Context() as ctx:
     # CHECK: index_bit_width: 32
     print(f"index_bit_width: {casted.index_bit_width}")
 
     created = st.EncodingAttr.get(casted.dim_level_types, casted.dim_ordering,
         16, 32)
     print(created)
     # CHECK: created_equal: True
     print(f"created_equal: {created == casted}")
+
+
+# CHECK-LABEL: TEST: testEncodingAttrOnTensor
+@run
+def testEncodingAttrOnTensor():
+  with Context() as ctx, Location.unknown():
+    encoding = st.EncodingAttr(Attribute.parse(
+      '#sparse_tensor.encoding<{ dimLevelType = [ "compressed" ], '
+      'pointerBitWidth = 16, indexBitWidth = 32 }>'))
+    tt = RankedTensorType.get((1024,), F32Type.get(), encoding=encoding)
+    # CHECK: tensor<1024xf32, #sparse_tensor
+    print(tt)
+    # CHECK: #sparse_tensor.encoding
+    print(tt.encoding)
+    assert tt.encoding == encoding

mlir/test/python/ir/builtin_types.py

try:
       tensor_invalid = RankedTensorType.get(shape, none)
     except ValueError as e:
       # CHECK: invalid 'Type(none)' and expected floating point, integer, vector
       # CHECK: or complex type.
       print(e)
     else:
       print("Exception not produced")
 
+    # Encoding should be None.
+    assert RankedTensorType.get(shape, f32).encoding is None
+
 
 # CHECK-LABEL: TEST: testUnrankedTensorType
 @run
 def testUnrankedTensorType():
   with Context(), Location.unknown():
     f32 = F32Type.get()
     loc = Location.unknown()
     unranked_tensor = UnrankedTensorType.get(f32)