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

[llvmlistbot at llvm.org]

Author: Valentin Clement (バレンタイン クレメン)
Date: 2026-09-04T16:04:34-07:00
New Revision: b6abc5ec3b4cf77853d534d67c61de0363c5e35e

URL: https://github.com/llvm/llvm-project/commit/b6abc5ec3b4cf77853d534d67c61de0363c5e35e
DIFF: https://github.com/llvm/llvm-project/commit/b6abc5ec3b4cf77853d534d67c61de0363c5e35e.diff

LOG: [mlir][openacc] Avoid per-routine module walks in ACCRoutineToGPUFunc (#221345)

ACCRoutineLowering uniquifies the specialized device copy (foo -> foo_0)
while acc.specialized_routine still names the original host function.
Moving those copies into gpu.func then called replaceAllSymbolUses on
the whole module once per routine, which is quadratic in the number of
specialized routines.

Collect the src->dest renames and rewrite symbol uses in a single walk
of the module body (still skipping nested symbol tables). Insert the
gpu.func ops under the final names afterward.

On a large TU this speeds up this pass by 3.5x

Added: 
    

Modified: 
    mlir/lib/Dialect/OpenACC/Transforms/ACCRoutineToGPUFunc.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/OpenACC/Transforms/ACCRoutineToGPUFunc.cpp b/mlir/lib/Dialect/OpenACC/Transforms/ACCRoutineToGPUFunc.cpp
index a83f4ac988408..6d9528ce57549 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,65 @@ 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());
+
+  // 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;
+  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;
@@ -244,23 +297,13 @@ 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())) {
+      deviceFuncOp.setName(specAttr.getFuncName());
       deviceFuncOp->setDiscardableAttr(getSpecializedRoutineAttrName(),
                                        specAttr);
+    }
 
     gpuSymTab.insert(deviceFuncOp);
   }
@@ -319,7 +362,7 @@ class ACCRoutineToGPUFunc
                                       materializedAccRoutines, funcsToClone)))
       return signalPassFailure();
 
-    if (failed(cloneFuncsToGPUModule(mod, accSupport, gpuSymTab, funcsToClone)))
+    if (failed(cloneFuncsToGPUModule(mod, gpuSymTab, funcsToClone)))
       return signalPassFailure();
 
     cleanupHostModule(funcsToClone);