[Mlir-commits] [mlir] [mlir][acc] Rewrite acc routine bind calls inside gpu.func (PR #204220)
Delaram Talaashrafi
llvmlistbot at llvm.org
Thu Jun 25 13:00:03 PDT 2026
https://github.com/delaram-talaashrafi updated https://github.com/llvm/llvm-project/pull/204220
>From 4a4abf592bbafd9a877a0873e959a5051a0a39f8 Mon Sep 17 00:00:00 2001
From: Delaram Talaashrafi <dtalaashrafi at rome5.pgi.net>
Date: Tue, 16 Jun 2026 10:56:32 -0700
Subject: [PATCH 1/5] [mlir][acc] Rewrite acc routine bind calls inside
gpu.func
Run `acc-bind-routine` on ModuleOp and rewrite calls to bound symbols in
offload regions and `gpu.func`. For string bind names, declare private
functions in the enclosing `gpu.module` symbol table when the call is inside
device code.
---
.../mlir/Dialect/OpenACC/Transforms/Passes.td | 2 +-
.../OpenACC/Transforms/ACCBindRoutine.cpp | 39 +++++++++++--------
.../Dialect/OpenACC/acc-bind-routine.mlir | 24 ++++++++++++
3 files changed, 48 insertions(+), 17 deletions(-)
diff --git a/mlir/include/mlir/Dialect/OpenACC/Transforms/Passes.td b/mlir/include/mlir/Dialect/OpenACC/Transforms/Passes.td
index 485f68dfc9338..6f5875f844846 100644
--- a/mlir/include/mlir/Dialect/OpenACC/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/OpenACC/Transforms/Passes.td
@@ -493,7 +493,7 @@ def ACCRoutineToGPUFunc : Pass<"acc-routine-to-gpu-func", "mlir::ModuleOp"> {
let options = [ AccDeviceTypeOption ];
}
-def ACCBindRoutine : Pass<"acc-bind-routine", "mlir::func::FuncOp"> {
+def ACCBindRoutine : Pass<"acc-bind-routine", "mlir::ModuleOp"> {
let summary = "Apply bind clause to function calls in ACC compute regions";
let description = [{
For calls inside offload regions that target a function with an
diff --git a/mlir/lib/Dialect/OpenACC/Transforms/ACCBindRoutine.cpp b/mlir/lib/Dialect/OpenACC/Transforms/ACCBindRoutine.cpp
index f4c8f18e5b04b..4ed4f4f0ffea0 100644
--- a/mlir/lib/Dialect/OpenACC/Transforms/ACCBindRoutine.cpp
+++ b/mlir/lib/Dialect/OpenACC/Transforms/ACCBindRoutine.cpp
@@ -81,24 +81,19 @@ class ACCBindRoutine : public acc::impl::ACCBindRoutineBase<ACCBindRoutine> {
using acc::impl::ACCBindRoutineBase<ACCBindRoutine>::ACCBindRoutineBase;
void runOnOperation() override {
- func::FuncOp func = getOperation();
- ModuleOp module = func->getParentOfType<ModuleOp>();
- if (!module)
- return;
-
+ ModuleOp module = getOperation();
+ OpenACCSupport &accSupport = getAnalysis<OpenACCSupport>();
SymbolTable symTab(module);
- auto cachedAnalysis =
- getCachedParentAnalysis<OpenACCSupport>(func->getParentOp());
- OpenACCSupport &accSupport =
- cachedAnalysis ? cachedAnalysis->get() : getAnalysis<OpenACCSupport>();
bool failed = false;
- func.walk([&](acc::OffloadRegionOpInterface offload) {
- Region ®ion = offload.getOffloadRegion();
- region.walk([&](CallOpInterface callOp) {
+ module.walk([&](FunctionOpInterface func) {
+ func.walk([&](CallOpInterface callOp) {
if (!callOp.getCallableForCallee())
return;
+ if (!callOp->getParentOfType<OffloadRegionOpInterface>() &&
+ !callOp->getParentOfType<gpu::GPUFuncOp>())
+ return;
SymbolRefAttr calleeSymbolRef =
dyn_cast<SymbolRefAttr>(callOp.getCallableForCallee());
if (!calleeSymbolRef)
@@ -131,14 +126,26 @@ class ACCBindRoutine : public acc::impl::ACCBindRoutineBase<ACCBindRoutine> {
bindNameOpt = routine.getBindNameValue();
if (!bindNameOpt)
return;
-
SymbolRefAttr calleeRef;
if (auto *symRef = std::get_if<SymbolRefAttr>(&*bindNameOpt)) {
calleeRef = *symRef;
} else {
- calleeRef = FlatSymbolRefAttr::get(
- callOp.getContext(),
- std::get<StringAttr>(*bindNameOpt).getValue());
+ StringRef bindName = std::get<StringAttr>(*bindNameOpt).getValue();
+ auto gpuMod = func->getParentOfType<gpu::GPUModuleOp>();
+ Operation *symbolTableOp =
+ gpuMod ? gpuMod.getOperation() : module.getOperation();
+ SymbolTable insertSymTab(symbolTableOp);
+ if (!insertSymTab.lookup(bindName)) {
+ OpBuilder builder(module.getContext());
+ Block *insertBlock = gpuMod ? gpuMod.getBody() : module.getBody();
+ builder.setInsertionPointToEnd(insertBlock);
+ auto funcType = cast<FunctionType>(callee.getFunctionType());
+ func::FuncOp bindFunc = func::FuncOp::create(
+ builder, callee.getLoc(), bindName, funcType);
+ bindFunc.setPrivate();
+ insertSymTab.insert(bindFunc);
+ }
+ calleeRef = FlatSymbolRefAttr::get(callOp.getContext(), bindName);
}
callOp.setCalleeFromCallable(calleeRef);
});
diff --git a/mlir/test/Dialect/OpenACC/acc-bind-routine.mlir b/mlir/test/Dialect/OpenACC/acc-bind-routine.mlir
index 52e1b9c675ceb..10771ed7e8828 100644
--- a/mlir/test/Dialect/OpenACC/acc-bind-routine.mlir
+++ b/mlir/test/Dialect/OpenACC/acc-bind-routine.mlir
@@ -89,3 +89,27 @@ module {
}
// CHECK: func.call_indirect %{{.*}}() : () -> ()
+
+// -----
+
+module {
+ acc.routine @acc_routine_0 func(@my_device_func) bind("__wrapper_my_device_func") seq
+ acc.routine @acc_routine_1 func(@my_device_sub) bind("__wrapper_my_device_sub") seq
+ func.func private @my_device_func(i32) -> i32 attributes {acc.routine_info = #acc.routine_info<[@acc_routine_0]>}
+ func.func private @my_device_sub(i32, memref<i32>) attributes {acc.routine_info = #acc.routine_info<[@acc_routine_1]>}
+ gpu.module @cuda_device_mod {
+ gpu.func @test(%arg0: i32, %arg1: memref<i32>) {
+ %0 = func.call @my_device_func(%arg0) : (i32) -> i32
+ func.call @my_device_sub(%arg0, %arg1) : (i32, memref<i32>) -> ()
+ gpu.return
+ }
+ func.func private @my_device_func(i32) -> i32 attributes {acc.routine_info = #acc.routine_info<[@acc_routine_0]>}
+ func.func private @my_device_sub(i32, memref<i32>) attributes {acc.routine_info = #acc.routine_info<[@acc_routine_1]>}
+ }
+}
+
+// CHECK-LABEL: gpu.func @test
+// CHECK: func.call @__wrapper_my_device_func
+// CHECK: func.call @__wrapper_my_device_sub
+// CHECK: func.func private @__wrapper_my_device_func
+// CHECK: func.func private @__wrapper_my_device_sub
>From e63d871a8f0d6df92377ab441c89a2cdbe70c1b2 Mon Sep 17 00:00:00 2001
From: Delaram Talaashrafi <dtalaashrafi at rome5.pgi.net>
Date: Thu, 25 Jun 2026 07:47:54 -0700
Subject: [PATCH 2/5] Change to func
---
.../mlir/Dialect/OpenACC/Transforms/Passes.td | 3 +-
.../OpenACC/Transforms/ACCBindRoutine.cpp | 42 +++++++++++--------
2 files changed, 26 insertions(+), 19 deletions(-)
diff --git a/mlir/include/mlir/Dialect/OpenACC/Transforms/Passes.td b/mlir/include/mlir/Dialect/OpenACC/Transforms/Passes.td
index 6f5875f844846..ba290331546f6 100644
--- a/mlir/include/mlir/Dialect/OpenACC/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/OpenACC/Transforms/Passes.td
@@ -493,7 +493,8 @@ def ACCRoutineToGPUFunc : Pass<"acc-routine-to-gpu-func", "mlir::ModuleOp"> {
let options = [ AccDeviceTypeOption ];
}
-def ACCBindRoutine : Pass<"acc-bind-routine", "mlir::ModuleOp"> {
+def ACCBindRoutine
+ : InterfacePass<"acc-bind-routine", "mlir::FunctionOpInterface"> {
let summary = "Apply bind clause to function calls in ACC compute regions";
let description = [{
For calls inside offload regions that target a function with an
diff --git a/mlir/lib/Dialect/OpenACC/Transforms/ACCBindRoutine.cpp b/mlir/lib/Dialect/OpenACC/Transforms/ACCBindRoutine.cpp
index 4ed4f4f0ffea0..44d17c113dbd8 100644
--- a/mlir/lib/Dialect/OpenACC/Transforms/ACCBindRoutine.cpp
+++ b/mlir/lib/Dialect/OpenACC/Transforms/ACCBindRoutine.cpp
@@ -13,10 +13,9 @@
//
// Overview:
// ---------
-// For each function, walk operations that implement OffloadRegionOpInterface.
-// For each call inside the offload region, if the callee is a function with
-// an acc routine that has bind(name), replace the call to use the bound
-// symbol.
+// For the current function, walk call operations inside offload regions (or
+// inside gpu.func). If the callee is a function with an acc routine that has
+// bind(name), replace the call to use the bound symbol.
//
// Requirements:
// -------------
@@ -35,6 +34,8 @@
#include "mlir/Dialect/OpenACC/Analysis/OpenACCSupport.h"
#include "mlir/Dialect/OpenACC/OpenACC.h"
#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/BuiltinOps.h"
+// #include "mlir/IR/FunctionInterfaces.h"
#include "mlir/IR/SymbolTable.h"
#include "mlir/Interfaces/CallInterfaces.h"
#include "llvm/Support/Debug.h"
@@ -81,23 +82,29 @@ class ACCBindRoutine : public acc::impl::ACCBindRoutineBase<ACCBindRoutine> {
using acc::impl::ACCBindRoutineBase<ACCBindRoutine>::ACCBindRoutineBase;
void runOnOperation() override {
- ModuleOp module = getOperation();
- OpenACCSupport &accSupport = getAnalysis<OpenACCSupport>();
+ FunctionOpInterface func = getOperation();
+ ModuleOp module = func->getParentOfType<ModuleOp>();
+ if (!module)
+ return;
+
+ auto cachedAnalysis =
+ getCachedParentAnalysis<OpenACCSupport>(func->getParentOp());
+ OpenACCSupport &accSupport = cachedAnalysis ? cachedAnalysis->get()
+ : getAnalysis<OpenACCSupport>();
SymbolTable symTab(module);
bool failed = false;
- module.walk([&](FunctionOpInterface func) {
- func.walk([&](CallOpInterface callOp) {
- if (!callOp.getCallableForCallee())
- return;
- if (!callOp->getParentOfType<OffloadRegionOpInterface>() &&
- !callOp->getParentOfType<gpu::GPUFuncOp>())
- return;
- SymbolRefAttr calleeSymbolRef =
- dyn_cast<SymbolRefAttr>(callOp.getCallableForCallee());
- if (!calleeSymbolRef)
- return;
+ func.walk([&](CallOpInterface callOp) {
+ if (!callOp.getCallableForCallee())
+ return;
+ if (!callOp->getParentOfType<OffloadRegionOpInterface>() &&
+ !callOp->getParentOfType<gpu::GPUFuncOp>())
+ return;
+ SymbolRefAttr calleeSymbolRef =
+ dyn_cast<SymbolRefAttr>(callOp.getCallableForCallee());
+ if (!calleeSymbolRef)
+ return;
FunctionOpInterface callee = symTab.lookup<FunctionOpInterface>(
calleeSymbolRef.getLeafReference());
@@ -148,7 +155,6 @@ class ACCBindRoutine : public acc::impl::ACCBindRoutineBase<ACCBindRoutine> {
calleeRef = FlatSymbolRefAttr::get(callOp.getContext(), bindName);
}
callOp.setCalleeFromCallable(calleeRef);
- });
});
if (failed)
>From da8e8c3d4522786121ecf4c6ad0b53743d1f311f Mon Sep 17 00:00:00 2001
From: Delaram Talaashrafi <dtalaashrafi at rome5.pgi.net>
Date: Thu, 25 Jun 2026 11:53:38 -0700
Subject: [PATCH 3/5] Fix test
---
mlir/test/Dialect/OpenACC/acc-bind-routine.mlir | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/mlir/test/Dialect/OpenACC/acc-bind-routine.mlir b/mlir/test/Dialect/OpenACC/acc-bind-routine.mlir
index 10771ed7e8828..6c49196249e28 100644
--- a/mlir/test/Dialect/OpenACC/acc-bind-routine.mlir
+++ b/mlir/test/Dialect/OpenACC/acc-bind-routine.mlir
@@ -1,4 +1,4 @@
-// RUN: mlir-opt %s -acc-bind-routine -split-input-file | FileCheck %s
+// RUN: mlir-opt %s --pass-pipeline='builtin.module(gpu.module(any(acc-bind-routine)), any(acc-bind-routine))' -split-input-file | FileCheck %s
// Call to routine with bind is rewritten to the bound symbol inside
// offload region.
>From 12b59940186810b454786489059e7ad512bf4c1e Mon Sep 17 00:00:00 2001
From: Delaram Talaashrafi <dtalaashrafi at rome5.pgi.net>
Date: Thu, 25 Jun 2026 12:01:50 -0700
Subject: [PATCH 4/5] Cleanup
---
mlir/lib/Dialect/OpenACC/Transforms/ACCBindRoutine.cpp | 5 ++---
1 file changed, 2 insertions(+), 3 deletions(-)
diff --git a/mlir/lib/Dialect/OpenACC/Transforms/ACCBindRoutine.cpp b/mlir/lib/Dialect/OpenACC/Transforms/ACCBindRoutine.cpp
index 44d17c113dbd8..c17fdc61702b5 100644
--- a/mlir/lib/Dialect/OpenACC/Transforms/ACCBindRoutine.cpp
+++ b/mlir/lib/Dialect/OpenACC/Transforms/ACCBindRoutine.cpp
@@ -13,8 +13,8 @@
//
// Overview:
// ---------
-// For the current function, walk call operations inside offload regions (or
-// inside gpu.func). If the callee is a function with an acc routine that has
+// For the current function, walk call operations inside offload regions, or
+// gpu.func). If the callee is a function with an acc routine that has
// bind(name), replace the call to use the bound symbol.
//
// Requirements:
@@ -35,7 +35,6 @@
#include "mlir/Dialect/OpenACC/OpenACC.h"
#include "mlir/IR/BuiltinAttributes.h"
#include "mlir/IR/BuiltinOps.h"
-// #include "mlir/IR/FunctionInterfaces.h"
#include "mlir/IR/SymbolTable.h"
#include "mlir/Interfaces/CallInterfaces.h"
#include "llvm/Support/Debug.h"
>From 39e62914472dbedd21c727e7c9fbc9365a6abff2 Mon Sep 17 00:00:00 2001
From: Delaram Talaashrafi <dtalaashrafi at rome5.pgi.net>
Date: Thu, 25 Jun 2026 12:57:32 -0700
Subject: [PATCH 5/5] Cleanup
---
.../OpenACC/Transforms/ACCBindRoutine.cpp | 90 +++++++++----------
1 file changed, 44 insertions(+), 46 deletions(-)
diff --git a/mlir/lib/Dialect/OpenACC/Transforms/ACCBindRoutine.cpp b/mlir/lib/Dialect/OpenACC/Transforms/ACCBindRoutine.cpp
index c17fdc61702b5..2892c369d57ad 100644
--- a/mlir/lib/Dialect/OpenACC/Transforms/ACCBindRoutine.cpp
+++ b/mlir/lib/Dialect/OpenACC/Transforms/ACCBindRoutine.cpp
@@ -34,7 +34,6 @@
#include "mlir/Dialect/OpenACC/Analysis/OpenACCSupport.h"
#include "mlir/Dialect/OpenACC/OpenACC.h"
#include "mlir/IR/BuiltinAttributes.h"
-#include "mlir/IR/BuiltinOps.h"
#include "mlir/IR/SymbolTable.h"
#include "mlir/Interfaces/CallInterfaces.h"
#include "llvm/Support/Debug.h"
@@ -88,8 +87,8 @@ class ACCBindRoutine : public acc::impl::ACCBindRoutineBase<ACCBindRoutine> {
auto cachedAnalysis =
getCachedParentAnalysis<OpenACCSupport>(func->getParentOp());
- OpenACCSupport &accSupport = cachedAnalysis ? cachedAnalysis->get()
- : getAnalysis<OpenACCSupport>();
+ OpenACCSupport &accSupport =
+ cachedAnalysis ? cachedAnalysis->get() : getAnalysis<OpenACCSupport>();
SymbolTable symTab(module);
bool failed = false;
@@ -105,55 +104,54 @@ class ACCBindRoutine : public acc::impl::ACCBindRoutineBase<ACCBindRoutine> {
if (!calleeSymbolRef)
return;
- FunctionOpInterface callee = symTab.lookup<FunctionOpInterface>(
- calleeSymbolRef.getLeafReference());
- if (!callee)
- return;
+ FunctionOpInterface callee = symTab.lookup<FunctionOpInterface>(
+ calleeSymbolRef.getLeafReference());
+ if (!callee)
+ return;
- if (!(isAccRoutine(callee) || isSpecializedAccRoutine(callee)))
- return;
+ if (!(isAccRoutine(callee) || isSpecializedAccRoutine(callee)))
+ return;
- if (auto routineInfo = callee->getAttrOfType<RoutineInfoAttr>(
- getRoutineInfoAttrName())) {
- if (routineInfo.getAccRoutines().size() > 1) {
- (void)accSupport.emitNYI(callOp.getLoc(),
- "multiple `acc routine`s");
- failed = true;
- return;
- }
+ if (auto routineInfo = callee->getAttrOfType<RoutineInfoAttr>(
+ getRoutineInfoAttrName())) {
+ if (routineInfo.getAccRoutines().size() > 1) {
+ (void)accSupport.emitNYI(callOp.getLoc(), "multiple `acc routine`s");
+ failed = true;
+ return;
}
+ }
- RoutineOp routine = getFirstAccRoutineOp(callee, symTab);
- if (!isACCRoutineBindDefaultOrDeviceType(routine, this->deviceType))
- return;
+ RoutineOp routine = getFirstAccRoutineOp(callee, symTab);
+ if (!isACCRoutineBindDefaultOrDeviceType(routine, this->deviceType))
+ return;
- auto bindNameOpt = routine.getBindNameValue(this->deviceType);
- if (!bindNameOpt)
- bindNameOpt = routine.getBindNameValue();
- if (!bindNameOpt)
- return;
- SymbolRefAttr calleeRef;
- if (auto *symRef = std::get_if<SymbolRefAttr>(&*bindNameOpt)) {
- calleeRef = *symRef;
- } else {
- StringRef bindName = std::get<StringAttr>(*bindNameOpt).getValue();
- auto gpuMod = func->getParentOfType<gpu::GPUModuleOp>();
- Operation *symbolTableOp =
- gpuMod ? gpuMod.getOperation() : module.getOperation();
- SymbolTable insertSymTab(symbolTableOp);
- if (!insertSymTab.lookup(bindName)) {
- OpBuilder builder(module.getContext());
- Block *insertBlock = gpuMod ? gpuMod.getBody() : module.getBody();
- builder.setInsertionPointToEnd(insertBlock);
- auto funcType = cast<FunctionType>(callee.getFunctionType());
- func::FuncOp bindFunc = func::FuncOp::create(
- builder, callee.getLoc(), bindName, funcType);
- bindFunc.setPrivate();
- insertSymTab.insert(bindFunc);
- }
- calleeRef = FlatSymbolRefAttr::get(callOp.getContext(), bindName);
+ auto bindNameOpt = routine.getBindNameValue(this->deviceType);
+ if (!bindNameOpt)
+ bindNameOpt = routine.getBindNameValue();
+ if (!bindNameOpt)
+ return;
+ SymbolRefAttr calleeRef;
+ if (auto *symRef = std::get_if<SymbolRefAttr>(&*bindNameOpt)) {
+ calleeRef = *symRef;
+ } else {
+ StringRef bindName = std::get<StringAttr>(*bindNameOpt).getValue();
+ auto gpuMod = func->getParentOfType<gpu::GPUModuleOp>();
+ Operation *symbolTableOp =
+ gpuMod ? gpuMod.getOperation() : module.getOperation();
+ SymbolTable insertSymTab(symbolTableOp);
+ if (!insertSymTab.lookup(bindName)) {
+ OpBuilder builder(module.getContext());
+ Block *insertBlock = gpuMod ? gpuMod.getBody() : module.getBody();
+ builder.setInsertionPointToEnd(insertBlock);
+ auto funcType = cast<FunctionType>(callee.getFunctionType());
+ func::FuncOp bindFunc = func::FuncOp::create(builder, callee.getLoc(),
+ bindName, funcType);
+ bindFunc.setPrivate();
+ insertSymTab.insert(bindFunc);
}
- callOp.setCalleeFromCallable(calleeRef);
+ calleeRef = FlatSymbolRefAttr::get(callOp.getContext(), bindName);
+ }
+ callOp.setCalleeFromCallable(calleeRef);
});
if (failed)
More information about the Mlir-commits
mailing list