Skip to content

Commit 7dbc664

Browse files
clementvalclementval
authored
[flang][cuda] Convert data transfer between scalar and arrays (#110180)
Add conversion of data transfer between scalars or between arrays. Scalar to array are not handled yet.
1 parent 432ba35 commit 7dbc664

File tree

2 files changed

+110
-18
lines changed

2 files changed

+110
-18
lines changed

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

Lines changed: 57 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -370,11 +370,6 @@ struct CufDataTransferOpConversion
370370
mlir::Type srcTy = fir::unwrapRefType(op.getSrc().getType());
371371
mlir::Type dstTy = fir::unwrapRefType(op.getDst().getType());
372372

373-
// Only convert cuf.data_transfer with at least one descripor.
374-
if (!mlir::isa<fir::BaseBoxType>(srcTy) &&
375-
!mlir::isa<fir::BaseBoxType>(dstTy))
376-
return failure();
377-
378373
unsigned mode;
379374
if (op.getTransferKind() == cuf::DataTransferKind::HostDevice) {
380375
mode = kHostToDevice;
@@ -387,7 +382,64 @@ struct CufDataTransferOpConversion
387382
auto mod = op->getParentOfType<mlir::ModuleOp>();
388383
fir::FirOpBuilder builder(rewriter, mod);
389384
mlir::Location loc = op.getLoc();
385+
fir::KindMapping kindMap{fir::getKindMapping(mod)};
386+
mlir::Value modeValue =
387+
builder.createIntegerConstant(loc, builder.getI32Type(), mode);
388+
389+
// Convert data transfer without any descriptor.
390+
if (!mlir::isa<fir::BaseBoxType>(srcTy) &&
391+
!mlir::isa<fir::BaseBoxType>(dstTy)) {
392+
393+
if (fir::isa_trivial(srcTy) && !fir::isa_trivial(dstTy)) {
394+
// TODO: scalar to array data transfer.
395+
mlir::emitError(loc,
396+
"not yet implemented: scalar to array data transfer\n");
397+
return mlir::failure();
398+
}
399+
400+
mlir::Type i64Ty = builder.getI64Type();
401+
mlir::Value nbElement;
402+
if (op.getShape()) {
403+
auto shapeOp =
404+
mlir::dyn_cast<fir::ShapeOp>(op.getShape().getDefiningOp());
405+
nbElement = rewriter.create<fir::ConvertOp>(loc, i64Ty,
406+
shapeOp.getExtents()[0]);
407+
for (unsigned i = 1; i < shapeOp.getExtents().size(); ++i) {
408+
auto operand = rewriter.create<fir::ConvertOp>(
409+
loc, i64Ty, shapeOp.getExtents()[i]);
410+
nbElement =
411+
rewriter.create<mlir::arith::MulIOp>(loc, nbElement, operand);
412+
}
413+
} else {
414+
if (auto seqTy = mlir::dyn_cast_or_null<fir::SequenceType>(dstTy))
415+
nbElement = builder.createIntegerConstant(
416+
loc, i64Ty, seqTy.getConstantArraySize());
417+
}
418+
int width = computeWidth(loc, dstTy, kindMap);
419+
mlir::Value widthValue = rewriter.create<mlir::arith::ConstantOp>(
420+
loc, i64Ty, rewriter.getIntegerAttr(i64Ty, width));
421+
mlir::Value bytes =
422+
nbElement
423+
? rewriter.create<mlir::arith::MulIOp>(loc, nbElement, widthValue)
424+
: widthValue;
425+
426+
mlir::func::FuncOp func =
427+
fir::runtime::getRuntimeFunc<mkRTKey(CUFDataTransferPtrPtr)>(loc,
428+
builder);
429+
auto fTy = func.getFunctionType();
430+
mlir::Value sourceFile = fir::factory::locationToFilename(builder, loc);
431+
mlir::Value sourceLine =
432+
fir::factory::locationToLineNo(builder, loc, fTy.getInput(5));
433+
434+
llvm::SmallVector<mlir::Value> args{fir::runtime::createArguments(
435+
builder, loc, fTy, op.getDst(), op.getSrc(), bytes, modeValue,
436+
sourceFile, sourceLine)};
437+
builder.create<fir::CallOp>(loc, func, args);
438+
rewriter.eraseOp(op);
439+
return mlir::success();
440+
}
390441

442+
// Conversion of data transfer involving at least one descriptor.
391443
if (mlir::isa<fir::BaseBoxType>(srcTy) &&
392444
mlir::isa<fir::BaseBoxType>(dstTy)) {
393445
// Transfer between two descriptor.
@@ -396,8 +448,6 @@ struct CufDataTransferOpConversion
396448
loc, builder);
397449

398450
auto fTy = func.getFunctionType();
399-
mlir::Value modeValue =
400-
builder.createIntegerConstant(loc, builder.getI32Type(), mode);
401451
mlir::Value sourceFile = fir::factory::locationToFilename(builder, loc);
402452
mlir::Value sourceLine =
403453
fir::factory::locationToLineNo(builder, loc, fTy.getInput(4));
@@ -430,8 +480,6 @@ struct CufDataTransferOpConversion
430480
builder.create<fir::CallOp>(loc, func, args);
431481
rewriter.eraseOp(op);
432482
} else {
433-
mlir::Value modeValue =
434-
builder.createIntegerConstant(loc, builder.getI32Type(), mode);
435483
// Type used to compute the width.
436484
mlir::Type computeType = dstTy;
437485
auto seqTy = mlir::dyn_cast<fir::SequenceType>(dstTy);
@@ -441,7 +489,6 @@ struct CufDataTransferOpConversion
441489
computeType = srcTy;
442490
seqTy = mlir::dyn_cast<fir::SequenceType>(srcTy);
443491
}
444-
fir::KindMapping kindMap{fir::getKindMapping(mod)};
445492
int width = computeWidth(loc, computeType, kindMap);
446493

447494
mlir::Value nbElement;
@@ -509,13 +556,6 @@ class CufOpConversion : public fir::impl::CufOpConversionBase<CufOpConversion> {
509556
fir::support::getOrSetDataLayout(module, /*allowDefaultLayout=*/false);
510557
fir::LLVMTypeConverter typeConverter(module, /*applyTBAA=*/false,
511558
/*forceUnifiedTBAATree=*/false, *dl);
512-
target.addDynamicallyLegalOp<cuf::DataTransferOp>(
513-
[](::cuf::DataTransferOp op) {
514-
mlir::Type srcTy = fir::unwrapRefType(op.getSrc().getType());
515-
mlir::Type dstTy = fir::unwrapRefType(op.getDst().getType());
516-
return !mlir::isa<fir::BaseBoxType>(srcTy) &&
517-
!mlir::isa<fir::BaseBoxType>(dstTy);
518-
});
519559
target.addLegalDialect<fir::FIROpsDialect, mlir::arith::ArithDialect>();
520560
cuf::populateCUFToFIRConversionPatterns(typeConverter, *dl, patterns);
521561
if (mlir::failed(mlir::applyPartialConversion(getOperation(), target,

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

Lines changed: 53 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -70,7 +70,6 @@ func.func @_QPsub4() {
7070
cuf.free %4#1 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xi32>>>> {data_attr = #cuf.cuda<device>}
7171
return
7272
}
73-
7473
// CHECK-LABEL: func.func @_QPsub4()
7574
// CHECK: %[[ADEV:.*]]:2 = hlfir.declare %{{.*}} {data_attr = #cuf.cuda<device>, fortran_attrs = #fir.var_attrs<allocatable>, uniq_name = "_QFsub4Eadev"} : (!fir.ref<!fir.box<!fir.heap<!fir.array<?xi32>>>>) -> (!fir.ref<!fir.box<!fir.heap<!fir.array<?xi32>>>>, !fir.ref<!fir.box<!fir.heap<!fir.array<?xi32>>>>)
7675
// CHECK: %[[AHOST:.*]]:2 = hlfir.declare %{{.*}}(%{{.*}}) {uniq_name = "_QFsub4Eahost"} : (!fir.ref<!fir.array<10xi32>>, !fir.shape<1>) -> (!fir.ref<!fir.array<10xi32>>, !fir.ref<!fir.array<10xi32>>)
@@ -137,4 +136,57 @@ func.func @_QPsub5(%arg0: !fir.ref<i32> {fir.bindc_name = "n"}) {
137136
// CHECK: %[[BYTES_CONV:.*]] = fir.convert %[[BYTES]] : (index) -> i64
138137
// CHECK: fir.call @_FortranACUFDataTransferPtrDesc(%[[AHOST_PTR]], %[[ADEV_BOX]], %[[BYTES_CONV]], %c1{{.*}}, %{{.*}}, %{{.*}}) : (!fir.llvm_ptr<i8>, !fir.box<none>, i64, i32, !fir.ref<i8>, i32) -> none
139138

139+
func.func @_QPsub6() {
140+
%0 = cuf.alloc i32 {bindc_name = "idev", data_attr = #cuf.cuda<device>, uniq_name = "_QFsub6Eidev"} -> !fir.ref<i32>
141+
%1:2 = hlfir.declare %0 {data_attr = #cuf.cuda<device>, uniq_name = "_QFsub6Eidev"} : (!fir.ref<i32>) -> (!fir.ref<i32>, !fir.ref<i32>)
142+
%2 = fir.alloca i32 {bindc_name = "ihost", uniq_name = "_QFsub6Eihost"}
143+
%3:2 = hlfir.declare %2 {uniq_name = "_QFsub6Eihost"} : (!fir.ref<i32>) -> (!fir.ref<i32>, !fir.ref<i32>)
144+
cuf.data_transfer %1#0 to %3#0 {transfer_kind = #cuf.cuda_transfer<device_host>} : !fir.ref<i32>, !fir.ref<i32>
145+
%4 = fir.load %3#0 : !fir.ref<i32>
146+
%5:3 = hlfir.associate %4 {uniq_name = ".cuf_host_tmp"} : (i32) -> (!fir.ref<i32>, !fir.ref<i32>, i1)
147+
cuf.data_transfer %5#0 to %1#0 {transfer_kind = #cuf.cuda_transfer<host_device>} : !fir.ref<i32>, !fir.ref<i32>
148+
hlfir.end_associate %5#1, %5#2 : !fir.ref<i32>, i1
149+
cuf.free %1#1 : !fir.ref<i32> {data_attr = #cuf.cuda<device>}
150+
return
151+
}
152+
153+
// CHECK-LABEL: func.func @_QPsub6()
154+
// CHECK: %[[IDEV:.*]]:2 = hlfir.declare %{{.*}} {data_attr = #cuf.cuda<device>, uniq_name = "_QFsub6Eidev"} : (!fir.ref<i32>) -> (!fir.ref<i32>, !fir.ref<i32>)
155+
// CHECK: %[[IHOST:.*]]:2 = hlfir.declare %{{.*}} {uniq_name = "_QFsub6Eihost"} : (!fir.ref<i32>) -> (!fir.ref<i32>, !fir.ref<i32>)
156+
// CHECK: %[[DST:.*]] = fir.convert %[[IHOST]]#0 : (!fir.ref<i32>) -> !fir.llvm_ptr<i8>
157+
// CHECK: %[[SRC:.*]] = fir.convert %[[IDEV]]#0 : (!fir.ref<i32>) -> !fir.llvm_ptr<i8>
158+
// CHECK: fir.call @_FortranACUFDataTransferPtrPtr(%[[DST]], %[[SRC]], %c4{{.*}}, %c1{{.*}}, %{{.*}}, %{{.*}}) : (!fir.llvm_ptr<i8>, !fir.llvm_ptr<i8>, i64, i32, !fir.ref<i8>, i32) -> none
159+
// CHECK: %[[LOAD:.*]] = fir.load %[[IHOST]]#0 : !fir.ref<i32>
160+
// CHECK: %[[ASSOC:.*]]:3 = hlfir.associate %[[LOAD]] {uniq_name = ".cuf_host_tmp"} : (i32) -> (!fir.ref<i32>, !fir.ref<i32>, i1)
161+
// CHECK: %[[DST:.*]] = fir.convert %[[IDEV]]#0 : (!fir.ref<i32>) -> !fir.llvm_ptr<i8>
162+
// CHECK: %[[SRC:.*]] = fir.convert %[[ASSOC]]#0 : (!fir.ref<i32>) -> !fir.llvm_ptr<i8>
163+
// CHECK: fir.call @_FortranACUFDataTransferPtrPtr(%[[DST]], %[[SRC]], %c4{{.*}}, %c0{{.*}}, %{{.*}}, %{{.*}}) : (!fir.llvm_ptr<i8>, !fir.llvm_ptr<i8>, i64, i32, !fir.ref<i8>, i32) -> none
164+
165+
func.func @_QPsub7() {
166+
%c10 = arith.constant 10 : index
167+
%0 = cuf.alloc !fir.array<10xi32> {bindc_name = "idev", data_attr = #cuf.cuda<device>, uniq_name = "_QFsub7Eidev"} -> !fir.ref<!fir.array<10xi32>>
168+
%1 = fir.shape %c10 : (index) -> !fir.shape<1>
169+
%2:2 = hlfir.declare %0(%1) {data_attr = #cuf.cuda<device>, uniq_name = "_QFsub7Eidev"} : (!fir.ref<!fir.array<10xi32>>, !fir.shape<1>) -> (!fir.ref<!fir.array<10xi32>>, !fir.ref<!fir.array<10xi32>>)
170+
%c10_0 = arith.constant 10 : index
171+
%3 = fir.alloca !fir.array<10xi32> {bindc_name = "ihost", uniq_name = "_QFsub7Eihost"}
172+
%4 = fir.shape %c10_0 : (index) -> !fir.shape<1>
173+
%5:2 = hlfir.declare %3(%4) {uniq_name = "_QFsub7Eihost"} : (!fir.ref<!fir.array<10xi32>>, !fir.shape<1>) -> (!fir.ref<!fir.array<10xi32>>, !fir.ref<!fir.array<10xi32>>)
174+
cuf.data_transfer %2#0 to %5#0 {transfer_kind = #cuf.cuda_transfer<device_host>} : !fir.ref<!fir.array<10xi32>>, !fir.ref<!fir.array<10xi32>>
175+
cuf.data_transfer %5#0 to %2#0 {transfer_kind = #cuf.cuda_transfer<host_device>} : !fir.ref<!fir.array<10xi32>>, !fir.ref<!fir.array<10xi32>>
176+
cuf.free %2#1 : !fir.ref<!fir.array<10xi32>> {data_attr = #cuf.cuda<device>}
177+
return
178+
}
179+
180+
// CHECK-LABEL: func.func @_QPsub7()
181+
// CHECK: %[[IDEV:.*]]:2 = hlfir.declare %{{.*}} {data_attr = #cuf.cuda<device>, uniq_name = "_QFsub7Eidev"} : (!fir.ref<!fir.array<10xi32>>, !fir.shape<1>) -> (!fir.ref<!fir.array<10xi32>>, !fir.ref<!fir.array<10xi32>>)
182+
// CHECK: %[[IHOST:.*]]:2 = hlfir.declare %{{.*}} {uniq_name = "_QFsub7Eihost"} : (!fir.ref<!fir.array<10xi32>>, !fir.shape<1>) -> (!fir.ref<!fir.array<10xi32>>, !fir.ref<!fir.array<10xi32>>)
183+
// CHECK: %[[BYTES:.*]] = arith.muli %c10{{.*}}, %c4{{.*}} : i64
184+
// CHECK: %[[DST:.*]] = fir.convert %[[IHOST]]#0 : (!fir.ref<!fir.array<10xi32>>) -> !fir.llvm_ptr<i8>
185+
// CHECK: %[[SRC:.*]] = fir.convert %[[IDEV]]#0 : (!fir.ref<!fir.array<10xi32>>) -> !fir.llvm_ptr<i8>
186+
// CHECK: fir.call @_FortranACUFDataTransferPtrPtr(%[[DST]], %[[SRC]], %[[BYTES]], %c1{{.*}}, %{{.*}}, %{{.*}}) : (!fir.llvm_ptr<i8>, !fir.llvm_ptr<i8>, i64, i32, !fir.ref<i8>, i32) -> none
187+
// CHECK: %[[BYTES:.*]] = arith.muli %c10{{.*}}, %c4{{.*}} : i64
188+
// CHECK: %[[DST:.*]] = fir.convert %[[IDEV]]#0 : (!fir.ref<!fir.array<10xi32>>) -> !fir.llvm_ptr<i8>
189+
// CHECK: %[[SRC:.*]] = fir.convert %[[IHOST]]#0 : (!fir.ref<!fir.array<10xi32>>) -> !fir.llvm_ptr<i8>
190+
// CHECK: fir.call @_FortranACUFDataTransferPtrPtr(%[[DST]], %[[SRC]], %[[BYTES]], %c0{{.*}}, %{{.*}}, %{{.*}}) : (!fir.llvm_ptr<i8>, !fir.llvm_ptr<i8>, i64, i32, !fir.ref<i8>, i32) -> none
191+
140192
} // end of module

0 commit comments

Comments
 (0)