Skip to content

Commit 0ec366c

Browse files
clementvalclementval
committed
[flang][cuda] Handle LHS with shape_shift in data transfer
1 parent 2baead0 commit 0ec366c

File tree

5 files changed

+62
-10
lines changed

5 files changed

+62
-10
lines changed

‎flang/include/flang/Optimizer/Dialect/CUF/CUFOps.td‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -162,7 +162,7 @@ def cuf_DataTransferOp : cuf_Op<"data_transfer", []> {
162162

163163
let arguments = (ins Arg<AnyType, "", [MemRead]>:$src,
164164
Arg<AnyRefOrBoxType, "", [MemWrite]>:$dst,
165-
Optional<fir_ShapeType>:$shape,
165+
Optional<AnyShapeOrShiftType>:$shape,
166166
cuf_DataTransferKindAttr:$transfer_kind);
167167

168168
let assemblyFormat = [{

‎flang/lib/Optimizer/Dialect/CUF/CUFOps.cpp‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -102,7 +102,7 @@ llvm::LogicalResult cuf::DataTransferOp::verify() {
102102
mlir::Type srcTy = getSrc().getType();
103103
mlir::Type dstTy = getDst().getType();
104104
if (getShape()) {
105-
if (!fir::isa_ref_type(srcTy) || !fir::isa_ref_type(dstTy))
105+
if (!fir::isa_ref_type(srcTy) && !fir::isa_ref_type(dstTy))
106106
return emitOpError()
107107
<< "shape can only be specified on data transfer with references";
108108
}

‎flang/lib/Optimizer/Transforms/CUFOpConversion.cpp‎

Lines changed: 15 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -570,13 +570,21 @@ struct CUFDataTransferOpConversion
570570
mlir::Type i64Ty = builder.getI64Type();
571571
mlir::Value nbElement;
572572
if (op.getShape()) {
573-
auto shapeOp =
574-
mlir::dyn_cast<fir::ShapeOp>(op.getShape().getDefiningOp());
575-
nbElement = rewriter.create<fir::ConvertOp>(loc, i64Ty,
576-
shapeOp.getExtents()[0]);
577-
for (unsigned i = 1; i < shapeOp.getExtents().size(); ++i) {
578-
auto operand = rewriter.create<fir::ConvertOp>(
579-
loc, i64Ty, shapeOp.getExtents()[i]);
573+
llvm::SmallVector<mlir::Value> extents;
574+
if (auto shapeOp =
575+
mlir::dyn_cast<fir::ShapeOp>(op.getShape().getDefiningOp())) {
576+
extents = shapeOp.getExtents();
577+
} else if (auto shapeShiftOp = mlir::dyn_cast<fir::ShapeShiftOp>(
578+
op.getShape().getDefiningOp())) {
579+
for (auto i : llvm::enumerate(shapeShiftOp.getPairs()))
580+
if (i.index() & 1)
581+
extents.push_back(i.value());
582+
}
583+
584+
nbElement = rewriter.create<fir::ConvertOp>(loc, i64Ty, extents[0]);
585+
for (unsigned i = 1; i < extents.size(); ++i) {
586+
auto operand =
587+
rewriter.create<fir::ConvertOp>(loc, i64Ty, extents[i]);
580588
nbElement =
581589
rewriter.create<mlir::arith::MulIOp>(loc, nbElement, operand);
582590
}

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

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -308,6 +308,7 @@ func.func @_QPtest_type() {
308308
// CHECK-LABEL: func.func @_QPtest_type()
309309
// CHECK: fir.call @_FortranACUFDataTransferPtrPtr(%{{.*}}, %{{.*}}, %c12{{.*}}, %c0{{.*}}, %{{.*}}, %{{.*}}) : (!fir.llvm_ptr<i8>, !fir.llvm_ptr<i8>, i64, i32, !fir.ref<i8>, i32) -> none
310310

311+
<<<<<<< HEAD
311312
func.func @_QPtest_array_type() {
312313
%c10 = arith.constant 10 : index
313314
%0 = cuf.alloc !fir.array<10x!fir.type<_QMbarTcmplx{id:i32,c:complex<f32>}>> {bindc_name = "a", data_attr = #cuf.cuda<device>, uniq_name = "_QFtest_array_typeEa"} -> !fir.ref<!fir.array<10x!fir.type<_QMbarTcmplx{id:i32,c:complex<f32>}>>>
@@ -326,5 +327,41 @@ func.func @_QPtest_array_type() {
326327
// CHECK: fir.call @_FortranACUFMemAlloc(%[[CONV_BYTES]], %c0{{.*}}, %{{.*}}, %{{.*}}) : (i64, i32, !fir.ref<i8>, i32) -> !fir.llvm_ptr<i8>
327328
// CHECK: %[[BYTES:.*]] = arith.muli %c10{{.*}}, %c12{{.*}} : i64
328329
// CHECK: fir.call @_FortranACUFDataTransferPtrPtr(%{{.*}}, %{{.*}}, %[[BYTES]], %c0{{.*}}, %{{.*}}, %{{.*}}) : (!fir.llvm_ptr<i8>, !fir.llvm_ptr<i8>, i64, i32, !fir.ref<i8>, i32) -> none
330+
=======
331+
func.func @_QPshape_shift() {
332+
%c0_i32 = arith.constant 0 : i32
333+
%c11 = arith.constant 11 : index
334+
%c10 = arith.constant 10 : index
335+
%0 = cuf.alloc !fir.array<10xi32> {bindc_name = "cdev", data_attr = #cuf.cuda<device>, uniq_name = "_QFshape_shiftEcdev"} -> !fir.ref<!fir.array<10xi32>>
336+
%1 = fir.shape_shift %c11, %c10 : (index, index) -> !fir.shapeshift<1>
337+
%2:2 = hlfir.declare %0(%1) {data_attr = #cuf.cuda<device>, uniq_name = "_QFshape_shiftEcdev"} : (!fir.ref<!fir.array<10xi32>>, !fir.shapeshift<1>) -> (!fir.box<!fir.array<10xi32>>, !fir.ref<!fir.array<10xi32>>)
338+
cuf.data_transfer %c0_i32 to %2#1, %1 : !fir.shapeshift<1> {transfer_kind = #cuf.cuda_transfer<host_device>} : i32, !fir.ref<!fir.array<10xi32>>
339+
cuf.free %2#1 : !fir.ref<!fir.array<10xi32>> {data_attr = #cuf.cuda<device>}
340+
return
341+
}
342+
343+
// CHECK-LABEL: func.func @_QPshape_shift()
344+
// CHECK: fir.call @_FortranACUFDataTransferDescDescNoRealloc
345+
346+
func.func @_QPshape_shift2() {
347+
%c11 = arith.constant 11 : index
348+
%c10 = arith.constant 10 : index
349+
%0 = fir.alloca !fir.array<10xi32> {bindc_name = "ahost", uniq_name = "_QFshape_shift2Eahost"}
350+
%1 = fir.shape %c10 : (index) -> !fir.shape<1>
351+
%2:2 = hlfir.declare %0(%1) {uniq_name = "_QFshape_shift2Eahost"} : (!fir.ref<!fir.array<10xi32>>, !fir.shape<1>) -> (!fir.ref<!fir.array<10xi32>>, !fir.ref<!fir.array<10xi32>>)
352+
%3 = cuf.alloc !fir.array<10xi32> {bindc_name = "cdev", data_attr = #cuf.cuda<device>, uniq_name = "_QFshape_shift2Ecdev"} -> !fir.ref<!fir.array<10xi32>>
353+
%4 = fir.shape_shift %c11, %c10 : (index, index) -> !fir.shapeshift<1>
354+
%5:2 = hlfir.declare %3(%4) {data_attr = #cuf.cuda<device>, uniq_name = "_QFshape_shift2Ecdev"} : (!fir.ref<!fir.array<10xi32>>, !fir.shapeshift<1>) -> (!fir.box<!fir.array<10xi32>>, !fir.ref<!fir.array<10xi32>>)
355+
cuf.data_transfer %2#0 to %5#1, %4 : !fir.shapeshift<1> {transfer_kind = #cuf.cuda_transfer<host_device>} : !fir.ref<!fir.array<10xi32>>, !fir.ref<!fir.array<10xi32>>
356+
cuf.free %5#1 : !fir.ref<!fir.array<10xi32>> {data_attr = #cuf.cuda<device>}
357+
return
358+
}
359+
360+
// CHECK-LABEL: func.func @_QPshape_shift2()
361+
// CHECK: %[[C10:.*]] = fir.convert %c10{{.*}} : (index) -> i64
362+
// CHECK: %[[BYTES:.*]] = arith.muli %[[C10]], %c4{{.*}} : i64
363+
// CHECK: fir.call @_FortranACUFDataTransferPtrPtr(%{{.*}}, %{{.*}}, %[[BYTES]], %{{.*}}, %{{.*}}, %{{.*}}) : (!fir.llvm_ptr<i8>, !fir.llvm_ptr<i8>, i64, i32, !fir.ref<i8>, i32) -> none
364+
365+
>>>>>>> 2480604fdb1c ([flang][cuda] Handle LHS with shape_shift in data transfer)
329366

330367
} // end of module

‎flang/test/Lower/CUDA/cuda-data-transfer.cuf‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,8 @@ module mod1
77
integer :: i
88
end type
99

10+
integer, device, dimension(11:20) :: cdev
11+
1012
contains
1113
function dev1(a)
1214
integer, device :: a(:)
@@ -16,6 +18,7 @@ contains
1618
end
1719

1820
subroutine sub1()
21+
use mod1
1922
integer, device :: m
2023
integer, device :: adev(10)
2124
integer :: i, ahost(10), bhost(10)
@@ -34,6 +37,8 @@ subroutine sub1()
3437

3538
adev = 10
3639

40+
cdev = 0
41+
3742
end
3843

3944
! CHECK-LABEL: func.func @_QPsub1()
@@ -70,6 +75,8 @@ end
7075

7176
! CHECK: cuf.data_transfer %c10{{.*}} to %[[ADEV]]#0 {transfer_kind = #cuf.cuda_transfer<host_device>} : i32, !fir.ref<!fir.array<10xi32>>
7277

78+
! CHECK: cuf.data_transfer %c0{{.*}} to %{{.*}}#1, %{{.*}} : !fir.shapeshift<1> {transfer_kind = #cuf.cuda_transfer<host_device>} : i32, !fir.ref<!fir.array<10xi32>>
79+
7380
subroutine sub2()
7481
integer, device :: m
7582
integer, device :: adev(10), bdev(10)
@@ -124,7 +131,7 @@ end
124131
! CHECK: %[[TMP:.*]] = fir.alloca !fir.type<_QMmod1Tt1{i:i32}> {bindc_name = ".tmp"}
125132
! CHECK: %[[AHOST:.*]]:2 = hlfir.declare %{{.*}}(%{{.*}}) {uniq_name = "_QFsub3Eahost"} : (!fir.ref<!fir.array<10xi32>>, !fir.shape<1>) -> (!fir.ref<!fir.array<10xi32>>, !fir.ref<!fir.array<10xi32>>)
126133
! CHECK: %[[BHOST:.*]]:2 = hlfir.declare %{{.*}}(%{{.*}}) {uniq_name = "_QFsub3Ebhost"} : (!fir.ref<!fir.array<10xi32>>, !fir.shape<1>) -> (!fir.ref<!fir.array<10xi32>>, !fir.ref<!fir.array<10xi32>>)
127-
! CHECK: %[[T:.*]]:2 = hlfir.declare %7 {data_attr = #cuf.cuda<device>, uniq_name = "_QFsub3Et"} : (!fir.ref<!fir.type<_QMmod1Tt1{i:i32}>>) -> (!fir.ref<!fir.type<_QMmod1Tt1{i:i32}>>, !fir.ref<!fir.type<_QMmod1Tt1{i:i32}>>)
134+
! CHECK: %[[T:.*]]:2 = hlfir.declare %{{.*}} {data_attr = #cuf.cuda<device>, uniq_name = "_QFsub3Et"} : (!fir.ref<!fir.type<_QMmod1Tt1{i:i32}>>) -> (!fir.ref<!fir.type<_QMmod1Tt1{i:i32}>>, !fir.ref<!fir.type<_QMmod1Tt1{i:i32}>>)
128135
! CHECK: %[[TMP_DECL:.*]]:2 = hlfir.declare %0 {uniq_name = ".tmp"} : (!fir.ref<!fir.type<_QMmod1Tt1{i:i32}>>) -> (!fir.ref<!fir.type<_QMmod1Tt1{i:i32}>>, !fir.ref<!fir.type<_QMmod1Tt1{i:i32}>>)
129136
! CHECK: cuf.data_transfer %[[T]]#1 to %[[TMP_DECL]]#0 {transfer_kind = #cuf.cuda_transfer<device_host>} : !fir.ref<!fir.type<_QMmod1Tt1{i:i32}>>, !fir.ref<!fir.type<_QMmod1Tt1{i:i32}>>
130137

0 commit comments

Comments
 (0)