[flang-commits] [flang] [flang][cuda] Fix device-to-device transfers involving scalar CUDA constants (PR #219577)

[Valentin Clement バレンタイン クレメン via flang-commits]

================
@@ -293,33 +293,86 @@ struct CUFDataTransferOpConversion
           return declareOp.getMemref().getDefiningOp<fir::AddrOfOp>();
         return {};
       };
+      // Address of the host shadow when val designates a scalar CUDA constant,
+      // null otherwise.
+      auto getShadowAddrOf = [&](mlir::Value val) -> fir::AddrOfOp {
+        fir::AddrOfOp addrOfOp = getAddrOf(val);
+        if (!addrOfOp)
+          return {};
+        auto global = symtab.lookup<fir::GlobalOp>(
+            addrOfOp.getSymbol().getRootReference().getValue());
+        if (!isScalarCudaConstantGlobal(global))
+          return {};
+        return addrOfOp;
+      };
       if (op.getTransferKind() == cuf::DataTransferKind::DeviceHost) {
-        if (fir::AddrOfOp addrOfOp = getAddrOf(src)) {
-          auto global = symtab.lookup<fir::GlobalOp>(
-              addrOfOp.getSymbol().getRootReference().getValue());
-          if (isScalarCudaConstantGlobal(global) &&
-              fir::isa_ref_type(dst.getType())) {
-            mlir::Value hostValue = fir::LoadOp::create(builder, loc, src);
-            hostValue = createConvertOp(rewriter, loc, dstTy, hostValue);
-            fir::StoreOp::create(builder, loc, hostValue, dst);
-            rewriter.eraseOp(op);
-            return mlir::success();
-          }
+        if (getShadowAddrOf(src) && fir::isa_ref_type(dst.getType())) {
+          mlir::Value hostValue = fir::LoadOp::create(builder, loc, src);
+          hostValue = createConvertOp(rewriter, loc, dstTy, hostValue);
+          fir::StoreOp::create(builder, loc, hostValue, dst);
+          rewriter.eraseOp(op);
+          return mlir::success();
         }
       }
       if (op.getTransferKind() == cuf::DataTransferKind::HostDevice) {
-        if (fir::AddrOfOp addrOfOp = getAddrOf(dst)) {
-          auto global = symtab.lookup<fir::GlobalOp>(
-              addrOfOp.getSymbol().getRootReference().getValue());
-          if (isScalarCudaConstantGlobal(global)) {
-            mlir::Value hostValue = src;
-            if (fir::isa_ref_type(src.getType()))
-              hostValue = fir::LoadOp::create(builder, loc, src);
-            hostValue = createConvertOp(rewriter, loc, dstTy, hostValue);
-            fir::StoreOp::create(builder, loc, hostValue, addrOfOp);
+        // A non-null shadow means the destination is a scalar constant. Keep
+        // its shadow up to date for later host reads, then aim the copy at the
+        // device symbol instead of the shadow.
+        if (fir::AddrOfOp addrOfOp = getShadowAddrOf(dst)) {
+          mlir::Value hostValue = src;
+          if (fir::isa_ref_type(src.getType()))
+            hostValue = fir::LoadOp::create(builder, loc, src);
+          hostValue = createConvertOp(rewriter, loc, dstTy, hostValue);
+          fir::StoreOp::create(builder, loc, hostValue, addrOfOp);
+          dst = cuf::DeviceAddressOp::create(rewriter, loc, dst.getType(),
+                                             addrOfOp.getSymbol());
+        }
+      }
+      if (op.getTransferKind() == cuf::DataTransferKind::DeviceDevice) {
+        // A scalar CUDA constant is designated by its host shadow, which lives
+        // in host memory and cannot take part in a device to device copy. Route
+        // the transfer through the shadow instead. There are three cases,
+        // depending on which side is a scalar constant.
----------------
clementval wrote:

We do have a global on the device side and it is registered. Is this a limitation because constant can only be initialized from the host? 

https://github.com/llvm/llvm-project/pull/219577