[flang-commits] [flang] [flang][cuda] Flatten memref descriptors in GPU kernel argument packing (PR #193651)

[Razvan Lupusoru via flang-commits]

================
@@ -253,3 +253,146 @@ module attributes {dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<!llvm.ptr<272>, d
 // CHECK-LABEL: llvm.func @_QQmain()
 // CHECK: %[[STREAM:.*]] = llvm.alloca %{{.*}} x i64 : (i64) -> !llvm.ptr
 // CHECK: llvm.call @_FortranACUFLaunchKernel(%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %[[STREAM]], %{{.*}}, %{{.*}}, %{{.*}})
+
+// -----
+
+// Verify that a memref kernel operand is unpacked into its descriptor fields
+// (allocatedPtr, alignedPtr, offset, sizes[...], strides[...]) so that the
+// host-side kernelParams match the NVVM-lowered device kernel signature.
+module attributes {dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<i1, dense<8> : vector<2xi64>>, #dlti.dl_entry<!llvm.ptr, dense<64> : vector<4xi64>>, #dlti.dl_entry<!llvm.ptr<270>, dense<32> : vector<4xi64>>, #dlti.dl_entry<!llvm.ptr<271>, dense<32> : vector<4xi64>>, #dlti.dl_entry<i8, dense<8> : vector<2xi64>>, #dlti.dl_entry<i16, dense<16> : vector<2xi64>>, #dlti.dl_entry<!llvm.ptr<272>, dense<64> : vector<4xi64>>, #dlti.dl_entry<i64, dense<64> : vector<2xi64>>, #dlti.dl_entry<i32, dense<32> : vector<2xi64>>, #dlti.dl_entry<f128, dense<128> : vector<2xi64>>, #dlti.dl_entry<i128, 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<"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.target_triple = "x86_64-unknown-linux-gnu"} {
+  llvm.func @_QMmod1Phost_sub() {
+    %c1_i64 = llvm.mlir.constant(1 : i64) : i64
+    %c0_i64 = llvm.mlir.constant(0 : i64) : i64
+    %c10000_i64 = llvm.mlir.constant(10000 : i64) : i64
+    %c128_i64 = llvm.mlir.constant(128 : i64) : i64
+    %c0_i32 = llvm.mlir.constant(0 : i32) : i32
+    %buf = llvm.alloca %c10000_i64 x f32 : (i64) -> !llvm.ptr
+    %u0 = llvm.mlir.undef : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+    %u1 = llvm.insertvalue %buf, %u0[0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+    %u2 = llvm.insertvalue %buf, %u1[1] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+    %u3 = llvm.insertvalue %c0_i64, %u2[2] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+    %u4 = llvm.insertvalue %c10000_i64, %u3[3, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+    %desc = llvm.insertvalue %c1_i64, %u4[4, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+    %memref = builtin.unrealized_conversion_cast %desc : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> to memref<10000xf32>
+    gpu.launch_func @cuda_device_mod::@_QMmod1Psub1 blocks in (%c1_i64, %c1_i64, %c1_i64) threads in (%c128_i64, %c1_i64, %c1_i64) : i64 dynamic_shared_memory_size %c0_i32 args(%memref : memref<10000xf32>) {cuf.proc_attr = #cuf.cuda_proc<global>}
----------------
razvanlupusoru wrote:

Will you also add testing for dynamic shapes (eg `memref<?x?xf32>`)?

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