[Mlir-commits] [mlir] 55d0610 - [mlir] make `parseCustomOperationName()` API check token type (#136306)

[llvmlistbot at llvm.org]

Author: Oleksandr "Alex" Zinenko
Date: 2025-04-30T09:28:21+02:00
New Revision: 55d0610e7bbea470c47985d9d3f638dff85411b5

URL: https://github.com/llvm/llvm-project/commit/55d0610e7bbea470c47985d9d3f638dff85411b5
DIFF: https://github.com/llvm/llvm-project/commit/55d0610e7bbea470c47985d9d3f638dff85411b5.diff

LOG: [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.

Added: 
    

Modified: 
    mlir/lib/AsmParser/Parser.cpp
    mlir/test/Dialect/Linalg/invalid.mlir
    mlir/test/IR/invalid.mlir
    mlir/test/IR/parser.mlir
    mlir/test/lib/Dialect/Test/TestOpsSyntax.cpp
    mlir/test/lib/Dialect/Test/TestOpsSyntax.td

Removed: 
    


################################################################################
diff  --git a/mlir/lib/AsmParser/Parser.cpp b/mlir/lib/AsmParser/Parser.cpp
index 6240fdc833501..756d3d01a4534 100644
--- a/mlir/lib/AsmParser/Parser.cpp
+++ b/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());

diff  --git a/mlir/test/Dialect/Linalg/invalid.mlir b/mlir/test/Dialect/Linalg/invalid.mlir
index 90ceadebbc1fa..79f2a0d7c7ae6 100644
--- a/mlir/test/Dialect/Linalg/invalid.mlir
+++ b/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>
+}

diff  --git a/mlir/test/IR/invalid.mlir b/mlir/test/IR/invalid.mlir
index d5b23fcf27950..d0bbf8669b63d 100644
--- a/mlir/test/IR/invalid.mlir
+++ b/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() {

diff  --git a/mlir/test/IR/parser.mlir b/mlir/test/IR/parser.mlir
index 8b192ff11d573..bb19cb63b65a5 100644
--- a/mlir/test/IR/parser.mlir
+++ b/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 
diff erent name scopes for regions isolated from above.
 
 // CHECK-LABEL: func @op_with_passthrough_region_args

diff  --git a/mlir/test/lib/Dialect/Test/TestOpsSyntax.cpp b/mlir/test/lib/Dialect/Test/TestOpsSyntax.cpp
index 7a9f07030215d..6d4e5e3cb4401 100644
--- a/mlir/test/lib/Dialect/Test/TestOpsSyntax.cpp
+++ b/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"
 

diff  --git a/mlir/test/lib/Dialect/Test/TestOpsSyntax.td b/mlir/test/lib/Dialect/Test/TestOpsSyntax.td
index d9003428e3746..9a96fe8b11c14 100644
--- a/mlir/test/lib/Dialect/Test/TestOpsSyntax.td
+++ b/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