From eb6c63cb0b6e6ead346f68e438f90ee0451906a3 Mon Sep 17 00:00:00 2001
From: Vladislav Vinogradov <vlad.vinogradov@intel.com>
Date: Wed, 7 Jul 2021 18:07:51 +0300
Subject: [PATCH] [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
---
 mlir/lib/IR/BuiltinTypes.cpp               | 2 ++
 mlir/test/IR/print-attr-type-aliases.mlir  | 4 ++++
 mlir/test/lib/Dialect/Test/TestDialect.cpp | 1 +
 3 files changed, 7 insertions(+)

diff --git a/mlir/lib/IR/BuiltinTypes.cpp b/mlir/lib/IR/BuiltinTypes.cpp
index 5dfb14b..24d5f53 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 7c343e6..8b4016d 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 02ad79a..18be25a 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();
-- 
2.7.4