diff --git a/flang/lib/Optimizer/Transforms/CUFOpConversion.cpp b/flang/lib/Optimizer/Transforms/CUFOpConversion.cpp index b14a1c338e087..c070c3de94cf8 100644 --- a/flang/lib/Optimizer/Transforms/CUFOpConversion.cpp +++ b/flang/lib/Optimizer/Transforms/CUFOpConversion.cpp @@ -640,15 +640,20 @@ struct CUFDataTransferOpConversion loc, builder); mlir::Value dst = op.getDst(); mlir::Value src = op.getSrc(); - if (!mlir::isa(srcTy)) { src = emboxSrc(rewriter, op, symtab); - } else if (mlir::isa(src.getDefiningOp())) { - // Materialize the box to memory to be able to call the runtime. - mlir::Value box = builder.createTemporary(loc, src.getType()); - builder.create(loc, src, box); - src = box; } + 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); auto fTy = func.getFunctionType(); mlir::Value sourceFile = fir::factory::locationToFilename(builder, loc); diff --git a/flang/test/Fir/CUDA/cuda-data-transfer.fir b/flang/test/Fir/CUDA/cuda-data-transfer.fir index 69baf7d15a7d0..ad392beed56d4 100644 --- a/flang/test/Fir/CUDA/cuda-data-transfer.fir +++ b/flang/test/Fir/CUDA/cuda-data-transfer.fir @@ -400,7 +400,6 @@ func.func @_QPdevmul(%arg0: !fir.ref> {fir.bindc_name = "b"} %9 = fir.convert %8 : (i32) -> index %12 = fir.shape %c1, %9 : (index, index) -> !fir.shape<2> %13 = fir.declare %arg0(%12) dummy_scope %0 {uniq_name = "_QFdevmulEb"} : (!fir.ref>, !fir.shape<2>, !fir.dscope) -> !fir.ref> - %24 = fir.load %7 : !fir.ref %25 = fir.convert %24 : (i32) -> index %26 = arith.cmpi sgt, %25, %c0 : index @@ -414,14 +413,19 @@ func.func @_QPdevmul(%arg0: !fir.ref> {fir.bindc_name = "b"} %34 = fir.slice %c1, %25, %c1, %c1, %29, %c1 : (index, index, index, index, index, index) -> !fir.slice<2> %35 = fir.embox %13(%12) [%34] : (!fir.ref>, !fir.shape<2>, !fir.slice<2>) -> !fir.box> cuf.data_transfer %35 to %6 {transfer_kind = #cuf.cuda_transfer} : !fir.box>, !fir.ref>>> + cuf.data_transfer %6 to %35 {transfer_kind = #cuf.cuda_transfer} : !fir.ref>>>, !fir.box> return } // CHECK-LABEL: func.func @_QPdevmul(%arg0: !fir.ref> {fir.bindc_name = "b"}, %arg1: !fir.ref {fir.bindc_name = "wa"}, %arg2: !fir.ref {fir.bindc_name = "wb"}) { -// CHECK: %[[ALLOCA:.*]] = fir.alloca !fir.box> +// CHECK: %[[ALLOCA0:.*]] = fir.alloca !fir.box> +// CHECK: %[[ALLOCA1:.*]] = fir.alloca !fir.box> // CHECK: %[[EMBOX:.*]] = fir.embox %{{.*}}(%{{.*}}) [%{{.*}}] : (!fir.ref>, !fir.shape<2>, !fir.slice<2>) -> !fir.box> -// CHECK: fir.store %[[EMBOX]] to %[[ALLOCA]] : !fir.ref>> -// CHECK: %[[SRC:.*]] = fir.convert %[[ALLOCA]] : (!fir.ref>>) -> !fir.ref> +// CHECK: fir.store %[[EMBOX]] to %[[ALLOCA1]] : !fir.ref>> +// CHECK: %[[SRC:.*]] = fir.convert %[[ALLOCA1]] : (!fir.ref>>) -> !fir.ref> // CHECK: fir.call @_FortranACUFDataTransferDescDesc(%{{.*}}, %[[SRC]], %{{.*}}, %{{.*}}, %{{.*}}) : (!fir.ref>, !fir.ref>, i32, !fir.ref, i32) -> none +// CHECK: fir.store %[[EMBOX]] to %[[ALLOCA0]] : !fir.ref>> +// CHECK: %[[DST:.*]] = fir.convert %[[ALLOCA0]] : (!fir.ref>>) -> !fir.ref> +// CHECK: fir.call @_FortranACUFDataTransferDescDesc(%[[DST]], %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}) : (!fir.ref>, !fir.ref>, i32, !fir.ref, i32) -> none } // end of module