[Mlir-commits] [mlir] 629460a - [mlir] Improve syntax of `distinct[n]<unit>`

[Markus Böck]

Author: Markus Böck
Date: 2023-07-14T08:26:54+02:00
New Revision: 629460a9b2e577a68d839b51f21351a03035d9ee

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

LOG: [mlir] Improve syntax of `distinct[n]<unit>`

In cases where memory is of less of a concern (e.g. small attributes where all instances have to be distinct by definition), using `DistinctAttr` with a unit attribute is a useful and conscious way of generating deterministic unique IDs.
The syntax as is however, makes them less useful to use, as it 1) always prints `<unit>` at the back and 2) always aliases them leading to not very useful `#distinct = distinct[n]<unit>` lines in the printer output.

This patch fixes that by special casing `UnitAttr` to simply elide the `unit` attribute in the back and not printing it as alias in that case.

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

Added: 
    

Modified: 
    mlir/lib/AsmParser/AttributeParser.cpp
    mlir/lib/IR/AsmPrinter.cpp
    mlir/lib/IR/BuiltinDialect.cpp
    mlir/test/IR/distinct-attr.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/AsmParser/AttributeParser.cpp b/mlir/lib/AsmParser/AttributeParser.cpp
index 20c89d342e13f2..3437ac9addc5ff 100644
--- a/mlir/lib/AsmParser/AttributeParser.cpp
+++ b/mlir/lib/AsmParser/AttributeParser.cpp
@@ -1244,10 +1244,20 @@ Attribute Parser::parseDistinctAttr(Type type) {
   if (parseToken(Token::r_square, "expected ']' to close distinct ID") ||
       parseToken(Token::less, "expected '<' after distinct ID"))
     return {};
-  Attribute referencedAttr = parseAttribute(type);
-  if (!referencedAttr) {
-    emitError("expected attribute");
-    return {};
+
+  Attribute referencedAttr;
+  if (getToken().is(Token::greater)) {
+    consumeToken();
+    referencedAttr = builder.getUnitAttr();
+  } else {
+    referencedAttr = parseAttribute(type);
+    if (!referencedAttr) {
+      emitError("expected attribute");
+      return {};
+    }
+
+    if (parseToken(Token::greater, "expected '>' to close distinct attribute"))
+      return {};
   }
 
   // Add the distinct attribute to the parser state, if it has not been parsed
@@ -1265,8 +1275,5 @@ Attribute Parser::parseDistinctAttr(Type type) {
     return {};
   }
 
-  if (parseToken(Token::greater, "expected '>' to close distinct attribute"))
-    return {};
-
   return it->getSecond();
 }

diff  --git a/mlir/lib/IR/AsmPrinter.cpp b/mlir/lib/IR/AsmPrinter.cpp
index 603e4143aae21d..78762082e05314 100644
--- a/mlir/lib/IR/AsmPrinter.cpp
+++ b/mlir/lib/IR/AsmPrinter.cpp
@@ -2141,7 +2141,9 @@ void AsmPrinter::Impl::printAttributeImpl(Attribute attr,
     return;
   } else if (auto distinctAttr = llvm::dyn_cast<DistinctAttr>(attr)) {
     os << "distinct[" << state.getDistinctState().getId(distinctAttr) << "]<";
-    printAttribute(distinctAttr.getReferencedAttr());
+    if (!llvm::isa<UnitAttr>(distinctAttr.getReferencedAttr())) {
+      printAttribute(distinctAttr.getReferencedAttr());
+    }
     os << '>';
     return;
   } else if (auto dictAttr = llvm::dyn_cast<DictionaryAttr>(attr)) {

diff  --git a/mlir/lib/IR/BuiltinDialect.cpp b/mlir/lib/IR/BuiltinDialect.cpp
index a0f6af389f19af..dcb1119fe52076 100644
--- a/mlir/lib/IR/BuiltinDialect.cpp
+++ b/mlir/lib/IR/BuiltinDialect.cpp
@@ -60,10 +60,11 @@ struct BuiltinOpAsmDialectInterface : public OpAsmDialectInterface {
       os << "loc";
       return AliasResult::OverridableAlias;
     }
-    if (llvm::isa<DistinctAttr>(attr)) {
-      os << "distinct";
-      return AliasResult::OverridableAlias;
-    }
+    if (auto distinct = llvm::dyn_cast<DistinctAttr>(attr))
+      if (!llvm::isa<UnitAttr>(distinct.getReferencedAttr())) {
+        os << "distinct";
+        return AliasResult::OverridableAlias;
+      }
     return AliasResult::NoAlias;
   }
 

diff  --git a/mlir/test/IR/distinct-attr.mlir b/mlir/test/IR/distinct-attr.mlir
index a9ec9372445971..e0727fec448913 100644
--- a/mlir/test/IR/distinct-attr.mlir
+++ b/mlir/test/IR/distinct-attr.mlir
@@ -20,3 +20,9 @@
 // CHECK:         distinct.attr = #[[DISTINCT2]]
 // CHECK-GENERIC: distinct.attr = distinct[2]<42 : i32>
 "test.op"() {distinct.attr = distinct[42]<42 : i32>} : () -> ()
+
+// CHECK: distinct.attr = distinct[3]<>
+"test.op"() {distinct.attr = distinct[3]<>} : () -> ()
+
+// CHECK: distinct.attr = distinct[4]<>
+"test.op"() {distinct.attr = distinct[4]<unit>} : () -> ()