[Mlir-commits] [mlir] e3fd612 - [mlir] Add fully dynamic constructor to StridedLayoutAttr bindings

[Denys Shabalin]

Author: Denys Shabalin
Date: 2022-10-04T13:02:55Z
New Revision: e3fd612e99b626860eff34a5eabf8036d6a9ca8c

URL: https://github.com/llvm/llvm-project/commit/e3fd612e99b626860eff34a5eabf8036d6a9ca8c
DIFF: https://github.com/llvm/llvm-project/commit/e3fd612e99b626860eff34a5eabf8036d6a9ca8c.diff

LOG: [mlir] Add fully dynamic constructor to StridedLayoutAttr bindings

Reviewed By: ftynse

Differential Revision: https://reviews.llvm.org/D135139

Added: 
    

Modified: 
    mlir/lib/Bindings/Python/IRAttributes.cpp
    mlir/test/python/ir/attributes.py

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Bindings/Python/IRAttributes.cpp b/mlir/lib/Bindings/Python/IRAttributes.cpp
index e62f1550c6dc6..0c8c9b8ba79dd 100644
--- a/mlir/lib/Bindings/Python/IRAttributes.cpp
+++ b/mlir/lib/Bindings/Python/IRAttributes.cpp
@@ -1050,6 +1050,19 @@ class PyStridedLayoutAttribute
         },
         py::arg("offset"), py::arg("strides"), py::arg("context") = py::none(),
         "Gets a strided layout attribute.");
+    c.def_static(
+        "get_fully_dynamic",
+        [](int64_t rank, DefaultingPyMlirContext ctx) {
+          auto dynamic = mlirShapedTypeGetDynamicStrideOrOffset();
+          std::vector<int64_t> strides(rank);
+          std::fill(strides.begin(), strides.end(), dynamic);
+          MlirAttribute attr = mlirStridedLayoutAttrGet(
+              ctx->get(), dynamic, strides.size(), strides.data());
+          return PyStridedLayoutAttribute(ctx->getRef(), attr);
+        },
+        py::arg("rank"), py::arg("context") = py::none(),
+        "Gets a strided layout attribute with dynamic offset and strides of a "
+        "given rank.");
     c.def_property_readonly(
         "offset",
         [](PyStridedLayoutAttribute &self) {

diff  --git a/mlir/test/python/ir/attributes.py b/mlir/test/python/ir/attributes.py
index e0960e3f1c456..684d52c3ae28f 100644
--- a/mlir/test/python/ir/attributes.py
+++ b/mlir/test/python/ir/attributes.py
@@ -542,3 +542,14 @@ def testStridedLayoutAttr():
     print(attr.strides[1])
     # CHECK: 13
     print(attr.strides[2])
+
+    attr = StridedLayoutAttr.get_fully_dynamic(3)
+    dynamic = ShapedType.get_dynamic_stride_or_offset()
+    # CHECK: strided<[?, ?, ?], offset: ?>
+    print(attr)
+    # CHECK: offset is dynamic: True
+    print(f"offset is dynamic: {attr.offset == dynamic}")
+    # CHECK: rank: 3
+    print(f"rank: {len(attr.strides)}")
+    # CHECK: strides are dynamic: [True, True, True]
+    print(f"strides are dynamic: {[s == dynamic for s in attr.strides]}")