[flang][cuda] Materialize box when needed

[clementval]

Materialize the box when the src comes from a embox or rebox operation. This was done in the case of transfer to a descriptor but not when transferring from a descriptor.

[llvmbot]

@llvm/pr-subscribers-flang-fir-hlfir

Author: Valentin Clement (バレンタイン クレメン) (clementval)

Changes

Materialize the box when the src comes from a embox or rebox operation. This was done in the case of transfer to a descriptor but not when transferring from a descriptor.

---

Full diff: https://github.com/llvm/llvm-project/pull/117810.diff

2 Files Affected:

• (modified) flang/lib/Optimizer/Transforms/CUFOpConversion.cpp (+14-12)
• (modified) flang/test/Fir/CUDA/cuda-data-transfer.fir (+23)

diff --git a/flang/lib/Optimizer/Transforms/CUFOpConversion.cpp b/flang/lib/Optimizer/Transforms/CUFOpConversion.cpp
index 3983336516db9e..1f9a7f17b35b66 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<fir::EmboxOp, fir::ReboxOp>(val.getDefiningOp())) {
+        // Materialize the box to memory to be able to call the runtime.
+        mlir::Value box = builder.createTemporary(loc, val.getType());
+        builder.create<fir::StoreOp>(loc, val, box);
+        return box;
+      }
+      return val;
+    };
+
     // Conversion of data transfer involving at least one descriptor.
     if (mlir::isa<fir::BaseBoxType>(dstTy)) {
       // Transfer to a descriptor.
@@ -682,15 +692,7 @@ struct CUFDataTransferOpConversion
           func = fir::runtime::getRuntimeFunc<mkRTKey(CUFDataTransferCstDesc)>(
               loc, builder);
       }
-      auto materializeBoxIfNeeded = [&](mlir::Value val) -> mlir::Value {
-        if (mlir::isa<fir::EmboxOp, fir::ReboxOp>(val.getDefiningOp())) {
-          // Materialize the box to memory to be able to call the runtime.
-          mlir::Value box = builder.createTemporary(loc, val.getType());
-          builder.create<fir::StoreOp>(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<mkRTKey(
           CUFDataTransferDescDescNoRealloc)>(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<mlir::Value> args{
-          fir::runtime::createArguments(builder, loc, fTy, dst, op.getSrc(),
-                                        modeValue, sourceFile, sourceLine)};
+      llvm::SmallVector<mlir::Value> args{fir::runtime::createArguments(
+          builder, loc, fTy, dst, src, modeValue, sourceFile, sourceLine)};
       builder.create<fir::CallOp>(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 0f9ca6e640a802..b371d397777280 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<f32>) -> !fir.llvm_ptr<i8>
 // CHECK: fir.call @_FortranACUFDataTransferPtrPtr(%{{.*}}, %[[CONV]], %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}) : (!fir.llvm_ptr<i8>, !fir.llvm_ptr<i8>, i64, i32, !fir.ref<i8>, i32) -> none
 
+func.func @_QPchecksums(%arg0: !fir.box<!fir.array<?xf64>> {cuf.data_attr = #cuf.cuda<device>, fir.bindc_name = "a"}, %arg1: !fir.ref<i32> {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<device>, uniq_name = "_QFchecksumsEa"} : (!fir.box<!fir.array<?xf64>>, !fir.dscope) -> !fir.box<!fir.array<?xf64>>
+  %2 = fir.rebox %1 : (!fir.box<!fir.array<?xf64>>) -> !fir.box<!fir.array<?xf64>>
+  %3 = fir.declare %arg1 dummy_scope %0 {uniq_name = "_QFchecksumsEn"} : (!fir.ref<i32>, !fir.dscope) -> !fir.ref<i32>
+  %4 = fir.load %3 : !fir.ref<i32>
+  %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<?xf64>, %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.array<?xf64>>, !fir.shape<1>) -> !fir.ref<!fir.array<?xf64>>
+  %11 = fir.embox %10(%9) : (!fir.ref<!fir.array<?xf64>>, !fir.shape<1>) -> !fir.box<!fir.array<?xf64>>
+  cuf.data_transfer %2 to %10, %9 : !fir.shape<1> {transfer_kind = #cuf.cuda_transfer<device_host>} : !fir.box<!fir.array<?xf64>>, !fir.ref<!fir.array<?xf64>>
+  return
+}
+
+// CHECK-LABEL: func.func @_QPchecksums
+// CHECK: %[[DST:.*]] = fir.convert %{{.*}} : (!fir.ref<!fir.box<!fir.array<?xf64>>>) -> !fir.ref<!fir.box<none>>
+// CHECK: %[[SRC:.*]] = fir.convert %{{.*}} : (!fir.ref<!fir.box<!fir.array<?xf64>>>) -> !fir.ref<!fir.box<none>>
+// CHECK: fir.call @_FortranACUFDataTransferDescDescNoRealloc(%[[DST]], %[[SRC]], %{{.*}}, %{{.*}}, %{{.*}}) : (!fir.ref<!fir.box<none>>, !fir.ref<!fir.box<none>>, i32, !fir.ref<i8>, i32) -> none
+
 } // end of module