[flang-commits] [flang] [flang][cuda] Recognize mangled on_device() names in CUFFunctionRewrite (PR #207298)
Zhen Wang via flang-commits
flang-commits at lists.llvm.org
Thu Jul 2 17:50:28 PDT 2026
https://github.com/wangzpgi updated https://github.com/llvm/llvm-project/pull/207298
>From d4876b68e08f4965558e607d97c1ae08a3274487 Mon Sep 17 00:00:00 2001
From: Zhen Wang <zhenw at nvidia.com>
Date: Wed, 1 Jul 2026 16:21:16 -0700
Subject: [PATCH 1/3] add on_device mapping
---
.../Transforms/CUDA/CUFFunctionRewrite.cpp | 21 ++++-
.../test/Fir/CUDA/cuda-function-rewrite.mlir | 86 +++++++++++++++++++
2 files changed, 106 insertions(+), 1 deletion(-)
diff --git a/flang/lib/Optimizer/Transforms/CUDA/CUFFunctionRewrite.cpp b/flang/lib/Optimizer/Transforms/CUDA/CUFFunctionRewrite.cpp
index bcbfb5294c56f..dc9d75d0e7e46 100644
--- a/flang/lib/Optimizer/Transforms/CUDA/CUFFunctionRewrite.cpp
+++ b/flang/lib/Optimizer/Transforms/CUDA/CUFFunctionRewrite.cpp
@@ -12,6 +12,7 @@
#include "flang/Optimizer/Dialect/FIRType.h"
#include "flang/Optimizer/Support/DataLayout.h"
#include "flang/Optimizer/Transforms/Passes.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/MLIRContext.h"
@@ -79,7 +80,9 @@ class CallConversion : public OpRewritePattern<fir::CallOp> {
}
const llvm::StringMap<genFunctionType> genMappings_ = {
- {"on_device", &genOnDevice}};
+ {"on_device", &genOnDevice},
+ {"_QPon_device", &genOnDevice},
+ {"on_device_", &genOnDevice}};
};
class CUFFunctionRewrite
@@ -96,6 +99,22 @@ class CUFFunctionRewrite
mlir::emitError(mlir::UnknownLoc::get(ctx),
"error in CUFFunctionRewrite op conversion\n");
signalPassFailure();
+ return;
+ }
+
+ static const llvm::StringRef rewrittenNames[] = {
+ "on_device", "_QPon_device", "on_device_"};
+ auto removeUnusedDecl = [](mlir::ModuleOp mod, llvm::StringRef name) {
+ if (auto func = mod.lookupSymbol<mlir::func::FuncOp>(name))
+ if (func.isDeclaration() && func.symbolKnownUseEmpty(mod))
+ func.erase();
+ };
+ if (auto mod = getOperation()->getParentOfType<mlir::ModuleOp>()) {
+ for (auto name : rewrittenNames)
+ removeUnusedDecl(mod, name);
+ } else if (auto mod = mlir::dyn_cast<mlir::ModuleOp>(getOperation())) {
+ for (auto name : rewrittenNames)
+ removeUnusedDecl(mod, name);
}
}
};
diff --git a/flang/test/Fir/CUDA/cuda-function-rewrite.mlir b/flang/test/Fir/CUDA/cuda-function-rewrite.mlir
index da1d601a2eb8b..721ba3aa1c320 100644
--- a/flang/test/Fir/CUDA/cuda-function-rewrite.mlir
+++ b/flang/test/Fir/CUDA/cuda-function-rewrite.mlir
@@ -42,3 +42,89 @@ func.func @_QMmtestsPdo3(%arg0: !fir.ref<i32> {cuf.data_attr = #cuf.cuda<device>
// CHECK-LABEL: func.func @_QMmtestsPdo3
// CHECK: fir.if %false
+
+// -----
+
+// Test on_device() with Fortran name mangling (_QPon_device) in device context
+gpu.module @acc_device_mod {
+ func.func @_QMmtestPsub_device() {
+ %c2_i32 = arith.constant 2 : i32
+ %c1_i32 = arith.constant 1 : i32
+ %0 = fir.alloca i32
+ %13 = fir.call @_QPon_device() fastmath<contract> : () -> !fir.logical<4>
+ %14 = fir.convert %13 : (!fir.logical<4>) -> i1
+ fir.if %14 {
+ fir.store %c1_i32 to %0 : !fir.ref<i32>
+ } else {
+ fir.store %c2_i32 to %0 : !fir.ref<i32>
+ }
+ return
+ }
+}
+
+// CHECK-LABEL: gpu.module @acc_device_mod
+// CHECK: func.func @_QMmtestPsub_device
+// CHECK: fir.if %true
+
+// -----
+
+// Test _QPon_device on host side
+func.func @_QMmtestPsub_host() {
+ %c2_i32 = arith.constant 2 : i32
+ %c1_i32 = arith.constant 1 : i32
+ %0 = fir.alloca i32
+ %13 = fir.call @_QPon_device() fastmath<contract> : () -> !fir.logical<4>
+ %14 = fir.convert %13 : (!fir.logical<4>) -> i1
+ fir.if %14 {
+ fir.store %c1_i32 to %0 : !fir.ref<i32>
+ } else {
+ fir.store %c2_i32 to %0 : !fir.ref<i32>
+ }
+ return
+}
+
+// CHECK-LABEL: func.func @_QMmtestPsub_host
+// CHECK: fir.if %false
+
+// -----
+
+// Test externally-mangled on_device_ (after ExternalNameConversion) in device context
+gpu.module @acc_extname_device_mod {
+ func.func @_QMmtestPsub_extname_device() {
+ %c2_i32 = arith.constant 2 : i32
+ %c1_i32 = arith.constant 1 : i32
+ %0 = fir.alloca i32
+ %13 = fir.call @on_device_() fastmath<contract> : () -> !fir.logical<4>
+ %14 = fir.convert %13 : (!fir.logical<4>) -> i1
+ fir.if %14 {
+ fir.store %c1_i32 to %0 : !fir.ref<i32>
+ } else {
+ fir.store %c2_i32 to %0 : !fir.ref<i32>
+ }
+ return
+ }
+}
+
+// CHECK-LABEL: gpu.module @acc_extname_device_mod
+// CHECK: func.func @_QMmtestPsub_extname_device
+// CHECK: fir.if %true
+
+// -----
+
+// Test on_device_ on host side
+func.func @_QMmtestPsub_extname_host() {
+ %c2_i32 = arith.constant 2 : i32
+ %c1_i32 = arith.constant 1 : i32
+ %0 = fir.alloca i32
+ %13 = fir.call @on_device_() fastmath<contract> : () -> !fir.logical<4>
+ %14 = fir.convert %13 : (!fir.logical<4>) -> i1
+ fir.if %14 {
+ fir.store %c1_i32 to %0 : !fir.ref<i32>
+ } else {
+ fir.store %c2_i32 to %0 : !fir.ref<i32>
+ }
+ return
+}
+
+// CHECK-LABEL: func.func @_QMmtestPsub_extname_host
+// CHECK: fir.if %false
>From 60db232b39c17ddb59459c3cc3bc01a6d5395b54 Mon Sep 17 00:00:00 2001
From: Zhen Wang <zhenw at nvidia.com>
Date: Thu, 2 Jul 2026 14:18:09 -0700
Subject: [PATCH 2/3] cleanup
---
.../Transforms/CUDA/CUFFunctionRewrite.cpp | 21 ++++---------------
1 file changed, 4 insertions(+), 17 deletions(-)
diff --git a/flang/lib/Optimizer/Transforms/CUDA/CUFFunctionRewrite.cpp b/flang/lib/Optimizer/Transforms/CUDA/CUFFunctionRewrite.cpp
index dc9d75d0e7e46..12a08626fd067 100644
--- a/flang/lib/Optimizer/Transforms/CUDA/CUFFunctionRewrite.cpp
+++ b/flang/lib/Optimizer/Transforms/CUDA/CUFFunctionRewrite.cpp
@@ -12,7 +12,6 @@
#include "flang/Optimizer/Dialect/FIRType.h"
#include "flang/Optimizer/Support/DataLayout.h"
#include "flang/Optimizer/Transforms/Passes.h"
-#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/MLIRContext.h"
@@ -79,6 +78,10 @@ class CallConversion : public OpRewritePattern<fir::CallOp> {
return fir::ConvertOp::create(rewriter, loc, op.getResult(0).getType(), t);
}
+ // on_device may reach this pass under different symbol names depending on how
+ // it was declared and when the pass runs: its bind(c) name ("on_device"), its
+ // Fortran internal name ("_QPon_device"), or its externally-mangled name after
+ // external-name conversion ("on_device_").
const llvm::StringMap<genFunctionType> genMappings_ = {
{"on_device", &genOnDevice},
{"_QPon_device", &genOnDevice},
@@ -99,22 +102,6 @@ class CUFFunctionRewrite
mlir::emitError(mlir::UnknownLoc::get(ctx),
"error in CUFFunctionRewrite op conversion\n");
signalPassFailure();
- return;
- }
-
- static const llvm::StringRef rewrittenNames[] = {
- "on_device", "_QPon_device", "on_device_"};
- auto removeUnusedDecl = [](mlir::ModuleOp mod, llvm::StringRef name) {
- if (auto func = mod.lookupSymbol<mlir::func::FuncOp>(name))
- if (func.isDeclaration() && func.symbolKnownUseEmpty(mod))
- func.erase();
- };
- if (auto mod = getOperation()->getParentOfType<mlir::ModuleOp>()) {
- for (auto name : rewrittenNames)
- removeUnusedDecl(mod, name);
- } else if (auto mod = mlir::dyn_cast<mlir::ModuleOp>(getOperation())) {
- for (auto name : rewrittenNames)
- removeUnusedDecl(mod, name);
}
}
};
>From 48f0e0e09a1336fc45db77db6059359efd010f5b Mon Sep 17 00:00:00 2001
From: Zhen Wang <zhenw at nvidia.com>
Date: Thu, 2 Jul 2026 17:50:07 -0700
Subject: [PATCH 3/3] format
---
flang/lib/Optimizer/Transforms/CUDA/CUFFunctionRewrite.cpp | 4 ++--
1 file changed, 2 insertions(+), 2 deletions(-)
diff --git a/flang/lib/Optimizer/Transforms/CUDA/CUFFunctionRewrite.cpp b/flang/lib/Optimizer/Transforms/CUDA/CUFFunctionRewrite.cpp
index 12a08626fd067..87bf105303dee 100644
--- a/flang/lib/Optimizer/Transforms/CUDA/CUFFunctionRewrite.cpp
+++ b/flang/lib/Optimizer/Transforms/CUDA/CUFFunctionRewrite.cpp
@@ -80,8 +80,8 @@ class CallConversion : public OpRewritePattern<fir::CallOp> {
// on_device may reach this pass under different symbol names depending on how
// it was declared and when the pass runs: its bind(c) name ("on_device"), its
- // Fortran internal name ("_QPon_device"), or its externally-mangled name after
- // external-name conversion ("on_device_").
+ // Fortran internal name ("_QPon_device"), or its externally-mangled name
+ // after external-name conversion ("on_device_").
const llvm::StringMap<genFunctionType> genMappings_ = {
{"on_device", &genOnDevice},
{"_QPon_device", &genOnDevice},
More information about the flang-commits
mailing list