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

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Jonas Rickert (jorickert)

<details>
<summary>Changes</summary>

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

---
Full diff: https://github.com/llvm/llvm-project/pull/105529.diff


3 Files Affected:

- (modified) mlir/lib/AsmParser/Parser.h (-1) 
- (modified) mlir/unittests/Parser/CMakeLists.txt (+2) 
- (modified) mlir/unittests/Parser/ParserTest.cpp (+81) 


``````````diff
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();
   }
 
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

``````````

</details>


https://github.com/llvm/llvm-project/pull/105529