[Mlir-commits] [mlir] e1249b4 - [mlir][affine] Fix parsing operands that specify result number (#215539)

llvmlistbot at llvm.org llvmlistbot at llvm.org
Tue Aug 25 04:12:11 PDT 2026


Author: Ivan R. Ivanov
Date: 2026-08-25T13:12:06+02:00
New Revision: e1249b4f571cbedd461ed13b7cc4d1b4f2973da3

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

LOG: [mlir][affine] Fix parsing operands that specify result number (#215539)

The check for a previously parsed operand was only checking the name of
the operand and not its result number, resulting in wrong parsing.

This patch restructures the parser to use the parseOperand function for
all occurances instead of hand rolling operand parsing, and then also
tracking the result number in the dim sym operand map.

Paths that parse exprs with predefined non-SSA symbols get an implicit
result number 0.

Added: 
    

Modified: 
    mlir/lib/AsmParser/AffineParser.cpp
    mlir/lib/AsmParser/Parser.cpp
    mlir/lib/AsmParser/Parser.h
    mlir/test/Dialect/Affine/ops.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/AsmParser/AffineParser.cpp b/mlir/lib/AsmParser/AffineParser.cpp
index 0d0c74b965a2b..405970d4fb15e 100644
--- a/mlir/lib/AsmParser/AffineParser.cpp
+++ b/mlir/lib/AsmParser/AffineParser.cpp
@@ -58,10 +58,14 @@ enum AffineHighPrecOp {
 /// bodies.
 class AffineParser : public Parser {
 public:
-  AffineParser(ParserState &state, bool allowParsingSSAIds = false,
-               function_ref<ParseResult(bool)> parseElement = nullptr)
+  AffineParser(
+      ParserState &state, bool allowParsingSSAIds = false,
+      function_ref<FailureOr<OpAsmParser::UnresolvedOperand>()> parseElement =
+          nullptr,
+      function_ref<void(bool, OpAsmParser::UnresolvedOperand)> addOperand =
+          nullptr)
       : Parser(state), allowParsingSSAIds(allowParsingSSAIds),
-        parseElement(parseElement) {}
+        parseElement(parseElement), addOperand(addOperand) {}
 
   ParseResult parseAffineMapRange(unsigned numDims, unsigned numSymbols,
                                   AffineMap &result);
@@ -107,10 +111,12 @@ class AffineParser : public Parser {
 
 private:
   bool allowParsingSSAIds;
-  function_ref<ParseResult(bool)> parseElement;
+  function_ref<FailureOr<OpAsmParser::UnresolvedOperand>()> parseElement;
+  function_ref<void(bool, OpAsmParser::UnresolvedOperand)> addOperand;
   unsigned numDimOperands = 0;
   unsigned numSymbolOperands = 0;
-  SmallVector<std::pair<StringRef, AffineExpr>, 4> dimsAndSymbols;
+  SmallVector<std::pair<std::pair<StringRef, unsigned>, AffineExpr>, 4>
+      dimsAndSymbols;
 };
 } // namespace
 
@@ -296,7 +302,7 @@ AffineExpr AffineParser::parseBareIdExpr() {
 
   StringRef sRef = getTokenSpelling();
   for (auto entry : dimsAndSymbols) {
-    if (entry.first == sRef) {
+    if (entry.first.first == sRef) {
       consumeToken();
       return entry.second;
     }
@@ -311,21 +317,21 @@ AffineExpr AffineParser::parseSSAIdExpr(bool isSymbol) {
     return emitWrongTokenError("unexpected ssa identifier"), nullptr;
   if (getToken().isNot(Token::percent_identifier))
     return emitWrongTokenError("expected ssa identifier"), nullptr;
-  auto name = getTokenSpelling();
+  FailureOr<OpAsmParser::UnresolvedOperand> operand = parseElement();
+  if (failed(operand))
+    return nullptr;
   // Check if we already parsed this SSA id.
-  for (auto entry : dimsAndSymbols) {
-    if (entry.first == name) {
-      consumeToken(Token::percent_identifier);
+  for (auto entry : dimsAndSymbols)
+    if (entry.first.first == operand->name &&
+        entry.first.second == operand->number)
       return entry.second;
-    }
-  }
-  // Parse the SSA id and add an AffineDim/SymbolExpr to represent it.
-  if (parseElement(isSymbol))
-    return nullptr;
+  // If we have not, register it in the dim/symbol operand list, and create an
+  // AffineDim/SymbolExpr to represent it.
+  addOperand(isSymbol, *operand);
   auto idExpr = isSymbol
                     ? getAffineSymbolExpr(numSymbolOperands++, getContext())
                     : getAffineDimExpr(numDimOperands++, getContext());
-  dimsAndSymbols.push_back({name, idExpr});
+  dimsAndSymbols.push_back({{operand->name, operand->number}, idExpr});
   return idExpr;
 }
 
@@ -496,12 +502,12 @@ ParseResult AffineParser::parseIdentifierDefinition(AffineExpr idExpr) {
 
   auto name = getTokenSpelling();
   for (auto entry : dimsAndSymbols) {
-    if (entry.first == name)
+    if (entry.first.first == name)
       return emitError("redefinition of identifier '" + name + "'");
   }
   consumeToken();
 
-  dimsAndSymbols.push_back({name, idExpr});
+  dimsAndSymbols.push_back({{name, 0}, idExpr});
   return success();
 }
 
@@ -559,7 +565,8 @@ ParseResult AffineParser::parseAffineMapOrIntegerSetInline(AffineMap &map,
 /// Parse an affine expresion definition inline, with given symbols.
 ParseResult AffineParser::parseAffineExprInline(
     ArrayRef<std::pair<StringRef, AffineExpr>> symbolSet, AffineExpr &expr) {
-  dimsAndSymbols.assign(symbolSet.begin(), symbolSet.end());
+  for (const auto &[name, expr] : symbolSet)
+    dimsAndSymbols.push_back({{name, 0}, expr});
   expr = parseAffineExpr();
   return success(expr != nullptr);
 }
@@ -750,20 +757,24 @@ ParseResult Parser::parseIntegerSetReference(IntegerSet &set) {
 
 /// Parse an AffineMap of SSA ids. The callback 'parseElement' is used to
 /// parse SSA value uses encountered while parsing affine expressions.
-ParseResult
-Parser::parseAffineMapOfSSAIds(AffineMap &map,
-                               function_ref<ParseResult(bool)> parseElement,
-                               OpAsmParser::Delimiter delimiter) {
-  return AffineParser(state, /*allowParsingSSAIds=*/true, parseElement)
+ParseResult Parser::parseAffineMapOfSSAIds(
+    AffineMap &map,
+    function_ref<FailureOr<OpAsmParser::UnresolvedOperand>()> parseElement,
+    function_ref<void(bool, OpAsmParser::UnresolvedOperand)> addOperand,
+    OpAsmParser::Delimiter delimiter) {
+  return AffineParser(state, /*allowParsingSSAIds=*/true, parseElement,
+                      addOperand)
       .parseAffineMapOfSSAIds(map, delimiter);
 }
 
-/// Parse an AffineExpr of SSA ids. The callback `parseElement` is used to parse
-/// SSA value uses encountered while parsing.
-ParseResult
-Parser::parseAffineExprOfSSAIds(AffineExpr &expr,
-                                function_ref<ParseResult(bool)> parseElement) {
-  return AffineParser(state, /*allowParsingSSAIds=*/true, parseElement)
+/// Parse an AffineExpr of SSA ids. The callback `parseElement` is used to
+/// parse SSA value uses encountered while parsing.
+ParseResult Parser::parseAffineExprOfSSAIds(
+    AffineExpr &expr,
+    function_ref<FailureOr<OpAsmParser::UnresolvedOperand>()> parseElement,
+    function_ref<void(bool, OpAsmParser::UnresolvedOperand)> addOperand) {
+  return AffineParser(state, /*allowParsingSSAIds=*/true, parseElement,
+                      addOperand)
       .parseAffineExprOfSSAIds(expr);
 }
 

diff  --git a/mlir/lib/AsmParser/Parser.cpp b/mlir/lib/AsmParser/Parser.cpp
index 5cfd2f793313b..22f01c7fdaf82 100644
--- a/mlir/lib/AsmParser/Parser.cpp
+++ b/mlir/lib/AsmParser/Parser.cpp
@@ -1820,19 +1820,21 @@ class CustomOpAsmParser : public AsmParserImpl<OpAsmParser> {
     SmallVector<UnresolvedOperand, 2> dimOperands;
     SmallVector<UnresolvedOperand, 1> symOperands;
 
-    auto parseElement = [&](bool isSymbol) -> ParseResult {
+    auto parseElement = [&]() -> FailureOr<UnresolvedOperand> {
       UnresolvedOperand operand;
       if (parseOperand(operand))
-        return failure();
+        return {};
+      return operand;
+    };
+    auto addOperand = [&](bool isSymbol, UnresolvedOperand operand) {
       if (isSymbol)
         symOperands.push_back(operand);
       else
         dimOperands.push_back(operand);
-      return success();
     };
 
     AffineMap map;
-    if (parser.parseAffineMapOfSSAIds(map, parseElement, delimiter))
+    if (parser.parseAffineMapOfSSAIds(map, parseElement, addOperand, delimiter))
       return failure();
     // Add AffineMap attribute.
     if (map) {
@@ -1849,20 +1851,22 @@ class CustomOpAsmParser : public AsmParserImpl<OpAsmParser> {
   /// Parse an AffineExpr of SSA ids.
   ParseResult
   parseAffineExprOfSSAIds(SmallVectorImpl<UnresolvedOperand> &dimOperands,
-                          SmallVectorImpl<UnresolvedOperand> &symbOperands,
+                          SmallVectorImpl<UnresolvedOperand> &symOperands,
                           AffineExpr &expr) override {
-    auto parseElement = [&](bool isSymbol) -> ParseResult {
+    auto parseElement = [&]() -> FailureOr<UnresolvedOperand> {
       UnresolvedOperand operand;
       if (parseOperand(operand))
-        return failure();
+        return {};
+      return operand;
+    };
+    auto addOperand = [&](bool isSymbol, UnresolvedOperand operand) {
       if (isSymbol)
-        symbOperands.push_back(operand);
+        symOperands.push_back(operand);
       else
         dimOperands.push_back(operand);
-      return success();
     };
 
-    return parser.parseAffineExprOfSSAIds(expr, parseElement);
+    return parser.parseAffineExprOfSSAIds(expr, parseElement, addOperand);
   }
 
   //===--------------------------------------------------------------------===//

diff  --git a/mlir/lib/AsmParser/Parser.h b/mlir/lib/AsmParser/Parser.h
index ecc128cf767b3..7e8a0fc8e51b0 100644
--- a/mlir/lib/AsmParser/Parser.h
+++ b/mlir/lib/AsmParser/Parser.h
@@ -332,15 +332,17 @@ class Parser {
   ParseResult parseIntegerSetReference(IntegerSet &set);
 
   /// Parse an AffineMap where the dim and symbol identifiers are SSA ids.
-  ParseResult
-  parseAffineMapOfSSAIds(AffineMap &map,
-                         function_ref<ParseResult(bool)> parseElement,
-                         Delimiter delimiter);
+  ParseResult parseAffineMapOfSSAIds(
+      AffineMap &map,
+      function_ref<FailureOr<OpAsmParser::UnresolvedOperand>()> parseElement,
+      function_ref<void(bool, OpAsmParser::UnresolvedOperand)> addOperand,
+      Delimiter delimiter);
 
   /// Parse an AffineExpr where dim and symbol identifiers are SSA ids.
-  ParseResult
-  parseAffineExprOfSSAIds(AffineExpr &expr,
-                          function_ref<ParseResult(bool)> parseElement);
+  ParseResult parseAffineExprOfSSAIds(
+      AffineExpr &expr,
+      function_ref<FailureOr<OpAsmParser::UnresolvedOperand>()> parseElement,
+      function_ref<void(bool, OpAsmParser::UnresolvedOperand)> addOperand);
 
   //===--------------------------------------------------------------------===//
   // Code Completion

diff  --git a/mlir/test/Dialect/Affine/ops.mlir b/mlir/test/Dialect/Affine/ops.mlir
index 148c130146f33..db85981942688 100644
--- a/mlir/test/Dialect/Affine/ops.mlir
+++ b/mlir/test/Dialect/Affine/ops.mlir
@@ -530,3 +530,15 @@ func.func @affine_vector_load_store_alignment(%memref: memref<16xi32>) {
   affine.vector_store %val, %memref[0] { alignment = 8 } : memref<16xi32>, vector<4xi32>
   return
 }
+
+// -----
+
+// CHECK-LABEL: func @symbol_multi_result_reuse
+func.func @symbol_multi_result_reuse(%mem: memref<?x?xf32>) {
+  %res:2 = "test.op"() : () -> (index, index)
+  // CHECK: %{{.*}} = affine.load %{{.*}}[symbol(%{{.*}}#0), symbol(%{{.*}}#0)] : memref<?x?xf32>
+  %v1 = affine.load %mem[symbol(%res#0), symbol(%res#0)] : memref<?x?xf32>
+  // CHECK: %{{.*}} = affine.load %{{.*}}[symbol(%{{.*}}#0), symbol(%{{.*}}#1)] : memref<?x?xf32>
+  %v2 = affine.load %mem[symbol(%res#0), symbol(%res#1)] : memref<?x?xf32>
+  return
+}


        


More information about the Mlir-commits mailing list