shylie/llvm-project
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

[flang][cuda] Compute size of derived type arrays (#115914)

Browse Source
This commit is contained in:
Valentin Clement (バレンタイン クレメン) 2024-11-12 21:23:58 -08:00 committed by GitHub
parent 4714215efb
commit 2583071fb4
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 31 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

17
flang/lib/Optimizer/Transforms/CUFOpConversion.cpp
View File

@ -321,9 +321,15 @@ struct CUFAllocOpConversion : public mlir::OpRewritePattern<cuf::AllocOp> {
builder.createIntegerConstant(loc, builder.getIndexType(), width);
} else if (auto seqTy = mlir::dyn_cast_or_null<fir::SequenceType>(
op.getInType())) {
mlir::Value width = builder.createIntegerConstant(
loc, builder.getIndexType(),
computeWidth(loc, seqTy.getEleTy(), kindMap));
std::size_t size = 0;
if (fir::isa_derived(seqTy.getEleTy())) {
mlir::Type structTy = typeConverter->convertType(seqTy.getEleTy());
size = dl->getTypeSizeInBits(structTy) / 8;
} else {
size = computeWidth(loc, seqTy.getEleTy(), kindMap);
}
mlir::Value width =
builder.createIntegerConstant(loc, builder.getIndexType(), size);
mlir::Value nbElem;
if (fir::sequenceWithNonConstantShape(seqTy)) {
assert(!op.getShape().empty() && "expect shape with dynamic arrays");
@ -580,8 +586,9 @@ struct CUFDataTransferOpConversion
loc, i64Ty, seqTy.getConstantArraySize());
}
unsigned width = 0;
if (fir::isa_derived(dstTy)) {
mlir::Type structTy = typeConverter->convertType(dstTy);
if (fir::isa_derived(fir::unwrapSequenceType(dstTy))) {
mlir::Type structTy =
typeConverter->convertType(fir::unwrapSequenceType(dstTy));
width = dl->getTypeSizeInBits(structTy) / 8;
} else {
width = computeWidth(loc, dstTy, kindMap);

19
flang/test/Fir/CUDA/cuda-data-transfer.fir
View File

@ -308,4 +308,23 @@ func.func @_QPtest_type() {
// 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
func.func @_QPtest_array_type() {
%c10 = arith.constant 10 : index
%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>}>>>
%1 = fir.shape %c10 : (index) -> !fir.shape<1>
%2 = fir.declare %0(%1) {data_attr = #cuf.cuda<device>, uniq_name = "_QFtest_array_typeEa"} : (!fir.ref<!fir.array<10x!fir.type<_QMbarTcmplx{id:i32,c:complex<f32>}>>>, !fir.shape<1>) -> !fir.ref<!fir.array<10x!fir.type<_QMbarTcmplx{id:i32,c:complex<f32>}>>>
%3 = fir.alloca !fir.array<10x!fir.type<_QMbarTcmplx{id:i32,c:complex<f32>}>> {bindc_name = "b", uniq_name = "_QFtest_array_typeEb"}
%4 = fir.declare %3(%1) {uniq_name = "_QFtest_array_typeEb"} : (!fir.ref<!fir.array<10x!fir.type<_QMbarTcmplx{id:i32,c:complex<f32>}>>>, !fir.shape<1>) -> !fir.ref<!fir.array<10x!fir.type<_QMbarTcmplx{id:i32,c:complex<f32>}>>>
cuf.data_transfer %4 to %2 {transfer_kind = #cuf.cuda_transfer<host_device>} : !fir.ref<!fir.array<10x!fir.type<_QMbarTcmplx{id:i32,c:complex<f32>}>>>, !fir.ref<!fir.array<10x!fir.type<_QMbarTcmplx{id:i32,c:complex<f32>}>>>
cuf.free %2 : !fir.ref<!fir.array<10x!fir.type<_QMbarTcmplx{id:i32,c:complex<f32>}>>> {data_attr = #cuf.cuda<device>}
return
}
// CHECK-LABEL: func.func @_QPtest_array_type()
// CHECK: %[[BYTES:.*]] = arith.muli %c10{{.*}}, %c12 : index
// CHECK: %[[CONV_BYTES:.*]] = fir.convert %[[BYTES]] : (index) -> i64
// CHECK: fir.call @_FortranACUFMemAlloc(%[[CONV_BYTES]], %c0{{.*}}, %{{.*}}, %{{.*}}) : (i64, i32, !fir.ref<i8>, i32) -> !fir.llvm_ptr<i8>
// CHECK: %[[BYTES:.*]] = arith.muli %c10{{.*}}, %c12{{.*}} : i64
// CHECK: fir.call @_FortranACUFDataTransferPtrPtr(%{{.*}}, %{{.*}}, %[[BYTES]], %c0{{.*}}, %{{.*}}, %{{.*}}) : (!fir.llvm_ptr<i8>, !fir.llvm_ptr<i8>, i64, i32, !fir.ref<i8>, i32) -> none
} // end of module