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

[zero9178]

Summary

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.

[mehdi_amini]

mlir/lib/AsmParser/AttributeParser.cpp

1249

getUnitAttr isn't "free", please do this in a else branch

1261

I would rather keep the empty <>, I'm worried of adding opportunities for parsing ambiguities.

mlir/test/IR/distinct-attr.mlir

28

The same CHECK-NOT should be here?

[zero9178]

Address review comments

[zero9178]

mlir/lib/AsmParser/AttributeParser.cpp

1261

Sure!

mlir/test/IR/distinct-attr.mlir

28

I am simply explicitly checking for <> now

mlir/lib/AsmParser/AttributeParser.cpp

if (!value) {
     emitError("expected an unsigned 64-bit integer");
     return {};
   }
 
   // Parse the referenced attribute.
   if (parseToken(Token::r_square, "expected ']' to close distinct ID") ||
       parseToken(Token::less, "expected '<' after distinct ID"))
     return {};
-  Attribute referencedAttr = parseAttribute(type);
+
+  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
   // before. Otherwise, check if the parsed reference attribute matches the one
   // found in the parser state.
   DenseMap<uint64_t, DistinctAttr> &distinctAttrs =
       state.symbols.distinctAttributes;
   auto it = distinctAttrs.find(*value);
   if (it == distinctAttrs.end()) {
     DistinctAttr distinctAttr = DistinctAttr::create(referencedAttr);
     it = distinctAttrs.try_emplace(*value, distinctAttr).first;
   } else if (it->getSecond().getReferencedAttr() != referencedAttr) {
     emitError("referenced attribute does not match previous definition: ")
         << it->getSecond().getReferencedAttr();
     return {};
   }
 
-  if (parseToken(Token::greater, "expected '>' to close distinct attribute"))
-    return {};
-
   return it->getSecond();
 }

mlir/lib/IR/AsmPrinter.cpp

void AsmPrinter::Impl::printAttributeImpl(Attribute attr,
   } else if (auto opaqueAttr = llvm::dyn_cast<OpaqueAttr>(attr)) {
     printDialectSymbol(os, "#", opaqueAttr.getDialectNamespace(),
                        opaqueAttr.getAttrData());
   } else if (llvm::isa<UnitAttr>(attr)) {
     os << "unit";
     return;
   } else if (auto distinctAttr = llvm::dyn_cast<DistinctAttr>(attr)) {
     os << "distinct[" << state.getDistinctState().getId(distinctAttr) << "]<";
+    if (!llvm::isa<UnitAttr>(distinctAttr.getReferencedAttr())) {
       printAttribute(distinctAttr.getReferencedAttr());
+    }
     os << '>';
     return;
   } else if (auto dictAttr = llvm::dyn_cast<DictionaryAttr>(attr)) {
     os << '{';
     interleaveComma(dictAttr.getValue(),
                     [&](NamedAttribute attr) { printNamedAttribute(attr); });
     os << '}';
 

mlir/lib/IR/BuiltinDialect.cpp

AliasResult getAlias(Attribute attr, raw_ostream &os) const override {
     if (llvm::isa<IntegerSetAttr>(attr)) {
       os << "set";
       return AliasResult::OverridableAlias;
     }
     if (llvm::isa<LocationAttr>(attr)) {
       os << "loc";
       return AliasResult::OverridableAlias;
     }
-    if (llvm::isa<DistinctAttr>(attr)) {
+    if (auto distinct = llvm::dyn_cast<DistinctAttr>(attr))
+      if (!llvm::isa<UnitAttr>(distinct.getReferencedAttr())) {
         os << "distinct";
         return AliasResult::OverridableAlias;
       }
     return AliasResult::NoAlias;
   }
 
   AliasResult getAlias(Type type, raw_ostream &os) const final {
     if (auto tupleType = llvm::dyn_cast<TupleType>(type)) {
       if (tupleType.size() > 16) {
         os << "tuple";
         return AliasResult::OverridableAlias;

mlir/test/IR/distinct-attr.mlir

 
 // CHECK:         distinct.attr = #[[DISTINCT0]]
 // CHECK-GENERIC: distinct.attr = distinct[0]<42 : i32>
 "test.op"() {distinct.attr = distinct[0]<42 : i32>} : () -> ()
 
 // 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>} : () -> ()