[Mlir-commits] [mlir] [mlir][openacc] Avoid per-routine module walks in ACCRoutineToGPUFunc (PR #221345)

[Valentin Clement バレンタイン クレメン]

https://github.com/clementval updated https://github.com/llvm/llvm-project/pull/221345

>From 1665c3a16ccdd1a5d61a1ef464052034411782d5 Mon Sep 17 00:00:00 2001
From: Valentin Clement <clementval at gmail.com>
Date: Fri, 4 Sep 2026 14:01:08 -0700
Subject: [PATCH 1/2] [mlir][openacc] Avoid per-routine module walks in
 ACCRoutineToGPUFunc

---
 .../Transforms/ACCRoutineToGPUFunc.cpp        | 71 ++++++++++++++-----
 1 file changed, 53 insertions(+), 18 deletions(-)

diff --git a/mlir/lib/Dialect/OpenACC/Transforms/ACCRoutineToGPUFunc.cpp b/mlir/lib/Dialect/OpenACC/Transforms/ACCRoutineToGPUFunc.cpp
index a83f4ac988408..4ce31cf059215 100644
--- a/mlir/lib/Dialect/OpenACC/Transforms/ACCRoutineToGPUFunc.cpp
+++ b/mlir/lib/Dialect/OpenACC/Transforms/ACCRoutineToGPUFunc.cpp
@@ -67,10 +67,12 @@
 #include "mlir/Dialect/GPU/IR/GPUDialect.h"
 #include "mlir/Dialect/OpenACC/Analysis/OpenACCSupport.h"
 #include "mlir/Dialect/OpenACC/OpenACC.h"
+#include "mlir/IR/AttrTypeSubElements.h"
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/IRMapping.h"
 #include "mlir/IR/SymbolTable.h"
 #include "mlir/Interfaces/CallInterfaces.h"
+#include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/SetVector.h"
 #include <string>
 
@@ -224,14 +226,48 @@ static LogicalResult processCallsInRoutines(
   return success();
 }
 
+/// Rewrite symbol uses of uniqued specialized names back to the original
+/// routine names. One walk of the module body instead of
+/// replaceAllSymbolUses per routine (which walks all of the IR each time).
+static void
+rewriteSpecializedSymbolUses(ModuleOp mod,
+                             const DenseMap<StringAttr, StringAttr> &renames) {
+  if (renames.empty())
+    return;
+
+  AttrTypeReplacer replacer;
+  replacer.addReplacement(
+      [&](SymbolRefAttr attr) -> std::pair<Attribute, WalkResult> {
+        auto it = renames.find(attr.getRootReference());
+        if (it == renames.end())
+          return {attr, WalkResult::skip()};
+        if (isa<FlatSymbolRefAttr>(attr))
+          return {FlatSymbolRefAttr::get(it->second), WalkResult::skip()};
+        return {SymbolRefAttr::get(it->second, attr.getNestedReferences()),
+                WalkResult::skip()};
+      });
+
+  // Match SymbolTable::replaceAllSymbolUses: walk the module body and do
+  // not enter nested symbol tables (e.g. gpu.module).
+  for (Region &region : mod->getRegions()) {
+    region.walk([&](Operation *op) {
+      if (op->hasTrait<OpTrait::SymbolTable>())
+        return WalkResult::skip();
+      replacer.replaceElementsIn(op);
+      return WalkResult::advance();
+    });
+  }
+}
+
 /// Clone each function in funcsToClone into the GPU module (declaration or
 /// full body). Fix up symbol names and specialized_routine attr for ACC
 /// routines.
 static LogicalResult cloneFuncsToGPUModule(
-    ModuleOp mod, OpenACCSupport &accSupport, SymbolTable &gpuSymTab,
+    ModuleOp mod, SymbolTable &gpuSymTab,
     const llvm::SmallSetVector<CloneCandidate, 4> &funcsToClone) {
-  MLIRContext *ctx = mod.getContext();
-  OpBuilder builder(ctx);
+  OpBuilder builder(mod.getContext());
+  DenseMap<StringAttr, StringAttr> specializedRenames;
+  SmallVector<gpu::GPUFuncOp> deviceFuncs;
 
   for (CloneCandidate candidate : funcsToClone) {
     func::FuncOp srcFunc = candidate.first;
@@ -244,26 +280,25 @@ static LogicalResult cloneFuncsToGPUModule(
 
     gpu::GPUFuncOp deviceFuncOp = createGPUFuncFromFunc(builder, srcFunc);
 
-    if (auto specRoutineAttr =
-            srcFunc->getDiscardableAttrOfType<SpecializedRoutineAttr>(
-                getSpecializedRoutineAttrName())) {
-      StringAttr funcName = specRoutineAttr.getFuncName();
-      if (failed(SymbolTable::replaceAllSymbolUses(
-              StringAttr::get(ctx, deviceFuncOp.getName()), funcName, mod))) {
-        accSupport.emitNYI(deviceFuncOp.getLoc(),
-                           "cannot replace symbol for acc routine");
-        return failure();
-      }
-      deviceFuncOp.setName(funcName);
-    }
     if (auto specAttr =
             srcFunc->getDiscardableAttrOfType<SpecializedRoutineAttr>(
-                getSpecializedRoutineAttrName()))
+                getSpecializedRoutineAttrName())) {
+      StringAttr destName = specAttr.getFuncName();
+      StringAttr srcName = deviceFuncOp.getNameAttr();
+      if (srcName != destName)
+        specializedRenames.try_emplace(srcName, destName);
+      deviceFuncOp.setName(destName);
       deviceFuncOp->setDiscardableAttr(getSpecializedRoutineAttrName(),
                                        specAttr);
+    }
 
-    gpuSymTab.insert(deviceFuncOp);
+    deviceFuncs.push_back(deviceFuncOp);
   }
+
+  rewriteSpecializedSymbolUses(mod, specializedRenames);
+
+  for (gpu::GPUFuncOp deviceFuncOp : deviceFuncs)
+    gpuSymTab.insert(deviceFuncOp);
   return success();
 }
 
@@ -319,7 +354,7 @@ class ACCRoutineToGPUFunc
                                       materializedAccRoutines, funcsToClone)))
       return signalPassFailure();
 
-    if (failed(cloneFuncsToGPUModule(mod, accSupport, gpuSymTab, funcsToClone)))
+    if (failed(cloneFuncsToGPUModule(mod, gpuSymTab, funcsToClone)))
       return signalPassFailure();
 
     cleanupHostModule(funcsToClone);

>From f2121074d93e1fb6780f95951bb3f09d6cf84eac Mon Sep 17 00:00:00 2001
From: Valentin Clement <clementval at gmail.com>
Date: Fri, 4 Sep 2026 14:54:10 -0700
Subject: [PATCH 2/2] Fix ordering

---
 .../Transforms/ACCRoutineToGPUFunc.cpp        | 32 ++++++++++++-------
 1 file changed, 20 insertions(+), 12 deletions(-)

diff --git a/mlir/lib/Dialect/OpenACC/Transforms/ACCRoutineToGPUFunc.cpp b/mlir/lib/Dialect/OpenACC/Transforms/ACCRoutineToGPUFunc.cpp
index 4ce31cf059215..6d9528ce57549 100644
--- a/mlir/lib/Dialect/OpenACC/Transforms/ACCRoutineToGPUFunc.cpp
+++ b/mlir/lib/Dialect/OpenACC/Transforms/ACCRoutineToGPUFunc.cpp
@@ -266,8 +266,25 @@ static LogicalResult cloneFuncsToGPUModule(
     ModuleOp mod, SymbolTable &gpuSymTab,
     const llvm::SmallSetVector<CloneCandidate, 4> &funcsToClone) {
   OpBuilder builder(mod.getContext());
+
+  // ACCRoutineLowering inserts the device copy next to the host function, so
+  // the symbol table uniqued its name while acc.specialized_routine still
+  // holds the original one. Collect all of those renames and apply them in a
+  // single walk below.
   DenseMap<StringAttr, StringAttr> specializedRenames;
-  SmallVector<gpu::GPUFuncOp> deviceFuncs;
+  for (CloneCandidate candidate : funcsToClone) {
+    func::FuncOp srcFunc = candidate.first;
+    if (srcFunc.isDeclaration())
+      continue;
+    if (auto specAttr =
+            srcFunc->getDiscardableAttrOfType<SpecializedRoutineAttr>(
+                getSpecializedRoutineAttrName())) {
+      StringAttr destName = specAttr.getFuncName();
+      if (srcFunc.getNameAttr() != destName)
+        specializedRenames.try_emplace(srcFunc.getNameAttr(), destName);
+    }
+  }
+  rewriteSpecializedSymbolUses(mod, specializedRenames);
 
   for (CloneCandidate candidate : funcsToClone) {
     func::FuncOp srcFunc = candidate.first;
@@ -283,22 +300,13 @@ static LogicalResult cloneFuncsToGPUModule(
     if (auto specAttr =
             srcFunc->getDiscardableAttrOfType<SpecializedRoutineAttr>(
                 getSpecializedRoutineAttrName())) {
-      StringAttr destName = specAttr.getFuncName();
-      StringAttr srcName = deviceFuncOp.getNameAttr();
-      if (srcName != destName)
-        specializedRenames.try_emplace(srcName, destName);
-      deviceFuncOp.setName(destName);
+      deviceFuncOp.setName(specAttr.getFuncName());
       deviceFuncOp->setDiscardableAttr(getSpecializedRoutineAttrName(),
                                        specAttr);
     }
 
-    deviceFuncs.push_back(deviceFuncOp);
-  }
-
-  rewriteSpecializedSymbolUses(mod, specializedRenames);
-
-  for (gpu::GPUFuncOp deviceFuncOp : deviceFuncs)
     gpuSymTab.insert(deviceFuncOp);
+  }
   return success();
 }