diff --git a/flang/lib/Optimizer/Transforms/CUFOpConversion.cpp b/flang/lib/Optimizer/Transforms/CUFOpConversion.cpp index 7ecb3b1a7bf27..2db2973212501 100644 --- a/flang/lib/Optimizer/Transforms/CUFOpConversion.cpp +++ b/flang/lib/Optimizer/Transforms/CUFOpConversion.cpp @@ -321,9 +321,15 @@ struct CUFAllocOpConversion : public mlir::OpRewritePattern { builder.createIntegerConstant(loc, builder.getIndexType(), width); } else if (auto seqTy = mlir::dyn_cast_or_null( 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 1a31c4c6d17a4..9a425cbc59624 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, !fir.llvm_ptr, i64, i32, !fir.ref, 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}>> {bindc_name = "a", data_attr = #cuf.cuda, uniq_name = "_QFtest_array_typeEa"} -> !fir.ref}>>> + %1 = fir.shape %c10 : (index) -> !fir.shape<1> + %2 = fir.declare %0(%1) {data_attr = #cuf.cuda, uniq_name = "_QFtest_array_typeEa"} : (!fir.ref}>>>, !fir.shape<1>) -> !fir.ref}>>> + %3 = fir.alloca !fir.array<10x!fir.type<_QMbarTcmplx{id:i32,c:complex}>> {bindc_name = "b", uniq_name = "_QFtest_array_typeEb"} + %4 = fir.declare %3(%1) {uniq_name = "_QFtest_array_typeEb"} : (!fir.ref}>>>, !fir.shape<1>) -> !fir.ref}>>> + cuf.data_transfer %4 to %2 {transfer_kind = #cuf.cuda_transfer} : !fir.ref}>>>, !fir.ref}>>> + cuf.free %2 : !fir.ref}>>> {data_attr = #cuf.cuda} + 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, i32) -> !fir.llvm_ptr +// CHECK: %[[BYTES:.*]] = arith.muli %c10{{.*}}, %c12{{.*}} : i64 +// CHECK: fir.call @_FortranACUFDataTransferPtrPtr(%{{.*}}, %{{.*}}, %[[BYTES]], %c0{{.*}}, %{{.*}}, %{{.*}}) : (!fir.llvm_ptr, !fir.llvm_ptr, i64, i32, !fir.ref, i32) -> none + } // end of module