[Mlir-commits] [mlir] [mlir][affine] Fix parsing operands that specify result number (PR #215539)
Ivan R. Ivanov
llvmlistbot at llvm.org
Tue Aug 11 05:25:44 PDT 2026
https://github.com/ivanradanov created https://github.com/llvm/llvm-project/pull/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.
>From 3aa6b6c0fd164cc395896fa6bbb922cf1da1c309 Mon Sep 17 00:00:00 2001
From: Ivan Radanov Ivanov <iivanov at nvidia.com>
Date: Tue, 11 Aug 2026 03:18:11 -0700
Subject: [PATCH] [mlir][affine] Fix parsing operands that specify result
number
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.
---
mlir/lib/AsmParser/AffineParser.cpp | 71 +++++++++++++++++------------
mlir/lib/AsmParser/Parser.cpp | 24 ++++++----
mlir/lib/AsmParser/Parser.h | 16 ++++---
mlir/test/Dialect/Affine/ops.mlir | 12 +++++
4 files changed, 76 insertions(+), 47 deletions(-)
diff --git a/mlir/lib/AsmParser/AffineParser.cpp b/mlir/lib/AsmParser/AffineParser.cpp
index 1797611858c06..b371df6629127 100644
--- a/mlir/lib/AsmParser/AffineParser.cpp
+++ b/mlir/lib/AsmParser/AffineParser.cpp
@@ -57,10 +57,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);
@@ -106,10 +110,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
@@ -295,7 +301,7 @@ AffineExpr AffineParser::parseBareIdExpr() {
StringRef sRef = getTokenSpelling();
for (auto entry : dimsAndSymbols) {
- if (entry.first == sRef) {
+ if (entry.first.first == sRef) {
consumeToken();
return entry.second;
}
@@ -310,21 +316,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;
}
@@ -488,12 +494,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();
}
@@ -551,7 +557,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);
}
@@ -742,20 +749,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 952d7e460c6e2..c9db1387dfed2 100644
--- a/mlir/lib/AsmParser/Parser.cpp
+++ b/mlir/lib/AsmParser/Parser.cpp
@@ -1812,19 +1812,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) {
@@ -1841,20 +1843,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