[flang-commits] [flang] [flang][cuda] Recognize mangled on_device() names in CUFFunctionRewrite (PR #207298)

[via flang-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-flang-fir-hlfir

Author: Zhen Wang (wangzpgi)

<details>
<summary>Changes</summary>

`CUFFunctionRewrite` folds `on_device()` to a compile-time constant: `true` inside a `gpu.module` (device), `false` on the host. It only matched the bind(c) name `on_device`, missing the mangled forms `_QPon_device` and `on_device_`. Extend the mapping to recognize all three.

---
Full diff: https://github.com/llvm/llvm-project/pull/207298.diff


2 Files Affected:

- (modified) flang/lib/Optimizer/Transforms/CUDA/CUFFunctionRewrite.cpp (+7-1) 
- (modified) flang/test/Fir/CUDA/cuda-function-rewrite.mlir (+86) 


``````````diff
diff --git a/flang/lib/Optimizer/Transforms/CUDA/CUFFunctionRewrite.cpp b/flang/lib/Optimizer/Transforms/CUDA/CUFFunctionRewrite.cpp
index bcbfb5294c56f..87bf105303dee 100644
--- a/flang/lib/Optimizer/Transforms/CUDA/CUFFunctionRewrite.cpp
+++ b/flang/lib/Optimizer/Transforms/CUDA/CUFFunctionRewrite.cpp
@@ -78,8 +78,14 @@ 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}};
+      {"on_device", &genOnDevice},
+      {"_QPon_device", &genOnDevice},
+      {"on_device_", &genOnDevice}};
 };
 
 class CUFFunctionRewrite
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

``````````

</details>


https://github.com/llvm/llvm-project/pull/207298