From 34f72d91252b92e956b80b97e0d586e1ddce5221 Mon Sep 17 00:00:00 2001 From: Alex Zinenko Date: Tue, 2 Nov 2021 16:44:37 +0100 Subject: [PATCH] [mlir][python] expose the shape property of shaped types This has been missing in the original definition of shaped types. Reviewed By: gysit Differential Revision: https://reviews.llvm.org/D113025 --- mlir/lib/Bindings/Python/IRTypes.cpp | 13 +++++++++++++ mlir/test/python/ir/builtin_types.py | 5 +++++ 2 files changed, 18 insertions(+) diff --git a/mlir/lib/Bindings/Python/IRTypes.cpp b/mlir/lib/Bindings/Python/IRTypes.cpp index 1cfd799..89fdb1f 100644 --- a/mlir/lib/Bindings/Python/IRTypes.cpp +++ b/mlir/lib/Bindings/Python/IRTypes.cpp @@ -284,6 +284,19 @@ public: }, "Returns whether the given value is used as a placeholder for dynamic " "strides and offsets in shaped types."); + c.def_property_readonly( + "shape", + [](PyShapedType &self) { + self.requireHasRank(); + + std::vector shape; + 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."); } private: diff --git a/mlir/test/python/ir/builtin_types.py b/mlir/test/python/ir/builtin_types.py index c5b32e8..7d881b9 100644 --- a/mlir/test/python/ir/builtin_types.py +++ b/mlir/test/python/ir/builtin_types.py @@ -315,6 +315,9 @@ def testRankedTensorType(): # Encoding should be None. assert RankedTensorType.get(shape, f32).encoding is None + tensor = RankedTensorType.get(shape, f32) + assert tensor.shape == shape + # CHECK-LABEL: TEST: testUnrankedTensorType @run @@ -396,6 +399,8 @@ def testMemRefType(): else: print("Exception not produced") + assert memref.shape == shape + # CHECK-LABEL: TEST: testUnrankedMemRefType @run -- 2.7.4