[flang][cuda] Run target rewrite in gpu.module

[clementval]

Apply signature conversion for func.func in the gpu.module. More work will need to be done for gpu.func op and implement the NVVM ABI for conversion in the gpu module.

[llvmbot]

@llvm/pr-subscribers-flang-codegen

Author: Valentin Clement (バレンタイン クレメン) (clementval)

Changes

Apply signature conversion for func.func in the gpu.module. More work will need to be done for gpu.func op and implement the NVVM ABI for conversion in the gpu module.

---

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

2 Files Affected:

• (modified) flang/lib/Optimizer/CodeGen/TargetRewrite.cpp (+6)
• (added) flang/test/Fir/CUDA/cuda-target-rewrite.mlir (+16)

diff --git a/flang/lib/Optimizer/CodeGen/TargetRewrite.cpp b/flang/lib/Optimizer/CodeGen/TargetRewrite.cpp
index ae6e7ce798d998..1b86d5241704b1 100644
--- a/flang/lib/Optimizer/CodeGen/TargetRewrite.cpp
+++ b/flang/lib/Optimizer/CodeGen/TargetRewrite.cpp
@@ -27,6 +27,7 @@
 #include "flang/Optimizer/Dialect/Support/FIRContext.h"
 #include "flang/Optimizer/Support/DataLayout.h"
 #include "mlir/Dialect/DLTI/DLTI.h"
+#include "mlir/Dialect/GPU/IR/GPUDialect.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Transforms/DialectConversion.h"
 #include "llvm/ADT/STLExtras.h"
@@ -720,6 +721,11 @@ class TargetRewrite : public fir::impl::TargetRewritePassBase<TargetRewrite> {
 
       convertSignature(fn);
     }
+
+    for (auto gpuMod : mod.getOps<mlir::gpu::GPUModuleOp>())
+      for (auto fn : gpuMod.getOps<mlir::func::FuncOp>())
+        convertSignature(fn);
+
     return mlir::success();
   }
 
diff --git a/flang/test/Fir/CUDA/cuda-target-rewrite.mlir b/flang/test/Fir/CUDA/cuda-target-rewrite.mlir
new file mode 100644
index 00000000000000..d85cca38870adb
--- /dev/null
+++ b/flang/test/Fir/CUDA/cuda-target-rewrite.mlir
@@ -0,0 +1,16 @@
+// RUN: fir-opt --target-rewrite %s | FileCheck %s
+
+gpu.module @testmod {
+  gpu.func @_QPvcpowdk(%arg0: !fir.ref<complex<f64>> {cuf.data_attr = #cuf.cuda<device>, fir.bindc_name = "a"}) attributes {cuf.proc_attr = #cuf.cuda_proc<global>} {
+    %0 = fir.alloca i64
+    %1 = fir.load %0 : !fir.ref<i64>
+    %2 = fir.load %arg0 : !fir.ref<complex<f64>>
+    %3 = fir.call @_FortranAzpowk(%2, %1) fastmath<contract> : (complex<f64>, i64) -> complex<f64>
+    gpu.return
+  }
+  func.func private @_FortranAzpowk(complex<f64>, i64) -> complex<f64> attributes {fir.bindc_name = "_FortranAzpowk", fir.runtime}
+}
+
+// CHECK-LABEL: gpu.func @_QPvcpowdk
+// CHECK: %{{.*}} = fir.call @_FortranAzpowk(%{{.*}}, %{{.*}}, %{{.*}}) : (f64, f64, i64) -> tuple<f64, f64>
+// CHECK: func.func private @_FortranAzpowk(f64, f64, i64) -> tuple<f64, f64> attributes {fir.bindc_name = "_FortranAzpowk", fir.runtime}

[llvmbot]

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

Author: Valentin Clement (バレンタイン クレメン) (clementval)

Changes

Apply signature conversion for func.func in the gpu.module. More work will need to be done for gpu.func op and implement the NVVM ABI for conversion in the gpu module.

---

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

2 Files Affected:

• (modified) flang/lib/Optimizer/CodeGen/TargetRewrite.cpp (+6)
• (added) flang/test/Fir/CUDA/cuda-target-rewrite.mlir (+16)

diff --git a/flang/lib/Optimizer/CodeGen/TargetRewrite.cpp b/flang/lib/Optimizer/CodeGen/TargetRewrite.cpp
index ae6e7ce798d998..1b86d5241704b1 100644
--- a/flang/lib/Optimizer/CodeGen/TargetRewrite.cpp
+++ b/flang/lib/Optimizer/CodeGen/TargetRewrite.cpp
@@ -27,6 +27,7 @@
 #include "flang/Optimizer/Dialect/Support/FIRContext.h"
 #include "flang/Optimizer/Support/DataLayout.h"
 #include "mlir/Dialect/DLTI/DLTI.h"
+#include "mlir/Dialect/GPU/IR/GPUDialect.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Transforms/DialectConversion.h"
 #include "llvm/ADT/STLExtras.h"
@@ -720,6 +721,11 @@ class TargetRewrite : public fir::impl::TargetRewritePassBase<TargetRewrite> {
 
       convertSignature(fn);
     }
+
+    for (auto gpuMod : mod.getOps<mlir::gpu::GPUModuleOp>())
+      for (auto fn : gpuMod.getOps<mlir::func::FuncOp>())
+        convertSignature(fn);
+
     return mlir::success();
   }
 
diff --git a/flang/test/Fir/CUDA/cuda-target-rewrite.mlir b/flang/test/Fir/CUDA/cuda-target-rewrite.mlir
new file mode 100644
index 00000000000000..d85cca38870adb
--- /dev/null
+++ b/flang/test/Fir/CUDA/cuda-target-rewrite.mlir
@@ -0,0 +1,16 @@
+// RUN: fir-opt --target-rewrite %s | FileCheck %s
+
+gpu.module @testmod {
+  gpu.func @_QPvcpowdk(%arg0: !fir.ref<complex<f64>> {cuf.data_attr = #cuf.cuda<device>, fir.bindc_name = "a"}) attributes {cuf.proc_attr = #cuf.cuda_proc<global>} {
+    %0 = fir.alloca i64
+    %1 = fir.load %0 : !fir.ref<i64>
+    %2 = fir.load %arg0 : !fir.ref<complex<f64>>
+    %3 = fir.call @_FortranAzpowk(%2, %1) fastmath<contract> : (complex<f64>, i64) -> complex<f64>
+    gpu.return
+  }
+  func.func private @_FortranAzpowk(complex<f64>, i64) -> complex<f64> attributes {fir.bindc_name = "_FortranAzpowk", fir.runtime}
+}
+
+// CHECK-LABEL: gpu.func @_QPvcpowdk
+// CHECK: %{{.*}} = fir.call @_FortranAzpowk(%{{.*}}, %{{.*}}, %{{.*}}) : (f64, f64, i64) -> tuple<f64, f64>
+// CHECK: func.func private @_FortranAzpowk(f64, f64, i64) -> tuple<f64, f64> attributes {fir.bindc_name = "_FortranAzpowk", fir.runtime}