[Mlir-commits] [mlir] [mlir] make parseCustomOperationNameAPI check token type (PR #136306)

[Oleksandr Alex Zinenko]

https://github.com/ftynse created https://github.com/llvm/llvm-project/pull/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.

>From e9da71e60ebe8b581f54bd24526551697c92d18a Mon Sep 17 00:00:00 2001
From: Alex Zinenko <git at ozinenko.com>
Date: Fri, 18 Apr 2025 15:40:29 +0200
Subject: [PATCH] [mlir] make parseCustomOperationNameAPI check token type

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.
---
 mlir/lib/AsmParser/Parser.cpp                |  5 +++++
 mlir/test/Dialect/Linalg/invalid.mlir        |  9 +++++++++
 mlir/test/IR/invalid.mlir                    | 10 ++++++++++
 mlir/test/IR/parser.mlir                     |  9 +++++++++
 mlir/test/lib/Dialect/Test/TestOpsSyntax.cpp | 19 +++++++++++++++++++
 mlir/test/lib/Dialect/Test/TestOpsSyntax.td  |  6 ++++++
 6 files changed, 58 insertions(+)

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 different 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