[mlir][AsmParser] Improve parse{Attribute,Type} error handling

[rkayaith]

Summary

Currently these functions report errors directly to stderr, this updates
them to use diagnostics instead. This also makes partially-consumed
strings an error if the numRead parameter isn't provided (the
docstrings already claimed this happened, but it didn't.)

While here I also tried to reduce the number of overloads by switching
to using default parameters.

[rkayaith]

mlir/lib/AsmParser/DialectSymbolParser.cpp

317

switched this to true since I found the parser reads past the end of the buffer if it isn't null-terminated (e.g. StringRef("999", 1) gets parsed as 999 instead of 9)

mlir/test/Dialect/Linalg/transform-op-pad.mlir

121–126

"foo" is still an error, but it emits an extra diagnostic that can't be checked through expected-error:

foo:1:1: error: expected attribute value

(previously that error was just ending up on stderr)

[rriddle]

mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-yaml-gen.cpp

320–323

Do you need the temp variable here? (after the RequiresNullTerminator change?)

[rkayaith]

mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-yaml-gen.cpp

320–323

Yea, all RequiresNullTerminator seems to do is enable an assert that checks that the input is null-terminated. Which ends up triggering here without the std::string copy.

[rriddle]

mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-yaml-gen.cpp

320–323

Hmrm, we should likely change the API to better differentiate cases where the user knows there is a null-terminator, vs where one is needed. We have a bug related to this as well.

[bixia]

This change breaks tensorflow test github.com/tensorflow/tensorflow/blob/master/tensorflow/compiler/mlir/tensorflow/tests/compile_mlir_util/mlir-module-serialized-str-attr.mlir

In particular, the parser returns error ":16:1: error: found trailing characters: ': return [[IDENTITY]] : tensor<?xi32> loc(unknown)" because it doesn't finish reading the trailing comments in the test input.
I couldn't figure out which part of this PR may cause this change in the parser.

I temporary fix the test by this change:
tensorflow/compiler/mlir/tensorflow/utils/tf_xla_mlir_translate.cc ====
388c388,391

< mlir::Attribute attr = mlir::parseAttribute(input, context);

// When the parser doesn't read all the input chars, it issues an error unless
// an output parameter is provided for returning the number of chars read.
size_t numRead;
mlir::Attribute attr = mlir::parseAttribute(input, context, {}, &numRead);

[rkayaith]

In D144804#4163284, @bixia wrote:

This change breaks tensorflow test github.com/tensorflow/tensorflow/blob/master/tensorflow/compiler/mlir/tensorflow/tests/compile_mlir_util/mlir-module-serialized-str-attr.mlir

In particular, the parser returns error ":16:1: error: found trailing characters: ': return [[IDENTITY]] : tensor<?xi32> loc(unknown)" because it doesn't finish reading the trailing comments in the test input.
I couldn't figure out which part of this PR may cause this change in the parser.

I temporary fix the test by this change:
tensorflow/compiler/mlir/tensorflow/utils/tf_xla_mlir_translate.cc ====
388c388,391

< mlir::Attribute attr = mlir::parseAttribute(input, context);

// When the parser doesn't read all the input chars, it issues an error unless
// an output parameter is provided for returning the number of chars read.
size_t numRead;
mlir::Attribute attr = mlir::parseAttribute(input, context, {}, &numRead);

If I understand correctly, the error here is an intentional change (previously the extra input would be silently ignored, which could hide potential errors). If you know you only want to parse part of the string, you can either trim the StringRef before passing it to parseAttribute, or pass numRead as you did in your fix.

[bixia]

The input contains the "meaningfule string for string attr" and followed by comment lines. Do you know why the parser stop in the middle of a comment line?

[rkayaith]

In D144804#4165109, @bixia wrote:

The input contains the "meaningfule string for string attr" and followed by comment lines. Do you know why the parser stop in the middle of a comment line?

I'm not familiar with the parser internals, but that sounds like a bug to me. It would've existed before this change though, just hidden.

[kuhar]

Hi @rkayaith,

Seems like this does not allow for passing trimmed StringRefs to mlir::parseType anymore. These are not null-terminated Was this intentional?

[rkayaith]

In D144804#4165327, @kuhar wrote:

Hi @rkayaith,

Seems like this does not allow for passing trimmed StringRefs to mlir::parseType anymore. These are not null-terminated Was this intentional?

From what I could tell, this is an existing requirement of the parser: https://reviews.llvm.org/D144804#inline-1398522

So you'd have to trim + copy to a null terminated string, e.g. std::string

[rkayaith]

In D144804#4165413, @rkayaith wrote:

In D144804#4165327, @kuhar wrote:

Hi @rkayaith,

Seems like this does not allow for passing trimmed StringRefs to mlir::parseType anymore. These are not null-terminated Was this intentional?

From what I could tell, this is an existing requirement of the parser: https://reviews.llvm.org/D144804#inline-1398522

So you'd have to trim + copy to a null terminated string, e.g. std::string

Actually it's probably better if parseType just does this copy by default, patch for that here: https://reviews.llvm.org/D145182

mlir/include/mlir/AsmParser/AsmParser.h

 /// used to signal tracking of a code completion event (generally only ever
 /// useful for LSP or other high level language tooling).
 LogicalResult
 parseAsmSourceFile(const llvm::SourceMgr &sourceMgr, Block *block,
                    const ParserConfig &config,
                    AsmParserState *asmState = nullptr,
                    AsmParserCodeCompleteContext *codeCompleteContext = nullptr);
 
-/// This parses a single MLIR attribute to an MLIR context if it was valid.  If
-/// not, an error message is emitted through a new SourceMgrDiagnosticHandler
-/// constructed from a new SourceMgr with a single a MemoryBuffer wrapping
-/// `attrStr`. If the passed `attrStr` has additional tokens that were not part
-/// of the type, an error is emitted.
-// TODO: Improve diagnostic reporting.
-Attribute parseAttribute(llvm::StringRef attrStr, MLIRContext *context);
-Attribute parseAttribute(llvm::StringRef attrStr, Type type);
-
-/// This parses a single MLIR attribute to an MLIR context if it was valid.  If
-/// not, an error message is emitted through a new SourceMgrDiagnosticHandler
-/// constructed from a new SourceMgr with a single a MemoryBuffer wrapping
-/// `attrStr`. The number of characters of `attrStr` parsed in the process is
-/// returned in `numRead`.
+/// This parses a single MLIR attribute to an MLIR context if it was valid. If
+/// not, an error diagnostic is emitted to the context and a null value is
+/// returned.
+/// If `numRead` is provided, it is set to the number of consumed characters on
+/// succesful parse. Otherwise, parsing fails if the entire string is not
+/// consumed.
 Attribute parseAttribute(llvm::StringRef attrStr, MLIRContext *context,
-                         size_t &numRead);
-Attribute parseAttribute(llvm::StringRef attrStr, Type type, size_t &numRead);
-
-/// This parses a single MLIR type to an MLIR context if it was valid.  If not,
-/// an error message is emitted through a new SourceMgrDiagnosticHandler
-/// constructed from a new SourceMgr with a single a MemoryBuffer wrapping
-/// `typeStr`. If the passed `typeStr` has additional tokens that were not part
-/// of the type, an error is emitted.
-// TODO: Improve diagnostic reporting.
-Type parseType(llvm::StringRef typeStr, MLIRContext *context);
+                         Type type = {}, size_t *numRead = nullptr);
 
-/// This parses a single MLIR type to an MLIR context if it was valid.  If not,
-/// an error message is emitted through a new SourceMgrDiagnosticHandler
-/// constructed from a new SourceMgr with a single a MemoryBuffer wrapping
-/// `typeStr`. The number of characters of `typeStr` parsed in the process is
-/// returned in `numRead`.
-Type parseType(llvm::StringRef typeStr, MLIRContext *context, size_t &numRead);
+/// This parses a single MLIR type to an MLIR context if it was valid. If not,
+/// an error diagnostic is emitted to the context.
+/// If `numRead` is provided, it is set to the number of consumed characters on
+/// succesful parse. Otherwise, parsing fails if the entire string is not
+/// consumed.
+Type parseType(llvm::StringRef typeStr, MLIRContext *context,
+               size_t *numRead = nullptr);
 
 /// This parses a single IntegerSet/AffineMap to an MLIR context if it was
 /// valid. If not, an error message is emitted through a new
 /// SourceMgrDiagnosticHandler constructed from a new SourceMgr with a single
 /// MemoryBuffer wrapping `str`. If the passed `str` has additional tokens that
 /// were not part of the IntegerSet/AffineMap, a failure is returned.
 AffineMap parseAffineMap(llvm::StringRef str, MLIRContext *context);
 IntegerSet parseIntegerSet(llvm::StringRef str, MLIRContext *context);
 
 } // namespace mlir
 
 #endif // MLIR_ASMPARSER_ASMPARSER_H

mlir/lib/AsmParser/DialectSymbolParser.cpp

 //===----------------------------------------------------------------------===//
 // mlir::parseAttribute/parseType
 //===----------------------------------------------------------------------===//
 
 /// Parses a symbol, of type 'T', and returns it if parsing was successful. If
 /// parsing failed, nullptr is returned. The number of bytes read from the input
 /// string is returned in 'numRead'.
 template <typename T, typename ParserFn>
-static T parseSymbol(StringRef inputStr, MLIRContext *context, size_t &numRead,
-                     ParserFn &&parserFn) {
+static T parseSymbol(StringRef inputStr, MLIRContext *context,
+                     size_t *numReadOut, ParserFn &&parserFn) {
+  // Set the buffer name to the string being parsed, so that it appears in error
+  // diagnostics.
+  auto memBuffer = MemoryBuffer::getMemBuffer(inputStr, /*BufferName=*/inputStr,
+                                              /*RequiresNullTerminator=*/true);
   SourceMgr sourceMgr;
-  auto memBuffer = MemoryBuffer::getMemBuffer(
-      inputStr, /*BufferName=*/"<mlir_parser_buffer>",
-      /*RequiresNullTerminator=*/false);
   sourceMgr.AddNewSourceBuffer(std::move(memBuffer), SMLoc());
   SymbolState aliasState;
   ParserConfig config(context);
   ParserState state(sourceMgr, config, aliasState, /*asmState=*/nullptr,
                     /*codeCompleteContext=*/nullptr);
   Parser parser(state);
 
-  SourceMgrDiagnosticHandler handler(
-      const_cast<llvm::SourceMgr &>(parser.getSourceMgr()),
-      parser.getContext());
   Token startTok = parser.getToken();
   T symbol = parserFn(parser);
   if (!symbol)
     return T();
 
   // Provide the number of bytes that were read.
   Token endTok = parser.getToken();
-  numRead = static_cast<size_t>(endTok.getLoc().getPointer() -
-                                startTok.getLoc().getPointer());
-  return symbol;
+  size_t numRead =
+      endTok.getLoc().getPointer() - startTok.getLoc().getPointer();
+  if (numReadOut) {
+    *numReadOut = numRead;
+  } else if (numRead != inputStr.size()) {
+    parser.emitError(endTok.getLoc()) << "found trailing characters: '"
+                                      << inputStr.drop_front(numRead) << "'";
+    return T();
   }
-
-Attribute mlir::parseAttribute(StringRef attrStr, MLIRContext *context) {
-  size_t numRead = 0;
-  return parseAttribute(attrStr, context, numRead);
-}
-Attribute mlir::parseAttribute(StringRef attrStr, Type type) {
-  size_t numRead = 0;
-  return parseAttribute(attrStr, type, numRead);
+  return symbol;
 }
 
 Attribute mlir::parseAttribute(StringRef attrStr, MLIRContext *context,
-                               size_t &numRead) {
-  return parseSymbol<Attribute>(attrStr, context, numRead, [](Parser &parser) {
-    return parser.parseAttribute();
-  });
-}
-Attribute mlir::parseAttribute(StringRef attrStr, Type type, size_t &numRead) {
+                               Type type, size_t *numRead) {
   return parseSymbol<Attribute>(
-      attrStr, type.getContext(), numRead,
+      attrStr, context, numRead,
       [type](Parser &parser) { return parser.parseAttribute(type); });
 }
-
-Type mlir::parseType(StringRef typeStr, MLIRContext *context) {
-  size_t numRead = 0;
-  return parseType(typeStr, context, numRead);
-}
-
-Type mlir::parseType(StringRef typeStr, MLIRContext *context, size_t &numRead) {
+Type mlir::parseType(StringRef typeStr, MLIRContext *context, size_t *numRead) {
   return parseSymbol<Type>(typeStr, context, numRead,
                            [](Parser &parser) { return parser.parseType(); });
 }

mlir/lib/Bytecode/Reader/BytecodeReader.cpp

LogicalResult AttrTypeReader::parseAsmEntry(T &result, EncodingReader &reader,
   StringRef asmStr;
   if (failed(reader.parseNullTerminatedString(asmStr)))
     return failure();
 
   // Invoke the MLIR assembly parser to parse the entry text.
   size_t numRead = 0;
   MLIRContext *context = fileLoc->getContext();
   if constexpr (std::is_same_v<T, Type>)
-    result = ::parseType(asmStr, context, numRead);
+    result = ::parseType(asmStr, context, &numRead);
   else
-    result = ::parseAttribute(asmStr, context, numRead);
+    result = ::parseAttribute(asmStr, context, Type(), &numRead);
   if (!result)
     return failure();
 
   // Ensure there weren't dangling characters after the entry.
   if (numRead != asmStr.size()) {
     return reader.emitError("trailing characters found after ", entryType,
                             " assembly format: ", asmStr.drop_front(numRead));
   }

mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp

for (auto const &it :
     auto attr = std::get<0>(it).dyn_cast<TypedAttr>();
     if (!attr) {
       emitOpError("expects padding values to be typed attributes");
       return DiagnosedSilenceableFailure::definiteFailure();
     }
     Type elementType = getElementTypeOrSelf(std::get<1>(it));
     // Try to parse string attributes to obtain an attribute of element type.
     if (auto stringAttr = attr.dyn_cast<StringAttr>()) {
-      paddingValues.push_back(
-          parseAttribute(attr.cast<StringAttr>(), elementType));
-      if (!paddingValues.back()) {
+      auto parsedAttr = dyn_cast_if_present<TypedAttr>(
+          parseAttribute(stringAttr, getContext(), elementType));
+      if (!parsedAttr || parsedAttr.getType() != elementType) {
         auto diag = this->emitOpError("expects a padding that parses to ")
                     << elementType << ", got " << std::get<0>(it);
         diag.attachNote(target.getLoc()) << "when applied to this op";
         return DiagnosedSilenceableFailure::definiteFailure();
       }
+      paddingValues.push_back(parsedAttr);
       continue;
     }
     // Otherwise, add the attribute directly.
     if (attr.getType() != elementType) {
       auto diag = this->emitOpError("expects a padding value of type ")
                   << elementType << ", got " << attr;
       diag.attachNote(target.getLoc()) << "when applied to this op";
       return DiagnosedSilenceableFailure::definiteFailure();

mlir/test/Dialect/Linalg/transform-op-pad.mlir

func.func @pad(%arg0: tensor<24x12xf32>,
   // expected-note @below {{when applied to this op}}
   %0 = linalg.matmul ins(%arg0, %arg1 : tensor<24x12xf32>, tensor<12x25xf32>) outs(%arg2 : tensor<24x25xf32>) -> tensor<24x25xf32>
   func.return %0 : tensor<24x25xf32>
 }
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !pdl.operation):
   %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!pdl.operation) -> !pdl.operation
-  // expected-error @below {{expects a padding that parses to 'f32', got "foo"}}
+  // expected-error @below {{expects a padding that parses to 'f32', got "{foo}"}}
   %1 = transform.structured.pad %0 {
-    padding_values=["foo", 0.0 : f32, 0.0 : f32],
+    padding_values=["{foo}", 0.0 : f32, 0.0 : f32],
     padding_dimensions=[0, 1, 2],
     pack_paddings=[1, 1, 0]
   }
 }
 
 // -----
 
 func.func @pad(%arg0: tensor<24x12xf32>,
                %arg1: tensor<12x25xf32>,
                %arg2: tensor<24x25xf32>) -> tensor<24x25xf32> {
   // This is attached to an error that is silenceable and is not reported by this transform
   //   {{when applied to this op}}

mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-yaml-gen.cpp

static void output(const SerializedAffineMap &value, void *rawYamlContext,
                      raw_ostream &out) {
     assert(value.affineMapAttr);
     value.affineMapAttr.print(out);
   }
   static StringRef input(StringRef scalar, void *rawYamlContext,
                          SerializedAffineMap &value) {
     assert(rawYamlContext);
     auto *yamlContext = static_cast<LinalgYAMLContext *>(rawYamlContext);
-    if (auto attr = mlir::parseAttribute(scalar, yamlContext->mlirContext)
-                        .dyn_cast_or_null<AffineMapAttr>())
+    std::string nullTerminatedScalar(scalar);
+    if (auto attr =
+            mlir::parseAttribute(nullTerminatedScalar, yamlContext->mlirContext)
       value.affineMapAttr = attr;
     else if (!value.affineMapAttr || !value.affineMapAttr.isa<AffineMapAttr>())
       return "could not parse as an affine map attribute";
     return StringRef();
   }
   static QuotingType mustQuote(StringRef) { return QuotingType::None; }
 };
 

mlir/unittests/Parser/ParserTest.cpp

 //===- ParserTest.cpp -----------------------------------------------------===//
 //
 // This file is licensed under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Parser/Parser.h"
+#include "mlir/AsmParser/AsmParser.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/Verifier.h"
 
 #include "gmock/gmock.h"
 
 using namespace mlir;
 
 namespace {
TEST(MLIRParser, ParseAtEnd) {
   LogicalResult secondParse =
       parseSourceString(secondModuleStr, &block, &context);
   EXPECT_TRUE(succeeded(secondParse));
 
   // Check the we parse at the end.
   EXPECT_EQ(block.front().getName().getStringRef(), "test.first");
   EXPECT_EQ(block.back().getName().getStringRef(), "test.second");
 }
+
+TEST(MLIRParser, ParseAttr) {
+  using namespace testing;
+  MLIRContext context;
+  Builder b(&context);
+  { // Successful parse
+    StringLiteral attrAsm = "array<i64: 1, 2, 3>";
+    size_t numRead = 0;
+    Attribute attr = parseAttribute(attrAsm, &context, Type(), &numRead);
+    EXPECT_EQ(attr, b.getDenseI64ArrayAttr({1, 2, 3}));
+    EXPECT_EQ(numRead, attrAsm.size());
+  }
+  { // Failed parse
+    std::vector<std::string> diagnostics;
+    ScopedDiagnosticHandler handler(&context, [&](Diagnostic &d) {
+      llvm::raw_string_ostream(diagnostics.emplace_back())
+          << d.getLocation() << ": " << d;
+    });
+    size_t numRead = 0;
+    EXPECT_FALSE(parseAttribute("dense<>", &context, Type(), &numRead));
+    EXPECT_THAT(diagnostics, ElementsAre("loc(\"dense<>\":1:7): expected ':'"));
+    EXPECT_EQ(numRead, size_t(0));
+  }
+  { // Parse with trailing characters
+    std::vector<std::string> diagnostics;
+    ScopedDiagnosticHandler handler(&context, [&](Diagnostic &d) {
+      llvm::raw_string_ostream(diagnostics.emplace_back())
+          << d.getLocation() << ": " << d;
+    });
+    EXPECT_FALSE(parseAttribute("10  foo", &context));
+    EXPECT_THAT(
+        diagnostics,
+        ElementsAre("loc(\"10  foo\":1:5): found trailing characters: 'foo'"));
+
+    size_t numRead = 0;
+    EXPECT_EQ(parseAttribute("10  foo", &context, Type(), &numRead),
+              b.getI64IntegerAttr(10));
+    EXPECT_EQ(numRead, size_t(4)); // includes trailing whitespace
+  }
+}
 } // namespace