[flang-commits] [flang] 9c6a5f0 - [flang][cuda] Set kernel intent(in) as const __restrict__ (#203652) (#203966)

[via flang-commits]

Author: Valentin Clement (バレンタイン クレメン)
Date: 2026-06-16T08:30:07-07:00
New Revision: 9c6a5f064611b12fd38f7282c0a3b600eb30950d

URL: https://github.com/llvm/llvm-project/commit/9c6a5f064611b12fd38f7282c0a3b600eb30950d
DIFF: https://github.com/llvm/llvm-project/commit/9c6a5f064611b12fd38f7282c0a3b600eb30950d.diff

LOG: [flang][cuda] Set kernel intent(in) as const __restrict__ (#203652) (#203966)

Added: 
    

Modified: 
    flang/lib/Optimizer/Transforms/CUDA/CUFDeviceFuncTransform.cpp
    flang/test/Fir/CUDA/cuda-device-func-transform.mlir

Removed: 
    


################################################################################
diff  --git a/flang/lib/Optimizer/Transforms/CUDA/CUFDeviceFuncTransform.cpp b/flang/lib/Optimizer/Transforms/CUDA/CUFDeviceFuncTransform.cpp
index 88382cafcc3bd..97540820ef415 100644
--- a/flang/lib/Optimizer/Transforms/CUDA/CUFDeviceFuncTransform.cpp
+++ b/flang/lib/Optimizer/Transforms/CUDA/CUFDeviceFuncTransform.cpp
@@ -12,6 +12,7 @@
 #include "flang/Optimizer/Dialect/FIRAttr.h"
 #include "flang/Optimizer/Dialect/FIRDialect.h"
 #include "flang/Optimizer/Dialect/FIROpsSupport.h"
+#include "flang/Optimizer/Dialect/FIRType.h"
 #include "flang/Optimizer/Support/InternalNames.h"
 #include "flang/Optimizer/Transforms/Passes.h"
 #include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"
@@ -41,6 +42,52 @@ class CUFDeviceFuncTransform
   using CUFDeviceFuncTransformBase<
       CUFDeviceFuncTransform>::CUFDeviceFuncTransformBase;
 
+  static bool isPointerLikeKernelArg(mlir::Type argType) {
+    return mlir::isa<fir::ReferenceType, fir::BaseBoxType>(argType);
+  }
+
+  // Decorate INTENT(IN) kernel arguments like C "const __restrict__". NVPTX
+  // only tags loads as invariant (and lowers them to ld.global.nc) when a
+  // kernel pointer parameter is both readonly and noalias; see
+  // NVPTXTagInvariantLoads.
+  static void setIntentInKernelArgAttrs(mlir::func::FuncOp funcOp,
+                                        gpu::GPUFuncOp deviceFuncOp) {
+    mlir::UnitAttr unitAttr = mlir::UnitAttr::get(funcOp.getContext());
+
+    auto markArg = [&](unsigned argIndex) {
+      if (argIndex >= deviceFuncOp.getNumArguments())
+        return;
+      mlir::Type argType = deviceFuncOp.getArgumentTypes()[argIndex];
+      if (!isPointerLikeKernelArg(argType))
+        return;
+      deviceFuncOp.setArgAttr(
+          argIndex, mlir::LLVM::LLVMDialect::getReadonlyAttrName(), unitAttr);
+      deviceFuncOp.setArgAttr(
+          argIndex, mlir::LLVM::LLVMDialect::getNoAliasAttrName(), unitAttr);
+    };
+
+    funcOp.walk([&](fir::DeclareOp declareOp) {
+      auto var =
+          mlir::cast<fir::FortranVariableOpInterface>(declareOp.getOperation());
+      if (!var.isIntentIn())
+        return;
+      if (auto attrs = var.getFortranAttrs())
+        if (fir::bitEnumContainsAny(*attrs,
+                                    fir::FortranVariableFlagsEnum::value))
+          return;
+      if (std::optional<uint32_t> dummyArgNo = declareOp.getDummyArgNo()) {
+        // Dummy argument numbers are 1-based in FIR.
+        markArg(*dummyArgNo - 1);
+        return;
+      }
+      if (auto blockArg =
+              mlir::dyn_cast<mlir::BlockArgument>(declareOp.getMemref()))
+        if (blockArg.getOwner()->isEntryBlock() &&
+            blockArg.getOwner()->getParentOp() == funcOp)
+          markArg(blockArg.getArgNumber());
+    });
+  }
+
   static gpu::GPUFuncOp createGPUFuncOp(mlir::func::FuncOp funcOp,
                                         bool isGlobal, int computeCap) {
     mlir::OpBuilder builder(funcOp.getContext());
@@ -67,6 +114,9 @@ class CUFDeviceFuncTransform
     auto deviceFuncOp =
         gpu::GPUFuncOp::create(builder, loc, funcOp.getName(), type,
                                mlir::TypeRange{}, mlir::TypeRange{});
+    if (mlir::ArrayAttr argAttrs = funcOp.getAllArgAttrs())
+      deviceFuncOp.setAllArgAttrs(argAttrs);
+    setIntentInKernelArgAttrs(funcOp, deviceFuncOp);
     if (isGlobal)
       deviceFuncOp.setKernel(true);
 

diff  --git a/flang/test/Fir/CUDA/cuda-device-func-transform.mlir b/flang/test/Fir/CUDA/cuda-device-func-transform.mlir
index cc183f0144e82..6af5e021c33df 100644
--- a/flang/test/Fir/CUDA/cuda-device-func-transform.mlir
+++ b/flang/test/Fir/CUDA/cuda-device-func-transform.mlir
@@ -34,7 +34,7 @@ func.func private @_QMmod1Psub1(!fir.ref<!fir.array<10xi32>> {cuf.data_attr = #c
 
 // CHECK: gpu.func @_QPsub_device1()
 
-// CHECK: gpu.func @_QPsub_device2(%[[ARG0:.*]]: !fir.ref<f32>) {
+// CHECK: gpu.func @_QPsub_device2(%[[ARG0:.*]]: !fir.ref<f32>
 // CHECK:   %[[DECL:.*]] = fir.declare %[[ARG0]] {uniq_name = "_QFsub1Ei"} : (!fir.ref<f32>) -> !fir.ref<f32>
 // CHECK:   %[[CST:.*]] = arith.constant 2.000000e+00 : f32
 // CHECK:   fir.store %[[CST]] to %[[DECL]] : !fir.ref<f32>
@@ -151,12 +151,39 @@ func.func @_QPpartialsumshflshflr8(%arg0: !fir.ref<!fir.array<?xf64>> {cuf.data_
 }
 
 // CHECK-LABEL: gpu.module @cuda_device_mod
-// CHECK: gpu.func @_QPpartialsumshflshflr8(%arg0: !fir.ref<!fir.array<?xf64>>, %arg1: i32) kernel
+// CHECK: gpu.func @_QPpartialsumshflshflr8(%{{.*}}) kernel
       
 // CHECK: func.func @_QPpartialsumshflshflr8
 
 // -----
 
+func.func @_QPldg_attrs(%arg0: !fir.ref<!fir.array<?xf32>> {fir.bindc_name = "a"}, %arg1: !fir.ref<!fir.array<?xf32>> {fir.bindc_name = "b"}) attributes {cuf.proc_attr = #cuf.cuda_proc<global>} {
+  %c10 = arith.constant 10 : index
+  %scope = fir.dummy_scope : !fir.dscope
+  %shape = fir.shape %c10 : (index) -> !fir.shape<1>
+  %0 = fir.declare %arg0(%shape) dummy_scope %scope arg 1 {fortran_attrs = #fir.var_attrs<intent_out>, uniq_name = "_QFldg_attrsEa"} : (!fir.ref<!fir.array<?xf32>>, !fir.shape<1>, !fir.dscope) -> !fir.ref<!fir.array<?xf32>>
+  %1 = fir.declare %arg1(%shape) dummy_scope %scope arg 2 {fortran_attrs = #fir.var_attrs<intent_in>, uniq_name = "_QFldg_attrsEb"} : (!fir.ref<!fir.array<?xf32>>, !fir.shape<1>, !fir.dscope) -> !fir.ref<!fir.array<?xf32>>
+  return
+}
+
+// CHECK-LABEL: gpu.module @cuda_device_mod
+// CHECK: gpu.func @_QPldg_attrs(%{{.*}}: !fir.ref<!fir.array<?xf32>>{{.*}}, %{{.*}}: !fir.ref<!fir.array<?xf32>> {{{.*}}llvm.noalias, llvm.readonly}) kernel
+
+// -----
+
+func.func @_QPvalue_in(%arg0: i32 {fir.bindc_name = "n"}) attributes {cuf.proc_attr = #cuf.cuda_proc<global>} {
+  %scope = fir.dummy_scope : !fir.dscope
+  %0 = fir.alloca i32
+  fir.store %arg0 to %0 : !fir.ref<i32>
+  %1 = fir.declare %0 dummy_scope %scope {fortran_attrs = #fir.var_attrs<intent_in, value>, uniq_name = "_QPvalue_inEn"} : (!fir.ref<i32>, !fir.dscope) -> !fir.ref<i32>
+  return
+}
+
+// CHECK-LABEL: gpu.module @cuda_device_mod
+// CHECK: gpu.func @_QPvalue_in(%{{.*}}: i32 {fir.bindc_name = "n"}) kernel
+
+// -----
+
 func.func @_QPsub_maxtnid() attributes {cuf.launch_bounds = #cuf.launch_bounds<maxTPB = 256 : i64, minBPM = 2 : i64, upperBoundClusterSize = 3 : i64>, cuf.proc_attr = #cuf.cuda_proc<global>} {
   %cst = arith.constant 2.000000e+00 : f32
   return