[mlir] make parseCustomOperationName() API check token type (#136306)

Previously, this parser API call would accept any token and interpret
its spelling as operation name, including tokens that are are not valid
operation names. Make it accept only bare identifiers and keywords. The
latter is questionable but consistent with current practices upstream.

Fixes #132889.

Author: ftynse

mlir/lib/AsmParser/Parser.cpp

@@ -2005,6 +2005,11 @@ class CustomOpAsmParser : public AsmParserImpl<OpAsmParser> {
 
 FailureOr<OperationName> OperationParser::parseCustomOperationName() {
   Token nameTok = getToken();
+  // Accept keywords here as they may be interpreted as a shortened operation
+  // name, e.g., `dialect.keyword` can be spelled as just `keyword` within a
+  // region of an operation from `dialect`.
+  if (nameTok.getKind() != Token::bare_identifier && !nameTok.isKeyword())
+    return emitError("expected bare identifier or keyword");
   StringRef opName = nameTok.getSpelling();
   if (opName.empty())
     return (emitError("empty operation name is invalid"), failure());

mlir/test/Dialect/Linalg/invalid.mlir

@@ -1666,3 +1666,12 @@ func.func @unpack_static_inner_tile_size_and_dynamic_output_shape(
   %0 = linalg.unpack %input inner_dims_pos = [0, 1] inner_tiles = [8, 4] into %output : tensor<?x?x?x4xf32> -> tensor<?x?xf32>
   return %0 : tensor<?x?xf32>
 }
+
+// -----
+
+func.func @reduce_non_operation_name(%arg0: tensor<4xf32>, %arg1: tensor<f32>) -> tensor<f32> {
+  // expected-error @below {{expected bare identifier or keyword}}
+  %0 = linalg.reduce {@reduce_fusion_elementwise} ins(
+    %arg0: tensor<4xf32>) outs(%arg1: tensor<f32>) dimensions = [0]
+  return %0 : tensor<f32>
+}

mlir/test/IR/invalid.mlir

@@ -642,6 +642,16 @@ func.func @invalid_region_dominance_with_dominance_free_regions() {
 
 // -----
 
+// expected-error @below {{expected bare identifier or keyword}}
+test.parse_custom_operation_name_api(@foo) {}
+
+// -----
+
+// expected-error @below {{expected bare identifier or keyword}}
+test.parse_custom_operation_name_api(42) {}
+
+// -----
+
 // This makes sure we emit an error at the end of the correct line, the : is
 // expected at the end of foo, not on the return line.
 func.func @error_at_end_of_line() {

mlir/test/IR/parser.mlir

@@ -1169,6 +1169,15 @@ func.func @op_with_region_args() {
   return
 }
 
+// Test parsing an operation name from within another op custom syntax.
+
+// CHECK-LABEL: @custom_name_api
+func.func @custom_name_api() {
+  // CHECK: test.parse_custom_operation_name_api(builtin.module)
+  test.parse_custom_operation_name_api(builtin.module)
+  return
+}
+
 // Test allowing different name scopes for regions isolated from above.
 
 // CHECK-LABEL: func @op_with_passthrough_region_args

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

@@ -487,6 +487,25 @@ static void printOptionalLoc(OpAsmPrinter &p, Operation *op, Attribute loc) {
   p.printOptionalLocationSpecifier(cast<LocationAttr>(loc));
 }
 
+//===----------------------------------------------------------------------===//
+// ParseCustomOperationNameAPI
+//===----------------------------------------------------------------------===//
+
+static ParseResult parseCustomOperationNameEntry(OpAsmParser &p,
+                                                 Attribute &name) {
+  FailureOr<OperationName> opName = p.parseCustomOperationName();
+  if (failed(opName))
+    return ParseResult::failure();
+
+  name = p.getBuilder().getStringAttr(opName->getStringRef());
+  return ParseResult::success();
+}
+
+static void printCustomOperationNameEntry(OpAsmPrinter &p, Operation *op,
+                                          Attribute name) {
+  p << cast<StringAttr>(name).getValue();
+}
+
 #define GET_OP_CLASSES
 #include "TestOpsSyntax.cpp.inc"
 

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

@@ -69,6 +69,12 @@ def TestAttrWithLoc : TEST_Op<"attr_with_loc"> {
   let assemblyFormat = "`(` $value `` custom<OptionalLoc>($loc) `)` attr-dict";
 }
 
+def ParseCustomOperationNameAPI : TEST_Op<"parse_custom_operation_name_api"> {
+  let summary = "noop that exercises the parseCustomOperationName API";
+  let arguments = (ins StrAttr:$name);
+  let assemblyFormat = "`(` custom<CustomOperationNameEntry>($name) `)` attr-dict";
+}
+
 // -----
 
 // This is used to test that the fallback for a custom op's parser and printer