Added PoC name retention for parsing and printing

Wheest · Jan 21, 2024 · 8b07f22 · 8b07f22
1 parent 01ba627
commit 8b07f22
Show file tree

Hide file tree

Showing 4 changed files with 29 additions and 7 deletions.
diff --git a/mlir/include/mlir/IR/MLIRContext.h b/mlir/include/mlir/IR/MLIRContext.h
@@ -12,6 +12,7 @@
 #include "mlir/Support/LLVM.h"
 #include "mlir/Support/TypeID.h"
 #include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/DenseMap.h"
 #include <functional>
 #include <memory>
 #include <vector>
@@ -34,6 +35,8 @@ class MLIRContextImpl;
 class RegisteredOperationName;
 class StorageUniquer;
 class IRUnit;
+class Value;
+class Operation;
 
 /// MLIRContext is the top-level object for a collection of MLIR operations. It
 /// holds immortal uniqued objects like types, and the tables used to unique
@@ -240,6 +243,11 @@ class MLIRContext {
   /// (attributes, operations, types, etc.).
   llvm::hash_code getRegistryHash();
 
+  ///===--------------------------------------------------------------------===//
+  llvm::DenseMap<Operation *, llvm::StringRef> valueToAlias;
+  void setAliasName(Value *value, llvm::StringRef aliasName);
+  llvm::StringRef getAliasName(Value *value);
+
   //===--------------------------------------------------------------------===//
   // Action API
   //===--------------------------------------------------------------------===//

diff --git a/mlir/lib/AsmParser/Parser.cpp b/mlir/lib/AsmParser/Parser.cpp
@@ -938,6 +938,11 @@ ParseResult OperationParser::addDefinition(UnresolvedOperand useInfo,
   /// Record this definition for the current scope.
   entries[useInfo.number] = {value, useInfo.location};
   recordDefinition(useInfo.name);
+
+  /// Register the alias name for use in printing
+  llvm::StringRef modifiedName = useInfo.name.drop_front(1);
+  getContext()->setAliasName(&value, modifiedName);
+
   return success();
 }
 

diff --git a/mlir/lib/IR/AsmPrinter.cpp b/mlir/lib/IR/AsmPrinter.cpp
@@ -72,13 +72,12 @@ 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()); });
-  }
-
-
+  return parseCommaSeparatedList(
+      [&]() { return parseType(result.emplace_back()); });
+}
 
 //===----------------------------------------------------------------------===//
 // DialectAsmPrinter
@@ -1417,7 +1416,9 @@ void SSANameState::printValueID(Value value, bool printResultNo,
   }
 
   stream << '%';
-  if (it->second != NameSentinel) {
+  if (!value.getContext()->getAliasName(&value).empty()) {
+    stream << value.getContext()->getAliasName(&value);
+  } else if (it->second != NameSentinel) {
     stream << it->second;
   } else {
     auto nameIt = valueNames.find(lookupValue);

diff --git a/mlir/lib/IR/MLIRContext.cpp b/mlir/lib/IR/MLIRContext.cpp
@@ -392,6 +392,14 @@ void MLIRContext::executeActionInternal(function_ref<void()> actionFn,
 
 bool MLIRContext::hasActionHandler() { return (bool)getImpl().actionHandler; }
 
+void MLIRContext::setAliasName(Value *value, llvm::StringRef aliasName) {
+  valueToAlias[value->getDefiningOp()] = aliasName;
+}
+llvm::StringRef MLIRContext::getAliasName(Value *value) {
+  auto found = valueToAlias.find(value->getDefiningOp());
+  return (found != valueToAlias.end()) ? found->second : llvm::StringRef();
+}
+
 //===----------------------------------------------------------------------===//
 // Diagnostic Handlers
 //===----------------------------------------------------------------------===//