[mlir] Add Python bindings for StridedLayoutAttr

[ftynse]

mlir/lib/Bindings/Python/IRAttributes.cpp

1060

Is the value that represents dynamic stride/offset exposed somewhere?

[shabalin]

Make magic constant bindings on ShapedType public.

mlir/include/mlir-c/BuiltinAttributes.h

 //===----------------------------------------------------------------------===//
 // Strided layout attribute.
 //===----------------------------------------------------------------------===//
 
 // Checks wheather the given attribute is a strided layout attribute.
 MLIR_CAPI_EXPORTED bool mlirAttributeIsAStridedLayout(MlirAttribute attr);
 
 // Creates a strided layout attribute from given strides and offset.
-MLIR_CAPI_EXPORTED MlirAttribute mlirStridedLayoutAttrGet(MlirContext ctx,
-                                                          int64_t offset,
-                                                          intptr_t numStrides,
-                                                          int64_t *strides);
+MLIR_CAPI_EXPORTED MlirAttribute
+mlirStridedLayoutAttrGet(MlirContext ctx, int64_t offset, intptr_t numStrides,
+                         const int64_t *strides);
 
 // Returns the offset in the given strided layout layout attribute.
 MLIR_CAPI_EXPORTED int64_t mlirStridedLayoutAttrGetOffset(MlirAttribute attr);
 
 // Returns the number of strides in the given strided layout attribute.
 MLIR_CAPI_EXPORTED intptr_t
 mlirStridedLayoutAttrGetNumStrides(MlirAttribute attr);
 

mlir/lib/Bindings/Python/IRAttributes.cpp

c.def_static(
         [](DefaultingPyMlirContext context) {
           return PyUnitAttribute(context->getRef(),
                                  mlirUnitAttrGet(context->get()));
         },
         py::arg("context") = py::none(), "Create a Unit attribute.");
   }
 };
 
+/// Strided layout attribute subclass.
+class PyStridedLayoutAttribute
+    : public PyConcreteAttribute<PyStridedLayoutAttribute> {
+public:
+  static constexpr IsAFunctionTy isaFunction = mlirAttributeIsAStridedLayout;
+  static constexpr const char *pyClassName = "StridedLayoutAttr";
+  using PyConcreteAttribute::PyConcreteAttribute;
+
+  static void bindDerived(ClassTy &c) {
+    c.def_static(
+        "get",
+        [](int64_t offset, const std::vector<int64_t> strides,
+           DefaultingPyMlirContext ctx) {
+          MlirAttribute attr = mlirStridedLayoutAttrGet(
+              ctx->get(), offset, strides.size(), strides.data());
+          return PyStridedLayoutAttribute(ctx->getRef(), attr);
+        },
+        py::arg("offset"), py::arg("strides"), py::arg("context") = py::none(),
+        "Gets a strided layout attribute.");
+    c.def_property_readonly(
+        "offset",
+        [](PyStridedLayoutAttribute &self) {
+          return mlirStridedLayoutAttrGetOffset(self);
+        },
+        "Returns the value of the float point attribute");
+    c.def_property_readonly(
+        "strides",
+        [](PyStridedLayoutAttribute &self) {
+          intptr_t size = mlirStridedLayoutAttrGetNumStrides(self);
+          std::vector<int64_t> strides(size);
+          for (intptr_t i = 0; i < size; i++) {
+            strides[i] = mlirStridedLayoutAttrGetStride(self, i);
+          }
+          return strides;
+        },
+        "Returns the value of the float point attribute");
+  }
+};
+
 } // namespace
 
 void mlir::python::populateIRAttributes(py::module &m) {
   PyAffineMapAttribute::bind(m);
 
   PyDenseBoolArrayAttribute::bind(m);
   PyDenseBoolArrayAttribute::PyDenseArrayIterator::bind(m);
   PyDenseI8ArrayAttribute::bind(m);
void mlir::python::populateIRAttributes(py::module &m) {
   PyDictAttribute::bind(m);
   PyFlatSymbolRefAttribute::bind(m);
   PyOpaqueAttribute::bind(m);
   PyFloatAttribute::bind(m);
   PyIntegerAttribute::bind(m);
   PyStringAttribute::bind(m);
   PyTypeAttribute::bind(m);
   PyUnitAttribute::bind(m);
+
+  PyStridedLayoutAttribute::bind(m);
 }

mlir/lib/Bindings/Python/IRTypes.cpp

c.def_property_readonly(
           int64_t rank = mlirShapedTypeGetRank(self);
           shape.reserve(rank);
           for (int64_t i = 0; i < rank; ++i)
             shape.push_back(mlirShapedTypeGetDimSize(self, i));
           return shape;
         },
         "Returns the shape of the ranked shaped type as a list of integers.");
     c.def_static(
-        "_get_dynamic_size", []() { return mlirShapedTypeGetDynamicSize(); },
+        "get_dynamic_size", []() { return mlirShapedTypeGetDynamicSize(); },
         "Returns the value used to indicate dynamic dimensions in shaped "
         "types.");
     c.def_static(
-        "_get_dynamic_stride_or_offset",
+        "get_dynamic_stride_or_offset",
         []() { return mlirShapedTypeGetDynamicStrideOrOffset(); },
         "Returns the value used to indicate dynamic strides or offsets in "
         "shaped types.");
   }
 
 private:
   void requireHasRank() {
     if (!mlirShapedTypeHasRank(*this)) {

mlir/lib/CAPI/IR/BuiltinAttributes.cpp

 // Strided layout attribute.
 //===----------------------------------------------------------------------===//
 
 bool mlirAttributeIsAStridedLayout(MlirAttribute attr) {
   return unwrap(attr).isa<StridedLayoutAttr>();
 }
 
 MlirAttribute mlirStridedLayoutAttrGet(MlirContext ctx, int64_t offset,
-                                       intptr_t numStrides, int64_t *strides) {
+                                       intptr_t numStrides,
+                                       const int64_t *strides) {
   return wrap(StridedLayoutAttr::get(unwrap(ctx), offset,
                                      ArrayRef<int64_t>(strides, numStrides)));
 }
 
 int64_t mlirStridedLayoutAttrGetOffset(MlirAttribute attr) {
   return unwrap(attr).cast<StridedLayoutAttr>().getOffset();
 }
 
 intptr_t mlirStridedLayoutAttrGetNumStrides(MlirAttribute attr) {
   return static_cast<intptr_t>(
       unwrap(attr).cast<StridedLayoutAttr>().getStrides().size());
 }
 
 int64_t mlirStridedLayoutAttrGetStride(MlirAttribute attr, intptr_t pos) {
   return unwrap(attr).cast<StridedLayoutAttr>().getStrides()[pos];
 }

mlir/python/mlir/dialects/_structured_transform_ops_ext.py

def __init__(self,
     dimension = _get_int64_attr(dimension)
     if isinstance(split_point, int):
       split_point = _get_int64_attr(split_point)
 
     if isinstance(split_point, Attribute):
       static_split_point = split_point
       dynamic_split_point = None
     else:
-      static_split_point = _get_int64_attr(ShapedType._get_dynamic_size())
+      static_split_point = _get_int64_attr(ShapedType.get_dynamic_size())
       dynamic_split_point = _get_op_result_or_value(split_point)
 
     pdl_operation_type = pdl.OperationType.get()
     super().__init__(
         pdl_operation_type,
         pdl_operation_type,
         _get_op_result_or_value(target),
         dimension=dimension,
def __init__(self,
     else:
       for size in sizes:
         if isinstance(size, int):
           static_sizes.append(IntegerAttr.get(i64_type, size))
         elif isinstance(size, IntegerAttr):
           static_sizes.append(size)
         else:
           static_sizes.append(
-              IntegerAttr.get(i64_type, ShapedType._get_dynamic_size()))
+              IntegerAttr.get(i64_type, ShapedType.get_dynamic_size()))
           dynamic_sizes.append(_get_op_result_or_value(size))
       sizes_attr = ArrayAttr.get(static_sizes)
 
     num_loops = sum(
         v if v == 0 else 1 for v in self.__extract_values(sizes_attr))
     super().__init__(
         pdl_operation_type, [pdl_operation_type] * num_loops,
         _get_op_result_or_value(target),

mlir/test/python/ir/attributes.py

except RuntimeError as e:
       # CHECK: Error: Invalid attribute when attempting to create an ArrayAttribute
       print("Error: ", e)
 
   with Context():
     array = ArrayAttr.get([StringAttr.get("a"), StringAttr.get("b")])
     array = array + [StringAttr.get("c")]
     # CHECK: concat: ["a", "b", "c"]
     print("concat: ", array)
+
+
+# CHECK-LABEL: TEST: testStridedLayoutAttr
+@run
+def testStridedLayoutAttr():
+  with Context():
+    attr = StridedLayoutAttr.get(42, [5, 7, 13])
+    # CHECK: strided<[5, 7, 13], offset: 42>
+    print(attr)
+    # CHECK: 42
+    print(attr.offset)
+    # CHECK: 3
+    print(len(attr.strides))
+    # CHECK: 5
+    print(attr.strides[0])
+    # CHECK: 7
+    print(attr.strides[1])
+    # CHECK: 13
+    print(attr.strides[2])

mlir/test/python/ir/builtin_types.py

with Context() as ctx:
     ctx.allow_unregistered_dialects = True
     opaque = OpaqueType.get("dialect", "type")
     # CHECK: opaque type: !dialect.type
     print("opaque type:", opaque)
     # CHECK: dialect namespace: dialect
     print("dialect namespace:", opaque.dialect_namespace)
     # CHECK: data: type
     print("data:", opaque.data)
+
+
+# CHECK-LABEL: TEST: testShapedTypeConstants
+# Tests that ShapedType exposes magic value constants.
+@run
+def testShapedTypeConstants():
+  # CHECK: <class 'int'>
+  print(type(ShapedType.get_dynamic_size()))
+  # CHECK: <class 'int'>
+  print(type(ShapedType.get_dynamic_stride_or_offset()))