Skip to content

Commit dddc57d

Browse files
committed
[flang][cuda] Get the descriptor pointer fot data transfer
1 parent fe00ab4 commit dddc57d

File tree

2 files changed

+37
-0
lines changed

2 files changed

+37
-0
lines changed

flang/lib/Optimizer/Transforms/CUFOpConversion.cpp

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -741,6 +741,9 @@ struct CUFDataTransferOpConversion
741741
fir::StoreOp::create(builder, loc, val, box);
742742
return box;
743743
}
744+
if (mlir::isa<fir::BaseBoxType>(val.getType()))
745+
if (auto loadOp = mlir::dyn_cast<fir::LoadOp>(val.getDefiningOp()))
746+
return loadOp.getMemref();
744747
return val;
745748
};
746749

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

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -691,5 +691,39 @@ func.func @_QPtesti4(%arg0: !fir.ref<i32> {fir.bindc_name = "n1"}, %arg1: !fir.r
691691
// CHECK-LABEL: func.func @_QPtesti4
692692
// CHECK: fir.call @_FortranACUFDataTransferCstDesc
693693

694+
// -----
695+
696+
func.func @_QQmain() attributes {fir.bindc_name = "T"} {
697+
%c2 = arith.constant 2 : index
698+
%c1 = arith.constant 1 : index
699+
%c80 = arith.constant 80 : index
700+
%c0 = arith.constant 0 : index
701+
%0 = fir.dummy_scope : !fir.dscope
702+
%1 = cuf.alloc !fir.box<!fir.heap<!fir.array<?x?x?xf16>>> {bindc_name = "a", data_attr = #cuf.cuda<device>, uniq_name = "_QFEa"} -> !fir.ref<!fir.box<!fir.heap<!fir.array<?x?x?xf16>>>>
703+
%2 = fir.zero_bits !fir.heap<!fir.array<?x?x?xf16>>
704+
%3 = fir.shape %c0, %c0, %c0 : (index, index, index) -> !fir.shape<3>
705+
%4 = fir.embox %2(%3) {allocator_idx = 2 : i32} : (!fir.heap<!fir.array<?x?x?xf16>>, !fir.shape<3>) -> !fir.box<!fir.heap<!fir.array<?x?x?xf16>>>
706+
fir.store %4 to %1 : !fir.ref<!fir.box<!fir.heap<!fir.array<?x?x?xf16>>>>
707+
%5 = fir.declare %1 {data_attr = #cuf.cuda<device>, fortran_attrs = #fir.var_attrs<allocatable>, uniq_name = "_QFEa"} : (!fir.ref<!fir.box<!fir.heap<!fir.array<?x?x?xf16>>>>) -> !fir.ref<!fir.box<!fir.heap<!fir.array<?x?x?xf16>>>>
708+
%6 = fir.address_of(@_QFEha) : !fir.ref<!fir.array<80x80x80xf32>>
709+
%7 = fir.shape %c80, %c80, %c80 : (index, index, index) -> !fir.shape<3>
710+
%8 = fir.declare %6(%7) {uniq_name = "_QFEha"} : (!fir.ref<!fir.array<80x80x80xf32>>, !fir.shape<3>) -> !fir.ref<!fir.array<80x80x80xf32>>
711+
%9 = fir.address_of(@_QFECn) : !fir.ref<i32>
712+
%10 = fir.declare %9 {fortran_attrs = #fir.var_attrs<parameter>, uniq_name = "_QFECn"} : (!fir.ref<i32>) -> !fir.ref<i32>
713+
%11 = fir.load %5 : !fir.ref<!fir.box<!fir.heap<!fir.array<?x?x?xf16>>>>
714+
%12:3 = fir.box_dims %11, %c0 : (!fir.box<!fir.heap<!fir.array<?x?x?xf16>>>, index) -> (index, index, index)
715+
%13:3 = fir.box_dims %11, %c1 : (!fir.box<!fir.heap<!fir.array<?x?x?xf16>>>, index) -> (index, index, index)
716+
%14:3 = fir.box_dims %11, %c2 : (!fir.box<!fir.heap<!fir.array<?x?x?xf16>>>, index) -> (index, index, index)
717+
%15 = fir.shape %12#1, %13#1, %14#1 : (index, index, index) -> !fir.shape<3>
718+
%16 = fir.allocmem !fir.array<?x?x?xf16>, %12#1, %13#1, %14#1 {bindc_name = ".tmp", uniq_name = ""}
719+
%17 = fir.declare %16(%15) {uniq_name = ".tmp"} : (!fir.heap<!fir.array<?x?x?xf16>>, !fir.shape<3>) -> !fir.heap<!fir.array<?x?x?xf16>>
720+
%18 = fir.embox %17(%15) : (!fir.heap<!fir.array<?x?x?xf16>>, !fir.shape<3>) -> !fir.box<!fir.array<?x?x?xf16>>
721+
cuf.data_transfer %11 to %18 {transfer_kind = #cuf.cuda_transfer<device_host>} : !fir.box<!fir.heap<!fir.array<?x?x?xf16>>>, !fir.box<!fir.array<?x?x?xf16>>
722+
return
723+
}
724+
725+
// CHECK-LABEL: func.func @_QQmain()
726+
// CHECK: fir.call @_FortranACUFDataTransferDescDesc
727+
694728
} // end of module
695729

0 commit comments

Comments
 (0)