Skip to content

Commit e3fd612

Browse files
committed
[mlir] Add fully dynamic constructor to StridedLayoutAttr bindings
Reviewed By: ftynse Differential Revision: https://reviews.llvm.org/D135139
1 parent ddff376 commit e3fd612

File tree

2 files changed

+24
-0
lines changed

2 files changed

+24
-0
lines changed

mlir/lib/Bindings/Python/IRAttributes.cpp

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1050,6 +1050,19 @@ class PyStridedLayoutAttribute
10501050
},
10511051
py::arg("offset"), py::arg("strides"), py::arg("context") = py::none(),
10521052
"Gets a strided layout attribute.");
1053+
c.def_static(
1054+
"get_fully_dynamic",
1055+
[](int64_t rank, DefaultingPyMlirContext ctx) {
1056+
auto dynamic = mlirShapedTypeGetDynamicStrideOrOffset();
1057+
std::vector<int64_t> strides(rank);
1058+
std::fill(strides.begin(), strides.end(), dynamic);
1059+
MlirAttribute attr = mlirStridedLayoutAttrGet(
1060+
ctx->get(), dynamic, strides.size(), strides.data());
1061+
return PyStridedLayoutAttribute(ctx->getRef(), attr);
1062+
},
1063+
py::arg("rank"), py::arg("context") = py::none(),
1064+
"Gets a strided layout attribute with dynamic offset and strides of a "
1065+
"given rank.");
10531066
c.def_property_readonly(
10541067
"offset",
10551068
[](PyStridedLayoutAttribute &self) {

mlir/test/python/ir/attributes.py

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -542,3 +542,14 @@ def testStridedLayoutAttr():
542542
print(attr.strides[1])
543543
# CHECK: 13
544544
print(attr.strides[2])
545+
546+
attr = StridedLayoutAttr.get_fully_dynamic(3)
547+
dynamic = ShapedType.get_dynamic_stride_or_offset()
548+
# CHECK: strided<[?, ?, ?], offset: ?>
549+
print(attr)
550+
# CHECK: offset is dynamic: True
551+
print(f"offset is dynamic: {attr.offset == dynamic}")
552+
# CHECK: rank: 3
553+
print(f"rank: {len(attr.strides)}")
554+
# CHECK: strides are dynamic: [True, True, True]
555+
print(f"strides are dynamic: {[s == dynamic for s in attr.strides]}")

0 commit comments

Comments
 (0)