caio/clang-p2996
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

[flang][cuda] Inline this_warp() calls (#146134)

Browse Source
This commit is contained in:
Valentin Clement (バレンタイン クレメン)
2025-06-27 14:12:17 -07:00
committed by GitHub
parent 68239b76f1
commit b2f504ff15
4 changed files with 75 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
flang/include/flang/Optimizer/Builder/IntrinsicCall.h
View File

@@ -443,6 +443,7 @@ struct IntrinsicLibrary {
fir::ExtendedValue genTranspose(mlir::Type,
llvm::ArrayRef<fir::ExtendedValue>);
mlir::Value genThisGrid(mlir::Type, llvm::ArrayRef<mlir::Value>);
mlir::Value genThisWarp(mlir::Type, llvm::ArrayRef<mlir::Value>);
void genThreadFence(llvm::ArrayRef<fir::ExtendedValue>);
void genThreadFenceBlock(llvm::ArrayRef<fir::ExtendedValue>);
void genThreadFenceSystem(llvm::ArrayRef<fir::ExtendedValue>);

40
flang/lib/Optimizer/Builder/IntrinsicCall.cpp
View File

@@ -933,6 +933,7 @@ static constexpr IntrinsicHandler handlers[]{
/*isElemental=*/false},
{"tand", &I::genTand},
{"this_grid", &I::genThisGrid, {}, /*isElemental=*/false},
{"this_warp", &I::genThisWarp, {}, /*isElemental=*/false},
{"threadfence", &I::genThreadFence, {}, /*isElemental=*/false},
{"threadfence_block", &I::genThreadFenceBlock, {}, /*isElemental=*/false},
{"threadfence_system", &I::genThreadFenceSystem, {}, /*isElemental=*/false},
@@ -8194,6 +8195,45 @@ mlir::Value IntrinsicLibrary::genThisGrid(mlir::Type resultType,
return res;
}
// THIS_WARP
mlir::Value IntrinsicLibrary::genThisWarp(mlir::Type resultType,
llvm::ArrayRef<mlir::Value> args) {
assert(args.size() == 0);
auto recTy = mlir::cast<fir::RecordType>(resultType);
assert(recTy && "RecordType expepected");
mlir::Value res = builder.create<fir::AllocaOp>(loc, resultType);
mlir::Type i32Ty = builder.getI32Type();
// coalesced_group%size = 32
mlir::Value size = builder.createIntegerConstant(loc, i32Ty, 32);
auto sizeFieldName = recTy.getTypeList()[1].first;
mlir::Type sizeFieldTy = recTy.getTypeList()[1].second;
mlir::Type fieldIndexType = fir::FieldType::get(resultType.getContext());
mlir::Value sizeFieldIndex = builder.create<fir::FieldIndexOp>(
loc, fieldIndexType, sizeFieldName, recTy,
/*typeParams=*/mlir::ValueRange{});
mlir::Value sizeCoord = builder.create<fir::CoordinateOp>(
loc, builder.getRefType(sizeFieldTy), res, sizeFieldIndex);
builder.create<fir::StoreOp>(loc, size, sizeCoord);
// coalesced_group%rank = threadIdx.x & 31 + 1
mlir::Value threadIdX = builder.create<mlir::NVVM::ThreadIdXOp>(loc, i32Ty);
mlir::Value mask = builder.createIntegerConstant(loc, i32Ty, 31);
mlir::Value one = builder.createIntegerConstant(loc, i32Ty, 1);
mlir::Value masked =
builder.create<mlir::arith::AndIOp>(loc, threadIdX, mask);
mlir::Value rank = builder.create<mlir::arith::AddIOp>(loc, masked, one);
auto rankFieldName = recTy.getTypeList()[2].first;
mlir::Type rankFieldTy = recTy.getTypeList()[2].second;
mlir::Value rankFieldIndex = builder.create<fir::FieldIndexOp>(
loc, fieldIndexType, rankFieldName, recTy,
/*typeParams=*/mlir::ValueRange{});
mlir::Value rankCoord = builder.create<fir::CoordinateOp>(
loc, builder.getRefType(rankFieldTy), res, rankFieldIndex);
builder.create<fir::StoreOp>(loc, rank, rankCoord);
return res;
}
// TRAILZ
mlir::Value IntrinsicLibrary::genTrailz(mlir::Type resultType,
llvm::ArrayRef<mlir::Value> args) {

13
flang/module/cooperative_groups.f90
View File

@@ -20,6 +20,12 @@ type :: grid_group
integer(4) :: rank
end type grid_group
type :: coalesced_group
type(c_devptr), private :: handle
integer(4) :: size
integer(4) :: rank
end type coalesced_group
interface
attributes(device) function this_grid()
import
@@ -27,4 +33,11 @@ interface
end function
end interface
interface this_warp
attributes(device) function this_warp()
import
type(coalesced_group) :: this_warp
end function
end interface
end module

21
flang/test/Lower/CUDA/cuda-cooperative.cuf
View File

@@ -50,3 +50,24 @@ end subroutine
! CHECK: fir.store %[[SIZE]] to %[[COORD_SIZE]] : !fir.ref<i32>
! CHECK: %[[COORD_RANK:.*]] = fir.coordinate_of %[[RES]], rank : (!fir.ref<!fir.type<_QMcooperative_groupsTgrid_group{_QMcooperative_groupsTgrid_group.handle:!fir.type<_QM__fortran_builtinsT__builtin_c_devptr{cptr:!fir.type<_QM__fortran_builtinsT__builtin_c_ptr{__address:i64}>}>,size:i32,rank:i32}>>) -> !fir.ref<i32>
! CHECK: fir.store %[[RANK]] to %[[COORD_RANK]] : !fir.ref<i32>
attributes(grid_global) subroutine w1()
use cooperative_groups
type(coalesced_group) :: gg
gg = this_warp()
end subroutine
! CHECK: %[[WARPSIZE:.*]] = fir.alloca i32 {bindc_name = "__builtin_warpsize", uniq_name = "_QM__fortran_builtinsEC__builtin_warpsize"}
! CHECK: %[[WARPSIZE_DECL:.*]]:2 = hlfir.declare %[[WARPSIZE]] {uniq_name = "_QM__fortran_builtinsEC__builtin_warpsize"} : (!fir.ref<i32>) -> (!fir.ref<i32>, !fir.ref<i32>)
! CHECK: %[[COALESCED_GROUP:.*]] = fir.alloca !fir.type<_QMcooperative_groupsTcoalesced_group{_QMcooperative_groupsTcoalesced_group.handle:!fir.type<_QM__fortran_builtinsT__builtin_c_devptr{cptr:!fir.type<_QM__fortran_builtinsT__builtin_c_ptr{__address:i64}>}>,size:i32,rank:i32}>
! CHECK: %[[C32:.*]] = arith.constant 32 : i32
! CHECK: %[[SIZE_COORD:.*]] = fir.coordinate_of %[[COALESCED_GROUP]], size : (!fir.ref<!fir.type<_QMcooperative_groupsTcoalesced_group{_QMcooperative_groupsTcoalesced_group.handle:!fir.type<_QM__fortran_builtinsT__builtin_c_devptr{cptr:!fir.type<_QM__fortran_builtinsT__builtin_c_ptr{__address:i64}>}>,size:i32,rank:i32}>>) -> !fir.ref<i32>
! CHECK: fir.store %[[C32]] to %[[SIZE_COORD]] : !fir.ref<i32>
! CHECK: %[[THREAD_ID:.*]] = nvvm.read.ptx.sreg.tid.x : i32
! CHECK: %[[C31:.*]] = arith.constant 31 : i32
! CHECK: %[[C1:.*]] = arith.constant 1 : i32
! CHECK: %[[AND:.*]] = arith.andi %[[THREAD_ID]], %[[C31]] : i32
! CHECK: %[[RANK:.*]] = arith.addi %[[AND]], %[[C1]] : i32
! CHECK: %[[RANK_COORD:.*]] = fir.coordinate_of %{{.*}}, rank : (!fir.ref<!fir.type<_QMcooperative_groupsTcoalesced_group{_QMcooperative_groupsTcoalesced_group.handle:!fir.type<_QM__fortran_builtinsT__builtin_c_devptr{cptr:!fir.type<_QM__fortran_builtinsT__builtin_c_ptr{__address:i64}>}>,size:i32,rank:i32}>>) -> !fir.ref<i32>
! CHECK: fir.store %[[RANK]] to %[[RANK_COORD]] : !fir.ref<i32>