[Mlir-commits] [mlir] 8249d08 - [mlir][acc] Rewrite acc routine bind calls inside gpu.func (#204220)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Fri Jun 26 07:36:19 PDT 2026
Author: Delaram Talaashrafi
Date: 2026-06-26T10:36:15-04:00
New Revision: 8249d087039e7e8e2967f0799902cb7ee8278a94
URL: https://github.com/llvm/llvm-project/commit/8249d087039e7e8e2967f0799902cb7ee8278a94
DIFF: https://github.com/llvm/llvm-project/commit/8249d087039e7e8e2967f0799902cb7ee8278a94.diff
LOG: [mlir][acc] Rewrite acc routine bind calls inside gpu.func (#204220)
Run `acc-bind-routine` on `FunctionOpInterface` 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.
Added:
Modified:
mlir/include/mlir/Dialect/OpenACC/Transforms/Passes.td
mlir/lib/Dialect/OpenACC/Transforms/ACCBindRoutine.cpp
mlir/test/Dialect/OpenACC/acc-bind-routine.mlir
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/OpenACC/Transforms/Passes.td b/mlir/include/mlir/Dialect/OpenACC/Transforms/Passes.td
index b52dee12d390a..3cf4241dbeb5b 100644
--- a/mlir/include/mlir/Dialect/OpenACC/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/OpenACC/Transforms/Passes.td
@@ -524,7 +524,8 @@ def ACCRoutineToGPUFunc : Pass<"acc-routine-to-gpu-func", "mlir::ModuleOp"> {
let options = [ AccDeviceTypeOption ];
}
-def ACCBindRoutine : Pass<"acc-bind-routine", "mlir::func::FuncOp"> {
+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 f4c8f18e5b04b..81b3cfb627952 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
+// 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:
// -------------
@@ -81,67 +80,77 @@ class ACCBindRoutine : public acc::impl::ACCBindRoutineBase<ACCBindRoutine> {
using acc::impl::ACCBindRoutineBase<ACCBindRoutine>::ACCBindRoutineBase;
void runOnOperation() override {
- func::FuncOp func = getOperation();
+ FunctionOpInterface func = getOperation();
ModuleOp module = func->getParentOfType<ModuleOp>();
if (!module)
return;
- SymbolTable symTab(module);
auto cachedAnalysis =
getCachedParentAnalysis<OpenACCSupport>(func->getParentOp());
OpenACCSupport &accSupport =
cachedAnalysis ? cachedAnalysis->get() : getAnalysis<OpenACCSupport>();
+ SymbolTable symTab(module);
bool failed = false;
- func.walk([&](acc::OffloadRegionOpInterface offload) {
- Region ®ion = offload.getOffloadRegion();
- region.walk([&](CallOpInterface callOp) {
- if (!callOp.getCallableForCallee())
- return;
- SymbolRefAttr calleeSymbolRef =
- dyn_cast<SymbolRefAttr>(callOp.getCallableForCallee());
- if (!calleeSymbolRef)
- return;
-
- FunctionOpInterface callee = symTab.lookup<FunctionOpInterface>(
- calleeSymbolRef.getLeafReference());
- if (!callee)
- return;
-
- if (!(isAccRoutine(callee) || isSpecializedAccRoutine(callee)))
+ 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());
+ if (!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;
-
- 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 {
- calleeRef = FlatSymbolRefAttr::get(
- callOp.getContext(),
- std::get<StringAttr>(*bindNameOpt).getValue());
+ }
+
+ 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);
}
- callOp.setCalleeFromCallable(calleeRef);
- });
+ calleeRef = FlatSymbolRefAttr::get(callOp.getContext(), bindName);
+ }
+ callOp.setCalleeFromCallable(calleeRef);
});
if (failed)
diff --git a/mlir/test/Dialect/OpenACC/acc-bind-routine.mlir b/mlir/test/Dialect/OpenACC/acc-bind-routine.mlir
index 52e1b9c675ceb..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.
@@ -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
More information about the Mlir-commits
mailing list