[Mlir-commits] [mlir] [mlir] Do not set lastToken in AsmParser's resetToken function and add a unit test for AsmParsers's locations (PR #105529)

Jonas Rickert llvmlistbot at llvm.org
Wed Aug 21 06:56:27 PDT 2024


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

>From cc73c18bb0fda041641bb53114065e66c17cb7e6 Mon Sep 17 00:00:00 2001
From: "Rickert, Jonas" <Jonas.Rickert at amd.com>
Date: Tue, 20 Aug 2024 12:45:42 +0100
Subject: [PATCH 1/2] [mlir] Do not set lastToken in AsmParser's resetToken
 function

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
---
 mlir/lib/AsmParser/Parser.h | 1 -
 1 file changed, 1 deletion(-)

diff --git a/mlir/lib/AsmParser/Parser.h b/mlir/lib/AsmParser/Parser.h
index 4caab499e1a0e4..3c556b19870770 100644
--- a/mlir/lib/AsmParser/Parser.h
+++ b/mlir/lib/AsmParser/Parser.h
@@ -133,7 +133,6 @@ class Parser {
   /// Reset the parser to the given lexer position.
   void resetToken(const char *tokPos) {
     state.lex.resetPointer(tokPos);
-    state.lastToken = state.curToken;
     state.curToken = state.lex.lexToken();
   }
 

>From 504d64611835978d9a8581df7514d0fda9f87f64 Mon Sep 17 00:00:00 2001
From: "Rickert, Jonas" <Jonas.Rickert at amd.com>
Date: Wed, 21 Aug 2024 10:56:50 +0100
Subject: [PATCH 2/2] [mlir] Add a unit test for AsmParser's file locations

This test checks the parsed ops' locations and scope locations.
It uses the !llvm.array type in the test to check that locations are correct for
ops that contain an extended-type.
The parser for extended-types calls resetToken, which caused wrong locations in
the past.
---
 mlir/unittests/Parser/CMakeLists.txt |  2 +
 mlir/unittests/Parser/ParserTest.cpp | 81 ++++++++++++++++++++++++++++
 2 files changed, 83 insertions(+)

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