From e3fd612e99b626860eff34a5eabf8036d6a9ca8c Mon Sep 17 00:00:00 2001
From: Denys Shabalin <shabalin@google.com>
Date: Tue, 4 Oct 2022 10:58:38 +0000
Subject: [PATCH] [mlir] Add fully dynamic constructor to StridedLayoutAttr
 bindings

Reviewed By: ftynse

Differential Revision: https://reviews.llvm.org/D135139
---
 mlir/lib/Bindings/Python/IRAttributes.cpp | 13 +++++++++++++
 mlir/test/python/ir/attributes.py         | 11 +++++++++++
 2 files changed, 24 insertions(+)

diff --git a/mlir/lib/Bindings/Python/IRAttributes.cpp b/mlir/lib/Bindings/Python/IRAttributes.cpp
index e62f155..0c8c9b8 100644
--- a/mlir/lib/Bindings/Python/IRAttributes.cpp
+++ b/mlir/lib/Bindings/Python/IRAttributes.cpp
@@ -1050,6 +1050,19 @@ public:
         },
         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 e0960e3..684d52c 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]}")
-- 
2.7.4