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

[llvmlistbot at llvm.org]

Author: Jeff Niu
Date: 2024-01-23T16:29:57-08:00
New Revision: df1e01b316b174b5aa10727bd2ee0ec1082a9f65

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

LOG: [mlir] Add example of `printAlias` to test dialect (NFC) (#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

Added: 
    

Modified: 
    mlir/test/IR/print-attr-type-aliases.mlir
    mlir/test/lib/Dialect/Test/TestAttrDefs.td
    mlir/test/lib/Dialect/Test/TestAttributes.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/test/IR/print-attr-type-aliases.mlir b/mlir/test/IR/print-attr-type-aliases.mlir
index a3db1f06009d88..162eacd0022832 100644
--- a/mlir/test/IR/print-attr-type-aliases.mlir
+++ b/mlir/test/IR/print-attr-type-aliases.mlir
@@ -54,3 +54,17 @@
 // 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])} : () -> ()
+
+// -----
+
+#unalias_me = "goodbye"
+#keep_aliased = "alias_test:dot_in_name"
+
+// 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">} : () -> ()
+// CHECK: #test.conditional_alias<goodbye>
+"test.op"() {attr = #test.conditional_alias<#unalias_me>} : () -> ()
+// CHECK-NEXT: #test.conditional_alias<#test2Ealias>
+"test.op"() {attr = #test.conditional_alias<#keep_aliased>} : () -> ()

diff  --git a/mlir/test/lib/Dialect/Test/TestAttrDefs.td b/mlir/test/lib/Dialect/Test/TestAttrDefs.td
index 945c54c04d47ce..40f035a3e3a4e5 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 c240354e5d9904..8819f6cbe94e2c 100644
--- a/mlir/test/lib/Dialect/Test/TestAttributes.cpp
+++ b/mlir/test/lib/Dialect/Test/TestAttributes.cpp
@@ -215,6 +215,30 @@ 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
 //===----------------------------------------------------------------------===//