[mlir] Print 0 element DenseElementsAttr as dense<> to fix parser bugs with expected shape.

[rriddle]

Summary

Depending on where the 0 dimension is within the shape, the parser will currently reject .mlir generated by the printer.

[GMNGeoffrey]

Not sure if my approval counts, but...

Thanks River :-D

[jpienaar]

Thanks (even removes need for regex in testing :-))

mlir/lib/IR/AsmPrinter.cpp

void ModulePrinter::printAttribute(Attribute attr,
   case StandardAttributes::SparseElements: {
     auto elementsAttr = attr.cast<SparseElementsAttr>();
     if (printerFlags.shouldElideElementsAttr(elementsAttr.getIndices()) ||
         printerFlags.shouldElideElementsAttr(elementsAttr.getValues())) {
       printElidedElementsAttr(os);
       break;
     }
     os << "sparse<";
-    printDenseIntOrFPElementsAttr(elementsAttr.getIndices(),
-                                  /*allowHex=*/false);
+    DenseIntElementsAttr indices = elementsAttr.getIndices();
+    if (indices.getNumElements() != 0) {
+      printDenseIntOrFPElementsAttr(indices, /*allowHex=*/false);
       os << ", ";
       printDenseElementsAttr(elementsAttr.getValues(), /*allowHex=*/true);
+    }
     os << '>';
     break;
   }
 
   // Location attributes.
   case StandardAttributes::CallSiteLocation:
   case StandardAttributes::FileLineColLocation:
   case StandardAttributes::FusedLocation:

 printDenseElementsAttrImpl(bool isSplat, ShapedType type, raw_ostream &os,
                            function_ref<void(unsigned)> printEltFn) {
   // Special case for 0-d and splat tensors.
   if (isSplat)
     return printEltFn(0);
 
   // Special case for degenerate tensors.
   auto numElements = type.getNumElements();
-  int64_t rank = type.getRank();
-  if (numElements == 0) {
-    for (int i = 0; i < rank; ++i)
-      os << '[';
-    for (int i = 0; i < rank; ++i)
-      os << ']';
+  if (numElements == 0)
     return;
-  }
 
   // We use a mixed-radix counter to iterate through the shape. When we bump a
   // non-least-significant digit, we emit a close bracket. When we next emit an
   // element we re-open all closed brackets.
 
   // The mixed-radix counter, with radices in 'shape'.
+  int64_t rank = type.getRank();
   SmallVector<unsigned, 4> counter(rank, 0);
   // The number of brackets that have been opened and not closed.
   unsigned openBrackets = 0;
 
   auto shape = type.getShape();
   auto bumpCounter = [&] {
     // Bump the least significant digit.
     ++counter[rank - 1];

mlir/lib/Parser/AttributeParser.cpp

 //===----------------------------------------------------------------------===//
 
 /// Parse a dense elements attribute.
 Attribute Parser::parseDenseElementsAttr(Type attrType) {
   consumeToken(Token::kw_dense);
   if (parseToken(Token::less, "expected '<' after 'dense'"))
     return nullptr;
 
-  // Parse the literal data.
+  // Parse the literal data if necessary.
   TensorLiteralParser literalParser(*this);
-  if (literalParser.parse(/*allowHex=*/true))
-    return nullptr;
-
-  if (parseToken(Token::greater, "expected '>'"))
+  if (!consumeIf(Token::greater)) {
+    if (literalParser.parse(/*allowHex=*/true) ||
+        parseToken(Token::greater, "expected '>'"))
       return nullptr;
+  }
 
   auto typeLoc = getToken().getLoc();
   auto type = parseElementsLiteralType(attrType);
   if (!type)
     return nullptr;
   return literalParser.getAttr(typeLoc, type);
 }
 

 }
 
 /// Parse a sparse elements attribute.
 Attribute Parser::parseSparseElementsAttr(Type attrType) {
   consumeToken(Token::kw_sparse);
   if (parseToken(Token::less, "Expected '<' after 'sparse'"))
     return nullptr;
 
+  // Check for the case where all elements are sparse. The indices are
+  // represented by a 2-dimensional shape where the second dimension is the rank
+  // of the type.
+  Type indiceEltType = builder.getIntegerType(64);
+  if (consumeIf(Token::greater)) {
+    ShapedType type = parseElementsLiteralType(attrType);
+    if (!type)
+      return nullptr;
+
+    // Construct the sparse elements attr using zero element indice/value
+    // attributes.
+    ShapedType indicesType =
+        RankedTensorType::get({0, type.getRank()}, indiceEltType);
+    ShapedType valuesType = RankedTensorType::get({0}, type.getElementType());
+    return SparseElementsAttr::get(
+        type, DenseElementsAttr::get(indicesType, ArrayRef<Attribute>()),
+        DenseElementsAttr::get(valuesType, ArrayRef<Attribute>()));
+  }
+
   /// Parse the indices. We don't allow hex values here as we may need to use
   /// the inferred shape.
   auto indicesLoc = getToken().getLoc();
   TensorLiteralParser indiceParser(*this);
   if (indiceParser.parse(/*allowHex=*/false))
     return nullptr;
 
   if (parseToken(Token::comma, "expected ','"))
Attribute Parser::parseSparseElementsAttr(Type attrType) {
   if (!type)
     return nullptr;
 
   // If the indices are a splat, i.e. the literal parser parsed an element and
   // not a list, we set the shape explicitly. The indices are represented by a
   // 2-dimensional shape where the second dimension is the rank of the type.
   // Given that the parsed indices is a splat, we know that we only have one
   // indice and thus one for the first dimension.
-  auto indiceEltType = builder.getIntegerType(64);
   ShapedType indicesType;
   if (indiceParser.getShape().empty()) {
     indicesType = RankedTensorType::get({1, type.getRank()}, indiceEltType);
   } else {
     // Otherwise, set the shape to the one parsed by the literal parser.
     indicesType = RankedTensorType::get(indiceParser.getShape(), indiceEltType);
   }
   auto indices = indiceParser.getAttr(indicesLoc, indicesType);

mlir/test/IR/parser.mlir

// CHECK: "fooi53"() {bar = dense<{{\[\[\[}}1, -2, 1, 2]], {{\[\[}}0, 2, -1, 2]]]> : tensor<2x1x4xi53>} : () -> ()
   "fooi53"(){bar = dense<[[[1, -2, 1, 2]], [[0, 2, -1, 2]]]> : tensor<2x1x4xi53>} : () -> ()
 // CHECK: "fooi64"() {bar = dense<{{\[\[\[}}1, -2, 1, 2]], {{\[\[}}0, 3, -1, 2]]]> : tensor<2x1x4xi64>} : () -> ()
   "fooi64"(){bar = dense<[[[1, -2, 1, 2]], [[0, 3, -1, 2]]]> : tensor<2x1x4xi64>} : () -> ()
 // CHECK: "fooi64"() {bar = dense<-5> : tensor<1x1x1xi64>} : () -> ()
   "fooi64"(){bar = dense<[[[-5]]]> : tensor<1x1x1xi64>} : () -> ()
 // CHECK: "fooi67"() {bar = dense<{{\[\[\[}}-5, 4, 6, 2]]]> : vector<1x1x4xi67>} : () -> ()
   "fooi67"(){bar = dense<[[[-5, 4, 6, 2]]]> : vector<1x1x4xi67>} : () -> ()
 
-// CHECK: "foo2"() {bar = dense<[]> : tensor<0xi32>} : () -> ()
-  "foo2"(){bar = dense<[]> : tensor<0xi32>} : () -> ()
-// CHECK: "foo2"() {bar = dense<{{\[\[}}]]> : tensor<1x0xi32>} : () -> ()
-  "foo2"(){bar = dense<[[]]> : tensor<1x0xi32>} : () -> ()
+// CHECK: "foo2"() {bar = dense<> : tensor<0xi32>} : () -> ()
+  "foo2"(){bar = dense<> : tensor<0xi32>} : () -> ()
+// CHECK: "foo2"() {bar = dense<> : tensor<1x0xi32>} : () -> ()
+  "foo2"(){bar = dense<> : tensor<1x0xi32>} : () -> ()
+// CHECK: dense<> : tensor<0x512x512xi32>
+  "foo2"(){bar = dense<> : tensor<0x512x512xi32>} : () -> ()
 // CHECK: "foo3"() {bar = dense<{{\[\[\[}}5, -6, 1, 2]], {{\[\[}}7, 8, 3, 4]]]> : tensor<2x1x4xi32>} : () -> ()
   "foo3"(){bar = dense<[[[5, -6, 1, 2]], [[7, 8, 3, 4]]]> : tensor<2x1x4xi32>} : () -> ()
 
 // CHECK: "float1"() {bar = dense<5.000000e+00> : tensor<1x1x1xf32>} : () -> ()
   "float1"(){bar = dense<[[[5.0]]]> : tensor<1x1x1xf32>} : () -> ()
-// CHECK: "float2"() {bar = dense<[]> : tensor<0xf32>} : () -> ()
-  "float2"(){bar = dense<[]> : tensor<0xf32>} : () -> ()
-// CHECK: "float2"() {bar = dense<{{\[\[}}]]> : tensor<1x0xf32>} : () -> ()
-  "float2"(){bar = dense<[[]]> : tensor<1x0xf32>} : () -> ()
+// CHECK: "float2"() {bar = dense<> : tensor<0xf32>} : () -> ()
+  "float2"(){bar = dense<> : tensor<0xf32>} : () -> ()
+// CHECK: "float2"() {bar = dense<> : tensor<1x0xf32>} : () -> ()
+  "float2"(){bar = dense<> : tensor<1x0xf32>} : () -> ()
 
 // CHECK: "bfloat16"() {bar = dense<{{\[\[\[}}-5.000000e+00, 6.000000e+00, 1.000000e+00, 2.000000e+00]], {{\[\[}}7.000000e+00, -8.000000e+00, 3.000000e+00, 4.000000e+00]]]> : tensor<2x1x4xbf16>} : () -> ()
   "bfloat16"(){bar = dense<[[[-5.0, 6.0, 1.0, 2.0]], [[7.0, -8.0, 3.0, 4.0]]]> : tensor<2x1x4xbf16>} : () -> ()
 // CHECK: "float16"() {bar = dense<{{\[\[\[}}-5.000000e+00, 6.000000e+00, 1.000000e+00, 2.000000e+00]], {{\[\[}}7.000000e+00, -8.000000e+00, 3.000000e+00, 4.000000e+00]]]> : tensor<2x1x4xf16>} : () -> ()
   "float16"(){bar = dense<[[[-5.0, 6.0, 1.0, 2.0]], [[7.0, -8.0, 3.0, 4.0]]]> : tensor<2x1x4xf16>} : () -> ()
 // CHECK: "float32"() {bar = dense<{{\[\[\[}}-5.000000e+00, 6.000000e+00, 1.000000e+00, 2.000000e+00]], {{\[\[}}7.000000e+00, -8.000000e+00, 3.000000e+00, 4.000000e+00]]]> : tensor<2x1x4xf32>} : () -> ()
   "float32"(){bar = dense<[[[-5.0, 6.0, 1.0, 2.0]], [[7.0, -8.0, 3.0, 4.0]]]> : tensor<2x1x4xf32>} : () -> ()
 // CHECK: "float64"() {bar = dense<{{\[\[\[}}-5.000000e+00, 6.000000e+00, 1.000000e+00, 2.000000e+00]], {{\[\[}}7.000000e+00, -8.000000e+00, 3.000000e+00, 4.000000e+00]]]> : tensor<2x1x4xf64>} : () -> ()

 // CHECK-LABEL: func @sparsetensorattr
 func @sparsetensorattr() -> () {
 ^bb0:
 // NOTE: The {{\[\[}} syntax is because "[[" confuses FileCheck.
 // CHECK: "fooi8"() {bar = sparse<0, -2> : tensor<1x1x1xi8>} : () -> ()
   "fooi8"(){bar = sparse<0, -2> : tensor<1x1x1xi8>} : () -> ()
 // CHECK: "fooi16"() {bar = sparse<{{\[\[}}1, 1, 0], {{\[}}0, 1, 0], {{\[}}0, 0, 1]], {{\[}}2, -1, 5]> : tensor<2x2x2xi16>} : () -> ()
   "fooi16"(){bar = sparse<[[1, 1, 0], [0, 1, 0], [0, 0, 1]], [2, -1, 5]> : tensor<2x2x2xi16>} : () -> ()
-// CHECK: "fooi32"() {bar = sparse<{{\[}}], {{\[}}]> : tensor<1x1xi32>} : () -> ()
-  "fooi32"(){bar = sparse<[], []> : tensor<1x1xi32>} : () -> ()
+// CHECK: "fooi32"() {bar = sparse<> : tensor<1x1xi32>} : () -> ()
+  "fooi32"(){bar = sparse<> : tensor<1x1xi32>} : () -> ()
 // CHECK: "fooi64"() {bar = sparse<0, -1> : tensor<1xi64>} : () -> ()
   "fooi64"(){bar = sparse<[[0]], [-1]> : tensor<1xi64>} : () -> ()
-// CHECK: "foo2"() {bar = sparse<{{\[}}], {{\[}}]> : tensor<0xi32>} : () -> ()
-  "foo2"(){bar = sparse<[], []> : tensor<0xi32>} : () -> ()
-// CHECK: "foo3"() {bar = sparse<{{\[}}], {{\[}}]> : tensor<i32>} : () -> ()
-  "foo3"(){bar = sparse<[], []> : tensor<i32>} : () -> ()
+// CHECK: "foo2"() {bar = sparse<> : tensor<0xi32>} : () -> ()
+  "foo2"(){bar = sparse<> : tensor<0xi32>} : () -> ()
+// CHECK: "foo3"() {bar = sparse<> : tensor<i32>} : () -> ()
+  "foo3"(){bar = sparse<> : tensor<i32>} : () -> ()
 
 // CHECK: "foof16"() {bar = sparse<0, -2.000000e+00> : tensor<1x1x1xf16>} : () -> ()
   "foof16"(){bar = sparse<0, -2.0> : tensor<1x1x1xf16>} : () -> ()
 // CHECK: "foobf16"() {bar = sparse<{{\[\[}}1, 1, 0], {{\[}}0, 1, 0], {{\[}}0, 0, 1]], {{\[}}2.000000e+00, -1.000000e+00, 5.000000e+00]> : tensor<2x2x2xbf16>} : () -> ()
   "foobf16"(){bar = sparse<[[1, 1, 0], [0, 1, 0], [0, 0, 1]], [2.0, -1.0, 5.0]> : tensor<2x2x2xbf16>} : () -> ()
-// CHECK: "foof32"() {bar = sparse<{{\[}}], {{\[}}]> : tensor<1x0x1xf32>} : () -> ()
-  "foof32"(){bar = sparse<[], []> : tensor<1x0x1xf32>} : () -> ()
+// CHECK: "foof32"() {bar = sparse<> : tensor<1x0x1xf32>} : () -> ()
+  "foof32"(){bar = sparse<> : tensor<1x0x1xf32>} : () -> ()
 // CHECK:  "foof64"() {bar = sparse<0, -1.000000e+00> : tensor<1xf64>} : () -> ()
   "foof64"(){bar = sparse<[[0]], [-1.0]> : tensor<1xf64>} : () -> ()
-// CHECK: "foof320"() {bar = sparse<{{\[}}], {{\[}}]> : tensor<0xf32>} : () -> ()
-  "foof320"(){bar = sparse<[], []> : tensor<0xf32>} : () -> ()
-// CHECK: "foof321"() {bar = sparse<{{\[}}], {{\[}}]> : tensor<f32>} : () -> ()
-  "foof321"(){bar = sparse<[], []> : tensor<f32>} : () -> ()
+// CHECK: "foof320"() {bar = sparse<> : tensor<0xf32>} : () -> ()
+  "foof320"(){bar = sparse<> : tensor<0xf32>} : () -> ()
+// CHECK: "foof321"() {bar = sparse<> : tensor<f32>} : () -> ()
+  "foof321"(){bar = sparse<> : tensor<f32>} : () -> ()
 
 // CHECK: "foostr"() {bar = sparse<0, "foo"> : tensor<1x1x1x!unknown<"">>} : () -> ()
   "foostr"(){bar = sparse<0, "foo"> : tensor<1x1x1x!unknown<"">>} : () -> ()
 // CHECK: "foostr"() {bar = sparse<{{\[\[}}1, 1, 0], {{\[}}0, 1, 0], {{\[}}0, 0, 1]], {{\[}}"a", "b", "c"]> : tensor<2x2x2x!unknown<"">>} : () -> ()
   "foostr"(){bar = sparse<[[1, 1, 0], [0, 1, 0], [0, 0, 1]], ["a", "b", "c"]> : tensor<2x2x2x!unknown<"">>} : () -> ()
   return
 }
 
 // CHECK-LABEL: func @sparsevectorattr
 func @sparsevectorattr() -> () {
 ^bb0:
 // NOTE: The {{\[\[}} syntax is because "[[" confuses FileCheck.
 // CHECK: "fooi8"() {bar = sparse<0, -2> : vector<1x1x1xi8>} : () -> ()
   "fooi8"(){bar = sparse<0, -2> : vector<1x1x1xi8>} : () -> ()
 // CHECK: "fooi16"() {bar = sparse<{{\[\[}}1, 1, 0], {{\[}}0, 1, 0], {{\[}}0, 0, 1]], {{\[}}2, -1, 5]> : vector<2x2x2xi16>} : () -> ()
   "fooi16"(){bar = sparse<[[1, 1, 0], [0, 1, 0], [0, 0, 1]], [2, -1, 5]> : vector<2x2x2xi16>} : () -> ()
-// CHECK: "fooi32"() {bar = sparse<{{\[}}], {{\[}}]> : vector<1x1xi32>} : () -> ()
-  "fooi32"(){bar = sparse<[], []> : vector<1x1xi32>} : () -> ()
+// CHECK: "fooi32"() {bar = sparse<> : vector<1x1xi32>} : () -> ()
+  "fooi32"(){bar = sparse<> : vector<1x1xi32>} : () -> ()
 // CHECK: "fooi64"() {bar = sparse<0, -1> : vector<1xi64>} : () -> ()
   "fooi64"(){bar = sparse<[[0]], [-1]> : vector<1xi64>} : () -> ()
 
 // CHECK: "foof16"() {bar = sparse<0, -2.000000e+00> : vector<1x1x1xf16>} : () -> ()
   "foof16"(){bar = sparse<0, -2.0> : vector<1x1x1xf16>} : () -> ()
 // CHECK: "foobf16"() {bar = sparse<{{\[\[}}1, 1, 0], {{\[}}0, 1, 0], {{\[}}0, 0, 1]], {{\[}}2.000000e+00, -1.000000e+00, 5.000000e+00]> : vector<2x2x2xbf16>} : () -> ()
   "foobf16"(){bar = sparse<[[1, 1, 0], [0, 1, 0], [0, 0, 1]], [2.0, -1.0, 5.0]> : vector<2x2x2xbf16>} : () -> ()
 // CHECK:  "foof64"() {bar = sparse<0, -1.000000e+00> : vector<1xf64>} : () -> ()