[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

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();
 }

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)) {

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;
   }
 

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>} : () -> ()