[flang][cuda] Compute size of derived type arrays (#115914)

2 files changed, 31 insertions, 5 deletions

diff --git a/flang/lib/Optimizer/Transforms/CUFOpConversion.cpp b/flang/lib/Optimizer/Transforms/CUFOpConversion.cpp
index 7ecb3b1..2db29732 100644
--- a/flang/lib/Optimizer/Transforms/CUFOpConversion.cpp
+++ b/flang/lib/Optimizer/Transforms/CUFOpConversion.cpp
@@ -321,9 +321,15 @@ struct CUFAllocOpConversion : public mlir::OpRewritePattern<cuf::AllocOp> {
             builder.createIntegerConstant(loc, builder.getIndexType(), width);
       } else if (auto seqTy = mlir::dyn_cast_or_null<fir::SequenceType>(
                      op.getInType())) {
-        mlir::Value width = builder.createIntegerConstant(
-            loc, builder.getIndexType(),
-            computeWidth(loc, seqTy.getEleTy(), kindMap));
+        std::size_t size = 0;
+        if (fir::isa_derived(seqTy.getEleTy())) {
+          mlir::Type structTy = typeConverter->convertType(seqTy.getEleTy());
+          size = dl->getTypeSizeInBits(structTy) / 8;
+        } else {
+          size = computeWidth(loc, seqTy.getEleTy(), kindMap);
+        }
+        mlir::Value width =
+            builder.createIntegerConstant(loc, builder.getIndexType(), size);
         mlir::Value nbElem;
         if (fir::sequenceWithNonConstantShape(seqTy)) {
           assert(!op.getShape().empty() && "expect shape with dynamic arrays");
@@ -580,8 +586,9 @@ struct CUFDataTransferOpConversion
               loc, i64Ty, seqTy.getConstantArraySize());
       }
       unsigned width = 0;
-      if (fir::isa_derived(dstTy)) {
-        mlir::Type structTy = typeConverter->convertType(dstTy);
+      if (fir::isa_derived(fir::unwrapSequenceType(dstTy))) {
+        mlir::Type structTy =
+            typeConverter->convertType(fir::unwrapSequenceType(dstTy));
         width = dl->getTypeSizeInBits(structTy) / 8;
       } else {
         width = computeWidth(loc, dstTy, kindMap);
diff --git a/flang/test/Fir/CUDA/cuda-data-transfer.fir b/flang/test/Fir/CUDA/cuda-data-transfer.fir
index 1a31c4c..9a425cb 100644
--- a/flang/test/Fir/CUDA/cuda-data-transfer.fir
+++ b/flang/test/Fir/CUDA/cuda-data-transfer.fir
@@ -308,4 +308,23 @@ func.func @_QPtest_type() {
 // CHECK-LABEL: func.func @_QPtest_type()
 // CHECK: fir.call @_FortranACUFDataTransferPtrPtr(%{{.*}}, %{{.*}}, %c12{{.*}}, %c0{{.*}}, %{{.*}}, %{{.*}}) : (!fir.llvm_ptr<i8>, !fir.llvm_ptr<i8>, i64, i32, !fir.ref<i8>, i32) -> none
 
+func.func @_QPtest_array_type() {
+  %c10 = arith.constant 10 : index
+  %0 = cuf.alloc !fir.array<10x!fir.type<_QMbarTcmplx{id:i32,c:complex<f32>}>> {bindc_name = "a", data_attr = #cuf.cuda<device>, uniq_name = "_QFtest_array_typeEa"} -> !fir.ref<!fir.array<10x!fir.type<_QMbarTcmplx{id:i32,c:complex<f32>}>>>
+  %1 = fir.shape %c10 : (index) -> !fir.shape<1>
+  %2 = fir.declare %0(%1) {data_attr = #cuf.cuda<device>, uniq_name = "_QFtest_array_typeEa"} : (!fir.ref<!fir.array<10x!fir.type<_QMbarTcmplx{id:i32,c:complex<f32>}>>>, !fir.shape<1>) -> !fir.ref<!fir.array<10x!fir.type<_QMbarTcmplx{id:i32,c:complex<f32>}>>>
+  %3 = fir.alloca !fir.array<10x!fir.type<_QMbarTcmplx{id:i32,c:complex<f32>}>> {bindc_name = "b", uniq_name = "_QFtest_array_typeEb"}
+  %4 = fir.declare %3(%1) {uniq_name = "_QFtest_array_typeEb"} : (!fir.ref<!fir.array<10x!fir.type<_QMbarTcmplx{id:i32,c:complex<f32>}>>>, !fir.shape<1>) -> !fir.ref<!fir.array<10x!fir.type<_QMbarTcmplx{id:i32,c:complex<f32>}>>>
+  cuf.data_transfer %4 to %2 {transfer_kind = #cuf.cuda_transfer<host_device>} : !fir.ref<!fir.array<10x!fir.type<_QMbarTcmplx{id:i32,c:complex<f32>}>>>, !fir.ref<!fir.array<10x!fir.type<_QMbarTcmplx{id:i32,c:complex<f32>}>>>
+  cuf.free %2 : !fir.ref<!fir.array<10x!fir.type<_QMbarTcmplx{id:i32,c:complex<f32>}>>> {data_attr = #cuf.cuda<device>}
+  return
+}
+
+// CHECK-LABEL: func.func @_QPtest_array_type()
+// CHECK: %[[BYTES:.*]] = arith.muli %c10{{.*}}, %c12 : index
+// CHECK: %[[CONV_BYTES:.*]] = fir.convert %[[BYTES]] : (index) -> i64
+// CHECK: fir.call @_FortranACUFMemAlloc(%[[CONV_BYTES]], %c0{{.*}}, %{{.*}}, %{{.*}}) : (i64, i32, !fir.ref<i8>, i32) -> !fir.llvm_ptr<i8>
+// CHECK: %[[BYTES:.*]] = arith.muli %c10{{.*}}, %c12{{.*}} : i64
+// CHECK: fir.call @_FortranACUFDataTransferPtrPtr(%{{.*}}, %{{.*}}, %[[BYTES]], %c0{{.*}}, %{{.*}}, %{{.*}}) : (!fir.llvm_ptr<i8>, !fir.llvm_ptr<i8>, i64, i32, !fir.ref<i8>, i32) -> none
+
 } // end of module