[mlir] Allow dense array to be parsed with type elision

[Mogball]

Summary

This patch makes parsing dense arrays with type elision work properly.
If a ranked tensor type is supplied to parseAttribute on a dense
array, the element type is skipped. Moreover, if type elision is set to
AttrTypeElision::Must, the element type is elided.

For example, this allows

memref.global @z : memref<3xi32> = array<1, 2, 3>

Fixes #57433

Depends on D132758

[rriddle]

mlir/lib/IR/AsmPrinter.cpp

1869

Why is this printing through the os again? Why not print through the printer.

[Mogball]

I added some tests. PTAL

[Mogball]

mlir/lib/IR/AsmPrinter.cpp

1869

Isn't this inside the printer already? Where stuff is being sent directly through the stream.

[rriddle]

mlir/lib/IR/AsmPrinter.cpp

1869

It's in the printer file, but calling operator<<(raw_ostream & is going to create a new nested AsmPrinter to print the type (i.e. it's going to call Type::print(os), which isn't what you want when you are already in the printer).

[Mogball]

mlir/lib/IR/AsmPrinter.cpp

1869

Right. I should be calling printType. Good catch

[rkayaith]

Thanks for the fix. Seems like this has a lot of asm forms now though, depending on how it gets printed (array<i8: 1, 2, 3>, array<1, 2, 3>, [1, 2, 3]).

[Mogball]

use printType

[Mogball]

rebase

mlir/lib/AsmParser/AttributeParser.cpp

if (p.consumeIf(Token::integer)) {
     return p.emitError("expected integer or floating point literal");
   }
 
   append(result->bitcastToAPInt());
   return success();
 }
 
 /// Parse a dense array attribute.
-Attribute Parser::parseDenseArrayAttr(Type type) {
+Attribute Parser::parseDenseArrayAttr(Type attrType) {
   consumeToken(Token::kw_array);
   if (parseToken(Token::less, "expected '<' after 'array'"))
     return {};
 
-  // Only bool or integer and floating point elements divisible by bytes are
-  // supported.
   SMLoc typeLoc = getToken().getLoc();
-  if (!type && !(type = parseType()))
+  Type eltType;
+  // If an attribute type was provided, use its element type.
+  if (attrType) {
+    auto tensorType = attrType.dyn_cast<RankedTensorType>();
+    if (!tensorType) {
+      emitError(typeLoc, "dense array attribute expected ranked tensor type");
+      return {};
+    }
+    eltType = tensorType.getElementType();
+
+    // Otherwise, parse a type.
+  } else if (!(eltType = parseType())) {
     return {};
-  if (!type.isIntOrIndexOrFloat()) {
-    emitError(typeLoc, "expected integer or float type, got: ") << type;
+  }
+
+  // Only bool or integer and floating point elements divisible by bytes are
+  // supported.
+  if (!eltType.isIntOrIndexOrFloat()) {
+    emitError(typeLoc, "expected integer or float type, got: ") << eltType;
     return {};
   }
-  if (!type.isInteger(1) && type.getIntOrFloatBitWidth() % 8 != 0) {
+  if (!eltType.isInteger(1) && eltType.getIntOrFloatBitWidth() % 8 != 0) {
     emitError(typeLoc, "element type bitwidth must be a multiple of 8");
     return {};
   }
 
+  // If a type was provided, check that it matches the parsed type.
+  auto checkProvidedType = [&](DenseArrayAttr result) -> Attribute {
+    if (attrType && result.getType() != attrType) {
+      emitError(typeLoc, "expected attribute type ")
+          << attrType << " does not match parsed type " << result.getType();
+      return {};
+    }
+    return result;
+  };
+
   // Check for empty list.
-  if (consumeIf(Token::greater))
-    return DenseArrayAttr::get(RankedTensorType::get(0, type), {});
-  if (parseToken(Token::colon, "expected ':' after dense array type"))
+  if (consumeIf(Token::greater)) {
+    return checkProvidedType(
+        DenseArrayAttr::get(RankedTensorType::get(0, eltType), {}));
+  }
+  if (!attrType &&
+      parseToken(Token::colon, "expected ':' after dense array type"))
     return {};
 
-  DenseArrayElementParser eltParser(type);
-  if (type.isIntOrIndex()) {
+  DenseArrayElementParser eltParser(eltType);
+  if (eltType.isIntOrIndex()) {
     if (parseCommaSeparatedList(
             [&] { return eltParser.parseIntegerElement(*this); }))
       return {};
   } else {
     if (parseCommaSeparatedList(
             [&] { return eltParser.parseFloatElement(*this); }))
       return {};
   }
   if (parseToken(Token::greater, "expected '>' to close an array attribute"))
     return {};
-  return eltParser.getAttr();
+  return checkProvidedType(eltParser.getAttr());
 }
 
 /// Parse a dense elements attribute.
 Attribute Parser::parseDenseElementsAttr(Type attrType) {
   auto attribLoc = getToken().getLoc();
   consumeToken(Token::kw_dense);
   if (parseToken(Token::less, "expected '<' after 'dense'"))
     return nullptr;

mlir/lib/IR/AsmPrinter.cpp

if (printerFlags.shouldElideElementsAttr(sparseEltAttr.getIndices()) ||
         os << ", ";
         printDenseElementsAttr(sparseEltAttr.getValues(), /*allowHex=*/true);
       }
       os << '>';
     }
   } else if (auto stridedLayoutAttr = attr.dyn_cast<StridedLayoutAttr>()) {
     stridedLayoutAttr.print(os);
   } else if (auto denseArrayAttr = attr.dyn_cast<DenseArrayAttr>()) {
-    typeElision = AttrTypeElision::Must;
-    os << "array<" << denseArrayAttr.getType().getElementType();
+    os << "array<";
+    if (typeElision != AttrTypeElision::Must)
+      printType(denseArrayAttr.getType().getElementType());
     if (!denseArrayAttr.empty()) {
+      if (typeElision != AttrTypeElision::Must)
         os << ": ";
       printDenseArrayAttr(denseArrayAttr);
     }
     os << ">";
+    return;
   } else if (auto resourceAttr = attr.dyn_cast<DenseResourceElementsAttr>()) {
     os << "dense_resource<";
     printResourceHandle(resourceAttr.getRawHandle());
     os << ">";
   } else if (auto locAttr = attr.dyn_cast<LocationAttr>()) {
     printLocation(locAttr);
   } else {
     llvm::report_fatal_error("Unknown builtin attribute");

mlir/test/IR/attribute.mlir

"array.sizes"() {
     x1_ui0 = array<ui0: 0, 0>,
     x2_si0 = array<si0: 0, 0>,
     x3_i24 = array<i24: -42, 42, 8388607>,
     x4_ui24 = array<ui24: 16777215>,
     x5_si24 = array<si24: -8388608>,
     x6_bf16 = array<bf16: 1.2, 3.4>,
     x7_f16 = array<f16: 1., 3.>
   }: () -> ()
+
+  // CHECK: test.typed_attr tensor<4xi32> = array<1, 2, 3, 4>
+  test.typed_attr tensor<4xi32> = array<1, 2, 3, 4>
   return
 }
 
 // -----
 
 //===----------------------------------------------------------------------===//
 // Test SymbolRefAttr
 //===----------------------------------------------------------------------===//

mlir/test/IR/invalid-builtin-attributes.mlir

 
 // expected-error@below {{expected ':' after dense array type}}
 #attr = array<i8)
 
 // -----
 
 // expected-error@below {{expected '>' to close an array attribute}}
 #attr = array<i8: 1)
+
+// -----
+
+// expected-error@below {{dense array attribute expected ranked tensor type}}
+test.typed_attr i32 = array<1>
+
+// -----
+
+// expected-error@below {{does not match parsed type}}
+test.typed_attr tensor<1xi32> = array<>
+
+// -----
+
+// expected-error@below {{does not match parsed type}}
+test.typed_attr tensor<0xi32> = array<1>

mlir/test/lib/Dialect/Test/TestDialect.cpp

if (opName == "test.dialect_custom_format_fallback") {
     return [](Operation *op, OpAsmPrinter &printer) {
       printer.getStream() << " custom_format_fallback";
     };
   }
   return {};
 }
 
 //===----------------------------------------------------------------------===//
+// TypedAttrOp
+//===----------------------------------------------------------------------===//
+
+/// Parse an attribute with a given type.
+static ParseResult parseAttrElideType(AsmParser &parser, TypeAttr type,
+                                      Attribute &attr) {
+  return parser.parseAttribute(attr, type.getValue());
+}
+
+/// Print an attribute without its type.
+static void printAttrElideType(AsmPrinter &printer, Operation *op,
+                               TypeAttr type, Attribute attr) {
+  printer.printAttributeWithoutType(attr);
+}
+
+//===----------------------------------------------------------------------===//
 // TestBranchOp
 //===----------------------------------------------------------------------===//
 
 SuccessorOperands TestBranchOp::getSuccessorOperands(unsigned index) {
   assert(index == 0 && "invalid successor index");
   return SuccessorOperands(getTargetOperandsMutable());
 }
 

mlir/test/lib/Dialect/Test/TestOps.td

 }
 
 def StringElementsAttrOp : TEST_Op<"string_elements_attr"> {
   let arguments = (ins
       StringElementsAttr:$scalar_string_attr
   );
 }
 
+def TypedAttrOp : TEST_Op<"typed_attr"> {
+  let arguments = (ins TypeAttr:$type, AnyAttr:$attr);
+  let assemblyFormat = [{
+    attr-dict $type `=` custom<AttrElideType>(ref($type), $attr)
+  }];
+}
+
 def DenseArrayAttrOp : TEST_Op<"dense_array_attr"> {
   let arguments = (ins
     DenseBoolArrayAttr:$i1attr,
     DenseI8ArrayAttr:$i8attr,
     DenseI16ArrayAttr:$i16attr,
     DenseI32ArrayAttr:$i32attr,
     DenseI64ArrayAttr:$i64attr,
     DenseF32ArrayAttr:$f32attr,