diff --git a/flang/lib/Optimizer/Transforms/CUFOpConversion.cpp b/flang/lib/Optimizer/Transforms/CUFOpConversion.cpp index e5c5ba9082426..759e3a65dd24f 100644 --- a/flang/lib/Optimizer/Transforms/CUFOpConversion.cpp +++ b/flang/lib/Optimizer/Transforms/CUFOpConversion.cpp @@ -741,6 +741,9 @@ struct CUFDataTransferOpConversion fir::StoreOp::create(builder, loc, val, box); return box; } + if (mlir::isa(val.getType())) + if (auto loadOp = mlir::dyn_cast(val.getDefiningOp())) + return loadOp.getMemref(); return val; }; diff --git a/flang/test/Fir/CUDA/cuda-data-transfer.fir b/flang/test/Fir/CUDA/cuda-data-transfer.fir index 5d3215dd07fce..b247fce44df3d 100644 --- a/flang/test/Fir/CUDA/cuda-data-transfer.fir +++ b/flang/test/Fir/CUDA/cuda-data-transfer.fir @@ -691,5 +691,39 @@ func.func @_QPtesti4(%arg0: !fir.ref {fir.bindc_name = "n1"}, %arg1: !fir.r // CHECK-LABEL: func.func @_QPtesti4 // CHECK: fir.call @_FortranACUFDataTransferCstDesc +// ----- + +func.func @_QQmain() attributes {fir.bindc_name = "T"} { + %c2 = arith.constant 2 : index + %c1 = arith.constant 1 : index + %c80 = arith.constant 80 : index + %c0 = arith.constant 0 : index + %0 = fir.dummy_scope : !fir.dscope + %1 = cuf.alloc !fir.box>> {bindc_name = "a", data_attr = #cuf.cuda, uniq_name = "_QFEa"} -> !fir.ref>>> + %2 = fir.zero_bits !fir.heap> + %3 = fir.shape %c0, %c0, %c0 : (index, index, index) -> !fir.shape<3> + %4 = fir.embox %2(%3) {allocator_idx = 2 : i32} : (!fir.heap>, !fir.shape<3>) -> !fir.box>> + fir.store %4 to %1 : !fir.ref>>> + %5 = fir.declare %1 {data_attr = #cuf.cuda, fortran_attrs = #fir.var_attrs, uniq_name = "_QFEa"} : (!fir.ref>>>) -> !fir.ref>>> + %6 = fir.address_of(@_QFEha) : !fir.ref> + %7 = fir.shape %c80, %c80, %c80 : (index, index, index) -> !fir.shape<3> + %8 = fir.declare %6(%7) {uniq_name = "_QFEha"} : (!fir.ref>, !fir.shape<3>) -> !fir.ref> + %9 = fir.address_of(@_QFECn) : !fir.ref + %10 = fir.declare %9 {fortran_attrs = #fir.var_attrs, uniq_name = "_QFECn"} : (!fir.ref) -> !fir.ref + %11 = fir.load %5 : !fir.ref>>> + %12:3 = fir.box_dims %11, %c0 : (!fir.box>>, index) -> (index, index, index) + %13:3 = fir.box_dims %11, %c1 : (!fir.box>>, index) -> (index, index, index) + %14:3 = fir.box_dims %11, %c2 : (!fir.box>>, index) -> (index, index, index) + %15 = fir.shape %12#1, %13#1, %14#1 : (index, index, index) -> !fir.shape<3> + %16 = fir.allocmem !fir.array, %12#1, %13#1, %14#1 {bindc_name = ".tmp", uniq_name = ""} + %17 = fir.declare %16(%15) {uniq_name = ".tmp"} : (!fir.heap>, !fir.shape<3>) -> !fir.heap> + %18 = fir.embox %17(%15) : (!fir.heap>, !fir.shape<3>) -> !fir.box> + cuf.data_transfer %11 to %18 {transfer_kind = #cuf.cuda_transfer} : !fir.box>>, !fir.box> + return +} + +// CHECK-LABEL: func.func @_QQmain() +// CHECK: fir.call @_FortranACUFDataTransferDescDesc + } // end of module