diff --git a/mlir/lib/AsmParser/AttributeParser.cpp b/mlir/lib/AsmParser/AttributeParser.cpp --- a/mlir/lib/AsmParser/AttributeParser.cpp +++ b/mlir/lib/AsmParser/AttributeParser.cpp @@ -1244,10 +1244,20 @@ 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 @@ 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 --- a/mlir/lib/IR/AsmPrinter.cpp +++ b/mlir/lib/IR/AsmPrinter.cpp @@ -2141,7 +2141,9 @@ return; } else if (auto distinctAttr = llvm::dyn_cast(attr)) { os << "distinct[" << state.getDistinctState().getId(distinctAttr) << "]<"; - printAttribute(distinctAttr.getReferencedAttr()); + if (!llvm::isa(distinctAttr.getReferencedAttr())) { + printAttribute(distinctAttr.getReferencedAttr()); + } os << '>'; return; } else if (auto dictAttr = llvm::dyn_cast(attr)) { diff --git a/mlir/lib/IR/BuiltinDialect.cpp b/mlir/lib/IR/BuiltinDialect.cpp --- a/mlir/lib/IR/BuiltinDialect.cpp +++ b/mlir/lib/IR/BuiltinDialect.cpp @@ -60,10 +60,11 @@ os << "loc"; return AliasResult::OverridableAlias; } - if (llvm::isa(attr)) { - os << "distinct"; - return AliasResult::OverridableAlias; - } + if (auto distinct = llvm::dyn_cast(attr)) + if (!llvm::isa(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 --- 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]} : () -> ()