From efa03d69be450d124346f646f0909c232661bc9a Mon Sep 17 00:00:00 2001 From: Valentin Clement Date: Mon, 14 Jul 2025 15:29:42 -0700 Subject: [PATCH 1/2] [flang][cuda] Add cuf.set_allocator_idx for device component --- flang/include/flang/Semantics/tools.h | 2 + flang/lib/Lower/ConvertVariable.cpp | 76 +++++++++++++++++++- flang/test/Lower/CUDA/cuda-set-allocator.cuf | 21 ++++++ 3 files changed, 96 insertions(+), 3 deletions(-) create mode 100644 flang/test/Lower/CUDA/cuda-set-allocator.cuf diff --git a/flang/include/flang/Semantics/tools.h b/flang/include/flang/Semantics/tools.h index fb670528f3ce4..317b9357b4c1f 100644 --- a/flang/include/flang/Semantics/tools.h +++ b/flang/include/flang/Semantics/tools.h @@ -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()}) { diff --git a/flang/lib/Lower/ConvertVariable.cpp b/flang/lib/Lower/ConvertVariable.cpp index 44f534e7d569a..cb931895600df 100644 --- a/flang/lib/Lower/ConvertVariable.cpp +++ b/flang/lib/Lower/ConvertVariable.cpp @@ -771,9 +771,79 @@ static mlir::Value createNewLocal(Fortran::lower::AbstractConverter &converter, return builder.create(loc, ty, nm, symNm, lenParams, indices); - if (!cuf::isCUDADeviceContext(builder.getRegion())) - return builder.create(loc, ty, nm, symNm, dataAttr, - lenParams, indices); + if (!cuf::isCUDADeviceContext(builder.getRegion())) { + mlir::Value alloc = builder.create( + loc, ty, nm, symNm, dataAttr, lenParams, indices); + if (const auto *details{ + ultimateSymbol + .detailsIf()}) { + 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(ty); + + mlir::Type fieldTy; + llvm::SmallVector coordinates; + for (const auto &sym : components) { + if (Fortran::semantics::IsDeviceAllocatable(sym)) { + unsigned fieldIdx = recTy.getFieldIndex(sym.name().ToString()); + mlir::Type fieldTy; + std::vector coordinates; + + if (fieldIdx != std::numeric_limits::max()) { + // Field found in the base record type. + auto fieldName = recTy.getTypeList()[fieldIdx].first; + fieldTy = recTy.getTypeList()[fieldIdx].second; + mlir::Value fieldIndex = builder.create( + 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(component.second)) { + fieldIdx = childRecTy.getFieldIndex(sym.name().ToString()); + if (fieldIdx != std::numeric_limits::max()) { + mlir::Value parentFieldIndex = + builder.create( + 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( + 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( + loc, builder.getRefType(fieldTy), alloc, coordinates); + cuf::DataAttributeAttr dataAttr = + Fortran::lower::translateSymbolCUFDataAttribute( + builder.getContext(), sym); + builder.create(loc, comp, dataAttr); + } + } + } + } + return alloc; + } } // Let the builder do all the heavy lifting. diff --git a/flang/test/Lower/CUDA/cuda-set-allocator.cuf b/flang/test/Lower/CUDA/cuda-set-allocator.cuf new file mode 100644 index 0000000000000..bf74e012a639d --- /dev/null +++ b/flang/test/Lower/CUDA/cuda-set-allocator.cuf @@ -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>>,y:i32,z:!fir.box>>}> {bindc_name = "a", data_attr = #cuf.cuda, uniq_name = "_QMm1Fsub1Ea"} -> !fir.ref>>,y:i32,z:!fir.box>>}>> +! CHECK: %[[X:.*]] = fir.coordinate_of %[[DT]], x : (!fir.ref>>,y:i32,z:!fir.box>>}>>) -> !fir.ref>>> +! CHECK: cuf.set_allocator_idx %[[X]] : !fir.ref>>> {data_attr = #cuf.cuda} +! CHECK: %[[Z:.*]] = fir.coordinate_of %[[DT]], z : (!fir.ref>>,y:i32,z:!fir.box>>}>>) -> !fir.ref>>> +! CHECK: cuf.set_allocator_idx %[[Z]] : !fir.ref>>> {data_attr = #cuf.cuda} + +end module From 1ad5cc881765c8c3ce37c6cea04f344037835d8e Mon Sep 17 00:00:00 2001 From: Valentin Clement Date: Mon, 14 Jul 2025 18:18:29 -0700 Subject: [PATCH 2/2] clang-format --- flang/lib/Lower/ConvertVariable.cpp | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/flang/lib/Lower/ConvertVariable.cpp b/flang/lib/Lower/ConvertVariable.cpp index cb931895600df..4ccaa59f503a7 100644 --- a/flang/lib/Lower/ConvertVariable.cpp +++ b/flang/lib/Lower/ConvertVariable.cpp @@ -830,7 +830,8 @@ static mlir::Value createNewLocal(Fortran::lower::AbstractConverter &converter, } if (coordinates.empty()) - TODO(loc, "device resident component in complex derived-type hierarchy"); + TODO(loc, "device resident component in complex derived-type " + "hierarchy"); mlir::Value comp = builder.create( loc, builder.getRefType(fieldTy), alloc, coordinates);