[Mlir-commits] [mlir] [mlir] Retain original identifier names for debugging v2 (PR #119944)

[Maksim Levental]

https://github.com/makslevental updated https://github.com/llvm/llvm-project/pull/119944

>From 150b74657e45d1f19ade9891b0c5b878d21fe8e1 Mon Sep 17 00:00:00 2001
From: max <maksim.levental at gmail.com>
Date: Fri, 13 Dec 2024 23:39:19 -0500
Subject: [PATCH] [mlir] retain identifier names

---
 mlir/lib/AsmParser/Parser.cpp | 14 ++++++++++++++
 mlir/lib/IR/AsmPrinter.cpp    | 18 ++++++++++++++----
 2 files changed, 28 insertions(+), 4 deletions(-)

diff --git a/mlir/lib/AsmParser/Parser.cpp b/mlir/lib/AsmParser/Parser.cpp
index e3db248164672c..6623ddeb01adc0 100644
--- a/mlir/lib/AsmParser/Parser.cpp
+++ b/mlir/lib/AsmParser/Parser.cpp
@@ -2071,6 +2071,14 @@ OperationParser::parseCustomOperation(ArrayRef<ResultRecord> resultIDs) {
 
   // Get location information for the operation.
   auto srcLocation = getEncodedSourceLocation(opLoc);
+  if (resultIDs.size() == 1) {
+    std::tuple<StringRef, unsigned, SMLoc> resultID = resultIDs.front();
+    auto id = std::get<0>(resultID);
+    srcLocation = NameLoc::get(
+        StringAttr::get(state.config.getContext(), id.drop_front(1)),
+        srcLocation);
+  }
+
   OperationState opState(srcLocation, *opNameInfo);
 
   // If we are populating the parser state, start a new operation definition.
@@ -2235,6 +2243,9 @@ ParseResult OperationParser::parseRegionBody(Region &region, SMLoc startLoc,
       Location loc = entryArg.sourceLoc.has_value()
                          ? *entryArg.sourceLoc
                          : getEncodedSourceLocation(argInfo.location);
+      loc = NameLoc::get(StringAttr::get(state.config.getContext(),
+                                         entryArg.ssaName.name.drop_front(1)),
+                         loc);
       BlockArgument arg = block->addArgument(entryArg.type, loc);
 
       // Add a definition of this arg to the assembly state if provided.
@@ -2415,6 +2426,9 @@ ParseResult OperationParser::parseOptionalBlockArgList(Block *owner) {
               return emitError("argument and block argument type mismatch");
           } else {
             auto loc = getEncodedSourceLocation(useInfo.location);
+            loc = NameLoc::get(StringAttr::get(state.config.getContext(),
+                                               useInfo.name.drop_front(1)),
+                               loc);
             arg = owner->addArgument(type, loc);
           }
 
diff --git a/mlir/lib/IR/AsmPrinter.cpp b/mlir/lib/IR/AsmPrinter.cpp
index 61b90bc9b0a7bb..e51a7305b43794 100644
--- a/mlir/lib/IR/AsmPrinter.cpp
+++ b/mlir/lib/IR/AsmPrinter.cpp
@@ -73,7 +73,8 @@ OpAsmParser::~OpAsmParser() = default;
 MLIRContext *AsmParser::getContext() const { return getBuilder().getContext(); }
 
 /// Parse a type list.
-/// This is out-of-line to work-around https://github.com/llvm/llvm-project/issues/62918
+/// This is out-of-line to work-around
+/// https://github.com/llvm/llvm-project/issues/62918
 ParseResult AsmParser::parseTypeList(SmallVectorImpl<Type> &result) {
   return parseCommaSeparatedList(
       [&]() { return parseType(result.emplace_back()); });
@@ -1511,7 +1512,10 @@ void SSANameState::numberValuesInRegion(Region &region) {
     assert(!valueIDs.count(arg) && "arg numbered multiple times");
     assert(llvm::cast<BlockArgument>(arg).getOwner()->getParent() == &region &&
            "arg not defined in current region");
-    setValueName(arg, name);
+    if (auto nameLoc = dyn_cast<NameLoc>(arg.getLoc()))
+      setValueName(arg, nameLoc.getName());
+    else
+      setValueName(arg, name);
   };
 
   if (!printerFlags.shouldPrintGenericOpForm()) {
@@ -1553,7 +1557,10 @@ void SSANameState::numberValuesInBlock(Block &block) {
       specialNameBuffer.resize(strlen("arg"));
       specialName << nextArgumentID++;
     }
-    setValueName(arg, specialName.str());
+    if (auto nameLoc = dyn_cast<NameLoc>(arg.getLoc()))
+      setValueName(arg, nameLoc.getName());
+    else
+      setValueName(arg, specialName.str());
   }
 
   // Number the operations in this block.
@@ -1567,7 +1574,10 @@ void SSANameState::numberValuesInOp(Operation &op) {
   auto setResultNameFn = [&](Value result, StringRef name) {
     assert(!valueIDs.count(result) && "result numbered multiple times");
     assert(result.getDefiningOp() == &op && "result not defined by 'op'");
-    setValueName(result, name);
+    if (auto nameLoc = dyn_cast<NameLoc>(result.getLoc()))
+      setValueName(result, nameLoc.getName());
+    else
+      setValueName(result, name);
 
     // Record the result number for groups not anchored at 0.
     if (int resultNo = llvm::cast<OpResult>(result).getResultNumber())