[Mlir-commits] [mlir] [mlir] Add example of `printAlias` to test dialect (NFC) (PR #79232)

[Jeff Niu]

https://github.com/Mogball created https://github.com/llvm/llvm-project/pull/79232

Follow-up from previous pull request. Motivate the API change with an attribute that decides between sugaring a sub-attribute or using an alias

>From cee91f0f7e591ffe85ea6fd84f4799a205638385 Mon Sep 17 00:00:00 2001
From: Mogball <jeff at modular.com>
Date: Tue, 23 Jan 2024 23:23:21 +0000
Subject: [PATCH] [mlir] Add example of `printAlias` to test dialect (NFC)

Follow-up from previous pull request.
---
 mlir/test/IR/print-attr-type-aliases.mlir     |  7 ++++++
 mlir/test/lib/Dialect/Test/TestAttrDefs.td    |  9 ++++++--
 mlir/test/lib/Dialect/Test/TestAttributes.cpp | 22 +++++++++++++++++++
 3 files changed, 36 insertions(+), 2 deletions(-)

diff --git a/mlir/test/IR/print-attr-type-aliases.mlir b/mlir/test/IR/print-attr-type-aliases.mlir
index a3db1f06009d88e..8b6fc470c799017 100644
--- a/mlir/test/IR/print-attr-type-aliases.mlir
+++ b/mlir/test/IR/print-attr-type-aliases.mlir
@@ -54,3 +54,10 @@
 // CHECK: #loc1 = loc(fused<memref<1xi32>
 // CHECK-NOT: #map
 "test.op"() {alias_test = loc(fused<memref<1xi32, affine_map<(d0) -> (d0)>>>["test.mlir":10:8])} : () -> ()
+
+// -----
+
+// CHECK: #test.conditional_alias<hello>
+"test.op"() {attr = #test.conditional_alias<"hello">} : () -> ()
+// CHECK-NEXT: #test.conditional_alias<#test_encoding>
+"test.op"() {attr = #test.conditional_alias<"alias_test:tensor_encoding">} : () -> ()
diff --git a/mlir/test/lib/Dialect/Test/TestAttrDefs.td b/mlir/test/lib/Dialect/Test/TestAttrDefs.td
index 945c54c04d47ce8..40f035a3e3a4e5e 100644
--- a/mlir/test/lib/Dialect/Test/TestAttrDefs.td
+++ b/mlir/test/lib/Dialect/Test/TestAttrDefs.td
@@ -332,7 +332,12 @@ def TestCopyCount : Test_Attr<"TestCopyCount"> {
   let assemblyFormat = "`<` $copy_count `>`";
 }
 
-
-
+def TestConditionalAliasAttr : Test_Attr<"TestConditionalAlias"> {
+  let mnemonic = "conditional_alias";
+  let parameters = (ins "mlir::StringAttr":$value);
+  let assemblyFormat = [{
+    `<` custom<ConditionalAlias>($value) `>`
+  }];
+}
 
 #endif // TEST_ATTRDEFS
diff --git a/mlir/test/lib/Dialect/Test/TestAttributes.cpp b/mlir/test/lib/Dialect/Test/TestAttributes.cpp
index c240354e5d99044..1628da46de521fb 100644
--- a/mlir/test/lib/Dialect/Test/TestAttributes.cpp
+++ b/mlir/test/lib/Dialect/Test/TestAttributes.cpp
@@ -215,6 +215,28 @@ llvm::hash_code hash_value(const test::CopyCount &copyCount) {
   return llvm::hash_value(copyCount.value);
 }
 } // namespace test
+
+//===----------------------------------------------------------------------===//
+// TestConditionalAliasAttr
+//===----------------------------------------------------------------------===//
+
+/// Attempt to parse the conditionally-aliased string attribute as a keyword or string, else try to parse an alias.
+static ParseResult parseConditionalAlias(AsmParser &p, StringAttr &value) {
+  std::string str;
+  if (succeeded(p.parseOptionalKeywordOrString(&str))) {
+    value = StringAttr::get(p.getContext(), str);
+    return success();
+  }
+  return p.parseAttribute(value);
+}
+
+/// Print the string attribute as an alias if it has one, otherwise print it as a keyword if possible.
+static void printConditionalAlias(AsmPrinter &p, StringAttr value) {
+   if (succeeded(p.printAlias(value)))
+      return;
+   p.printKeywordOrString(value);
+}
+
 //===----------------------------------------------------------------------===//
 // Tablegen Generated Definitions
 //===----------------------------------------------------------------------===//