diff --git a/flang/lib/Optimizer/Transforms/CUFOpConversion.cpp b/flang/lib/Optimizer/Transforms/CUFOpConversion.cpp index 3983336516db9..1f9a7f17b35b6 100644 --- a/flang/lib/Optimizer/Transforms/CUFOpConversion.cpp +++ b/flang/lib/Optimizer/Transforms/CUFOpConversion.cpp @@ -665,6 +665,16 @@ struct CUFDataTransferOpConversion return mlir::success(); } + auto materializeBoxIfNeeded = [&](mlir::Value val) -> mlir::Value { + if (mlir::isa(val.getDefiningOp())) { + // Materialize the box to memory to be able to call the runtime. + mlir::Value box = builder.createTemporary(loc, val.getType()); + builder.create(loc, val, box); + return box; + } + return val; + }; + // Conversion of data transfer involving at least one descriptor. if (mlir::isa(dstTy)) { // Transfer to a descriptor. @@ -682,15 +692,7 @@ struct CUFDataTransferOpConversion func = fir::runtime::getRuntimeFunc( loc, builder); } - auto materializeBoxIfNeeded = [&](mlir::Value val) -> mlir::Value { - if (mlir::isa(val.getDefiningOp())) { - // Materialize the box to memory to be able to call the runtime. - mlir::Value box = builder.createTemporary(loc, val.getType()); - builder.create(loc, val, box); - return box; - } - return val; - }; + src = materializeBoxIfNeeded(src); dst = materializeBoxIfNeeded(dst); @@ -705,6 +707,7 @@ struct CUFDataTransferOpConversion } else { // Transfer from a descriptor. mlir::Value dst = emboxDst(rewriter, op, symtab); + mlir::Value src = materializeBoxIfNeeded(op.getSrc()); mlir::func::FuncOp func = fir::runtime::getRuntimeFunc(loc, builder); @@ -713,9 +716,8 @@ struct CUFDataTransferOpConversion mlir::Value sourceFile = fir::factory::locationToFilename(builder, loc); mlir::Value sourceLine = fir::factory::locationToLineNo(builder, loc, fTy.getInput(4)); - llvm::SmallVector args{ - fir::runtime::createArguments(builder, loc, fTy, dst, op.getSrc(), - modeValue, sourceFile, sourceLine)}; + llvm::SmallVector args{fir::runtime::createArguments( + builder, loc, fTy, dst, src, modeValue, sourceFile, sourceLine)}; builder.create(loc, func, args); rewriter.eraseOp(op); } diff --git a/flang/test/Fir/CUDA/cuda-data-transfer.fir b/flang/test/Fir/CUDA/cuda-data-transfer.fir index 0f9ca6e640a80..b371d39777728 100644 --- a/flang/test/Fir/CUDA/cuda-data-transfer.fir +++ b/flang/test/Fir/CUDA/cuda-data-transfer.fir @@ -553,4 +553,27 @@ func.func @_QPsrc_cst() { // CHECK: %[[CONV:.*]] = fir.convert %[[ALLOCA]] : (!fir.ref) -> !fir.llvm_ptr // CHECK: fir.call @_FortranACUFDataTransferPtrPtr(%{{.*}}, %[[CONV]], %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}) : (!fir.llvm_ptr, !fir.llvm_ptr, i64, i32, !fir.ref, i32) -> none +func.func @_QPchecksums(%arg0: !fir.box> {cuf.data_attr = #cuf.cuda, fir.bindc_name = "a"}, %arg1: !fir.ref {fir.bindc_name = "n"}) { + %c0 = arith.constant 0 : index + %0 = fir.dummy_scope : !fir.dscope + %1 = fir.declare %arg0 dummy_scope %0 {data_attr = #cuf.cuda, uniq_name = "_QFchecksumsEa"} : (!fir.box>, !fir.dscope) -> !fir.box> + %2 = fir.rebox %1 : (!fir.box>) -> !fir.box> + %3 = fir.declare %arg1 dummy_scope %0 {uniq_name = "_QFchecksumsEn"} : (!fir.ref, !fir.dscope) -> !fir.ref + %4 = fir.load %3 : !fir.ref + %5 = fir.convert %4 : (i32) -> index + %6 = arith.cmpi sgt, %5, %c0 : index + %7 = arith.select %6, %5, %c0 : index + %8 = fir.alloca !fir.array, %7 {bindc_name = "hosttmp", uniq_name = "_QFchecksumsEhosttmp"} + %9 = fir.shape %7 : (index) -> !fir.shape<1> + %10 = fir.declare %8(%9) {uniq_name = "_QFchecksumsEhosttmp"} : (!fir.ref>, !fir.shape<1>) -> !fir.ref> + %11 = fir.embox %10(%9) : (!fir.ref>, !fir.shape<1>) -> !fir.box> + cuf.data_transfer %2 to %10, %9 : !fir.shape<1> {transfer_kind = #cuf.cuda_transfer} : !fir.box>, !fir.ref> + return +} + +// CHECK-LABEL: func.func @_QPchecksums +// CHECK: %[[DST:.*]] = fir.convert %{{.*}} : (!fir.ref>>) -> !fir.ref> +// CHECK: %[[SRC:.*]] = fir.convert %{{.*}} : (!fir.ref>>) -> !fir.ref> +// CHECK: fir.call @_FortranACUFDataTransferDescDescNoRealloc(%[[DST]], %[[SRC]], %{{.*}}, %{{.*}}, %{{.*}}) : (!fir.ref>, !fir.ref>, i32, !fir.ref, i32) -> none + } // end of module