280 lines
8.2 KiB
C++
280 lines
8.2 KiB
C++
|
//===-- AMDGPULowerKernelAttributes.cpp ------------------------------------------===//
|
||
|
|
//
|
||
|
|
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
|
||
|
|
// See https://llvm.org/LICENSE.txt for license information.
|
||
|
|
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
||
|
|
//
|
||
|
|
//===----------------------------------------------------------------------===//
|
||
|
|
//
|
||
|
|
/// \file This pass does attempts to make use of reqd_work_group_size metadata
|
||
|
|
/// to eliminate loads from the dispatch packet and to constant fold OpenCL
|
||
|
|
/// get_local_size-like functions.
|
||
|
|
//
|
||
|
|
//===----------------------------------------------------------------------===//
|
||
|
|
|
||
|
|
#include "AMDGPU.h"
|
||
|
|
#include "llvm/Analysis/ValueTracking.h"
|
||
|
|
#include "llvm/CodeGen/Passes.h"
|
||
|
|
#include "llvm/CodeGen/TargetPassConfig.h"
|
||
|
|
#include "llvm/IR/Constants.h"
|
||
|
|
#include "llvm/IR/Function.h"
|
||
|
|
#include "llvm/IR/InstIterator.h"
|
||
|
|
#include "llvm/IR/Instructions.h"
|
||
|
|
#include "llvm/IR/IntrinsicsAMDGPU.h"
|
||
|
|
#include "llvm/IR/PatternMatch.h"
|
||
|
|
#include "llvm/Pass.h"
|
||
|
|
|
||
|
|
#define DEBUG_TYPE "amdgpu-lower-kernel-attributes"
|
||
|
|
|
||
|
|
using namespace llvm;
|
||
|
|
|
||
|
|
namespace {
|
||
|
|
|
||
|
|
// Field offsets in hsa_kernel_dispatch_packet_t.
|
||
|
|
enum DispatchPackedOffsets {
|
||
|
|
WORKGROUP_SIZE_X = 4,
|
||
|
|
WORKGROUP_SIZE_Y = 6,
|
||
|
|
WORKGROUP_SIZE_Z = 8,
|
||
|
|
|
||
|
|
GRID_SIZE_X = 12,
|
||
|
|
GRID_SIZE_Y = 16,
|
||
|
|
GRID_SIZE_Z = 20
|
||
|
|
};
|
||
|
|
|
||
|
|
class AMDGPULowerKernelAttributes : public ModulePass {
|
||
|
|
public:
|
||
|
|
static char ID;
|
||
|
|
|
||
|
|
AMDGPULowerKernelAttributes() : ModulePass(ID) {}
|
||
|
|
|
||
|
|
bool runOnModule(Module &M) override;
|
||
|
|
|
||
|
|
StringRef getPassName() const override {
|
||
|
|
return "AMDGPU Kernel Attributes";
|
||
|
|
}
|
||
|
|
|
||
|
|
void getAnalysisUsage(AnalysisUsage &AU) const override {
|
||
|
|
AU.setPreservesAll();
|
||
|
|
}
|
||
|
|
};
|
||
|
|
|
||
|
|
} // end anonymous namespace
|
||
|
|
|
||
|
|
static bool processUse(CallInst *CI) {
|
||
|
|
Function *F = CI->getParent()->getParent();
|
||
|
|
|
||
|
|
auto MD = F->getMetadata("reqd_work_group_size");
|
||
|
|
const bool HasReqdWorkGroupSize = MD && MD->getNumOperands() == 3;
|
||
|
|
|
||
|
|
const bool HasUniformWorkGroupSize =
|
||
|
|
F->getFnAttribute("uniform-work-group-size").getValueAsString() == "true";
|
||
|
|
|
||
|
|
if (!HasReqdWorkGroupSize && !HasUniformWorkGroupSize)
|
||
|
|
return false;
|
||
|
|
|
||
|
|
Value *WorkGroupSizeX = nullptr;
|
||
|
|
Value *WorkGroupSizeY = nullptr;
|
||
|
|
Value *WorkGroupSizeZ = nullptr;
|
||
|
|
|
||
|
|
Value *GridSizeX = nullptr;
|
||
|
|
Value *GridSizeY = nullptr;
|
||
|
|
Value *GridSizeZ = nullptr;
|
||
|
|
|
||
|
|
const DataLayout &DL = F->getParent()->getDataLayout();
|
||
|
|
|
||
|
|
// We expect to see several GEP users, casted to the appropriate type and
|
||
|
|
// loaded.
|
||
|
|
for (User *U : CI->users()) {
|
||
|
|
if (!U->hasOneUse())
|
||
|
|
continue;
|
||
|
|
|
||
|
|
int64_t Offset = 0;
|
||
|
|
if (GetPointerBaseWithConstantOffset(U, Offset, DL) != CI)
|
||
|
|
continue;
|
||
|
|
|
||
|
|
auto *BCI = dyn_cast<BitCastInst>(*U->user_begin());
|
||
|
|
if (!BCI || !BCI->hasOneUse())
|
||
|
|
continue;
|
||
|
|
|
||
|
|
auto *Load = dyn_cast<LoadInst>(*BCI->user_begin());
|
||
|
|
if (!Load || !Load->isSimple())
|
||
|
|
continue;
|
||
|
|
|
||
|
|
unsigned LoadSize = DL.getTypeStoreSize(Load->getType());
|
||
|
|
|
||
|
|
// TODO: Handle merged loads.
|
||
|
|
switch (Offset) {
|
||
|
|
case WORKGROUP_SIZE_X:
|
||
|
|
if (LoadSize == 2)
|
||
|
|
WorkGroupSizeX = Load;
|
||
|
|
break;
|
||
|
|
case WORKGROUP_SIZE_Y:
|
||
|
|
if (LoadSize == 2)
|
||
|
|
WorkGroupSizeY = Load;
|
||
|
|
break;
|
||
|
|
case WORKGROUP_SIZE_Z:
|
||
|
|
if (LoadSize == 2)
|
||
|
|
WorkGroupSizeZ = Load;
|
||
|
|
break;
|
||
|
|
case GRID_SIZE_X:
|
||
|
|
if (LoadSize == 4)
|
||
|
|
GridSizeX = Load;
|
||
|
|
break;
|
||
|
|
case GRID_SIZE_Y:
|
||
|
|
if (LoadSize == 4)
|
||
|
|
GridSizeY = Load;
|
||
|
|
break;
|
||
|
|
case GRID_SIZE_Z:
|
||
|
|
if (LoadSize == 4)
|
||
|
|
GridSizeZ = Load;
|
||
|
|
break;
|
||
|
|
default:
|
||
|
|
break;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// Pattern match the code used to handle partial workgroup dispatches in the
|
||
|
|
// library implementation of get_local_size, so the entire function can be
|
||
|
|
// constant folded with a known group size.
|
||
|
|
//
|
||
|
|
// uint r = grid_size - group_id * group_size;
|
||
|
|
// get_local_size = (r < group_size) ? r : group_size;
|
||
|
|
//
|
||
|
|
// If we have uniform-work-group-size (which is the default in OpenCL 1.2),
|
||
|
|
// the grid_size is required to be a multiple of group_size). In this case:
|
||
|
|
//
|
||
|
|
// grid_size - (group_id * group_size) < group_size
|
||
|
|
// ->
|
||
|
|
// grid_size < group_size + (group_id * group_size)
|
||
|
|
//
|
||
|
|
// (grid_size / group_size) < 1 + group_id
|
||
|
|
//
|
||
|
|
// grid_size / group_size is at least 1, so we can conclude the select
|
||
|
|
// condition is false (except for group_id == 0, where the select result is
|
||
|
|
// the same).
|
||
|
|
|
||
|
|
bool MadeChange = false;
|
||
|
|
Value *WorkGroupSizes[3] = { WorkGroupSizeX, WorkGroupSizeY, WorkGroupSizeZ };
|
||
|
|
Value *GridSizes[3] = { GridSizeX, GridSizeY, GridSizeZ };
|
||
|
|
|
||
|
|
for (int I = 0; HasUniformWorkGroupSize && I < 3; ++I) {
|
||
|
|
Value *GroupSize = WorkGroupSizes[I];
|
||
|
|
Value *GridSize = GridSizes[I];
|
||
|
|
if (!GroupSize || !GridSize)
|
||
|
|
continue;
|
||
|
|
|
||
|
|
for (User *U : GroupSize->users()) {
|
||
|
|
auto *ZextGroupSize = dyn_cast<ZExtInst>(U);
|
||
|
|
if (!ZextGroupSize)
|
||
|
|
continue;
|
||
|
|
|
||
|
|
for (User *ZextUser : ZextGroupSize->users()) {
|
||
|
|
auto *SI = dyn_cast<SelectInst>(ZextUser);
|
||
|
|
if (!SI)
|
||
|
|
continue;
|
||
|
|
|
||
|
|
using namespace llvm::PatternMatch;
|
||
|
|
auto GroupIDIntrin = I == 0 ?
|
||
|
|
m_Intrinsic<Intrinsic::amdgcn_workgroup_id_x>() :
|
||
|
|
(I == 1 ? m_Intrinsic<Intrinsic::amdgcn_workgroup_id_y>() :
|
||
|
|
m_Intrinsic<Intrinsic::amdgcn_workgroup_id_z>());
|
||
|
|
|
||
|
|
auto SubExpr = m_Sub(m_Specific(GridSize),
|
||
|
|
m_Mul(GroupIDIntrin, m_Specific(ZextGroupSize)));
|
||
|
|
|
||
|
|
ICmpInst::Predicate Pred;
|
||
|
|
if (match(SI,
|
||
|
|
m_Select(m_ICmp(Pred, SubExpr, m_Specific(ZextGroupSize)),
|
||
|
|
SubExpr,
|
||
|
|
m_Specific(ZextGroupSize))) &&
|
||
|
|
Pred == ICmpInst::ICMP_ULT) {
|
||
|
|
if (HasReqdWorkGroupSize) {
|
||
|
|
ConstantInt *KnownSize
|
||
|
|
= mdconst::extract<ConstantInt>(MD->getOperand(I));
|
||
|
|
SI->replaceAllUsesWith(ConstantExpr::getIntegerCast(KnownSize,
|
||
|
|
SI->getType(),
|
||
|
|
false));
|
||
|
|
} else {
|
||
|
|
SI->replaceAllUsesWith(ZextGroupSize);
|
||
|
|
}
|
||
|
|
|
||
|
|
MadeChange = true;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
if (!HasReqdWorkGroupSize)
|
||
|
|
return MadeChange;
|
||
|
|
|
||
|
|
// Eliminate any other loads we can from the dispatch packet.
|
||
|
|
for (int I = 0; I < 3; ++I) {
|
||
|
|
Value *GroupSize = WorkGroupSizes[I];
|
||
|
|
if (!GroupSize)
|
||
|
|
continue;
|
||
|
|
|
||
|
|
ConstantInt *KnownSize = mdconst::extract<ConstantInt>(MD->getOperand(I));
|
||
|
|
GroupSize->replaceAllUsesWith(
|
||
|
|
ConstantExpr::getIntegerCast(KnownSize,
|
||
|
|
GroupSize->getType(),
|
||
|
|
false));
|
||
|
|
MadeChange = true;
|
||
|
|
}
|
||
|
|
|
||
|
|
return MadeChange;
|
||
|
|
}
|
||
|
|
|
||
|
|
// TODO: Move makeLIDRangeMetadata usage into here. Seem to not get
|
||
|
|
// TargetPassConfig for subtarget.
|
||
|
|
bool AMDGPULowerKernelAttributes::runOnModule(Module &M) {
|
||
|
|
StringRef DispatchPtrName
|
||
|
|
= Intrinsic::getName(Intrinsic::amdgcn_dispatch_ptr);
|
||
|
|
|
||
|
|
Function *DispatchPtr = M.getFunction(DispatchPtrName);
|
||
|
|
if (!DispatchPtr) // Dispatch ptr not used.
|
||
|
|
return false;
|
||
|
|
|
||
|
|
bool MadeChange = false;
|
||
|
|
|
||
|
|
SmallPtrSet<Instruction *, 4> HandledUses;
|
||
|
|
for (auto *U : DispatchPtr->users()) {
|
||
|
|
CallInst *CI = cast<CallInst>(U);
|
||
|
|
if (HandledUses.insert(CI).second) {
|
||
|
|
if (processUse(CI))
|
||
|
|
MadeChange = true;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
return MadeChange;
|
||
|
|
}
|
||
|
|
|
||
|
|
INITIALIZE_PASS_BEGIN(AMDGPULowerKernelAttributes, DEBUG_TYPE,
|
||
|
|
"AMDGPU IR optimizations", false, false)
|
||
|
|
INITIALIZE_PASS_END(AMDGPULowerKernelAttributes, DEBUG_TYPE, "AMDGPU IR optimizations",
|
||
|
|
false, false)
|
||
|
|
|
||
|
|
char AMDGPULowerKernelAttributes::ID = 0;
|
||
|
|
|
||
|
|
ModulePass *llvm::createAMDGPULowerKernelAttributesPass() {
|
||
|
|
return new AMDGPULowerKernelAttributes();
|
||
|
|
}
|
||
|
|
|
||
|
|
PreservedAnalyses
|
||
|
|
AMDGPULowerKernelAttributesPass::run(Function &F, FunctionAnalysisManager &AM) {
|
||
|
|
StringRef DispatchPtrName =
|
||
|
|
Intrinsic::getName(Intrinsic::amdgcn_dispatch_ptr);
|
||
|
|
|
||
|
|
Function *DispatchPtr = F.getParent()->getFunction(DispatchPtrName);
|
||
|
|
if (!DispatchPtr) // Dispatch ptr not used.
|
||
|
|
return PreservedAnalyses::all();
|
||
|
|
|
||
|
|
for (Instruction &I : instructions(F)) {
|
||
|
|
if (CallInst *CI = dyn_cast<CallInst>(&I)) {
|
||
|
|
if (CI->getCalledFunction() == DispatchPtr)
|
||
|
|
processUse(CI);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
return PreservedAnalyses::all();
|
||
|
|
}
|