From a1ffab8786a66644f7254faad0755f0dae2107a1 Mon Sep 17 00:00:00 2001
From: Mehdi Amini <joker.eph@gmail.com>
Date: Mon, 28 Apr 2025 11:43:52 -0700
Subject: [PATCH] [mlir] Fix `--mlir-print-ir-tree-dir` when the symbol name
 contains `/`

Fixes #137622
---
 mlir/lib/Pass/IRPrinting.cpp              | 15 ++++++++++-----
 mlir/test/Pass/ir-printing-file-tree.mlir | 10 ++++++++--
 2 files changed, 18 insertions(+), 7 deletions(-)

diff --git a/mlir/lib/Pass/IRPrinting.cpp b/mlir/lib/Pass/IRPrinting.cpp
index 99e0601e4749c..c11dbc627c0be 100644
--- a/mlir/lib/Pass/IRPrinting.cpp
+++ b/mlir/lib/Pass/IRPrinting.cpp
@@ -10,6 +10,7 @@
 #include "mlir/IR/SymbolTable.h"
 #include "mlir/Pass/PassManager.h"
 #include "mlir/Support/FileUtilities.h"
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/Support/FileSystem.h"
 #include "llvm/Support/FormatVariadic.h"
@@ -210,22 +211,26 @@ struct BasicIRPrinterConfig : public PassManager::IRPrinterConfig {
 /// `op` first, `op` last). This information is used to construct the directory
 /// tree for the `FileTreeIRPrinterConfig` below.
 /// The counter for `op` will be incremented by this call.
-static std::pair<SmallVector<std::pair<std::string, StringRef>>, std::string>
+static std::pair<SmallVector<std::pair<std::string, std::string>>, std::string>
 getOpAndSymbolNames(Operation *op, StringRef passName,
                     llvm::DenseMap<Operation *, unsigned> &counters) {
-  SmallVector<std::pair<std::string, StringRef>> pathElements;
+  SmallVector<std::pair<std::string, std::string>> pathElements;
   SmallVector<unsigned> countPrefix;
 
   Operation *iter = op;
   ++counters.try_emplace(op, -1).first->second;
   while (iter) {
     countPrefix.push_back(counters[iter]);
-    StringAttr symbolName =
+    StringAttr symbolNameAttr =
         iter->getAttrOfType<StringAttr>(SymbolTable::getSymbolAttrName());
+    std::string symbolName =
+        symbolNameAttr ? symbolNameAttr.str() : "no-symbol-name";
+    llvm::replace(symbolName, '/', '_');
+    llvm::replace(symbolName, '\\', '_');
+
     std::string opName =
         llvm::join(llvm::split(iter->getName().getStringRef().str(), '.'), "_");
-    pathElements.emplace_back(opName, symbolName ? symbolName.strref()
-                                                 : "no-symbol-name");
+    pathElements.emplace_back(std::move(opName), std::move(symbolName));
     iter = iter->getParentOp();
   }
   // Return in the order of top level (module) down to `op`.
diff --git a/mlir/test/Pass/ir-printing-file-tree.mlir b/mlir/test/Pass/ir-printing-file-tree.mlir
index b00d77db2c603..46bf6a236cc31 100644
--- a/mlir/test/Pass/ir-printing-file-tree.mlir
+++ b/mlir/test/Pass/ir-printing-file-tree.mlir
@@ -15,7 +15,8 @@
 // RUN:   -mlir-print-ir-after-all
 // RUN: test -f %t/builtin_module_outer/0_canonicalize.mlir
 // RUN: test -f %t/builtin_module_outer/1_canonicalize.mlir
-// RUN: test -f %t/builtin_module_outer/func_func_symA/1_0_cse.mlir
+// RUN: test -f %t/builtin_module_outer/func_func_sym_A/1_0_cse.mlir
+// RUN: test -f %t/builtin_module_outer/func_func_sym_backslash/1_0_cse.mlir
 // RUN: test -f %t/builtin_module_outer/builtin_module_inner/1_0_canonicalize.mlir
 // RUN: test -f %t/builtin_module_outer/builtin_module_inner/func_func_symB/1_0_0_cse.mlir
 // RUN: test -f %t/builtin_module_outer/builtin_module_inner/func_func_symB/1_0_1_canonicalize.mlir
@@ -26,7 +27,10 @@
 
 builtin.module @outer {
 
-  func.func @symA() {
+  func.func @"sym/A"() {
+    return
+  }
+  func.func @"sym\\backslash"() {
     return
   }
 
@@ -38,4 +42,6 @@ builtin.module @outer {
       return
     }
   }
+
+
 }