[flang][cuda] Support derived type in cuf.data_transfer conversion (#115557)

Support derived type in `cuf.data_transfer` conversion by computing
their size in bytes.

Author: clementval

flang/lib/Optimizer/Transforms/CUFOpConversion.cpp

@@ -507,8 +507,11 @@ struct CUFDataTransferOpConversion
   using OpRewritePattern::OpRewritePattern;
 
   CUFDataTransferOpConversion(mlir::MLIRContext *context,
-                              const mlir::SymbolTable &symtab)
-      : OpRewritePattern(context), symtab{symtab} {}
+                              const mlir::SymbolTable &symtab,
+                              mlir::DataLayout *dl,
+                              const fir::LLVMTypeConverter *typeConverter)
+      : OpRewritePattern(context), symtab{symtab}, dl{dl},
+        typeConverter{typeConverter} {}
 
   mlir::LogicalResult
   matchAndRewrite(cuf::DataTransferOp op,
@@ -576,7 +579,13 @@ struct CUFDataTransferOpConversion
           nbElement = builder.createIntegerConstant(
               loc, i64Ty, seqTy.getConstantArraySize());
       }
-      int width = computeWidth(loc, dstTy, kindMap);
+      unsigned width = 0;
+      if (fir::isa_derived(dstTy)) {
+        mlir::Type structTy = typeConverter->convertType(dstTy);
+        width = dl->getTypeSizeInBits(structTy) / 8;
+      } else {
+        width = computeWidth(loc, dstTy, kindMap);
+      }
       mlir::Value widthValue = rewriter.create<mlir::arith::ConstantOp>(
           loc, i64Ty, rewriter.getIntegerAttr(i64Ty, width));
       mlir::Value bytes =
@@ -647,6 +656,8 @@ struct CUFDataTransferOpConversion
 
 private:
   const mlir::SymbolTable &symtab;
+  mlir::DataLayout *dl;
+  const fir::LLVMTypeConverter *typeConverter;
 };
 
 struct CUFLaunchOpConversion
@@ -749,6 +760,7 @@ void cuf::populateCUFToFIRConversionPatterns(
   patterns.insert<CUFAllocOpConversion>(patterns.getContext(), &dl, &converter);
   patterns.insert<CUFAllocateOpConversion, CUFDeallocateOpConversion,
                   CUFFreeOpConversion>(patterns.getContext());
-  patterns.insert<CUFDataTransferOpConversion, CUFLaunchOpConversion>(
-      patterns.getContext(), symtab);
+  patterns.insert<CUFDataTransferOpConversion>(patterns.getContext(), symtab,
+                                               &dl, &converter);
+  patterns.insert<CUFLaunchOpConversion>(patterns.getContext(), symtab);
 }

flang/test/Fir/CUDA/cuda-data-transfer.fir

@@ -295,4 +295,17 @@ func.func @_QPscalar_to_array() {
 // CHECK-LABEL: func.func @_QPscalar_to_array()
 // CHECK: _FortranACUFDataTransferDescDescNoRealloc
 
+func.func @_QPtest_type() {
+  %0 = cuf.alloc !fir.type<_QMbarTcmplx{id:i32,c:complex<f32>}> {bindc_name = "a", data_attr = #cuf.cuda<device>, uniq_name = "_QFtest_typeEa"} -> !fir.ref<!fir.type<_QMbarTcmplx{id:i32,c:complex<f32>}>>
+  %1 = fir.declare %0 {data_attr = #cuf.cuda<device>, uniq_name = "_QFtest_typeEa"} : (!fir.ref<!fir.type<_QMbarTcmplx{id:i32,c:complex<f32>}>>) -> !fir.ref<!fir.type<_QMbarTcmplx{id:i32,c:complex<f32>}>>
+  %2 = fir.alloca !fir.type<_QMbarTcmplx{id:i32,c:complex<f32>}> {bindc_name = "b", uniq_name = "_QFtest_typeEb"}
+  %3 = fir.declare %2 {uniq_name = "_QFtest_typeEb"} : (!fir.ref<!fir.type<_QMbarTcmplx{id:i32,c:complex<f32>}>>) -> !fir.ref<!fir.type<_QMbarTcmplx{id:i32,c:complex<f32>}>>
+  cuf.data_transfer %3 to %1 {transfer_kind = #cuf.cuda_transfer<host_device>} : !fir.ref<!fir.type<_QMbarTcmplx{id:i32,c:complex<f32>}>>, !fir.ref<!fir.type<_QMbarTcmplx{id:i32,c:complex<f32>}>>
+  cuf.free %1 : !fir.ref<!fir.type<_QMbarTcmplx{id:i32,c:complex<f32>}>> {data_attr = #cuf.cuda<device>}
+  return
+}
+
+// CHECK-LABEL: func.func @_QPtest_type()
+// CHECK: fir.call @_FortranACUFDataTransferPtrPtr(%{{.*}}, %{{.*}}, %c12{{.*}}, %c0{{.*}}, %{{.*}}, %{{.*}}) : (!fir.llvm_ptr<i8>, !fir.llvm_ptr<i8>, i64, i32, !fir.ref<i8>, i32) -> none
+
 } // end of module