Skip to content

Commit

Permalink
[flang][cuda] Lower attribute for module variables (llvm#81226)
Browse files Browse the repository at this point in the history
Propagate the CUDA attribute to fir.global operation for simple module
variables.
  • Loading branch information
clementval committed Feb 9, 2024
1 parent 0329c1b commit 314ef96
Show file tree
Hide file tree
Showing 5 changed files with 61 additions and 38 deletions.
6 changes: 4 additions & 2 deletions flang/include/flang/Optimizer/Builder/FIRBuilder.h
Original file line number Diff line number Diff line change
Expand Up @@ -230,12 +230,14 @@ class FirOpBuilder : public mlir::OpBuilder, public mlir::OpBuilder::Listener {
llvm::StringRef name,
mlir::StringAttr linkage = {},
mlir::Attribute value = {}, bool isConst = false,
bool isTarget = false);
bool isTarget = false,
fir::CUDAAttributeAttr cudaAttr = {});

fir::GlobalOp createGlobal(mlir::Location loc, mlir::Type type,
llvm::StringRef name, bool isConst, bool isTarget,
std::function<void(FirOpBuilder &)> bodyBuilder,
mlir::StringAttr linkage = {});
mlir::StringAttr linkage = {},
fir::CUDAAttributeAttr cudaAttr = {});

/// Create a global constant (read-only) value.
fir::GlobalOp createGlobalConstant(mlir::Location loc, mlir::Type type,
Expand Down
3 changes: 2 additions & 1 deletion flang/include/flang/Optimizer/Dialect/FIROps.td
Original file line number Diff line number Diff line change
Expand Up @@ -2737,7 +2737,8 @@ def fir_GlobalOp : fir_Op<"global", [IsolatedFromAbove, Symbol]> {
OptionalAttr<AnyAttr>:$initVal,
OptionalAttr<UnitAttr>:$constant,
OptionalAttr<UnitAttr>:$target,
OptionalAttr<StrAttr>:$linkName
OptionalAttr<StrAttr>:$linkName,
OptionalAttr<fir_CUDAAttributeAttr>:$cuda_attr
);

let regions = (region AtMostRegion<1>:$region);
Expand Down
16 changes: 11 additions & 5 deletions flang/lib/Lower/ConvertVariable.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -138,7 +138,8 @@ static bool isConstant(const Fortran::semantics::Symbol &sym) {
static fir::GlobalOp defineGlobal(Fortran::lower::AbstractConverter &converter,
const Fortran::lower::pft::Variable &var,
llvm::StringRef globalName,
mlir::StringAttr linkage);
mlir::StringAttr linkage,
fir::CUDAAttributeAttr cudaAttr = {});

static mlir::Location genLocation(Fortran::lower::AbstractConverter &converter,
const Fortran::semantics::Symbol &sym) {
Expand Down Expand Up @@ -462,7 +463,8 @@ void Fortran::lower::createGlobalInitialization(
static fir::GlobalOp defineGlobal(Fortran::lower::AbstractConverter &converter,
const Fortran::lower::pft::Variable &var,
llvm::StringRef globalName,
mlir::StringAttr linkage) {
mlir::StringAttr linkage,
fir::CUDAAttributeAttr cudaAttr) {
fir::FirOpBuilder &builder = converter.getFirOpBuilder();
const Fortran::semantics::Symbol &sym = var.getSymbol();
mlir::Location loc = genLocation(converter, sym);
Expand Down Expand Up @@ -500,8 +502,9 @@ static fir::GlobalOp defineGlobal(Fortran::lower::AbstractConverter &converter,
}
}
if (!global)
global = builder.createGlobal(loc, symTy, globalName, linkage,
mlir::Attribute{}, isConst, var.isTarget());
global =
builder.createGlobal(loc, symTy, globalName, linkage, mlir::Attribute{},
isConst, var.isTarget(), cudaAttr);
if (Fortran::semantics::IsAllocatableOrPointer(sym) &&
!Fortran::semantics::IsProcedure(sym)) {
const auto *details =
Expand Down Expand Up @@ -2219,7 +2222,10 @@ void Fortran::lower::defineModuleVariable(
// Do nothing. Mapping will be done on user side.
} else {
std::string globalName = converter.mangleName(sym);
defineGlobal(converter, var, globalName, linkage);
fir::CUDAAttributeAttr cudaAttr =
Fortran::lower::translateSymbolCUDAAttribute(
converter.getFirOpBuilder().getContext(), sym);
defineGlobal(converter, var, globalName, linkage, cudaAttr);
}
}

Expand Down
23 changes: 14 additions & 9 deletions flang/lib/Optimizer/Builder/FIRBuilder.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -271,27 +271,32 @@ mlir::Value fir::FirOpBuilder::createHeapTemporary(

/// Create a global variable in the (read-only) data section. A global variable
/// must have a unique name to identify and reference it.
fir::GlobalOp fir::FirOpBuilder::createGlobal(mlir::Location loc,
mlir::Type type,
llvm::StringRef name,
mlir::StringAttr linkage,
mlir::Attribute value,
bool isConst, bool isTarget) {
fir::GlobalOp fir::FirOpBuilder::createGlobal(
mlir::Location loc, mlir::Type type, llvm::StringRef name,
mlir::StringAttr linkage, mlir::Attribute value, bool isConst,
bool isTarget, fir::CUDAAttributeAttr cudaAttr) {
auto module = getModule();
auto insertPt = saveInsertionPoint();
if (auto glob = module.lookupSymbol<fir::GlobalOp>(name))
return glob;
setInsertionPoint(module.getBody(), module.getBody()->end());
auto glob =
create<fir::GlobalOp>(loc, name, isConst, isTarget, type, value, linkage);
llvm::SmallVector<mlir::NamedAttribute> attrs;
if (cudaAttr) {
auto globalOpName = mlir::OperationName(fir::GlobalOp::getOperationName(),
module.getContext());
attrs.push_back(mlir::NamedAttribute(
fir::GlobalOp::getCudaAttrAttrName(globalOpName), cudaAttr));
}
auto glob = create<fir::GlobalOp>(loc, name, isConst, isTarget, type, value,
linkage, attrs);
restoreInsertionPoint(insertPt);
return glob;
}

fir::GlobalOp fir::FirOpBuilder::createGlobal(
mlir::Location loc, mlir::Type type, llvm::StringRef name, bool isConst,
bool isTarget, std::function<void(FirOpBuilder &)> bodyBuilder,
mlir::StringAttr linkage) {
mlir::StringAttr linkage, fir::CUDAAttributeAttr cudaAttr) {
auto module = getModule();
auto insertPt = saveInsertionPoint();
if (auto glob = module.lookupSymbol<fir::GlobalOp>(name))
Expand Down
51 changes: 30 additions & 21 deletions flang/test/Lower/CUDA/cuda-data-attribute.cuf
Original file line number Diff line number Diff line change
Expand Up @@ -3,53 +3,62 @@

! Test lowering of CUDA attribute on variables.

module cuda_var
real, constant :: mod_a_rc
! CHECK: fir.global @_QMcuda_varEmod_a_rc {cuda_attr = #fir.cuda<constant>} : f32
real, device :: mod_b_ra
! CHECK: fir.global @_QMcuda_varEmod_b_ra {cuda_attr = #fir.cuda<device>} : f32
real, allocatable, managed :: mod_c_rm
! CHECK: fir.global @_QMcuda_varEmod_c_rm {cuda_attr = #fir.cuda<managed>} : !fir.box<!fir.heap<f32>>
real, allocatable, pinned :: mod_d_rp
! CHECK: fir.global @_QMcuda_varEmod_d_rp {cuda_attr = #fir.cuda<pinned>} : !fir.box<!fir.heap<f32>>

contains

subroutine local_var_attrs
real, constant :: rc
real, device :: rd
real, allocatable, managed :: rm
real, allocatable, pinned :: rp
end subroutine

! CHECK-LABEL: func.func @_QPlocal_var_attrs()
! CHECK: %{{.*}}:2 = hlfir.declare %{{.*}} {cuda_attr = #fir.cuda<constant>, uniq_name = "_QFlocal_var_attrsErc"} : (!fir.ref<f32>) -> (!fir.ref<f32>, !fir.ref<f32>)
! CHECK: %{{.*}}:2 = hlfir.declare %{{.*}} {cuda_attr = #fir.cuda<device>, uniq_name = "_QFlocal_var_attrsErd"} : (!fir.ref<f32>) -> (!fir.ref<f32>, !fir.ref<f32>)
! CHECK: %{{.*}}:2 = hlfir.declare %{{.*}} {cuda_attr = #fir.cuda<managed>, fortran_attrs = #fir.var_attrs<allocatable>, uniq_name = "_QFlocal_var_attrsErm"} : (!fir.ref<!fir.box<!fir.heap<f32>>>) -> (!fir.ref<!fir.box<!fir.heap<f32>>>, !fir.ref<!fir.box<!fir.heap<f32>>>)
! CHECK: %{{.*}}:2 = hlfir.declare %{{.*}} {cuda_attr = #fir.cuda<pinned>, fortran_attrs = #fir.var_attrs<allocatable>, uniq_name = "_QFlocal_var_attrsErp"} : (!fir.ref<!fir.box<!fir.heap<f32>>>) -> (!fir.ref<!fir.box<!fir.heap<f32>>>, !fir.ref<!fir.box<!fir.heap<f32>>>)
! CHECK-LABEL: func.func @_QMcuda_varPlocal_var_attrs()
! CHECK: %{{.*}}:2 = hlfir.declare %{{.*}} {cuda_attr = #fir.cuda<constant>, uniq_name = "_QMcuda_varFlocal_var_attrsErc"} : (!fir.ref<f32>) -> (!fir.ref<f32>, !fir.ref<f32>)
! CHECK: %{{.*}}:2 = hlfir.declare %{{.*}} {cuda_attr = #fir.cuda<device>, uniq_name = "_QMcuda_varFlocal_var_attrsErd"} : (!fir.ref<f32>) -> (!fir.ref<f32>, !fir.ref<f32>)
! CHECK: %{{.*}}:2 = hlfir.declare %{{.*}} {cuda_attr = #fir.cuda<managed>, fortran_attrs = #fir.var_attrs<allocatable>, uniq_name = "_QMcuda_varFlocal_var_attrsErm"} : (!fir.ref<!fir.box<!fir.heap<f32>>>) -> (!fir.ref<!fir.box<!fir.heap<f32>>>, !fir.ref<!fir.box<!fir.heap<f32>>>)
! CHECK: %{{.*}}:2 = hlfir.declare %{{.*}} {cuda_attr = #fir.cuda<pinned>, fortran_attrs = #fir.var_attrs<allocatable>, uniq_name = "_QMcuda_varFlocal_var_attrsErp"} : (!fir.ref<!fir.box<!fir.heap<f32>>>) -> (!fir.ref<!fir.box<!fir.heap<f32>>>, !fir.ref<!fir.box<!fir.heap<f32>>>)

! FIR: %{{.*}} = fir.declare %{{.*}} {cuda_attr = #fir.cuda<constant>, uniq_name = "_QFlocal_var_attrsErc"} : (!fir.ref<f32>) -> !fir.ref<f32>
! FIR: %{{.*}} = fir.declare %{{.*}} {cuda_attr = #fir.cuda<device>, uniq_name = "_QFlocal_var_attrsErd"} : (!fir.ref<f32>) -> !fir.ref<f32>
! FIR: %{{.*}} = fir.declare %{{.*}} {cuda_attr = #fir.cuda<managed>, fortran_attrs = #fir.var_attrs<allocatable>, uniq_name = "_QFlocal_var_attrsErm"} : (!fir.ref<!fir.box<!fir.heap<f32>>>) -> !fir.ref<!fir.box<!fir.heap<f32>>>
! FIR: %{{.*}} = fir.declare %{{.*}} {cuda_attr = #fir.cuda<pinned>, fortran_attrs = #fir.var_attrs<allocatable>, uniq_name = "_QFlocal_var_attrsErp"} : (!fir.ref<!fir.box<!fir.heap<f32>>>) -> !fir.ref<!fir.box<!fir.heap<f32>>>
! FIR: %{{.*}} = fir.declare %{{.*}} {cuda_attr = #fir.cuda<constant>, uniq_name = "_QMcuda_varFlocal_var_attrsErc"} : (!fir.ref<f32>) -> !fir.ref<f32>
! FIR: %{{.*}} = fir.declare %{{.*}} {cuda_attr = #fir.cuda<device>, uniq_name = "_QMcuda_varFlocal_var_attrsErd"} : (!fir.ref<f32>) -> !fir.ref<f32>
! FIR: %{{.*}} = fir.declare %{{.*}} {cuda_attr = #fir.cuda<managed>, fortran_attrs = #fir.var_attrs<allocatable>, uniq_name = "_QMcuda_varFlocal_var_attrsErm"} : (!fir.ref<!fir.box<!fir.heap<f32>>>) -> !fir.ref<!fir.box<!fir.heap<f32>>>
! FIR: %{{.*}} = fir.declare %{{.*}} {cuda_attr = #fir.cuda<pinned>, fortran_attrs = #fir.var_attrs<allocatable>, uniq_name = "_QMcuda_varFlocal_var_attrsErp"} : (!fir.ref<!fir.box<!fir.heap<f32>>>) -> !fir.ref<!fir.box<!fir.heap<f32>>>

subroutine dummy_arg_constant(dc)
real, constant :: dc
end subroutine
! CHECK-LABEL: func.func @_QPdummy_arg_constant(
! CHECK-LABEL: func.func @_QMcuda_varPdummy_arg_constant(
! CHECK-SAME: %[[ARG0:.*]]: !fir.ref<f32> {fir.bindc_name = "dc", fir.cuda_attr = #fir.cuda<constant>}
! CHECK: %{{.*}}:2 = hlfir.declare %[[ARG0]] {cuda_attr = #fir.cuda<constant>, uniq_name = "_QFdummy_arg_constantEdc"} : (!fir.ref<f32>) -> (!fir.ref<f32>, !fir.ref<f32>)
! CHECK: %{{.*}}:2 = hlfir.declare %[[ARG0]] {cuda_attr = #fir.cuda<constant>, uniq_name = "_QMcuda_varFdummy_arg_constantEdc"} : (!fir.ref<f32>) -> (!fir.ref<f32>, !fir.ref<f32>)

subroutine dummy_arg_device(dd)
real, device :: dd
end subroutine
! CHECK-LABEL: func.func @_QPdummy_arg_device(
! CHECK-LABEL: func.func @_QMcuda_varPdummy_arg_device(
! CHECK-SAME: %[[ARG0:.*]]: !fir.ref<f32> {fir.bindc_name = "dd", fir.cuda_attr = #fir.cuda<device>}) {
! CHECK: %{{.*}}:2 = hlfir.declare %[[ARG0]] {cuda_attr = #fir.cuda<device>, uniq_name = "_QFdummy_arg_deviceEdd"} : (!fir.ref<f32>) -> (!fir.ref<f32>, !fir.ref<f32>)
! CHECK: %{{.*}}:2 = hlfir.declare %[[ARG0]] {cuda_attr = #fir.cuda<device>, uniq_name = "_QMcuda_varFdummy_arg_deviceEdd"} : (!fir.ref<f32>) -> (!fir.ref<f32>, !fir.ref<f32>)

subroutine dummy_arg_managed(dm)
real, allocatable, managed :: dm
end subroutine
! CHECK-LABEL: func.func @_QPdummy_arg_managed(
! CHECK-LABEL: func.func @_QMcuda_varPdummy_arg_managed(
! CHECK-SAME: %[[ARG0:.*]]: !fir.ref<!fir.box<!fir.heap<f32>>> {fir.bindc_name = "dm", fir.cuda_attr = #fir.cuda<managed>}) {
! CHECK: %{{.*}}:2 = hlfir.declare %[[ARG0]] {cuda_attr = #fir.cuda<managed>, fortran_attrs = #fir.var_attrs<allocatable>, uniq_name = "_QFdummy_arg_managedEdm"} : (!fir.ref<!fir.box<!fir.heap<f32>>>) -> (!fir.ref<!fir.box<!fir.heap<f32>>>, !fir.ref<!fir.box<!fir.heap<f32>>>)
! CHECK: %{{.*}}:2 = hlfir.declare %[[ARG0]] {cuda_attr = #fir.cuda<managed>, fortran_attrs = #fir.var_attrs<allocatable>, uniq_name = "_QMcuda_varFdummy_arg_managedEdm"} : (!fir.ref<!fir.box<!fir.heap<f32>>>) -> (!fir.ref<!fir.box<!fir.heap<f32>>>, !fir.ref<!fir.box<!fir.heap<f32>>>)

subroutine dummy_arg_pinned(dp)
real, allocatable, pinned :: dp
end subroutine
! CHECK-LABEL: func.func @_QPdummy_arg_pinned(
! CHECK-LABEL: func.func @_QMcuda_varPdummy_arg_pinned(
! CHECK-SAME: %[[ARG0:.*]]: !fir.ref<!fir.box<!fir.heap<f32>>> {fir.bindc_name = "dp", fir.cuda_attr = #fir.cuda<pinned>}) {
! CHECK: %{{.*}}:2 = hlfir.declare %[[ARG0]] {cuda_attr = #fir.cuda<pinned>, fortran_attrs = #fir.var_attrs<allocatable>, uniq_name = "_QFdummy_arg_pinnedEdp"} : (!fir.ref<!fir.box<!fir.heap<f32>>>) -> (!fir.ref<!fir.box<!fir.heap<f32>>>, !fir.ref<!fir.box<!fir.heap<f32>>>)




! CHECK: %{{.*}}:2 = hlfir.declare %[[ARG0]] {cuda_attr = #fir.cuda<pinned>, fortran_attrs = #fir.var_attrs<allocatable>, uniq_name = "_QMcuda_varFdummy_arg_pinnedEdp"} : (!fir.ref<!fir.box<!fir.heap<f32>>>) -> (!fir.ref<!fir.box<!fir.heap<f32>>>, !fir.ref<!fir.box<!fir.heap<f32>>>)

end module

0 comments on commit 314ef96

Please sign in to comment.