[flang][cuda] Allow static and dynamic shared memory in a single kernel

[clementval]

No description provided.

[llvmbot]

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

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

Changes

---

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

2 Files Affected:

• (modified) flang/lib/Optimizer/Transforms/CUDA/CUFComputeSharedMemoryOffsetsAndSize.cpp (+7-12)
• (modified) flang/test/Fir/CUDA/cuda-shared-offset.mlir (+33)

diff --git a/flang/lib/Optimizer/Transforms/CUDA/CUFComputeSharedMemoryOffsetsAndSize.cpp b/flang/lib/Optimizer/Transforms/CUDA/CUFComputeSharedMemoryOffsetsAndSize.cpp
index 18537f677661a..26245e39d255b 100644
--- a/flang/lib/Optimizer/Transforms/CUDA/CUFComputeSharedMemoryOffsetsAndSize.cpp
+++ b/flang/lib/Optimizer/Transforms/CUDA/CUFComputeSharedMemoryOffsetsAndSize.cpp
@@ -159,18 +159,13 @@ struct CUFComputeSharedMemoryOffsetsAndSize
       if (nbDynamicSharedVariables == 0 && nbStaticSharedVariables == 0)
         continue;
 
-      if (nbDynamicSharedVariables > 0 && nbStaticSharedVariables > 0)
-        mlir::emitError(
-            funcOp.getLoc(),
-            "static and dynamic shared variables in a single kernel");
-
-      if (nbStaticSharedVariables > 0)
-        continue;
-
-      auto sharedMemType = fir::SequenceType::get(sharedMemSize, i8Ty);
-      createSharedMemoryGlobal(builder, funcOp.getLoc(), funcOp.getName(), "",
-                               gpuMod, sharedMemType, sharedMemSize, alignment,
-                               /*isDynamic=*/true);
+      if (nbDynamicSharedVariables > 0) {
+        auto sharedMemType = fir::SequenceType::get(sharedMemSize, i8Ty);
+        createSharedMemoryGlobal(builder, funcOp.getLoc(), funcOp.getName(), "",
+                                 gpuMod, sharedMemType, sharedMemSize,
+                                 alignment,
+                                 /*isDynamic=*/true);
+      }
     }
   }
 };
diff --git a/flang/test/Fir/CUDA/cuda-shared-offset.mlir b/flang/test/Fir/CUDA/cuda-shared-offset.mlir
index a33c89b6d9af2..68c31356a278b 100644
--- a/flang/test/Fir/CUDA/cuda-shared-offset.mlir
+++ b/flang/test/Fir/CUDA/cuda-shared-offset.mlir
@@ -202,3 +202,36 @@ module attributes {dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<!llvm.ptr, dense<
 // CHECK:     cuf.shared_memory[%c0{{.*}} : i32] i32 align 4 {bindc_name = "_lro_private_1", isStatic, uniq_name = "_lro_private_1"} -> !fir.ref<i32>
 // CHECK: fir.global internal @_QPreduce_kernel__shared_mem___lro_private_0 {alignment = 4 : i64, data_attr = #cuf.cuda<shared>} : !fir.array<4xi8>
 // CHECK: fir.global internal @_QPreduce_kernel__shared_mem___lro_private_1 {alignment = 4 : i64, data_attr = #cuf.cuda<shared>} : !fir.array<4xi8>
+
+// -----
+
+module attributes {dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<!llvm.ptr, dense<64> : vector<4xi64>>, #dlti.dl_entry<!llvm.ptr<271>, dense<32> : vector<4xi64>>, #dlti.dl_entry<!llvm.ptr<270>, dense<32> : vector<4xi64>>, #dlti.dl_entry<f128, dense<128> : vector<2xi64>>, #dlti.dl_entry<f64, dense<64> : vector<2xi64>>, #dlti.dl_entry<f80, dense<128> : vector<2xi64>>, #dlti.dl_entry<f16, dense<16> : vector<2xi64>>, #dlti.dl_entry<i32, dense<32> : vector<2xi64>>, #dlti.dl_entry<i16, dense<16> : vector<2xi64>>, #dlti.dl_entry<i128, dense<128> : vector<2xi64>>, #dlti.dl_entry<i8, dense<8> : vector<2xi64>>, #dlti.dl_entry<!llvm.ptr<272>, dense<64> : vector<4xi64>>, #dlti.dl_entry<i64, dense<64> : vector<2xi64>>, #dlti.dl_entry<i1, dense<8> : vector<2xi64>>, #dlti.dl_entry<"dlti.endianness", "little">, #dlti.dl_entry<"dlti.stack_alignment", 128 : i64>>, fir.defaultkind = "a1c4d8i4l4r4", fir.kindmap = "", gpu.container_module, llvm.data_layout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-i128:128-f80:128-n8:16:32:64-S128", llvm.ident = "flang version 20.0.0 (https://github.com/llvm/llvm-project.git cae351f3453a0a26ec8eb2ddaf773c24a29d929e)", llvm.target_triple = "x86_64-unknown-linux-gnu"} {
+  gpu.module @cuda_device_mod {
+    gpu.func @_QPt1() kernel {
+      %c256 = arith.constant 256 : index
+      %0 = fir.dummy_scope : !fir.dscope
+      %1 = fir.address_of(@_QM__fortran_builtinsE__builtin_blockdim) : !fir.ref<!fir.type<_QM__fortran_builtinsT__builtin_dim3{x:i32,y:i32,z:i32}>>
+      %2:2 = hlfir.declare %1 {uniq_name = "_QM__fortran_builtinsE__builtin_blockdim"} : (!fir.ref<!fir.type<_QM__fortran_builtinsT__builtin_dim3{x:i32,y:i32,z:i32}>>) -> (!fir.ref<!fir.type<_QM__fortran_builtinsT__builtin_dim3{x:i32,y:i32,z:i32}>>, !fir.ref<!fir.type<_QM__fortran_builtinsT__builtin_dim3{x:i32,y:i32,z:i32}>>)
+      %3 = fir.address_of(@_QM__fortran_builtinsE__builtin_blockidx) : !fir.ref<!fir.type<_QM__fortran_builtinsT__builtin_dim3{x:i32,y:i32,z:i32}>>
+      %4:2 = hlfir.declare %3 {uniq_name = "_QM__fortran_builtinsE__builtin_blockidx"} : (!fir.ref<!fir.type<_QM__fortran_builtinsT__builtin_dim3{x:i32,y:i32,z:i32}>>) -> (!fir.ref<!fir.type<_QM__fortran_builtinsT__builtin_dim3{x:i32,y:i32,z:i32}>>, !fir.ref<!fir.type<_QM__fortran_builtinsT__builtin_dim3{x:i32,y:i32,z:i32}>>)
+      %5 = fir.assumed_size_extent : index
+      %6 = cuf.shared_memory !fir.array<?xf32>, %5 : index {bindc_name = "d", uniq_name = "_QFt1Ed"} -> !fir.ref<!fir.array<?xf32>>
+      %7 = fir.shape %5 : (index) -> !fir.shape<1>
+      %8:2 = hlfir.declare %6(%7) {data_attr = #cuf.cuda<shared>, uniq_name = "_QFt1Ed"} : (!fir.ref<!fir.array<?xf32>>, !fir.shape<1>) -> (!fir.box<!fir.array<?xf32>>, !fir.ref<!fir.array<?xf32>>)
+      %9 = fir.address_of(@_QM__fortran_builtinsE__builtin_griddim) : !fir.ref<!fir.type<_QM__fortran_builtinsT__builtin_dim3{x:i32,y:i32,z:i32}>>
+      %10:2 = hlfir.declare %9 {uniq_name = "_QM__fortran_builtinsE__builtin_griddim"} : (!fir.ref<!fir.type<_QM__fortran_builtinsT__builtin_dim3{x:i32,y:i32,z:i32}>>) -> (!fir.ref<!fir.type<_QM__fortran_builtinsT__builtin_dim3{x:i32,y:i32,z:i32}>>, !fir.ref<!fir.type<_QM__fortran_builtinsT__builtin_dim3{x:i32,y:i32,z:i32}>>)
+      %11 = cuf.shared_memory !fir.array<256xf32> {bindc_name = "s", uniq_name = "_QFt1Es"} -> !fir.ref<!fir.array<256xf32>>
+      %12 = fir.shape %c256 : (index) -> !fir.shape<1>
+      %13:2 = hlfir.declare %11(%12) {data_attr = #cuf.cuda<shared>, uniq_name = "_QFt1Es"} : (!fir.ref<!fir.array<256xf32>>, !fir.shape<1>) -> (!fir.ref<!fir.array<256xf32>>, !fir.ref<!fir.array<256xf32>>)
+      %14 = fir.address_of(@_QM__fortran_builtinsE__builtin_threadidx) : !fir.ref<!fir.type<_QM__fortran_builtinsT__builtin_dim3{x:i32,y:i32,z:i32}>>
+      %15:2 = hlfir.declare %14 {uniq_name = "_QM__fortran_builtinsE__builtin_threadidx"} : (!fir.ref<!fir.type<_QM__fortran_builtinsT__builtin_dim3{x:i32,y:i32,z:i32}>>) -> (!fir.ref<!fir.type<_QM__fortran_builtinsT__builtin_dim3{x:i32,y:i32,z:i32}>>, !fir.ref<!fir.type<_QM__fortran_builtinsT__builtin_dim3{x:i32,y:i32,z:i32}>>)
+      %16 = fir.alloca i32 {bindc_name = "__builtin_warpsize", uniq_name = "_QM__fortran_builtinsEC__builtin_warpsize"}
+      %17:2 = hlfir.declare %16 {uniq_name = "_QM__fortran_builtinsEC__builtin_warpsize"} : (!fir.ref<i32>) -> (!fir.ref<i32>, !fir.ref<i32>)
+      gpu.return
+    }
+  }
+}
+
+// CHECK-LABEL: gpu.func @_QPt1()
+// CHECK: fir.global internal @_QPt1__shared_mem__s {alignment = 4 : i64, data_attr = #cuf.cuda<shared>} : !fir.array<1024xi8>
+// CHECK: fir.global external @_QPt1__shared_mem__ {alignment = 4 : i64, data_attr = #cuf.cuda<shared>} : !fir.array<0xi8>