[Mlir-commits] [mlir] aa21ce4 - [mlir] Do not set lastToken in AsmParser's resetToken function and add a unit test for AsmParsers's locations (#105529)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Sep 9 02:48:18 PDT 2024
Author: Jonas Rickert
Date: 2024-09-09T11:48:13+02:00
New Revision: aa21ce4a792c170074193c32e8ba8dd35e57c628
URL: https://github.com/llvm/llvm-project/commit/aa21ce4a792c170074193c32e8ba8dd35e57c628
DIFF: https://github.com/llvm/llvm-project/commit/aa21ce4a792c170074193c32e8ba8dd35e57c628.diff
LOG: [mlir] Do not set lastToken in AsmParser's resetToken function and add a unit test for AsmParsers's locations (#105529)
This changes the function `resetToken` to not update `lastToken`.
The member `lastToken` is the last token that was consumed by the
parser.
Resetting the lexer position to a different position does not cause any
token to be consumed, so `lastToken` should not be updated.
Setting it to `curToken` can cause the scopeLoc.end location of
`OperationDefinition `to be off-by-one, pointing to the
first token after the operation.
An example for an operation for which the scopeLoc.end location was
wrong before is:
```
%0 = torch.vtensor.literal(dense_resource<__elided__> : tensor<768xbf16>) : !torch.vtensor<[768],bf16>
```
Here the scope end loc always pointed to the next token
This also adds a test for the Locations of `OperationDefinitions`.
Without the change to `resetToken` the test failes, with the scope end
location for `llvm.mlir.undef` pointing to the `func.return` in the next
line
Added:
Modified:
mlir/lib/AsmParser/Parser.h
mlir/unittests/Parser/CMakeLists.txt
mlir/unittests/Parser/ParserTest.cpp
Removed:
################################################################################
diff --git a/mlir/lib/AsmParser/Parser.h b/mlir/lib/AsmParser/Parser.h
index 4caab499e1a0e4..bf91831798056b 100644
--- a/mlir/lib/AsmParser/Parser.h
+++ b/mlir/lib/AsmParser/Parser.h
@@ -130,10 +130,14 @@ class Parser {
consumeToken();
}
- /// Reset the parser to the given lexer position.
+ /// Reset the parser to the given lexer position. Resetting the parser/lexer
+ /// position does not update 'state.lastToken'. 'state.lastToken' is the
+ /// last parsed token, and is used to provide the scope end location for
+ /// OperationDefinitions. To ensure the correctness of the end location, the
+ /// last consumed token of an OperationDefinition needs to be the last token
+ /// belonging to it.
void resetToken(const char *tokPos) {
state.lex.resetPointer(tokPos);
- state.lastToken = state.curToken;
state.curToken = state.lex.lexToken();
}
diff --git a/mlir/unittests/Parser/CMakeLists.txt b/mlir/unittests/Parser/CMakeLists.txt
index f5646ff3f2c345..a5e2da56ffb57e 100644
--- a/mlir/unittests/Parser/CMakeLists.txt
+++ b/mlir/unittests/Parser/CMakeLists.txt
@@ -8,6 +8,8 @@ add_mlir_unittest(MLIRParserTests
target_include_directories(MLIRParserTests PRIVATE "${MLIR_BINARY_DIR}/test/lib/Dialect/Test")
target_link_libraries(MLIRParserTests PRIVATE
+ MLIRFuncDialect
+ MLIRLLVMDialect
MLIRIR
MLIRParser
MLIRTestDialect
diff --git a/mlir/unittests/Parser/ParserTest.cpp b/mlir/unittests/Parser/ParserTest.cpp
index 62f609ecf80492..52b965bcc1326b 100644
--- a/mlir/unittests/Parser/ParserTest.cpp
+++ b/mlir/unittests/Parser/ParserTest.cpp
@@ -8,8 +8,12 @@
#include "mlir/Parser/Parser.h"
#include "mlir/AsmParser/AsmParser.h"
+#include "mlir/AsmParser/AsmParserState.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
#include "mlir/IR/BuiltinOps.h"
#include "mlir/IR/Verifier.h"
+#include "llvm/Support/SourceMgr.h"
#include "gmock/gmock.h"
@@ -101,4 +105,81 @@ TEST(MLIRParser, ParseAttr) {
EXPECT_EQ(attr, b.getI64IntegerAttr(9));
}
}
+
+TEST(MLIRParser, AsmParserLocations) {
+ std::string moduleStr = R"mlir(
+func.func @foo() -> !llvm.array<2 x f32> {
+ %0 = llvm.mlir.undef : !llvm.array<2 x f32>
+ func.return %0 : !llvm.array<2 x f32>
+}
+ )mlir";
+
+ DialectRegistry registry;
+ registry.insert<func::FuncDialect, LLVM::LLVMDialect>();
+ MLIRContext context(registry);
+
+ auto memBuffer =
+ llvm::MemoryBuffer::getMemBuffer(moduleStr, "AsmParserTest.mlir",
+ /*RequiresNullTerminator=*/false);
+ ASSERT_TRUE(memBuffer);
+
+ llvm::SourceMgr sourceMgr;
+ sourceMgr.AddNewSourceBuffer(std::move(memBuffer), llvm::SMLoc());
+
+ Block block;
+ AsmParserState parseState;
+ const LogicalResult parseResult =
+ parseAsmSourceFile(sourceMgr, &block, &context, &parseState);
+ ASSERT_TRUE(parseResult.succeeded());
+
+ auto funcOp = *block.getOps<func::FuncOp>().begin();
+ const AsmParserState::OperationDefinition *funcOpDefinition =
+ parseState.getOpDef(funcOp);
+ ASSERT_TRUE(funcOpDefinition);
+
+ const std::pair expectedStartFunc{2u, 1u};
+ const std::pair expectedEndFunc{2u, 10u};
+ const std::pair expectedScopeEndFunc{5u, 2u};
+ ASSERT_EQ(sourceMgr.getLineAndColumn(funcOpDefinition->loc.Start),
+ expectedStartFunc);
+ ASSERT_EQ(sourceMgr.getLineAndColumn(funcOpDefinition->loc.End),
+ expectedEndFunc);
+ ASSERT_EQ(funcOpDefinition->loc.Start, funcOpDefinition->scopeLoc.Start);
+ ASSERT_EQ(sourceMgr.getLineAndColumn(funcOpDefinition->scopeLoc.End),
+ expectedScopeEndFunc);
+
+ auto llvmUndef = *funcOp.getOps<LLVM::UndefOp>().begin();
+ const AsmParserState::OperationDefinition *llvmUndefDefinition =
+ parseState.getOpDef(llvmUndef);
+ ASSERT_TRUE(llvmUndefDefinition);
+
+ const std::pair expectedStartUndef{3u, 8u};
+ const std::pair expectedEndUndef{3u, 23u};
+ const std::pair expectedScopeEndUndef{3u, 46u};
+ ASSERT_EQ(sourceMgr.getLineAndColumn(llvmUndefDefinition->loc.Start),
+ expectedStartUndef);
+ ASSERT_EQ(sourceMgr.getLineAndColumn(llvmUndefDefinition->loc.End),
+ expectedEndUndef);
+ ASSERT_EQ(llvmUndefDefinition->loc.Start,
+ llvmUndefDefinition->scopeLoc.Start);
+ ASSERT_EQ(sourceMgr.getLineAndColumn(llvmUndefDefinition->scopeLoc.End),
+ expectedScopeEndUndef);
+
+ auto funcReturn = *funcOp.getOps<func::ReturnOp>().begin();
+ const AsmParserState::OperationDefinition *funcReturnDefinition =
+ parseState.getOpDef(funcReturn);
+ ASSERT_TRUE(funcReturnDefinition);
+
+ const std::pair expectedStartReturn{4u, 3u};
+ const std::pair expectedEndReturn{4u, 14u};
+ const std::pair expectedScopeEndReturn{4u, 40u};
+ ASSERT_EQ(sourceMgr.getLineAndColumn(funcReturnDefinition->loc.Start),
+ expectedStartReturn);
+ ASSERT_EQ(sourceMgr.getLineAndColumn(funcReturnDefinition->loc.End),
+ expectedEndReturn);
+ ASSERT_EQ(funcReturnDefinition->loc.Start,
+ funcReturnDefinition->scopeLoc.Start);
+ ASSERT_EQ(sourceMgr.getLineAndColumn(funcReturnDefinition->scopeLoc.End),
+ expectedScopeEndReturn);
+}
} // namespace
More information about the Mlir-commits
mailing list