blob: 11ba5c91dae9bb88dac54174667fccd6c3b4daa8 [file] [log] [blame]
//===-- AMDGPULowerModuleLDSPass.cpp ------------------------------*- C++ -*-=//
//
// 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
//
//===----------------------------------------------------------------------===//
//
// This pass eliminates local data store, LDS, uses from non-kernel functions.
// LDS is contiguous memory allocated per kernel execution.
//
// Background.
//
// The programming model is global variables, or equivalently function local
// static variables, accessible from kernels or other functions. For uses from
// kernels this is straightforward - assign an integer to the kernel for the
// memory required by all the variables combined, allocate them within that.
// For uses from functions there are performance tradeoffs to choose between.
//
// This model means the GPU runtime can specify the amount of memory allocated.
// If this is more than the kernel assumed, the excess can be made available
// using a language specific feature, which IR represents as a variable with
// no initializer. This feature is not yet implemented for non-kernel functions.
// This lowering could be extended to handle that use case, but would probably
// require closer integration with promoteAllocaToLDS.
//
// Consequences of this GPU feature:
// - memory is limited and exceeding it halts compilation
// - a global accessed by one kernel exists independent of other kernels
// - a global exists independent of simultaneous execution of the same kernel
// - the address of the global may be different from different kernels as they
// do not alias, which permits only allocating variables they use
// - if the address is allowed to differ, functions need help to find it
//
// Uses from kernels are implemented here by grouping them in a per-kernel
// struct instance. This duplicates the variables, accurately modelling their
// aliasing properties relative to a single global representation. It also
// permits control over alignment via padding.
//
// Uses from functions are more complicated and the primary purpose of this
// IR pass. Several different lowering are chosen between to meet requirements
// to avoid allocating any LDS where it is not necessary, as that impacts
// occupancy and may fail the compilation, while not imposing overhead on a
// feature whose primary advantage over global memory is performance. The basic
// design goal is to avoid one kernel imposing overhead on another.
//
// Implementation.
//
// LDS variables with constant annotation or non-undef initializer are passed
// through unchanged for simplification or error diagnostics in later passes.
// Non-undef initializers are not yet implemented for LDS.
//
// LDS variables that are always allocated at the same address can be found
// by lookup at that address. Otherwise runtime information/cost is required.
//
// The simplest strategy possible is to group all LDS variables in a single
// struct and allocate that struct in every kernel such that the original
// variables are always at the same address. LDS is however a limited resource
// so this strategy is unusable in practice. It is not implemented here.
//
// Strategy | Precise allocation | Zero runtime cost | General purpose |
// --------+--------------------+-------------------+-----------------+
// Module | No | Yes | Yes |
// Table | Yes | No | Yes |
// Kernel | Yes | Yes | No |
// Hybrid | Yes | Partial | Yes |
//
// Module spends LDS memory to save cycles. Table spends cycles and global
// memory to save LDS. Kernel is as fast as kernel allocation but only works
// for variables that are known reachable from a single kernel. Hybrid picks
// between all three. When forced to choose between LDS and cycles it minimises
// LDS use.
// The "module" lowering implemented here finds LDS variables which are used by
// non-kernel functions and creates a new struct with a field for each of those
// LDS variables. Variables that are only used from kernels are excluded.
// Kernels that do not use this struct are annoteated with the attribute
// amdgpu-elide-module-lds which allows the back end to elide the allocation.
//
// The "table" lowering implemented here has three components.
// First kernels are assigned a unique integer identifier which is available in
// functions it calls through the intrinsic amdgcn_lds_kernel_id. The integer
// is passed through a specific SGPR, thus works with indirect calls.
// Second, each kernel allocates LDS variables independent of other kernels and
// writes the addresses it chose for each variable into an array in consistent
// order. If the kernel does not allocate a given variable, it writes undef to
// the corresponding array location. These arrays are written to a constant
// table in the order matching the kernel unique integer identifier.
// Third, uses from non-kernel functions are replaced with a table lookup using
// the intrinsic function to find the address of the variable.
//
// "Kernel" lowering is only applicable for variables that are unambiguously
// reachable from exactly one kernel. For those cases, accesses to the variable
// can be lowered to ConstantExpr address of a struct instance specific to that
// one kernel. This is zero cost in space and in compute. It will raise a fatal
// error on any variable that might be reachable from multiple kernels and is
// thus most easily used as part of the hybrid lowering strategy.
//
// Hybrid lowering is a mixture of the above. It uses the zero cost kernel
// lowering where it can. It lowers the variable accessed by the greatest
// number of kernels using the module strategy as that is free for the first
// variable. Any futher variables that can be lowered with the module strategy
// without incurring LDS memory overhead are. The remaining ones are lowered
// via table.
//
// Consequences
// - No heuristics or user controlled magic numbers, hybrid is the right choice
// - Kernels that don't use functions (or have had them all inlined) are not
// affected by any lowering for kernels that do.
// - Kernels that don't make indirect function calls are not affected by those
// that do.
// - Variables which are used by lots of kernels, e.g. those injected by a
// language runtime in most kernels, are expected to have no overhead
// - Implementations that instantiate templates per-kernel where those templates
// use LDS are expected to hit the "Kernel" lowering strategy
// - The runtime properties impose a cost in compiler implementation complexity
//
//===----------------------------------------------------------------------===//
#include "AMDGPU.h"
#include "Utils/AMDGPUBaseInfo.h"
#include "Utils/AMDGPUMemoryUtils.h"
#include "llvm/ADT/BitVector.h"
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/DenseSet.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SetOperations.h"
#include "llvm/ADT/SetVector.h"
#include "llvm/Analysis/CallGraph.h"
#include "llvm/IR/Constants.h"
#include "llvm/IR/DerivedTypes.h"
#include "llvm/IR/IRBuilder.h"
#include "llvm/IR/InlineAsm.h"
#include "llvm/IR/Instructions.h"
#include "llvm/IR/IntrinsicsAMDGPU.h"
#include "llvm/IR/MDBuilder.h"
#include "llvm/InitializePasses.h"
#include "llvm/Pass.h"
#include "llvm/Support/CommandLine.h"
#include "llvm/Support/Debug.h"
#include "llvm/Support/OptimizedStructLayout.h"
#include "llvm/Transforms/Utils/BasicBlockUtils.h"
#include "llvm/Transforms/Utils/ModuleUtils.h"
#include <tuple>
#include <vector>
#include <cstdio>
#define DEBUG_TYPE "amdgpu-lower-module-lds"
using namespace llvm;
namespace {
cl::opt<bool> SuperAlignLDSGlobals(
"amdgpu-super-align-lds-globals",
cl::desc("Increase alignment of LDS if it is not on align boundary"),
cl::init(true), cl::Hidden);
enum class LoweringKind { module, table, kernel, hybrid };
cl::opt<LoweringKind> LoweringKindLoc(
"amdgpu-lower-module-lds-strategy",
cl::desc("Specify lowering strategy for function LDS access:"), cl::Hidden,
cl::init(LoweringKind::module),
cl::values(
clEnumValN(LoweringKind::table, "table", "Lower via table lookup"),
clEnumValN(LoweringKind::module, "module", "Lower via module struct"),
clEnumValN(
LoweringKind::kernel, "kernel",
"Lower variables reachable from one kernel, otherwise abort"),
clEnumValN(LoweringKind::hybrid, "hybrid",
"Lower via mixture of above strategies")));
bool isKernelLDS(const Function *F) {
// Some weirdness here. AMDGPU::isKernelCC does not call into
// AMDGPU::isKernel with the calling conv, it instead calls into
// isModuleEntryFunction which returns true for more calling conventions
// than AMDGPU::isKernel does. There's a FIXME on AMDGPU::isKernel.
// There's also a test that checks that the LDS lowering does not hit on
// a graphics shader, denoted amdgpu_ps, so stay with the limited case.
// Putting LDS in the name of the function to draw attention to this.
return AMDGPU::isKernel(F->getCallingConv());
}
class AMDGPULowerModuleLDS : public ModulePass {
static void
removeLocalVarsFromUsedLists(Module &M,
const DenseSet<GlobalVariable *> &LocalVars) {
// The verifier rejects used lists containing an inttoptr of a constant
// so remove the variables from these lists before replaceAllUsesWith
SmallPtrSet<Constant *, 8> LocalVarsSet;
for (GlobalVariable *LocalVar : LocalVars)
LocalVarsSet.insert(cast<Constant>(LocalVar->stripPointerCasts()));
removeFromUsedLists(
M, [&LocalVarsSet](Constant *C) { return LocalVarsSet.count(C); });
for (GlobalVariable *LocalVar : LocalVars)
LocalVar->removeDeadConstantUsers();
}
static void markUsedByKernel(IRBuilder<> &Builder, Function *Func,
GlobalVariable *SGV) {
// The llvm.amdgcn.module.lds instance is implicitly used by all kernels
// that might call a function which accesses a field within it. This is
// presently approximated to 'all kernels' if there are any such functions
// in the module. This implicit use is redefined as an explicit use here so
// that later passes, specifically PromoteAlloca, account for the required
// memory without any knowledge of this transform.
// An operand bundle on llvm.donothing works because the call instruction
// survives until after the last pass that needs to account for LDS. It is
// better than inline asm as the latter survives until the end of codegen. A
// totally robust solution would be a function with the same semantics as
// llvm.donothing that takes a pointer to the instance and is lowered to a
// no-op after LDS is allocated, but that is not presently necessary.
LLVMContext &Ctx = Func->getContext();
Builder.SetInsertPoint(Func->getEntryBlock().getFirstNonPHI());
FunctionType *FTy = FunctionType::get(Type::getVoidTy(Ctx), {});
Function *Decl =
Intrinsic::getDeclaration(Func->getParent(), Intrinsic::donothing, {});
Value *UseInstance[1] = {Builder.CreateInBoundsGEP(
SGV->getValueType(), SGV, ConstantInt::get(Type::getInt32Ty(Ctx), 0))};
Builder.CreateCall(FTy, Decl, {},
{OperandBundleDefT<Value *>("ExplicitUse", UseInstance)},
"");
}
static bool eliminateConstantExprUsesOfLDSFromAllInstructions(Module &M) {
// Constants are uniqued within LLVM. A ConstantExpr referring to a LDS
// global may have uses from multiple different functions as a result.
// This pass specialises LDS variables with respect to the kernel that
// allocates them.
// This is semantically equivalent to:
// for (auto &F : M.functions())
// for (auto &BB : F)
// for (auto &I : BB)
// for (Use &Op : I.operands())
// if (constantExprUsesLDS(Op))
// replaceConstantExprInFunction(I, Op);
bool Changed = false;
// Find all ConstantExpr that are direct users of an LDS global
SmallVector<ConstantExpr *> Stack;
for (auto &GV : M.globals())
if (AMDGPU::isLDSVariableToLower(GV))
for (User *U : GV.users())
if (ConstantExpr *C = dyn_cast<ConstantExpr>(U))
Stack.push_back(C);
// Expand to include constexpr users of direct users
SetVector<ConstantExpr *> ConstExprUsersOfLDS;
while (!Stack.empty()) {
ConstantExpr *V = Stack.pop_back_val();
if (ConstExprUsersOfLDS.contains(V))
continue;
ConstExprUsersOfLDS.insert(V);
for (auto *Nested : V->users())
if (ConstantExpr *CE = dyn_cast<ConstantExpr>(Nested))
Stack.push_back(CE);
}
// Find all instructions that use any of the ConstExpr users of LDS
SetVector<Instruction *> InstructionWorklist;
for (ConstantExpr *CE : ConstExprUsersOfLDS)
for (User *U : CE->users())
if (auto *I = dyn_cast<Instruction>(U))
InstructionWorklist.insert(I);
// Replace those ConstExpr operands with instructions
while (!InstructionWorklist.empty()) {
Instruction *I = InstructionWorklist.pop_back_val();
for (Use &U : I->operands()) {
auto *BI = I;
if (auto *Phi = dyn_cast<PHINode>(I)) {
BasicBlock *BB = Phi->getIncomingBlock(U);
BasicBlock::iterator It = BB->getFirstInsertionPt();
assert(It != BB->end() && "Unexpected empty basic block");
BI = &(*(It));
}
if (ConstantExpr *C = dyn_cast<ConstantExpr>(U.get())) {
if (ConstExprUsersOfLDS.contains(C)) {
Changed = true;
Instruction *NI = C->getAsInstruction(BI);
InstructionWorklist.insert(NI);
U.set(NI);
C->removeDeadConstantUsers();
}
}
}
}
return Changed;
}
public:
static char ID;
AMDGPULowerModuleLDS() : ModulePass(ID) {
initializeAMDGPULowerModuleLDSPass(*PassRegistry::getPassRegistry());
}
using FunctionVariableMap = DenseMap<Function *, DenseSet<GlobalVariable *>>;
using VariableFunctionMap = DenseMap<GlobalVariable *, DenseSet<Function *>>;
static void getUsesOfLDSByFunction(CallGraph const &CG, Module &M,
FunctionVariableMap &kernels,
FunctionVariableMap &functions) {
// Get uses from the current function, excluding uses by called functions
// Two output variables to avoid walking the globals list twice
for (auto &GV : M.globals()) {
if (!AMDGPU::isLDSVariableToLower(GV)) {
continue;
}
SmallVector<User *, 16> Stack(GV.users());
for (User *V : GV.users()) {
if (auto *I = dyn_cast<Instruction>(V)) {
Function *F = I->getFunction();
if (isKernelLDS(F)) {
kernels[F].insert(&GV);
} else {
functions[F].insert(&GV);
}
}
}
}
}
struct LDSUsesInfoTy {
FunctionVariableMap direct_access;
FunctionVariableMap indirect_access;
};
static LDSUsesInfoTy getTransitiveUsesOfLDS(CallGraph const &CG, Module &M) {
FunctionVariableMap direct_map_kernel;
FunctionVariableMap direct_map_function;
getUsesOfLDSByFunction(CG, M, direct_map_kernel, direct_map_function);
// Collect variables that are used by functions whose address has escaped
DenseSet<GlobalVariable *> VariablesReachableThroughFunctionPointer;
for (Function &F : M.functions()) {
if (!isKernelLDS(&F))
if (F.hasAddressTaken(nullptr,
/* IgnoreCallbackUses */ false,
/* IgnoreAssumeLikeCalls */ false,
/* IgnoreLLVMUsed */ true,
/* IgnoreArcAttachedCall */ false)) {
set_union(VariablesReachableThroughFunctionPointer,
direct_map_function[&F]);
}
}
auto functionMakesUnknownCall = [&](const Function *F) -> bool {
assert(!F->isDeclaration());
for (CallGraphNode::CallRecord R : *CG[F]) {
if (!R.second->getFunction()) {
return true;
}
}
return false;
};
// Work out which variables are reachable through function calls
FunctionVariableMap transitive_map_function = direct_map_function;
// If the function makes any unknown call, assume the worst case that it can
// access all variables accessed by functions whose address escaped
for (Function &F : M.functions()) {
if (!F.isDeclaration() && functionMakesUnknownCall(&F)) {
if (!isKernelLDS(&F)) {
set_union(transitive_map_function[&F],
VariablesReachableThroughFunctionPointer);
}
}
}
// Direct implementation of collecting all variables reachable from each
// function
for (Function &Func : M.functions()) {
if (Func.isDeclaration() || isKernelLDS(&Func))
continue;
DenseSet<Function *> seen; // catches cycles
SmallVector<Function *, 4> wip{&Func};
while (!wip.empty()) {
Function *F = wip.pop_back_val();
// Can accelerate this by referring to transitive map for functions that
// have already been computed, with more care than this
set_union(transitive_map_function[&Func], direct_map_function[F]);
for (CallGraphNode::CallRecord R : *CG[F]) {
Function *ith = R.second->getFunction();
if (ith) {
if (!seen.contains(ith)) {
seen.insert(ith);
wip.push_back(ith);
}
}
}
}
}
// direct_map_kernel lists which variables are used by the kernel
// find the variables which are used through a function call
FunctionVariableMap indirect_map_kernel;
for (Function &Func : M.functions()) {
if (Func.isDeclaration() || !isKernelLDS(&Func))
continue;
for (CallGraphNode::CallRecord R : *CG[&Func]) {
Function *ith = R.second->getFunction();
if (ith) {
set_union(indirect_map_kernel[&Func], transitive_map_function[ith]);
} else {
set_union(indirect_map_kernel[&Func],
VariablesReachableThroughFunctionPointer);
}
}
}
return {std::move(direct_map_kernel), std::move(indirect_map_kernel)};
}
struct LDSVariableReplacement {
GlobalVariable *SGV = nullptr;
DenseMap<GlobalVariable *, Constant *> LDSVarsToConstantGEP;
};
// remap from lds global to a constantexpr gep to where it has been moved to
// for each kernel
// an array with an element for each kernel containing where the corresponding
// variable was remapped to
static Constant *getAddressesOfVariablesInKernel(
LLVMContext &Ctx, ArrayRef<GlobalVariable *> Variables,
DenseMap<GlobalVariable *, Constant *> &LDSVarsToConstantGEP) {
// Create a ConstantArray containing the address of each Variable within the
// kernel corresponding to LDSVarsToConstantGEP, or poison if that kernel
// does not allocate it
// TODO: Drop the ptrtoint conversion
Type *I32 = Type::getInt32Ty(Ctx);
ArrayType *KernelOffsetsType = ArrayType::get(I32, Variables.size());
SmallVector<Constant *> Elements;
for (size_t i = 0; i < Variables.size(); i++) {
GlobalVariable *GV = Variables[i];
if (LDSVarsToConstantGEP.count(GV) != 0) {
auto elt = ConstantExpr::getPtrToInt(LDSVarsToConstantGEP[GV], I32);
Elements.push_back(elt);
} else {
Elements.push_back(PoisonValue::get(I32));
}
}
return ConstantArray::get(KernelOffsetsType, Elements);
}
static GlobalVariable *buildLookupTable(
Module &M, ArrayRef<GlobalVariable *> Variables,
ArrayRef<Function *> kernels,
DenseMap<Function *, LDSVariableReplacement> &KernelToReplacement) {
if (Variables.empty()) {
return nullptr;
}
LLVMContext &Ctx = M.getContext();
const size_t NumberVariables = Variables.size();
const size_t NumberKernels = kernels.size();
ArrayType *KernelOffsetsType =
ArrayType::get(Type::getInt32Ty(Ctx), NumberVariables);
ArrayType *AllKernelsOffsetsType =
ArrayType::get(KernelOffsetsType, NumberKernels);
std::vector<Constant *> overallConstantExprElts(NumberKernels);
for (size_t i = 0; i < NumberKernels; i++) {
LDSVariableReplacement Replacement = KernelToReplacement[kernels[i]];
overallConstantExprElts[i] = getAddressesOfVariablesInKernel(
Ctx, Variables, Replacement.LDSVarsToConstantGEP);
}
Constant *init =
ConstantArray::get(AllKernelsOffsetsType, overallConstantExprElts);
return new GlobalVariable(
M, AllKernelsOffsetsType, true, GlobalValue::InternalLinkage, init,
"llvm.amdgcn.lds.offset.table", nullptr, GlobalValue::NotThreadLocal,
AMDGPUAS::CONSTANT_ADDRESS);
}
void replaceUsesInInstructionsWithTableLookup(
Module &M, ArrayRef<GlobalVariable *> ModuleScopeVariables,
GlobalVariable *LookupTable) {
LLVMContext &Ctx = M.getContext();
IRBuilder<> Builder(Ctx);
Type *I32 = Type::getInt32Ty(Ctx);
// Accesses from a function use the amdgcn_lds_kernel_id intrinsic which
// lowers to a read from a live in register. Emit it once in the entry
// block to spare deduplicating it later.
DenseMap<Function *, Value *> tableKernelIndexCache;
auto getTableKernelIndex = [&](Function *F) -> Value * {
if (tableKernelIndexCache.count(F) == 0) {
LLVMContext &Ctx = M.getContext();
FunctionType *FTy = FunctionType::get(Type::getInt32Ty(Ctx), {});
Function *Decl =
Intrinsic::getDeclaration(&M, Intrinsic::amdgcn_lds_kernel_id, {});
BasicBlock::iterator it =
F->getEntryBlock().getFirstNonPHIOrDbgOrAlloca();
Instruction &i = *it;
Builder.SetInsertPoint(&i);
tableKernelIndexCache[F] = Builder.CreateCall(FTy, Decl, {});
}
return tableKernelIndexCache[F];
};
for (size_t Index = 0; Index < ModuleScopeVariables.size(); Index++) {
auto *GV = ModuleScopeVariables[Index];
for (Use &U : make_early_inc_range(GV->uses())) {
auto *I = dyn_cast<Instruction>(U.getUser());
if (!I)
continue;
Value *tableKernelIndex = getTableKernelIndex(I->getFunction());
// So if the phi uses this value multiple times, what does this look
// like?
if (auto *Phi = dyn_cast<PHINode>(I)) {
BasicBlock *BB = Phi->getIncomingBlock(U);
Builder.SetInsertPoint(&(*(BB->getFirstInsertionPt())));
} else {
Builder.SetInsertPoint(I);
}
Value *GEPIdx[3] = {
ConstantInt::get(I32, 0),
tableKernelIndex,
ConstantInt::get(I32, Index),
};
Value *Address = Builder.CreateInBoundsGEP(
LookupTable->getValueType(), LookupTable, GEPIdx, GV->getName());
Value *loaded = Builder.CreateLoad(I32, Address);
Value *replacement =
Builder.CreateIntToPtr(loaded, GV->getType(), GV->getName());
U.set(replacement);
}
}
}
static DenseSet<Function *> kernelsThatIndirectlyAccessAnyOfPassedVariables(
Module &M, LDSUsesInfoTy &LDSUsesInfo,
DenseSet<GlobalVariable *> const &VariableSet) {
DenseSet<Function *> KernelSet;
if (VariableSet.empty()) return KernelSet;
for (Function &Func : M.functions()) {
if (Func.isDeclaration() || !isKernelLDS(&Func))
continue;
for (GlobalVariable *GV : LDSUsesInfo.indirect_access[&Func]) {
if (VariableSet.contains(GV)) {
KernelSet.insert(&Func);
break;
}
}
}
return KernelSet;
}
static GlobalVariable *
chooseBestVariableForModuleStrategy(const DataLayout &DL,
VariableFunctionMap &LDSVars) {
// Find the global variable with the most indirect uses from kernels
struct CandidateTy {
GlobalVariable *GV = nullptr;
size_t UserCount = 0;
size_t Size = 0;
CandidateTy() = default;
CandidateTy(GlobalVariable *GV, uint64_t UserCount, uint64_t AllocSize)
: GV(GV), UserCount(UserCount), Size(AllocSize) {}
bool operator<(const CandidateTy &Other) const {
// Fewer users makes module scope variable less attractive
if (UserCount < Other.UserCount) {
return true;
}
if (UserCount > Other.UserCount) {
return false;
}
// Bigger makes module scope variable less attractive
if (Size < Other.Size) {
return false;
}
if (Size > Other.Size) {
return true;
}
// Arbitrary but consistent
return GV->getName() < Other.GV->getName();
}
};
CandidateTy MostUsed;
for (auto &K : LDSVars) {
GlobalVariable *GV = K.first;
if (K.second.size() <= 1) {
// A variable reachable by only one kernel is best lowered with kernel
// strategy
continue;
}
CandidateTy Candidate(GV, K.second.size(),
DL.getTypeAllocSize(GV->getValueType()).getFixedValue());
if (MostUsed < Candidate)
MostUsed = Candidate;
}
return MostUsed.GV;
}
bool runOnModule(Module &M) override {
LLVMContext &Ctx = M.getContext();
CallGraph CG = CallGraph(M);
bool Changed = superAlignLDSGlobals(M);
Changed |= eliminateConstantExprUsesOfLDSFromAllInstructions(M);
Changed = true; // todo: narrow this down
// For each kernel, what variables does it access directly or through
// callees
LDSUsesInfoTy LDSUsesInfo = getTransitiveUsesOfLDS(CG, M);
// For each variable accessed through callees, which kernels access it
VariableFunctionMap LDSToKernelsThatNeedToAccessItIndirectly;
for (auto &K : LDSUsesInfo.indirect_access) {
Function *F = K.first;
assert(isKernelLDS(F));
for (GlobalVariable *GV : K.second) {
LDSToKernelsThatNeedToAccessItIndirectly[GV].insert(F);
}
}
// Partition variables into the different strategies
DenseSet<GlobalVariable *> ModuleScopeVariables;
DenseSet<GlobalVariable *> TableLookupVariables;
DenseSet<GlobalVariable *> KernelAccessVariables;
{
GlobalVariable *HybridModuleRoot =
LoweringKindLoc != LoweringKind::hybrid
? nullptr
: chooseBestVariableForModuleStrategy(
M.getDataLayout(),
LDSToKernelsThatNeedToAccessItIndirectly);
DenseSet<Function *> const EmptySet;
DenseSet<Function *> const &HybridModuleRootKernels =
HybridModuleRoot
? LDSToKernelsThatNeedToAccessItIndirectly[HybridModuleRoot]
: EmptySet;
for (auto &K : LDSToKernelsThatNeedToAccessItIndirectly) {
// Each iteration of this loop assigns exactly one global variable to
// exactly one of the implementation strategies.
GlobalVariable *GV = K.first;
assert(AMDGPU::isLDSVariableToLower(*GV));
assert(K.second.size() != 0);
switch (LoweringKindLoc) {
case LoweringKind::module:
ModuleScopeVariables.insert(GV);
break;
case LoweringKind::table:
TableLookupVariables.insert(GV);
break;
case LoweringKind::kernel:
if (K.second.size() == 1) {
KernelAccessVariables.insert(GV);
} else {
report_fatal_error(
"cannot lower LDS '" + GV->getName() +
"' to kernel access as it is reachable from multiple kernels");
}
break;
case LoweringKind::hybrid: {
if (GV == HybridModuleRoot) {
assert(K.second.size() != 1);
ModuleScopeVariables.insert(GV);
} else if (K.second.size() == 1) {
KernelAccessVariables.insert(GV);
} else if (set_is_subset(K.second, HybridModuleRootKernels)) {
ModuleScopeVariables.insert(GV);
} else {
TableLookupVariables.insert(GV);
}
break;
}
}
}
assert(ModuleScopeVariables.size() + TableLookupVariables.size() +
KernelAccessVariables.size() ==
LDSToKernelsThatNeedToAccessItIndirectly.size());
} // Variables have now been partitioned into the three lowering strategies.
// If the kernel accesses a variable that is going to be stored in the
// module instance through a call then that kernel needs to allocate the
// module instance
DenseSet<Function *> KernelsThatAllocateModuleLDS =
kernelsThatIndirectlyAccessAnyOfPassedVariables(M, LDSUsesInfo,
ModuleScopeVariables);
DenseSet<Function *> KernelsThatAllocateTableLDS =
kernelsThatIndirectlyAccessAnyOfPassedVariables(M, LDSUsesInfo,
TableLookupVariables);
if (!ModuleScopeVariables.empty()) {
LDSVariableReplacement ModuleScopeReplacement =
createLDSVariableReplacement(M, "llvm.amdgcn.module.lds",
ModuleScopeVariables);
appendToCompilerUsed(M,
{static_cast<GlobalValue *>(
ConstantExpr::getPointerBitCastOrAddrSpaceCast(
cast<Constant>(ModuleScopeReplacement.SGV),
Type::getInt8PtrTy(Ctx)))});
// historic
removeLocalVarsFromUsedLists(M, ModuleScopeVariables);
// Replace all uses of module scope variable from non-kernel functions
replaceLDSVariablesWithStruct(
M, ModuleScopeVariables, ModuleScopeReplacement, [&](Use &U) {
Instruction *I = dyn_cast<Instruction>(U.getUser());
if (!I) {
return false;
}
Function *F = I->getFunction();
return !isKernelLDS(F);
});
// Replace uses of module scope variable from kernel functions that
// allocate the module scope variable, otherwise leave them unchanged
// Record on each kernel whether the module scope global is used by it
LLVMContext &Ctx = M.getContext();
IRBuilder<> Builder(Ctx);
for (Function &Func : M.functions()) {
if (Func.isDeclaration() || !isKernelLDS(&Func))
continue;
if (KernelsThatAllocateModuleLDS.contains(&Func)) {
replaceLDSVariablesWithStruct(
M, ModuleScopeVariables, ModuleScopeReplacement, [&](Use &U) {
Instruction *I = dyn_cast<Instruction>(U.getUser());
if (!I) {
return false;
}
Function *F = I->getFunction();
return F == &Func;
});
markUsedByKernel(Builder, &Func, ModuleScopeReplacement.SGV);
} else {
Func.addFnAttr("amdgpu-elide-module-lds");
}
}
}
// Create a struct for each kernel for the non-module-scope variables
DenseMap<Function *, LDSVariableReplacement> KernelToReplacement;
for (Function &Func : M.functions()) {
if (Func.isDeclaration() || !isKernelLDS(&Func))
continue;
DenseSet<GlobalVariable *> KernelUsedVariables;
for (auto &v : LDSUsesInfo.direct_access[&Func]) {
KernelUsedVariables.insert(v);
}
for (auto &v : LDSUsesInfo.indirect_access[&Func]) {
KernelUsedVariables.insert(v);
}
// Variables allocated in module lds must all resolve to that struct,
// not to the per-kernel instance.
if (KernelsThatAllocateModuleLDS.contains(&Func)) {
for (GlobalVariable *v : ModuleScopeVariables) {
KernelUsedVariables.erase(v);
}
}
if (KernelUsedVariables.empty()) {
// Either used no LDS, or all the LDS it used was also in module
continue;
}
// The association between kernel function and LDS struct is done by
// symbol name, which only works if the function in question has a
// name This is not expected to be a problem in practice as kernels
// are called by name making anonymous ones (which are named by the
// backend) difficult to use. This does mean that llvm test cases need
// to name the kernels.
if (!Func.hasName()) {
report_fatal_error("Anonymous kernels cannot use LDS variables");
}
std::string VarName =
(Twine("llvm.amdgcn.kernel.") + Func.getName() + ".lds").str();
auto Replacement =
createLDSVariableReplacement(M, VarName, KernelUsedVariables);
// remove preserves existing codegen
removeLocalVarsFromUsedLists(M, KernelUsedVariables);
KernelToReplacement[&Func] = Replacement;
// Rewrite uses within kernel to the new struct
replaceLDSVariablesWithStruct(
M, KernelUsedVariables, Replacement, [&Func](Use &U) {
Instruction *I = dyn_cast<Instruction>(U.getUser());
return I && I->getFunction() == &Func;
});
}
// Lower zero cost accesses to the kernel instances just created
for (auto &GV : KernelAccessVariables) {
auto &funcs = LDSToKernelsThatNeedToAccessItIndirectly[GV];
assert(funcs.size() == 1); // Only one kernel can access it
LDSVariableReplacement Replacement =
KernelToReplacement[*(funcs.begin())];
DenseSet<GlobalVariable *> Vec;
Vec.insert(GV);
replaceLDSVariablesWithStruct(M, Vec, Replacement, [](Use &U) {
return isa<Instruction>(U.getUser());
});
}
if (!KernelsThatAllocateTableLDS.empty()) {
// Collect the kernels that allocate table lookup LDS
std::vector<Function *> OrderedKernels;
{
for (Function &Func : M.functions()) {
if (Func.isDeclaration())
continue;
if (!isKernelLDS(&Func))
continue;
if (KernelsThatAllocateTableLDS.contains(&Func)) {
assert(Func.hasName()); // else fatal error earlier
OrderedKernels.push_back(&Func);
}
}
// Put them in an arbitrary but reproducible order
llvm::sort(OrderedKernels.begin(), OrderedKernels.end(),
[](const Function *lhs, const Function *rhs) -> bool {
return lhs->getName() < rhs->getName();
});
// Annotate the kernels with their order in this vector
LLVMContext &Ctx = M.getContext();
IRBuilder<> Builder(Ctx);
if (OrderedKernels.size() > UINT32_MAX) {
// 32 bit keeps it in one SGPR. > 2**32 kernels won't fit on the GPU
report_fatal_error("Unimplemented LDS lowering for > 2**32 kernels");
}
for (size_t i = 0; i < OrderedKernels.size(); i++) {
Metadata *AttrMDArgs[1] = {
ConstantAsMetadata::get(Builder.getInt32(i)),
};
OrderedKernels[i]->setMetadata("llvm.amdgcn.lds.kernel.id",
MDNode::get(Ctx, AttrMDArgs));
markUsedByKernel(Builder, OrderedKernels[i],
KernelToReplacement[OrderedKernels[i]].SGV);
}
}
// The order must be consistent between lookup table and accesses to
// lookup table
std::vector<GlobalVariable *> TableLookupVariablesOrdered(
TableLookupVariables.begin(), TableLookupVariables.end());
llvm::sort(TableLookupVariablesOrdered.begin(),
TableLookupVariablesOrdered.end(),
[](const GlobalVariable *lhs, const GlobalVariable *rhs) {
return lhs->getName() < rhs->getName();
});
GlobalVariable *LookupTable = buildLookupTable(
M, TableLookupVariablesOrdered, OrderedKernels, KernelToReplacement);
replaceUsesInInstructionsWithTableLookup(M, TableLookupVariablesOrdered,
LookupTable);
}
for (auto &GV : make_early_inc_range(M.globals()))
if (AMDGPU::isLDSVariableToLower(GV)) {
// probably want to remove from used lists
GV.removeDeadConstantUsers();
if (GV.use_empty())
GV.eraseFromParent();
}
return Changed;
}
private:
// Increase the alignment of LDS globals if necessary to maximise the chance
// that we can use aligned LDS instructions to access them.
static bool superAlignLDSGlobals(Module &M) {
const DataLayout &DL = M.getDataLayout();
bool Changed = false;
if (!SuperAlignLDSGlobals) {
return Changed;
}
for (auto &GV : M.globals()) {
if (GV.getType()->getPointerAddressSpace() != AMDGPUAS::LOCAL_ADDRESS) {
// Only changing alignment of LDS variables
continue;
}
if (!GV.hasInitializer()) {
// cuda/hip extern __shared__ variable, leave alignment alone
continue;
}
Align Alignment = AMDGPU::getAlign(DL, &GV);
TypeSize GVSize = DL.getTypeAllocSize(GV.getValueType());
if (GVSize > 8) {
// We might want to use a b96 or b128 load/store
Alignment = std::max(Alignment, Align(16));
} else if (GVSize > 4) {
// We might want to use a b64 load/store
Alignment = std::max(Alignment, Align(8));
} else if (GVSize > 2) {
// We might want to use a b32 load/store
Alignment = std::max(Alignment, Align(4));
} else if (GVSize > 1) {
// We might want to use a b16 load/store
Alignment = std::max(Alignment, Align(2));
}
if (Alignment != AMDGPU::getAlign(DL, &GV)) {
Changed = true;
GV.setAlignment(Alignment);
}
}
return Changed;
}
static LDSVariableReplacement createLDSVariableReplacement(
Module &M, std::string VarName,
DenseSet<GlobalVariable *> const &LDSVarsToTransform) {
// Create a struct instance containing LDSVarsToTransform and map from those
// variables to ConstantExprGEP
// Variables may be introduced to meet alignment requirements. No aliasing
// metadata is useful for these as they have no uses. Erased before return.
LLVMContext &Ctx = M.getContext();
const DataLayout &DL = M.getDataLayout();
assert(!LDSVarsToTransform.empty());
SmallVector<OptimizedStructLayoutField, 8> LayoutFields;
LayoutFields.reserve(LDSVarsToTransform.size());
{
// The order of fields in this struct depends on the order of
// varables in the argument which varies when changing how they
// are identified, leading to spurious test breakage.
std::vector<GlobalVariable *> Sorted(LDSVarsToTransform.begin(),
LDSVarsToTransform.end());
llvm::sort(Sorted.begin(), Sorted.end(),
[](const GlobalVariable *lhs, const GlobalVariable *rhs) {
return lhs->getName() < rhs->getName();
});
for (GlobalVariable *GV : Sorted) {
OptimizedStructLayoutField F(GV,
DL.getTypeAllocSize(GV->getValueType()),
AMDGPU::getAlign(DL, GV));
LayoutFields.emplace_back(F);
}
}
performOptimizedStructLayout(LayoutFields);
std::vector<GlobalVariable *> LocalVars;
BitVector IsPaddingField;
LocalVars.reserve(LDSVarsToTransform.size()); // will be at least this large
IsPaddingField.reserve(LDSVarsToTransform.size());
{
uint64_t CurrentOffset = 0;
for (size_t I = 0; I < LayoutFields.size(); I++) {
GlobalVariable *FGV = static_cast<GlobalVariable *>(
const_cast<void *>(LayoutFields[I].Id));
Align DataAlign = LayoutFields[I].Alignment;
uint64_t DataAlignV = DataAlign.value();
if (uint64_t Rem = CurrentOffset % DataAlignV) {
uint64_t Padding = DataAlignV - Rem;
// Append an array of padding bytes to meet alignment requested
// Note (o + (a - (o % a)) ) % a == 0
// (offset + Padding ) % align == 0
Type *ATy = ArrayType::get(Type::getInt8Ty(Ctx), Padding);
LocalVars.push_back(new GlobalVariable(
M, ATy, false, GlobalValue::InternalLinkage, UndefValue::get(ATy),
"", nullptr, GlobalValue::NotThreadLocal, AMDGPUAS::LOCAL_ADDRESS,
false));
IsPaddingField.push_back(true);
CurrentOffset += Padding;
}
LocalVars.push_back(FGV);
IsPaddingField.push_back(false);
CurrentOffset += LayoutFields[I].Size;
}
}
std::vector<Type *> LocalVarTypes;
LocalVarTypes.reserve(LocalVars.size());
std::transform(
LocalVars.cbegin(), LocalVars.cend(), std::back_inserter(LocalVarTypes),
[](const GlobalVariable *V) -> Type * { return V->getValueType(); });
StructType *LDSTy = StructType::create(Ctx, LocalVarTypes, VarName + ".t");
Align StructAlign = AMDGPU::getAlign(DL, LocalVars[0]);
GlobalVariable *SGV = new GlobalVariable(
M, LDSTy, false, GlobalValue::InternalLinkage, UndefValue::get(LDSTy),
VarName, nullptr, GlobalValue::NotThreadLocal, AMDGPUAS::LOCAL_ADDRESS,
false);
SGV->setAlignment(StructAlign);
DenseMap<GlobalVariable *, Constant *> Map;
Type *I32 = Type::getInt32Ty(Ctx);
for (size_t I = 0; I < LocalVars.size(); I++) {
GlobalVariable *GV = LocalVars[I];
Constant *GEPIdx[] = {ConstantInt::get(I32, 0), ConstantInt::get(I32, I)};
Constant *GEP = ConstantExpr::getGetElementPtr(LDSTy, SGV, GEPIdx, true);
if (IsPaddingField[I]) {
assert(GV->use_empty());
GV->eraseFromParent();
} else {
Map[GV] = GEP;
}
}
assert(Map.size() == LDSVarsToTransform.size());
return {SGV, std::move(Map)};
}
template <typename PredicateTy>
void replaceLDSVariablesWithStruct(
Module &M, DenseSet<GlobalVariable *> const &LDSVarsToTransformArg,
LDSVariableReplacement Replacement, PredicateTy Predicate) {
LLVMContext &Ctx = M.getContext();
const DataLayout &DL = M.getDataLayout();
// A hack... we need to insert the aliasing info in a predictable order for
// lit tests. Would like to have them in a stable order already, ideally the
// same order they get allocated, which might mean an ordered set container
std::vector<GlobalVariable *> LDSVarsToTransform(
LDSVarsToTransformArg.begin(), LDSVarsToTransformArg.end());
llvm::sort(LDSVarsToTransform.begin(), LDSVarsToTransform.end(),
[](const GlobalVariable *lhs, const GlobalVariable *rhs) {
return lhs->getName() < rhs->getName();
});
// Create alias.scope and their lists. Each field in the new structure
// does not alias with all other fields.
SmallVector<MDNode *> AliasScopes;
SmallVector<Metadata *> NoAliasList;
const size_t NumberVars = LDSVarsToTransform.size();
if (NumberVars > 1) {
MDBuilder MDB(Ctx);
AliasScopes.reserve(NumberVars);
MDNode *Domain = MDB.createAnonymousAliasScopeDomain();
for (size_t I = 0; I < NumberVars; I++) {
MDNode *Scope = MDB.createAnonymousAliasScope(Domain);
AliasScopes.push_back(Scope);
}
NoAliasList.append(&AliasScopes[1], AliasScopes.end());
}
// Replace uses of ith variable with a constantexpr to the corresponding
// field of the instance that will be allocated by AMDGPUMachineFunction
for (size_t I = 0; I < NumberVars; I++) {
GlobalVariable *GV = LDSVarsToTransform[I];
Constant *GEP = Replacement.LDSVarsToConstantGEP[GV];
GV->replaceUsesWithIf(GEP, Predicate);
APInt APOff(DL.getIndexTypeSizeInBits(GEP->getType()), 0);
GEP->stripAndAccumulateInBoundsConstantOffsets(DL, APOff);
uint64_t Offset = APOff.getZExtValue();
Align A =
commonAlignment(Replacement.SGV->getAlign().valueOrOne(), Offset);
if (I)
NoAliasList[I - 1] = AliasScopes[I - 1];
MDNode *NoAlias =
NoAliasList.empty() ? nullptr : MDNode::get(Ctx, NoAliasList);
MDNode *AliasScope =
AliasScopes.empty() ? nullptr : MDNode::get(Ctx, {AliasScopes[I]});
refineUsesAlignmentAndAA(GEP, A, DL, AliasScope, NoAlias);
}
}
void refineUsesAlignmentAndAA(Value *Ptr, Align A, const DataLayout &DL,
MDNode *AliasScope, MDNode *NoAlias,
unsigned MaxDepth = 5) {
if (!MaxDepth || (A == 1 && !AliasScope))
return;
for (User *U : Ptr->users()) {
if (auto *I = dyn_cast<Instruction>(U)) {
if (AliasScope && I->mayReadOrWriteMemory()) {
MDNode *AS = I->getMetadata(LLVMContext::MD_alias_scope);
AS = (AS ? MDNode::getMostGenericAliasScope(AS, AliasScope)
: AliasScope);
I->setMetadata(LLVMContext::MD_alias_scope, AS);
MDNode *NA = I->getMetadata(LLVMContext::MD_noalias);
NA = (NA ? MDNode::intersect(NA, NoAlias) : NoAlias);
I->setMetadata(LLVMContext::MD_noalias, NA);
}
}
if (auto *LI = dyn_cast<LoadInst>(U)) {
LI->setAlignment(std::max(A, LI->getAlign()));
continue;
}
if (auto *SI = dyn_cast<StoreInst>(U)) {
if (SI->getPointerOperand() == Ptr)
SI->setAlignment(std::max(A, SI->getAlign()));
continue;
}
if (auto *AI = dyn_cast<AtomicRMWInst>(U)) {
// None of atomicrmw operations can work on pointers, but let's
// check it anyway in case it will or we will process ConstantExpr.
if (AI->getPointerOperand() == Ptr)
AI->setAlignment(std::max(A, AI->getAlign()));
continue;
}
if (auto *AI = dyn_cast<AtomicCmpXchgInst>(U)) {
if (AI->getPointerOperand() == Ptr)
AI->setAlignment(std::max(A, AI->getAlign()));
continue;
}
if (auto *GEP = dyn_cast<GetElementPtrInst>(U)) {
unsigned BitWidth = DL.getIndexTypeSizeInBits(GEP->getType());
APInt Off(BitWidth, 0);
if (GEP->getPointerOperand() == Ptr) {
Align GA;
if (GEP->accumulateConstantOffset(DL, Off))
GA = commonAlignment(A, Off.getLimitedValue());
refineUsesAlignmentAndAA(GEP, GA, DL, AliasScope, NoAlias,
MaxDepth - 1);
}
continue;
}
if (auto *I = dyn_cast<Instruction>(U)) {
if (I->getOpcode() == Instruction::BitCast ||
I->getOpcode() == Instruction::AddrSpaceCast)
refineUsesAlignmentAndAA(I, A, DL, AliasScope, NoAlias, MaxDepth - 1);
}
}
}
};
} // namespace
char AMDGPULowerModuleLDS::ID = 0;
char &llvm::AMDGPULowerModuleLDSID = AMDGPULowerModuleLDS::ID;
INITIALIZE_PASS(AMDGPULowerModuleLDS, DEBUG_TYPE,
"Lower uses of LDS variables from non-kernel functions", false,
false)
ModulePass *llvm::createAMDGPULowerModuleLDSPass() {
return new AMDGPULowerModuleLDS();
}
PreservedAnalyses AMDGPULowerModuleLDSPass::run(Module &M,
ModuleAnalysisManager &) {
return AMDGPULowerModuleLDS().runOnModule(M) ? PreservedAnalyses::none()
: PreservedAnalyses::all();
}