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

[flang][cuda] Add cuf.set_allocator_idx for device component (#148750)

Browse Source
This commit is contained in:
Valentin Clement (バレンタイン クレメン) 2025-07-14 19:31:44 -07:00 committed by GitHub
parent 7cde974233
commit 90ef114a33
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 97 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
flang/include/flang/Semantics/tools.h
View File

@ -199,6 +199,8 @@ bool IsPolymorphic(const Symbol &);
bool IsUnlimitedPolymorphic(const Symbol &);
bool IsPolymorphicAllocatable(const Symbol &);
bool IsDeviceAllocatable(const Symbol &symbol);
inline bool IsCUDADeviceContext(const Scope *scope) {
if (scope) {
if (const Symbol * symbol{scope->symbol()}) {

77
flang/lib/Lower/ConvertVariable.cpp
View File

@ -771,9 +771,80 @@ static mlir::Value createNewLocal(Fortran::lower::AbstractConverter &converter,
return builder.create<cuf::SharedMemoryOp>(loc, ty, nm, symNm, lenParams,
indices);
if (!cuf::isCUDADeviceContext(builder.getRegion()))
return builder.create<cuf::AllocOp>(loc, ty, nm, symNm, dataAttr,
lenParams, indices);
if (!cuf::isCUDADeviceContext(builder.getRegion())) {
mlir::Value alloc = builder.create<cuf::AllocOp>(
loc, ty, nm, symNm, dataAttr, lenParams, indices);
if (const auto *details{
ultimateSymbol
.detailsIf<Fortran::semantics::ObjectEntityDetails>()}) {
const Fortran::semantics::DeclTypeSpec *type{details->type()};
const Fortran::semantics::DerivedTypeSpec *derived{
type ? type->AsDerived() : nullptr};
if (derived) {
Fortran::semantics::UltimateComponentIterator components{*derived};
auto recTy = mlir::dyn_cast<fir::RecordType>(ty);
mlir::Type fieldTy;
llvm::SmallVector<mlir::Value> coordinates;
for (const auto &sym : components) {
if (Fortran::semantics::IsDeviceAllocatable(sym)) {
unsigned fieldIdx = recTy.getFieldIndex(sym.name().ToString());
mlir::Type fieldTy;
std::vector<mlir::Value> coordinates;
if (fieldIdx != std::numeric_limits<unsigned>::max()) {
// Field found in the base record type.
auto fieldName = recTy.getTypeList()[fieldIdx].first;
fieldTy = recTy.getTypeList()[fieldIdx].second;
mlir::Value fieldIndex = builder.create<fir::FieldIndexOp>(
loc, fir::FieldType::get(fieldTy.getContext()), fieldName,
recTy,
/*typeParams=*/mlir::ValueRange{});
coordinates.push_back(fieldIndex);
} else {
// Field not found in base record type, search in potential
// record type components.
for (auto component : recTy.getTypeList()) {
if (auto childRecTy =
mlir::dyn_cast<fir::RecordType>(component.second)) {
fieldIdx = childRecTy.getFieldIndex(sym.name().ToString());
if (fieldIdx != std::numeric_limits<unsigned>::max()) {
mlir::Value parentFieldIndex =
builder.create<fir::FieldIndexOp>(
loc, fir::FieldType::get(childRecTy.getContext()),
component.first, recTy,
/*typeParams=*/mlir::ValueRange{});
coordinates.push_back(parentFieldIndex);
auto fieldName = childRecTy.getTypeList()[fieldIdx].first;
fieldTy = childRecTy.getTypeList()[fieldIdx].second;
mlir::Value childFieldIndex =
builder.create<fir::FieldIndexOp>(
loc, fir::FieldType::get(fieldTy.getContext()),
fieldName, childRecTy,
/*typeParams=*/mlir::ValueRange{});
coordinates.push_back(childFieldIndex);
break;
}
}
}
}
if (coordinates.empty())
TODO(loc, "device resident component in complex derived-type "
"hierarchy");
mlir::Value comp = builder.create<fir::CoordinateOp>(
loc, builder.getRefType(fieldTy), alloc, coordinates);
cuf::DataAttributeAttr dataAttr =
Fortran::lower::translateSymbolCUFDataAttribute(
builder.getContext(), sym);
builder.create<cuf::SetAllocatorIndexOp>(loc, comp, dataAttr);
}
}
}
}
return alloc;
}
}
// Let the builder do all the heavy lifting.

21
flang/test/Lower/CUDA/cuda-set-allocator.cuf Normal file
View File

@ -0,0 +1,21 @@
! RUN: bbc -emit-hlfir -fcuda %s -o - | FileCheck %s
module m1
type ty_device
integer, device, allocatable, dimension(:) :: x
integer :: y
integer, device, allocatable, dimension(:) :: z
end type
contains
subroutine sub1()
type(ty_device) :: a
end subroutine
! CHECK-LABEL: func.func @_QMm1Psub1()
! CHECK: %[[DT:.*]] = cuf.alloc !fir.type<_QMm1Tty_device{x:!fir.box<!fir.heap<!fir.array<?xi32>>>,y:i32,z:!fir.box<!fir.heap<!fir.array<?xi32>>>}> {bindc_name = "a", data_attr = #cuf.cuda<managed>, uniq_name = "_QMm1Fsub1Ea"} -> !fir.ref<!fir.type<_QMm1Tty_device{x:!fir.box<!fir.heap<!fir.array<?xi32>>>,y:i32,z:!fir.box<!fir.heap<!fir.array<?xi32>>>}>>
! CHECK: %[[X:.*]] = fir.coordinate_of %[[DT]], x : (!fir.ref<!fir.type<_QMm1Tty_device{x:!fir.box<!fir.heap<!fir.array<?xi32>>>,y:i32,z:!fir.box<!fir.heap<!fir.array<?xi32>>>}>>) -> !fir.ref<!fir.box<!fir.heap<!fir.array<?xi32>>>>
! CHECK: cuf.set_allocator_idx %[[X]] : !fir.ref<!fir.box<!fir.heap<!fir.array<?xi32>>>> {data_attr = #cuf.cuda<device>}
! CHECK: %[[Z:.*]] = fir.coordinate_of %[[DT]], z : (!fir.ref<!fir.type<_QMm1Tty_device{x:!fir.box<!fir.heap<!fir.array<?xi32>>>,y:i32,z:!fir.box<!fir.heap<!fir.array<?xi32>>>}>>) -> !fir.ref<!fir.box<!fir.heap<!fir.array<?xi32>>>>
! CHECK: cuf.set_allocator_idx %[[Z]] : !fir.ref<!fir.box<!fir.heap<!fir.array<?xi32>>>> {data_attr = #cuf.cuda<device>}
end module