[Mlir-commits] [mlir] eb6c63c - [mlir] Fix RankedTensorType::walkImmediateSubElements method

[Vladislav Vinogradov]

Author: Vladislav Vinogradov
Date: 2021-07-26T11:49:25+03:00
New Revision: eb6c63cb0b6e6ead346f68e438f90ee0451906a3

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

LOG: [mlir] Fix RankedTensorType::walkImmediateSubElements method

Add 'enconding' attribute visitor.
Without it ASM printer doesn't use attribute aliases for 'enconding'.

Reviewed By: rriddle

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

Added: 
    

Modified: 
    mlir/lib/IR/BuiltinTypes.cpp
    mlir/test/IR/print-attr-type-aliases.mlir
    mlir/test/lib/Dialect/Test/TestDialect.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/IR/BuiltinTypes.cpp b/mlir/lib/IR/BuiltinTypes.cpp
index 5dfb14b60abf8..24d5f53eaa5ab 100644
--- a/mlir/lib/IR/BuiltinTypes.cpp
+++ b/mlir/lib/IR/BuiltinTypes.cpp
@@ -515,6 +515,8 @@ void RankedTensorType::walkImmediateSubElements(
     function_ref<void(Attribute)> walkAttrsFn,
     function_ref<void(Type)> walkTypesFn) const {
   walkTypesFn(getElementType());
+  if (Attribute encoding = getEncoding())
+    walkAttrsFn(encoding);
 }
 
 //===----------------------------------------------------------------------===//

diff  --git a/mlir/test/IR/print-attr-type-aliases.mlir b/mlir/test/IR/print-attr-type-aliases.mlir
index 7c343e6463761..8b4016d49ad56 100644
--- a/mlir/test/IR/print-attr-type-aliases.mlir
+++ b/mlir/test/IR/print-attr-type-aliases.mlir
@@ -17,3 +17,7 @@
 
 // CHECK-DAG: !tuple = type tuple<i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32>
 "test.op"() {alias_test = "alias_test:large_tuple"} : () -> (tuple<i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32>)
+
+// CHECK-DAG: #test_encoding = "alias_test:tensor_encoding"
+// CHECK-DAG: tensor<32xf32, #test_encoding>
+"test.op"() : () -> tensor<32xf32, "alias_test:tensor_encoding">

diff  --git a/mlir/test/lib/Dialect/Test/TestDialect.cpp b/mlir/test/lib/Dialect/Test/TestDialect.cpp
index 02ad79ad085de..18be25a32cada 100644
--- a/mlir/test/lib/Dialect/Test/TestDialect.cpp
+++ b/mlir/test/lib/Dialect/Test/TestDialect.cpp
@@ -67,6 +67,7 @@ struct TestOpAsmInterface : public OpAsmDialectInterface {
                   StringRef("test_alias_conflict0"))
             .Case("alias_test:sanitize_conflict_b",
                   StringRef("test_alias_conflict0_"))
+            .Case("alias_test:tensor_encoding", StringRef("test_encoding"))
             .Default(llvm::None);
     if (!aliasName)
       return failure();