[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 Sep 4 02:31:27 PDT 2024


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

>From 585e04ad72315b685e5761a8304f36888dbbc18a 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 | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

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

>From 8e2f39e2ffcdd1f4538c23404dd00b92afc30549 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