Google Git
Sign in
llvm / llvm-project / 2d287f51eff2a5fbf84458a33f7fb2493cf67965 / . / llvm / lib / Target / AMDGPU / AMDGPUSwLowerLDS.cpp
blob: 9a5ffe1c2dad85415c439571d08da876a6c5d4c7 [file] [log] [blame]
//===-- AMDGPUSwLowerLDS.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
//
//===----------------------------------------------------------------------===//
//
// This pass lowers the local data store, LDS, uses in kernel and non-kernel
// functions in module to use dynamically allocated global memory.
// Packed LDS Layout is emulated in the global memory.
// The lowered memory instructions from LDS to global memory are then
// instrumented for address sanitizer, to catch addressing errors.
// This pass only work when address sanitizer has been enabled and has
// instrumented the IR. It identifies that IR has been instrumented using
// "nosanitize_address" module flag.
//
// Replacement of Kernel LDS accesses:
// For a kernel, LDS access can be static or dynamic which are direct
// (accessed within kernel) and indirect (accessed through non-kernels).
// All these LDS accesses corresponding to kernel will be packed together,
// where all static LDS accesses will be allocated first and then dynamic
// LDS follows. The total size with alignment is calculated. A new LDS global
// will be created for the kernel called "SW LDS" and it will have the
// attribute "amdgpu-lds-size" attached with value of the size calculated.
// All the LDS accesses in the module will be replaced by GEP with offset
// into the "Sw LDS".
// A new "llvm.amdgcn.<kernel>.dynlds" is created per kernel accessing
// the dynamic LDS. This will be marked used by kernel and will have
// MD_absolue_symbol metadata set to total static LDS size, Since dynamic
// LDS allocation starts after all static LDS allocation.
//
// A device global memory equal to the total LDS size will be allocated.
// At the prologue of the kernel, a single work-item from the
// work-group, does a "malloc" and stores the pointer of the
// allocation in "SW LDS".
//
// To store the offsets corresponding to all LDS accesses, another global
// variable is created which will be called "SW LDS metadata" in this pass.
// - SW LDS Global:
// It is LDS global of ptr type with name
// "llvm.amdgcn.sw.lds.<kernel-name>".
// - Metadata Global:
// It is of struct type, with n members. n equals the number of LDS
// globals accessed by the kernel(direct and indirect). Each member of
// struct is another struct of type {i32, i32, i32}. First member
// corresponds to offset, second member corresponds to size of LDS global
// being replaced and third represents the total aligned size. It will
// have name "llvm.amdgcn.sw.lds.<kernel-name>.md". This global will have
// an intializer with static LDS related offsets and sizes initialized.
// But for dynamic LDS related entries, offsets will be intialized to
// previous static LDS allocation end offset. Sizes for them will be zero
// initially. These dynamic LDS offset and size values will be updated
// within the kernel, since kernel can read the dynamic LDS size
// allocation done at runtime with query to "hidden_dynamic_lds_size"
// hidden kernel argument.
//
// At the epilogue of kernel, allocated memory would be made free by the same
// single work-item.
//
// Replacement of non-kernel LDS accesses:
// Multiple kernels can access the same non-kernel function.
// All the kernels accessing LDS through non-kernels are sorted and
// assigned a kernel-id. All the LDS globals accessed by non-kernels
// are sorted. This information is used to build two tables:
// - Base table:
// Base table will have single row, with elements of the row
// placed as per kernel ID. Each element in the row corresponds
// to ptr of "SW LDS" variable created for that kernel.
// - Offset table:
// Offset table will have multiple rows and columns.
// Rows are assumed to be from 0 to (n-1). n is total number
// of kernels accessing the LDS through non-kernels.
// Each row will have m elements. m is the total number of
// unique LDS globals accessed by all non-kernels.
// Each element in the row correspond to the ptr of
// the replacement of LDS global done by that particular kernel.
// A LDS variable in non-kernel will be replaced based on the information
// from base and offset tables. Based on kernel-id query, ptr of "SW
// LDS" for that corresponding kernel is obtained from base table.
// The Offset into the base "SW LDS" is obtained from
// corresponding element in offset table. With this information, replacement
// value is obtained.
//===----------------------------------------------------------------------===//
#include "AMDGPU.h"
#include "AMDGPUAsanInstrumentation.h"
#include "AMDGPUMemoryUtils.h"
#include "AMDGPUTargetMachine.h"
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/DenseSet.h"
#include "llvm/ADT/SetVector.h"
#include "llvm/ADT/StringExtras.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/Analysis/CallGraph.h"
#include "llvm/Analysis/DomTreeUpdater.h"
#include "llvm/CodeGen/TargetPassConfig.h"
#include "llvm/IR/Constants.h"
#include "llvm/IR/DIBuilder.h"
#include "llvm/IR/DebugInfo.h"
#include "llvm/IR/DebugInfoMetadata.h"
#include "llvm/IR/IRBuilder.h"
#include "llvm/IR/Instructions.h"
#include "llvm/IR/IntrinsicsAMDGPU.h"
#include "llvm/IR/MDBuilder.h"
#include "llvm/IR/ReplaceConstant.h"
#include "llvm/InitializePasses.h"
#include "llvm/Pass.h"
#include "llvm/Support/raw_ostream.h"
#include "llvm/Transforms/Instrumentation/AddressSanitizerCommon.h"
#include "llvm/Transforms/Utils/ModuleUtils.h"
#include <algorithm>
#define DEBUG_TYPE "amdgpu-sw-lower-lds"
#define COV5_HIDDEN_DYN_LDS_SIZE_ARG 15
using namespace llvm;
using namespace AMDGPU;
namespace {
cl::opt<bool>
AsanInstrumentLDS("amdgpu-asan-instrument-lds",
cl::desc("Run asan instrumentation on LDS instructions "
"lowered to global memory"),
cl::init(true), cl::Hidden);
using DomTreeCallback = function_ref<DominatorTree *(Function &F)>;
struct LDSAccessTypeInfo {
SetVector<GlobalVariable *> StaticLDSGlobals;
SetVector<GlobalVariable *> DynamicLDSGlobals;
};
// Struct to hold all the Metadata required for a kernel
// to replace a LDS global uses with corresponding offset
// in to device global memory.
struct KernelLDSParameters {
GlobalVariable *SwLDS = nullptr;
GlobalVariable *SwDynLDS = nullptr;
GlobalVariable *SwLDSMetadata = nullptr;
LDSAccessTypeInfo DirectAccess;
LDSAccessTypeInfo IndirectAccess;
DenseMap<GlobalVariable *, SmallVector<uint32_t, 3>>
LDSToReplacementIndicesMap;
uint32_t MallocSize = 0;
uint32_t LDSSize = 0;
SmallVector<std::pair<uint32_t, uint32_t>, 64> RedzoneOffsetAndSizeVector;
};
// Struct to store information for creation of offset table
// for all the non-kernel LDS accesses.
struct NonKernelLDSParameters {
GlobalVariable *LDSBaseTable = nullptr;
GlobalVariable *LDSOffsetTable = nullptr;
SetVector<Function *> OrderedKernels;
SetVector<GlobalVariable *> OrdereLDSGlobals;
};
struct AsanInstrumentInfo {
int Scale = 0;
uint32_t Offset = 0;
SetVector<Instruction *> Instructions;
};
struct FunctionsAndLDSAccess {
DenseMap<Function *, KernelLDSParameters> KernelToLDSParametersMap;
SetVector<Function *> KernelsWithIndirectLDSAccess;
SetVector<Function *> NonKernelsWithLDSArgument;
SetVector<GlobalVariable *> AllNonKernelLDSAccess;
FunctionVariableMap NonKernelToLDSAccessMap;
};
class AMDGPUSwLowerLDS {
public:
AMDGPUSwLowerLDS(Module &Mod, const AMDGPUTargetMachine &TM,
DomTreeCallback Callback)
: M(Mod), AMDGPUTM(TM), IRB(M.getContext()), DTCallback(Callback) {}
bool run();
void getUsesOfLDSByNonKernels();
void getNonKernelsWithLDSArguments(const CallGraph &CG);
SetVector<Function *>
getOrderedIndirectLDSAccessingKernels(SetVector<Function *> &Kernels);
SetVector<GlobalVariable *>
getOrderedNonKernelAllLDSGlobals(SetVector<GlobalVariable *> &Variables);
void buildSwLDSGlobal(Function *Func);
void buildSwDynLDSGlobal(Function *Func);
void populateSwMetadataGlobal(Function *Func);
void populateSwLDSAttributeAndMetadata(Function *Func);
void populateLDSToReplacementIndicesMap(Function *Func);
void getLDSMemoryInstructions(Function *Func,
SetVector<Instruction *> &LDSInstructions);
void replaceKernelLDSAccesses(Function *Func);
Value *getTranslatedGlobalMemoryPtrOfLDS(Value *LoadMallocPtr, Value *LDSPtr);
void translateLDSMemoryOperationsToGlobalMemory(
Function *Func, Value *LoadMallocPtr,
SetVector<Instruction *> &LDSInstructions);
void poisonRedzones(Function *Func, Value *MallocPtr);
void lowerKernelLDSAccesses(Function *Func, DomTreeUpdater &DTU);
void buildNonKernelLDSOffsetTable(NonKernelLDSParameters &NKLDSParams);
void buildNonKernelLDSBaseTable(NonKernelLDSParameters &NKLDSParams);
Constant *
getAddressesOfVariablesInKernel(Function *Func,
SetVector<GlobalVariable *> &Variables);
void lowerNonKernelLDSAccesses(Function *Func,
SetVector<GlobalVariable *> &LDSGlobals,
NonKernelLDSParameters &NKLDSParams);
void
updateMallocSizeForDynamicLDS(Function *Func, Value **CurrMallocSize,
Value *HiddenDynLDSSize,
SetVector<GlobalVariable *> &DynamicLDSGlobals);
void initAsanInfo();
private:
Module &M;
const AMDGPUTargetMachine &AMDGPUTM;
IRBuilder<> IRB;
DomTreeCallback DTCallback;
FunctionsAndLDSAccess FuncLDSAccessInfo;
AsanInstrumentInfo AsanInfo;
};
template <typename T> SetVector<T> sortByName(std::vector<T> &&V) {
// Sort the vector of globals or Functions based on their name.
// Returns a SetVector of globals/Functions.
sort(V, [](const auto *L, const auto *R) {
return L->getName() < R->getName();
});
return {SetVector<T>(llvm::from_range, V)};
}
SetVector<GlobalVariable *> AMDGPUSwLowerLDS::getOrderedNonKernelAllLDSGlobals(
SetVector<GlobalVariable *> &Variables) {
// Sort all the non-kernel LDS accesses based on their name.
return sortByName(
std::vector<GlobalVariable *>(Variables.begin(), Variables.end()));
}
SetVector<Function *> AMDGPUSwLowerLDS::getOrderedIndirectLDSAccessingKernels(
SetVector<Function *> &Kernels) {
// Sort the non-kernels accessing LDS based on their name.
// Also assign a kernel ID metadata based on the sorted order.
LLVMContext &Ctx = M.getContext();
if (Kernels.size() > UINT32_MAX) {
report_fatal_error("Unimplemented SW LDS lowering for > 2**32 kernels");
}
SetVector<Function *> OrderedKernels =
sortByName(std::vector<Function *>(Kernels.begin(), Kernels.end()));
for (size_t i = 0; i < Kernels.size(); i++) {
Metadata *AttrMDArgs[1] = {
ConstantAsMetadata::get(IRB.getInt32(i)),
};
Function *Func = OrderedKernels[i];
Func->setMetadata("llvm.amdgcn.lds.kernel.id",
MDNode::get(Ctx, AttrMDArgs));
}
return OrderedKernels;
}
void AMDGPUSwLowerLDS::getNonKernelsWithLDSArguments(const CallGraph &CG) {
// Among the kernels accessing LDS, get list of
// Non-kernels to which a call is made and a ptr
// to addrspace(3) is passed as argument.
for (auto &K : FuncLDSAccessInfo.KernelToLDSParametersMap) {
Function *Func = K.first;
const CallGraphNode *CGN = CG[Func];
if (!CGN)
continue;
for (auto &I : *CGN) {
CallGraphNode *CallerCGN = I.second;
Function *CalledFunc = CallerCGN->getFunction();
if (!CalledFunc || CalledFunc->isDeclaration())
continue;
if (AMDGPU::isKernelLDS(CalledFunc))
continue;
for (auto AI = CalledFunc->arg_begin(), E = CalledFunc->arg_end();
AI != E; ++AI) {
Type *ArgTy = (*AI).getType();
if (!ArgTy->isPointerTy())
continue;
if (ArgTy->getPointerAddressSpace() != AMDGPUAS::LOCAL_ADDRESS)
continue;
FuncLDSAccessInfo.NonKernelsWithLDSArgument.insert(CalledFunc);
// Also add the Calling function to KernelsWithIndirectLDSAccess list
// so that base table of LDS is generated.
FuncLDSAccessInfo.KernelsWithIndirectLDSAccess.insert(Func);
}
}
}
}
void AMDGPUSwLowerLDS::getUsesOfLDSByNonKernels() {
for (GlobalVariable *GV : FuncLDSAccessInfo.AllNonKernelLDSAccess) {
if (!AMDGPU::isLDSVariableToLower(*GV))
continue;
for (User *V : GV->users()) {
if (auto *I = dyn_cast<Instruction>(V)) {
Function *F = I->getFunction();
if (!isKernelLDS(F) && !F->isDeclaration())
FuncLDSAccessInfo.NonKernelToLDSAccessMap[F].insert(GV);
}
}
}
}
static void recordLDSAbsoluteAddress(Module &M, GlobalVariable *GV,
uint32_t Address) {
// Write the specified address into metadata where it can be retrieved by
// the assembler. Format is a half open range, [Address Address+1)
LLVMContext &Ctx = M.getContext();
auto *IntTy = M.getDataLayout().getIntPtrType(Ctx, AMDGPUAS::LOCAL_ADDRESS);
MDBuilder MDB(Ctx);
MDNode *MetadataNode = MDB.createRange(ConstantInt::get(IntTy, Address),
ConstantInt::get(IntTy, Address + 1));
GV->setMetadata(LLVMContext::MD_absolute_symbol, MetadataNode);
}
static void addLDSSizeAttribute(Function *Func, uint32_t Offset,
bool IsDynLDS) {
if (Offset != 0) {
std::string Buffer;
raw_string_ostream SS{Buffer};
SS << Offset;
if (IsDynLDS)
SS << "," << Offset;
Func->addFnAttr("amdgpu-lds-size", Buffer);
}
}
static void markUsedByKernel(Function *Func, GlobalVariable *SGV) {
BasicBlock *Entry = &Func->getEntryBlock();
IRBuilder<> Builder(Entry, Entry->getFirstNonPHIIt());
Function *Decl = Intrinsic::getOrInsertDeclaration(Func->getParent(),
Intrinsic::donothing, {});
Value *UseInstance[1] = {
Builder.CreateConstInBoundsGEP1_32(SGV->getValueType(), SGV, 0)};
Builder.CreateCall(Decl, {},
{OperandBundleDefT<Value *>("ExplicitUse", UseInstance)});
}
void AMDGPUSwLowerLDS::buildSwLDSGlobal(Function *Func) {
// Create new LDS global required for each kernel to store
// device global memory pointer.
auto &LDSParams = FuncLDSAccessInfo.KernelToLDSParametersMap[Func];
// Create new global pointer variable
LDSParams.SwLDS = new GlobalVariable(
M, IRB.getPtrTy(), false, GlobalValue::InternalLinkage,
PoisonValue::get(IRB.getPtrTy()), "llvm.amdgcn.sw.lds." + Func->getName(),
nullptr, GlobalValue::NotThreadLocal, AMDGPUAS::LOCAL_ADDRESS, false);
GlobalValue::SanitizerMetadata MD;
MD.NoAddress = true;
LDSParams.SwLDS->setSanitizerMetadata(MD);
}
void AMDGPUSwLowerLDS::buildSwDynLDSGlobal(Function *Func) {
// Create new Dyn LDS global if kernel accesses dyn LDS.
auto &LDSParams = FuncLDSAccessInfo.KernelToLDSParametersMap[Func];
if (LDSParams.DirectAccess.DynamicLDSGlobals.empty() &&
LDSParams.IndirectAccess.DynamicLDSGlobals.empty())
return;
// Create new global pointer variable
auto *emptyCharArray = ArrayType::get(IRB.getInt8Ty(), 0);
LDSParams.SwDynLDS = new GlobalVariable(
M, emptyCharArray, false, GlobalValue::ExternalLinkage, nullptr,
"llvm.amdgcn." + Func->getName() + ".dynlds", nullptr,
GlobalValue::NotThreadLocal, AMDGPUAS::LOCAL_ADDRESS, false);
markUsedByKernel(Func, LDSParams.SwDynLDS);
GlobalValue::SanitizerMetadata MD;
MD.NoAddress = true;
LDSParams.SwDynLDS->setSanitizerMetadata(MD);
}
void AMDGPUSwLowerLDS::populateSwLDSAttributeAndMetadata(Function *Func) {
auto &LDSParams = FuncLDSAccessInfo.KernelToLDSParametersMap[Func];
bool IsDynLDSUsed = LDSParams.SwDynLDS;
uint32_t Offset = LDSParams.LDSSize;
recordLDSAbsoluteAddress(M, LDSParams.SwLDS, 0);
addLDSSizeAttribute(Func, Offset, IsDynLDSUsed);
if (LDSParams.SwDynLDS)
recordLDSAbsoluteAddress(M, LDSParams.SwDynLDS, Offset);
}
void AMDGPUSwLowerLDS::populateSwMetadataGlobal(Function *Func) {
// Create new metadata global for every kernel and initialize the
// start offsets and sizes corresponding to each LDS accesses.
auto &LDSParams = FuncLDSAccessInfo.KernelToLDSParametersMap[Func];
auto &Ctx = M.getContext();
auto &DL = M.getDataLayout();
std::vector<Type *> Items;
Type *Int32Ty = IRB.getInt32Ty();
std::vector<Constant *> Initializers;
Align MaxAlignment(1);
auto UpdateMaxAlignment = [&MaxAlignment, &DL](GlobalVariable *GV) {
Align GVAlign = AMDGPU::getAlign(DL, GV);
MaxAlignment = std::max(MaxAlignment, GVAlign);
};
for (GlobalVariable *GV : LDSParams.DirectAccess.StaticLDSGlobals)
UpdateMaxAlignment(GV);
for (GlobalVariable *GV : LDSParams.DirectAccess.DynamicLDSGlobals)
UpdateMaxAlignment(GV);
for (GlobalVariable *GV : LDSParams.IndirectAccess.StaticLDSGlobals)
UpdateMaxAlignment(GV);
for (GlobalVariable *GV : LDSParams.IndirectAccess.DynamicLDSGlobals)
UpdateMaxAlignment(GV);
//{StartOffset, AlignedSizeInBytes}
SmallString<128> MDItemStr;
raw_svector_ostream MDItemOS(MDItemStr);
MDItemOS << "llvm.amdgcn.sw.lds." << Func->getName() << ".md.item";
StructType *LDSItemTy =
StructType::create(Ctx, {Int32Ty, Int32Ty, Int32Ty}, MDItemOS.str());
uint32_t &MallocSize = LDSParams.MallocSize;
SetVector<GlobalVariable *> UniqueLDSGlobals;
int AsanScale = AsanInfo.Scale;
auto buildInitializerForSwLDSMD =
[&](SetVector<GlobalVariable *> &LDSGlobals) {
for (auto &GV : LDSGlobals) {
if (is_contained(UniqueLDSGlobals, GV))
continue;
UniqueLDSGlobals.insert(GV);
Type *Ty = GV->getValueType();
const uint64_t SizeInBytes = DL.getTypeAllocSize(Ty);
Items.push_back(LDSItemTy);
Constant *ItemStartOffset = ConstantInt::get(Int32Ty, MallocSize);
Constant *SizeInBytesConst = ConstantInt::get(Int32Ty, SizeInBytes);
// Get redzone size corresponding a size.
const uint64_t RightRedzoneSize =
AMDGPU::getRedzoneSizeForGlobal(AsanScale, SizeInBytes);
// Update MallocSize with current size and redzone size.
MallocSize += SizeInBytes;
if (!AMDGPU::isDynamicLDS(*GV))
LDSParams.RedzoneOffsetAndSizeVector.emplace_back(MallocSize,
RightRedzoneSize);
MallocSize += RightRedzoneSize;
// Align current size plus redzone.
uint64_t AlignedSize =
alignTo(SizeInBytes + RightRedzoneSize, MaxAlignment);
Constant *AlignedSizeInBytesConst =
ConstantInt::get(Int32Ty, AlignedSize);
// Align MallocSize
MallocSize = alignTo(MallocSize, MaxAlignment);
Constant *InitItem =
ConstantStruct::get(LDSItemTy, {ItemStartOffset, SizeInBytesConst,
AlignedSizeInBytesConst});
Initializers.push_back(InitItem);
}
};
SetVector<GlobalVariable *> SwLDSVector;
SwLDSVector.insert(LDSParams.SwLDS);
buildInitializerForSwLDSMD(SwLDSVector);
buildInitializerForSwLDSMD(LDSParams.DirectAccess.StaticLDSGlobals);
buildInitializerForSwLDSMD(LDSParams.IndirectAccess.StaticLDSGlobals);
buildInitializerForSwLDSMD(LDSParams.DirectAccess.DynamicLDSGlobals);
buildInitializerForSwLDSMD(LDSParams.IndirectAccess.DynamicLDSGlobals);
// Update the LDS size used by the kernel.
Type *Ty = LDSParams.SwLDS->getValueType();
const uint64_t SizeInBytes = DL.getTypeAllocSize(Ty);
uint64_t AlignedSize = alignTo(SizeInBytes, MaxAlignment);
LDSParams.LDSSize = AlignedSize;
SmallString<128> MDTypeStr;
raw_svector_ostream MDTypeOS(MDTypeStr);
MDTypeOS << "llvm.amdgcn.sw.lds." << Func->getName() << ".md.type";
StructType *MetadataStructType =
StructType::create(Ctx, Items, MDTypeOS.str());
SmallString<128> MDStr;
raw_svector_ostream MDOS(MDStr);
MDOS << "llvm.amdgcn.sw.lds." << Func->getName() << ".md";
LDSParams.SwLDSMetadata = new GlobalVariable(
M, MetadataStructType, false, GlobalValue::InternalLinkage,
PoisonValue::get(MetadataStructType), MDOS.str(), nullptr,
GlobalValue::NotThreadLocal, AMDGPUAS::GLOBAL_ADDRESS, false);
Constant *data = ConstantStruct::get(MetadataStructType, Initializers);
LDSParams.SwLDSMetadata->setInitializer(data);
assert(LDSParams.SwLDS);
// Set the alignment to MaxAlignment for SwLDS.
LDSParams.SwLDS->setAlignment(MaxAlignment);
if (LDSParams.SwDynLDS)
LDSParams.SwDynLDS->setAlignment(MaxAlignment);
GlobalValue::SanitizerMetadata MD;
MD.NoAddress = true;
LDSParams.SwLDSMetadata->setSanitizerMetadata(MD);
}
void AMDGPUSwLowerLDS::populateLDSToReplacementIndicesMap(Function *Func) {
// Fill the corresponding LDS replacement indices for each LDS access
// related to this kernel.
auto &LDSParams = FuncLDSAccessInfo.KernelToLDSParametersMap[Func];
SetVector<GlobalVariable *> UniqueLDSGlobals;
auto PopulateIndices = [&](SetVector<GlobalVariable *> &LDSGlobals,
uint32_t &Idx) {
for (auto &GV : LDSGlobals) {
if (is_contained(UniqueLDSGlobals, GV))
continue;
UniqueLDSGlobals.insert(GV);
LDSParams.LDSToReplacementIndicesMap[GV] = {0, Idx, 0};
++Idx;
}
};
uint32_t Idx = 0;
SetVector<GlobalVariable *> SwLDSVector;
SwLDSVector.insert(LDSParams.SwLDS);
PopulateIndices(SwLDSVector, Idx);
PopulateIndices(LDSParams.DirectAccess.StaticLDSGlobals, Idx);
PopulateIndices(LDSParams.IndirectAccess.StaticLDSGlobals, Idx);
PopulateIndices(LDSParams.DirectAccess.DynamicLDSGlobals, Idx);
PopulateIndices(LDSParams.IndirectAccess.DynamicLDSGlobals, Idx);
}
static void replacesUsesOfGlobalInFunction(Function *Func, GlobalVariable *GV,
Value *Replacement) {
// Replace all uses of LDS global in this Function with a Replacement.
auto ReplaceUsesLambda = [Func](const Use &U) -> bool {
auto *V = U.getUser();
if (auto *Inst = dyn_cast<Instruction>(V)) {
auto *Func1 = Inst->getParent()->getParent();
if (Func == Func1)
return true;
}
return false;
};
GV->replaceUsesWithIf(Replacement, ReplaceUsesLambda);
}
void AMDGPUSwLowerLDS::replaceKernelLDSAccesses(Function *Func) {
auto &LDSParams = FuncLDSAccessInfo.KernelToLDSParametersMap[Func];
GlobalVariable *SwLDS = LDSParams.SwLDS;
assert(SwLDS);
GlobalVariable *SwLDSMetadata = LDSParams.SwLDSMetadata;
assert(SwLDSMetadata);
StructType *SwLDSMetadataStructType =
cast<StructType>(SwLDSMetadata->getValueType());
Type *Int32Ty = IRB.getInt32Ty();
auto &IndirectAccess = LDSParams.IndirectAccess;
auto &DirectAccess = LDSParams.DirectAccess;
// Replace all uses of LDS global in this Function with a Replacement.
SetVector<GlobalVariable *> UniqueLDSGlobals;
auto ReplaceLDSGlobalUses = [&](SetVector<GlobalVariable *> &LDSGlobals) {
for (auto &GV : LDSGlobals) {
// Do not generate instructions if LDS access is in non-kernel
// i.e indirect-access.
if ((IndirectAccess.StaticLDSGlobals.contains(GV) ||
IndirectAccess.DynamicLDSGlobals.contains(GV)) &&
(!DirectAccess.StaticLDSGlobals.contains(GV) &&
!DirectAccess.DynamicLDSGlobals.contains(GV)))
continue;
if (is_contained(UniqueLDSGlobals, GV))
continue;
UniqueLDSGlobals.insert(GV);
auto &Indices = LDSParams.LDSToReplacementIndicesMap[GV];
assert(Indices.size() == 3);
Constant *GEPIdx[] = {ConstantInt::get(Int32Ty, Indices[0]),
ConstantInt::get(Int32Ty, Indices[1]),
ConstantInt::get(Int32Ty, Indices[2])};
Constant *GEP = ConstantExpr::getGetElementPtr(
SwLDSMetadataStructType, SwLDSMetadata, GEPIdx, true);
Value *Offset = IRB.CreateLoad(Int32Ty, GEP);
Value *BasePlusOffset =
IRB.CreateInBoundsGEP(IRB.getInt8Ty(), SwLDS, {Offset});
LLVM_DEBUG(GV->printAsOperand(dbgs() << "Sw LDS Lowering, Replacing LDS ",
false));
replacesUsesOfGlobalInFunction(Func, GV, BasePlusOffset);
}
};
ReplaceLDSGlobalUses(DirectAccess.StaticLDSGlobals);
ReplaceLDSGlobalUses(IndirectAccess.StaticLDSGlobals);
ReplaceLDSGlobalUses(DirectAccess.DynamicLDSGlobals);
ReplaceLDSGlobalUses(IndirectAccess.DynamicLDSGlobals);
}
void AMDGPUSwLowerLDS::updateMallocSizeForDynamicLDS(
Function *Func, Value **CurrMallocSize, Value *HiddenDynLDSSize,
SetVector<GlobalVariable *> &DynamicLDSGlobals) {
auto &LDSParams = FuncLDSAccessInfo.KernelToLDSParametersMap[Func];
Type *Int32Ty = IRB.getInt32Ty();
GlobalVariable *SwLDS = LDSParams.SwLDS;
GlobalVariable *SwLDSMetadata = LDSParams.SwLDSMetadata;
assert(SwLDS && SwLDSMetadata);
StructType *MetadataStructType =
cast<StructType>(SwLDSMetadata->getValueType());
unsigned MaxAlignment = SwLDS->getAlignment();
Value *MaxAlignValue = IRB.getInt32(MaxAlignment);
Value *MaxAlignValueMinusOne = IRB.getInt32(MaxAlignment - 1);
for (GlobalVariable *DynGV : DynamicLDSGlobals) {
auto &Indices = LDSParams.LDSToReplacementIndicesMap[DynGV];
// Update the Offset metadata.
Constant *Index0 = ConstantInt::get(Int32Ty, 0);
Constant *Index1 = ConstantInt::get(Int32Ty, Indices[1]);
Constant *Index2Offset = ConstantInt::get(Int32Ty, 0);
auto *GEPForOffset = IRB.CreateInBoundsGEP(
MetadataStructType, SwLDSMetadata, {Index0, Index1, Index2Offset});
IRB.CreateStore(*CurrMallocSize, GEPForOffset);
// Update the size and Aligned Size metadata.
Constant *Index2Size = ConstantInt::get(Int32Ty, 1);
auto *GEPForSize = IRB.CreateInBoundsGEP(MetadataStructType, SwLDSMetadata,
{Index0, Index1, Index2Size});
Value *CurrDynLDSSize = IRB.CreateLoad(Int32Ty, HiddenDynLDSSize);
IRB.CreateStore(CurrDynLDSSize, GEPForSize);
Constant *Index2AlignedSize = ConstantInt::get(Int32Ty, 2);
auto *GEPForAlignedSize = IRB.CreateInBoundsGEP(
MetadataStructType, SwLDSMetadata, {Index0, Index1, Index2AlignedSize});
Value *AlignedDynLDSSize =
IRB.CreateAdd(CurrDynLDSSize, MaxAlignValueMinusOne);
AlignedDynLDSSize = IRB.CreateUDiv(AlignedDynLDSSize, MaxAlignValue);
AlignedDynLDSSize = IRB.CreateMul(AlignedDynLDSSize, MaxAlignValue);
IRB.CreateStore(AlignedDynLDSSize, GEPForAlignedSize);
// Update the Current Malloc Size
*CurrMallocSize = IRB.CreateAdd(*CurrMallocSize, AlignedDynLDSSize);
}
}
static DebugLoc getOrCreateDebugLoc(const Instruction *InsertBefore,
DISubprogram *SP) {
assert(InsertBefore);
if (InsertBefore->getDebugLoc())
return InsertBefore->getDebugLoc();
if (SP)
return DILocation::get(SP->getContext(), SP->getLine(), 1, SP);
return DebugLoc();
}
void AMDGPUSwLowerLDS::getLDSMemoryInstructions(
Function *Func, SetVector<Instruction *> &LDSInstructions) {
for (BasicBlock &BB : *Func) {
for (Instruction &Inst : BB) {
if (LoadInst *LI = dyn_cast<LoadInst>(&Inst)) {
if (LI->getPointerAddressSpace() == AMDGPUAS::LOCAL_ADDRESS)
LDSInstructions.insert(&Inst);
} else if (StoreInst *SI = dyn_cast<StoreInst>(&Inst)) {
if (SI->getPointerAddressSpace() == AMDGPUAS::LOCAL_ADDRESS)
LDSInstructions.insert(&Inst);
} else if (AtomicRMWInst *RMW = dyn_cast<AtomicRMWInst>(&Inst)) {
if (RMW->getPointerAddressSpace() == AMDGPUAS::LOCAL_ADDRESS)
LDSInstructions.insert(&Inst);
} else if (AtomicCmpXchgInst *XCHG = dyn_cast<AtomicCmpXchgInst>(&Inst)) {
if (XCHG->getPointerAddressSpace() == AMDGPUAS::LOCAL_ADDRESS)
LDSInstructions.insert(&Inst);
} else if (AddrSpaceCastInst *ASC = dyn_cast<AddrSpaceCastInst>(&Inst)) {
if (ASC->getSrcAddressSpace() == AMDGPUAS::LOCAL_ADDRESS &&
ASC->getDestAddressSpace() == AMDGPUAS::FLAT_ADDRESS)
LDSInstructions.insert(&Inst);
} else
continue;
}
}
}
Value *AMDGPUSwLowerLDS::getTranslatedGlobalMemoryPtrOfLDS(Value *LoadMallocPtr,
Value *LDSPtr) {
assert(LDSPtr && "Invalid LDS pointer operand");
Type *LDSPtrType = LDSPtr->getType();
LLVMContext &Ctx = M.getContext();
const DataLayout &DL = M.getDataLayout();
Type *IntTy = DL.getIntPtrType(Ctx, AMDGPUAS::LOCAL_ADDRESS);
if (auto *VecPtrTy = dyn_cast<VectorType>(LDSPtrType)) {
// Handle vector of pointers
ElementCount NumElements = VecPtrTy->getElementCount();
IntTy = VectorType::get(IntTy, NumElements);
}
Value *GepIndex = IRB.CreatePtrToInt(LDSPtr, IntTy);
return IRB.CreateInBoundsGEP(IRB.getInt8Ty(), LoadMallocPtr, {GepIndex});
}
void AMDGPUSwLowerLDS::translateLDSMemoryOperationsToGlobalMemory(
Function *Func, Value *LoadMallocPtr,
SetVector<Instruction *> &LDSInstructions) {
LLVM_DEBUG(dbgs() << "Translating LDS memory operations to global memory : "
<< Func->getName());
for (Instruction *Inst : LDSInstructions) {
IRB.SetInsertPoint(Inst);
if (LoadInst *LI = dyn_cast<LoadInst>(Inst)) {
Value *LIOperand = LI->getPointerOperand();
Value *Replacement =
getTranslatedGlobalMemoryPtrOfLDS(LoadMallocPtr, LIOperand);
LoadInst *NewLI = IRB.CreateAlignedLoad(LI->getType(), Replacement,
LI->getAlign(), LI->isVolatile());
NewLI->setAtomic(LI->getOrdering(), LI->getSyncScopeID());
AsanInfo.Instructions.insert(NewLI);
LI->replaceAllUsesWith(NewLI);
LI->eraseFromParent();
} else if (StoreInst *SI = dyn_cast<StoreInst>(Inst)) {
Value *SIOperand = SI->getPointerOperand();
Value *Replacement =
getTranslatedGlobalMemoryPtrOfLDS(LoadMallocPtr, SIOperand);
StoreInst *NewSI = IRB.CreateAlignedStore(
SI->getValueOperand(), Replacement, SI->getAlign(), SI->isVolatile());
NewSI->setAtomic(SI->getOrdering(), SI->getSyncScopeID());
AsanInfo.Instructions.insert(NewSI);
SI->replaceAllUsesWith(NewSI);
SI->eraseFromParent();
} else if (AtomicRMWInst *RMW = dyn_cast<AtomicRMWInst>(Inst)) {
Value *RMWPtrOperand = RMW->getPointerOperand();
Value *RMWValOperand = RMW->getValOperand();
Value *Replacement =
getTranslatedGlobalMemoryPtrOfLDS(LoadMallocPtr, RMWPtrOperand);
AtomicRMWInst *NewRMW = IRB.CreateAtomicRMW(
RMW->getOperation(), Replacement, RMWValOperand, RMW->getAlign(),
RMW->getOrdering(), RMW->getSyncScopeID());
NewRMW->setVolatile(RMW->isVolatile());
AsanInfo.Instructions.insert(NewRMW);
RMW->replaceAllUsesWith(NewRMW);
RMW->eraseFromParent();
} else if (AtomicCmpXchgInst *XCHG = dyn_cast<AtomicCmpXchgInst>(Inst)) {
Value *XCHGPtrOperand = XCHG->getPointerOperand();
Value *Replacement =
getTranslatedGlobalMemoryPtrOfLDS(LoadMallocPtr, XCHGPtrOperand);
AtomicCmpXchgInst *NewXCHG = IRB.CreateAtomicCmpXchg(
Replacement, XCHG->getCompareOperand(), XCHG->getNewValOperand(),
XCHG->getAlign(), XCHG->getSuccessOrdering(),
XCHG->getFailureOrdering(), XCHG->getSyncScopeID());
NewXCHG->setVolatile(XCHG->isVolatile());
AsanInfo.Instructions.insert(NewXCHG);
XCHG->replaceAllUsesWith(NewXCHG);
XCHG->eraseFromParent();
} else if (AddrSpaceCastInst *ASC = dyn_cast<AddrSpaceCastInst>(Inst)) {
Value *AIOperand = ASC->getPointerOperand();
Value *Replacement =
getTranslatedGlobalMemoryPtrOfLDS(LoadMallocPtr, AIOperand);
Value *NewAI = IRB.CreateAddrSpaceCast(Replacement, ASC->getType());
// Note: No need to add the instruction to AsanInfo instructions to be
// instrumented list. FLAT_ADDRESS ptr would have been already
// instrumented by asan pass prior to this pass.
ASC->replaceAllUsesWith(NewAI);
ASC->eraseFromParent();
} else
report_fatal_error("Unimplemented LDS lowering instruction");
}
}
void AMDGPUSwLowerLDS::poisonRedzones(Function *Func, Value *MallocPtr) {
auto &LDSParams = FuncLDSAccessInfo.KernelToLDSParametersMap[Func];
Type *Int64Ty = IRB.getInt64Ty();
Type *VoidTy = IRB.getVoidTy();
FunctionCallee AsanPoisonRegion = M.getOrInsertFunction(
"__asan_poison_region",
FunctionType::get(VoidTy, {Int64Ty, Int64Ty}, false));
auto RedzonesVec = LDSParams.RedzoneOffsetAndSizeVector;
size_t VecSize = RedzonesVec.size();
for (unsigned i = 0; i < VecSize; i++) {
auto &RedzonePair = RedzonesVec[i];
uint64_t RedzoneOffset = RedzonePair.first;
uint64_t RedzoneSize = RedzonePair.second;
Value *RedzoneAddrOffset = IRB.CreateInBoundsGEP(
IRB.getInt8Ty(), MallocPtr, {IRB.getInt64(RedzoneOffset)});
Value *RedzoneAddress = IRB.CreatePtrToInt(RedzoneAddrOffset, Int64Ty);
IRB.CreateCall(AsanPoisonRegion,
{RedzoneAddress, IRB.getInt64(RedzoneSize)});
}
}
void AMDGPUSwLowerLDS::lowerKernelLDSAccesses(Function *Func,
DomTreeUpdater &DTU) {
LLVM_DEBUG(dbgs() << "Sw Lowering Kernel LDS for : " << Func->getName());
auto &LDSParams = FuncLDSAccessInfo.KernelToLDSParametersMap[Func];
auto &Ctx = M.getContext();
auto *PrevEntryBlock = &Func->getEntryBlock();
SetVector<Instruction *> LDSInstructions;
getLDSMemoryInstructions(Func, LDSInstructions);
// Create malloc block.
auto *MallocBlock = BasicBlock::Create(Ctx, "Malloc", Func, PrevEntryBlock);
// Create WIdBlock block which has instructions related to selection of
// {0,0,0} indiex work item in the work group.
auto *WIdBlock = BasicBlock::Create(Ctx, "WId", Func, MallocBlock);
IRB.SetInsertPoint(WIdBlock, WIdBlock->begin());
DebugLoc FirstDL =
getOrCreateDebugLoc(&*PrevEntryBlock->begin(), Func->getSubprogram());
IRB.SetCurrentDebugLocation(FirstDL);
Value *WIdx = IRB.CreateIntrinsic(Intrinsic::amdgcn_workitem_id_x, {});
Value *WIdy = IRB.CreateIntrinsic(Intrinsic::amdgcn_workitem_id_y, {});
Value *WIdz = IRB.CreateIntrinsic(Intrinsic::amdgcn_workitem_id_z, {});
Value *XYOr = IRB.CreateOr(WIdx, WIdy);
Value *XYZOr = IRB.CreateOr(XYOr, WIdz);
Value *WIdzCond = IRB.CreateICmpEQ(XYZOr, IRB.getInt32(0));
// All work items will branch to PrevEntryBlock except {0,0,0} index
// work item which will branch to malloc block.
IRB.CreateCondBr(WIdzCond, MallocBlock, PrevEntryBlock);
// Malloc block
IRB.SetInsertPoint(MallocBlock, MallocBlock->begin());
// If Dynamic LDS globals are accessed by the kernel,
// Get the size of dyn lds from hidden dyn_lds_size kernel arg.
// Update the corresponding metadata global entries for this dyn lds global.
GlobalVariable *SwLDS = LDSParams.SwLDS;
GlobalVariable *SwLDSMetadata = LDSParams.SwLDSMetadata;
assert(SwLDS && SwLDSMetadata);
StructType *MetadataStructType =
cast<StructType>(SwLDSMetadata->getValueType());
uint32_t MallocSize = 0;
Value *CurrMallocSize;
Type *Int32Ty = IRB.getInt32Ty();
Type *Int64Ty = IRB.getInt64Ty();
SetVector<GlobalVariable *> UniqueLDSGlobals;
auto GetUniqueLDSGlobals = [&](SetVector<GlobalVariable *> &LDSGlobals) {
for (auto &GV : LDSGlobals) {
if (is_contained(UniqueLDSGlobals, GV))
continue;
UniqueLDSGlobals.insert(GV);
}
};
GetUniqueLDSGlobals(LDSParams.DirectAccess.StaticLDSGlobals);
GetUniqueLDSGlobals(LDSParams.IndirectAccess.StaticLDSGlobals);
unsigned NumStaticLDS = 1 + UniqueLDSGlobals.size();
UniqueLDSGlobals.clear();
if (NumStaticLDS) {
auto *GEPForEndStaticLDSOffset =
IRB.CreateInBoundsGEP(MetadataStructType, SwLDSMetadata,
{ConstantInt::get(Int32Ty, 0),
ConstantInt::get(Int32Ty, NumStaticLDS - 1),
ConstantInt::get(Int32Ty, 0)});
auto *GEPForEndStaticLDSSize =
IRB.CreateInBoundsGEP(MetadataStructType, SwLDSMetadata,
{ConstantInt::get(Int32Ty, 0),
ConstantInt::get(Int32Ty, NumStaticLDS - 1),
ConstantInt::get(Int32Ty, 2)});
Value *EndStaticLDSOffset =
IRB.CreateLoad(Int32Ty, GEPForEndStaticLDSOffset);
Value *EndStaticLDSSize = IRB.CreateLoad(Int32Ty, GEPForEndStaticLDSSize);
CurrMallocSize = IRB.CreateAdd(EndStaticLDSOffset, EndStaticLDSSize);
} else
CurrMallocSize = IRB.getInt32(MallocSize);
if (LDSParams.SwDynLDS) {
if (!(AMDGPU::getAMDHSACodeObjectVersion(M) >= AMDGPU::AMDHSA_COV5))
report_fatal_error(
"Dynamic LDS size query is only supported for CO V5 and later.");
// Get size from hidden dyn_lds_size argument of kernel
Value *ImplicitArg =
IRB.CreateIntrinsic(Intrinsic::amdgcn_implicitarg_ptr, {});
Value *HiddenDynLDSSize = IRB.CreateInBoundsGEP(
ImplicitArg->getType(), ImplicitArg,
{ConstantInt::get(Int64Ty, COV5_HIDDEN_DYN_LDS_SIZE_ARG)});
UniqueLDSGlobals.clear();
GetUniqueLDSGlobals(LDSParams.DirectAccess.DynamicLDSGlobals);
GetUniqueLDSGlobals(LDSParams.IndirectAccess.DynamicLDSGlobals);
updateMallocSizeForDynamicLDS(Func, &CurrMallocSize, HiddenDynLDSSize,
UniqueLDSGlobals);
}
CurrMallocSize = IRB.CreateZExt(CurrMallocSize, Int64Ty);
// Create a call to malloc function which does device global memory allocation
// with size equals to all LDS global accesses size in this kernel.
Value *ReturnAddress =
IRB.CreateIntrinsic(Intrinsic::returnaddress, {IRB.getInt32(0)});
FunctionCallee MallocFunc = M.getOrInsertFunction(
StringRef("__asan_malloc_impl"),
FunctionType::get(Int64Ty, {Int64Ty, Int64Ty}, false));
Value *RAPtrToInt = IRB.CreatePtrToInt(ReturnAddress, Int64Ty);
Value *MallocCall = IRB.CreateCall(MallocFunc, {CurrMallocSize, RAPtrToInt});
Value *MallocPtr =
IRB.CreateIntToPtr(MallocCall, IRB.getPtrTy(AMDGPUAS::GLOBAL_ADDRESS));
// Create store of malloc to new global
IRB.CreateStore(MallocPtr, SwLDS);
// Create calls to __asan_poison_region to poison redzones.
poisonRedzones(Func, MallocPtr);
// Create branch to PrevEntryBlock
IRB.CreateBr(PrevEntryBlock);
// Create wave-group barrier at the starting of Previous entry block
Type *Int1Ty = IRB.getInt1Ty();
IRB.SetInsertPoint(PrevEntryBlock, PrevEntryBlock->begin());
auto *XYZCondPhi = IRB.CreatePHI(Int1Ty, 2, "xyzCond");
XYZCondPhi->addIncoming(IRB.getInt1(0), WIdBlock);
XYZCondPhi->addIncoming(IRB.getInt1(1), MallocBlock);
IRB.CreateIntrinsic(Intrinsic::amdgcn_s_barrier, {});
// Load malloc pointer from Sw LDS.
Value *LoadMallocPtr =
IRB.CreateLoad(IRB.getPtrTy(AMDGPUAS::GLOBAL_ADDRESS), SwLDS);
// Replace All uses of LDS globals with new LDS pointers.
replaceKernelLDSAccesses(Func);
// Replace Memory Operations on LDS with corresponding
// global memory pointers.
translateLDSMemoryOperationsToGlobalMemory(Func, LoadMallocPtr,
LDSInstructions);
auto *CondFreeBlock = BasicBlock::Create(Ctx, "CondFree", Func);
auto *FreeBlock = BasicBlock::Create(Ctx, "Free", Func);
auto *EndBlock = BasicBlock::Create(Ctx, "End", Func);
for (BasicBlock &BB : *Func) {
if (!BB.empty()) {
if (ReturnInst *RI = dyn_cast<ReturnInst>(&BB.back())) {
RI->eraseFromParent();
IRB.SetInsertPoint(&BB, BB.end());
IRB.CreateBr(CondFreeBlock);
}
}
}
// Cond Free Block
IRB.SetInsertPoint(CondFreeBlock, CondFreeBlock->begin());
IRB.CreateIntrinsic(Intrinsic::amdgcn_s_barrier, {});
IRB.CreateCondBr(XYZCondPhi, FreeBlock, EndBlock);
// Free Block
IRB.SetInsertPoint(FreeBlock, FreeBlock->begin());
// Free the previously allocate device global memory.
FunctionCallee AsanFreeFunc = M.getOrInsertFunction(
StringRef("__asan_free_impl"),
FunctionType::get(IRB.getVoidTy(), {Int64Ty, Int64Ty}, false));
Value *ReturnAddr =
IRB.CreateIntrinsic(Intrinsic::returnaddress, IRB.getInt32(0));
Value *RAPToInt = IRB.CreatePtrToInt(ReturnAddr, Int64Ty);
Value *MallocPtrToInt = IRB.CreatePtrToInt(LoadMallocPtr, Int64Ty);
IRB.CreateCall(AsanFreeFunc, {MallocPtrToInt, RAPToInt});
IRB.CreateBr(EndBlock);
// End Block
IRB.SetInsertPoint(EndBlock, EndBlock->begin());
IRB.CreateRetVoid();
// Update the DomTree with corresponding links to basic blocks.
DTU.applyUpdates({{DominatorTree::Insert, WIdBlock, MallocBlock},
{DominatorTree::Insert, MallocBlock, PrevEntryBlock},
{DominatorTree::Insert, CondFreeBlock, FreeBlock},
{DominatorTree::Insert, FreeBlock, EndBlock}});
}
Constant *AMDGPUSwLowerLDS::getAddressesOfVariablesInKernel(
Function *Func, SetVector<GlobalVariable *> &Variables) {
Type *Int32Ty = IRB.getInt32Ty();
auto &LDSParams = FuncLDSAccessInfo.KernelToLDSParametersMap[Func];
GlobalVariable *SwLDSMetadata = LDSParams.SwLDSMetadata;
assert(SwLDSMetadata);
auto *SwLDSMetadataStructType =
cast<StructType>(SwLDSMetadata->getValueType());
ArrayType *KernelOffsetsType =
ArrayType::get(IRB.getPtrTy(AMDGPUAS::GLOBAL_ADDRESS), Variables.size());
SmallVector<Constant *> Elements;
for (auto *GV : Variables) {
auto It = LDSParams.LDSToReplacementIndicesMap.find(GV);
if (It == LDSParams.LDSToReplacementIndicesMap.end()) {
Elements.push_back(
PoisonValue::get(IRB.getPtrTy(AMDGPUAS::GLOBAL_ADDRESS)));
continue;
}
auto &Indices = It->second;
Constant *GEPIdx[] = {ConstantInt::get(Int32Ty, Indices[0]),
ConstantInt::get(Int32Ty, Indices[1]),
ConstantInt::get(Int32Ty, Indices[2])};
Constant *GEP = ConstantExpr::getGetElementPtr(SwLDSMetadataStructType,
SwLDSMetadata, GEPIdx, true);
Elements.push_back(GEP);
}
return ConstantArray::get(KernelOffsetsType, Elements);
}
void AMDGPUSwLowerLDS::buildNonKernelLDSBaseTable(
NonKernelLDSParameters &NKLDSParams) {
// Base table will have single row, with elements of the row
// placed as per kernel ID. Each element in the row corresponds
// to addresss of "SW LDS" global of the kernel.
auto &Kernels = NKLDSParams.OrderedKernels;
if (Kernels.empty())
return;
Type *Int32Ty = IRB.getInt32Ty();
const size_t NumberKernels = Kernels.size();
ArrayType *AllKernelsOffsetsType =
ArrayType::get(IRB.getPtrTy(AMDGPUAS::LOCAL_ADDRESS), NumberKernels);
std::vector<Constant *> OverallConstantExprElts(NumberKernels);
for (size_t i = 0; i < NumberKernels; i++) {
Function *Func = Kernels[i];
auto &LDSParams = FuncLDSAccessInfo.KernelToLDSParametersMap[Func];
GlobalVariable *SwLDS = LDSParams.SwLDS;
assert(SwLDS);
Constant *GEPIdx[] = {ConstantInt::get(Int32Ty, 0)};
Constant *GEP =
ConstantExpr::getGetElementPtr(SwLDS->getType(), SwLDS, GEPIdx, true);
OverallConstantExprElts[i] = GEP;
}
Constant *init =
ConstantArray::get(AllKernelsOffsetsType, OverallConstantExprElts);
NKLDSParams.LDSBaseTable = new GlobalVariable(
M, AllKernelsOffsetsType, true, GlobalValue::InternalLinkage, init,
"llvm.amdgcn.sw.lds.base.table", nullptr, GlobalValue::NotThreadLocal,
AMDGPUAS::GLOBAL_ADDRESS);
GlobalValue::SanitizerMetadata MD;
MD.NoAddress = true;
NKLDSParams.LDSBaseTable->setSanitizerMetadata(MD);
}
void AMDGPUSwLowerLDS::buildNonKernelLDSOffsetTable(
NonKernelLDSParameters &NKLDSParams) {
// Offset table will have multiple rows and columns.
// Rows are assumed to be from 0 to (n-1). n is total number
// of kernels accessing the LDS through non-kernels.
// Each row will have m elements. m is the total number of
// unique LDS globals accessed by non-kernels.
// Each element in the row correspond to the address of
// the replacement of LDS global done by that particular kernel.
auto &Variables = NKLDSParams.OrdereLDSGlobals;
auto &Kernels = NKLDSParams.OrderedKernels;
if (Variables.empty() || Kernels.empty())
return;
const size_t NumberVariables = Variables.size();
const size_t NumberKernels = Kernels.size();
ArrayType *KernelOffsetsType =
ArrayType::get(IRB.getPtrTy(AMDGPUAS::GLOBAL_ADDRESS), NumberVariables);
ArrayType *AllKernelsOffsetsType =
ArrayType::get(KernelOffsetsType, NumberKernels);
std::vector<Constant *> overallConstantExprElts(NumberKernels);
for (size_t i = 0; i < NumberKernels; i++) {
Function *Func = Kernels[i];
overallConstantExprElts[i] =
getAddressesOfVariablesInKernel(Func, Variables);
}
Constant *Init =
ConstantArray::get(AllKernelsOffsetsType, overallConstantExprElts);
NKLDSParams.LDSOffsetTable = new GlobalVariable(
M, AllKernelsOffsetsType, true, GlobalValue::InternalLinkage, Init,
"llvm.amdgcn.sw.lds.offset.table", nullptr, GlobalValue::NotThreadLocal,
AMDGPUAS::GLOBAL_ADDRESS);
GlobalValue::SanitizerMetadata MD;
MD.NoAddress = true;
NKLDSParams.LDSOffsetTable->setSanitizerMetadata(MD);
}
void AMDGPUSwLowerLDS::lowerNonKernelLDSAccesses(
Function *Func, SetVector<GlobalVariable *> &LDSGlobals,
NonKernelLDSParameters &NKLDSParams) {
// Replace LDS access in non-kernel with replacement queried from
// Base table and offset from offset table.
LLVM_DEBUG(dbgs() << "Sw LDS lowering, lower non-kernel access for : "
<< Func->getName());
auto InsertAt = Func->getEntryBlock().getFirstNonPHIOrDbgOrAlloca();
IRB.SetInsertPoint(InsertAt);
// Get LDS memory instructions.
SetVector<Instruction *> LDSInstructions;
getLDSMemoryInstructions(Func, LDSInstructions);
auto *KernelId = IRB.CreateIntrinsic(Intrinsic::amdgcn_lds_kernel_id, {});
GlobalVariable *LDSBaseTable = NKLDSParams.LDSBaseTable;
GlobalVariable *LDSOffsetTable = NKLDSParams.LDSOffsetTable;
auto &OrdereLDSGlobals = NKLDSParams.OrdereLDSGlobals;
Value *BaseGEP = IRB.CreateInBoundsGEP(
LDSBaseTable->getValueType(), LDSBaseTable, {IRB.getInt32(0), KernelId});
Value *BaseLoad =
IRB.CreateLoad(IRB.getPtrTy(AMDGPUAS::LOCAL_ADDRESS), BaseGEP);
Value *LoadMallocPtr =
IRB.CreateLoad(IRB.getPtrTy(AMDGPUAS::GLOBAL_ADDRESS), BaseLoad);
for (GlobalVariable *GV : LDSGlobals) {
const auto *GVIt = llvm::find(OrdereLDSGlobals, GV);
assert(GVIt != OrdereLDSGlobals.end());
uint32_t GVOffset = std::distance(OrdereLDSGlobals.begin(), GVIt);
Value *OffsetGEP = IRB.CreateInBoundsGEP(
LDSOffsetTable->getValueType(), LDSOffsetTable,
{IRB.getInt32(0), KernelId, IRB.getInt32(GVOffset)});
Value *OffsetLoad =
IRB.CreateLoad(IRB.getPtrTy(AMDGPUAS::GLOBAL_ADDRESS), OffsetGEP);
Value *Offset = IRB.CreateLoad(IRB.getInt32Ty(), OffsetLoad);
Value *BasePlusOffset =
IRB.CreateInBoundsGEP(IRB.getInt8Ty(), BaseLoad, {Offset});
LLVM_DEBUG(dbgs() << "Sw LDS Lowering, Replace non-kernel LDS for "
<< GV->getName());
replacesUsesOfGlobalInFunction(Func, GV, BasePlusOffset);
}
translateLDSMemoryOperationsToGlobalMemory(Func, LoadMallocPtr,
LDSInstructions);
}
static void reorderStaticDynamicIndirectLDSSet(KernelLDSParameters &LDSParams) {
// Sort Static, dynamic LDS globals which are either
// direct or indirect access on basis of name.
auto &DirectAccess = LDSParams.DirectAccess;
auto &IndirectAccess = LDSParams.IndirectAccess;
LDSParams.DirectAccess.StaticLDSGlobals = sortByName(
std::vector<GlobalVariable *>(DirectAccess.StaticLDSGlobals.begin(),
DirectAccess.StaticLDSGlobals.end()));
LDSParams.DirectAccess.DynamicLDSGlobals = sortByName(
std::vector<GlobalVariable *>(DirectAccess.DynamicLDSGlobals.begin(),
DirectAccess.DynamicLDSGlobals.end()));
LDSParams.IndirectAccess.StaticLDSGlobals = sortByName(
std::vector<GlobalVariable *>(IndirectAccess.StaticLDSGlobals.begin(),
IndirectAccess.StaticLDSGlobals.end()));
LDSParams.IndirectAccess.DynamicLDSGlobals = sortByName(
std::vector<GlobalVariable *>(IndirectAccess.DynamicLDSGlobals.begin(),
IndirectAccess.DynamicLDSGlobals.end()));
}
void AMDGPUSwLowerLDS::initAsanInfo() {
// Get Shadow mapping scale and offset.
unsigned LongSize =
M.getDataLayout().getPointerSizeInBits(AMDGPUAS::GLOBAL_ADDRESS);
uint64_t Offset;
int Scale;
bool OrShadowOffset;
llvm::getAddressSanitizerParams(AMDGPUTM.getTargetTriple(), LongSize, false,
&Offset, &Scale, &OrShadowOffset);
AsanInfo.Scale = Scale;
AsanInfo.Offset = Offset;
}
static bool hasFnWithSanitizeAddressAttr(FunctionVariableMap &LDSAccesses) {
for (auto &K : LDSAccesses) {
Function *F = K.first;
if (!F)
continue;
if (F->hasFnAttribute(Attribute::SanitizeAddress))
return true;
}
return false;
}
bool AMDGPUSwLowerLDS::run() {
bool Changed = false;
CallGraph CG = CallGraph(M);
Changed |= eliminateConstantExprUsesOfLDSFromAllInstructions(M);
// Get all the direct and indirect access of LDS for all the kernels.
LDSUsesInfoTy LDSUsesInfo = getTransitiveUsesOfLDS(CG, M);
// Flag to decide whether to lower all the LDS accesses
// based on sanitize_address attribute.
bool LowerAllLDS = hasFnWithSanitizeAddressAttr(LDSUsesInfo.direct_access) ||
hasFnWithSanitizeAddressAttr(LDSUsesInfo.indirect_access);
if (!LowerAllLDS)
return Changed;
// Utility to group LDS access into direct, indirect, static and dynamic.
auto PopulateKernelStaticDynamicLDS = [&](FunctionVariableMap &LDSAccesses,
bool DirectAccess) {
for (auto &K : LDSAccesses) {
Function *F = K.first;
if (!F || K.second.empty())
continue;
assert(isKernelLDS(F));
// Only inserts if key isn't already in the map.
FuncLDSAccessInfo.KernelToLDSParametersMap.insert(
{F, KernelLDSParameters()});
auto &LDSParams = FuncLDSAccessInfo.KernelToLDSParametersMap[F];
if (!DirectAccess)
FuncLDSAccessInfo.KernelsWithIndirectLDSAccess.insert(F);
for (GlobalVariable *GV : K.second) {
if (!DirectAccess) {
if (AMDGPU::isDynamicLDS(*GV))
LDSParams.IndirectAccess.DynamicLDSGlobals.insert(GV);
else
LDSParams.IndirectAccess.StaticLDSGlobals.insert(GV);
FuncLDSAccessInfo.AllNonKernelLDSAccess.insert(GV);
} else {
if (AMDGPU::isDynamicLDS(*GV))
LDSParams.DirectAccess.DynamicLDSGlobals.insert(GV);
else
LDSParams.DirectAccess.StaticLDSGlobals.insert(GV);
}
}
}
};
PopulateKernelStaticDynamicLDS(LDSUsesInfo.direct_access, true);
PopulateKernelStaticDynamicLDS(LDSUsesInfo.indirect_access, false);
// Get address sanitizer scale.
initAsanInfo();
for (auto &K : FuncLDSAccessInfo.KernelToLDSParametersMap) {
Function *Func = K.first;
auto &LDSParams = FuncLDSAccessInfo.KernelToLDSParametersMap[Func];
if (LDSParams.DirectAccess.StaticLDSGlobals.empty() &&
LDSParams.DirectAccess.DynamicLDSGlobals.empty() &&
LDSParams.IndirectAccess.StaticLDSGlobals.empty() &&
LDSParams.IndirectAccess.DynamicLDSGlobals.empty()) {
Changed = false;
} else {
removeFnAttrFromReachable(
CG, Func,
{"amdgpu-no-workitem-id-x", "amdgpu-no-workitem-id-y",
"amdgpu-no-workitem-id-z", "amdgpu-no-heap-ptr"});
if (!LDSParams.IndirectAccess.StaticLDSGlobals.empty() ||
!LDSParams.IndirectAccess.DynamicLDSGlobals.empty())
removeFnAttrFromReachable(CG, Func, {"amdgpu-no-lds-kernel-id"});
reorderStaticDynamicIndirectLDSSet(LDSParams);
buildSwLDSGlobal(Func);
buildSwDynLDSGlobal(Func);
populateSwMetadataGlobal(Func);
populateSwLDSAttributeAndMetadata(Func);
populateLDSToReplacementIndicesMap(Func);
DomTreeUpdater DTU(DTCallback(*Func),
DomTreeUpdater::UpdateStrategy::Lazy);
lowerKernelLDSAccesses(Func, DTU);
Changed = true;
}
}
// Get the Uses of LDS from non-kernels.
getUsesOfLDSByNonKernels();
// Get non-kernels with LDS ptr as argument and called by kernels.
getNonKernelsWithLDSArguments(CG);
// Lower LDS accesses in non-kernels.
if (!FuncLDSAccessInfo.NonKernelToLDSAccessMap.empty() ||
!FuncLDSAccessInfo.NonKernelsWithLDSArgument.empty()) {
NonKernelLDSParameters NKLDSParams;
NKLDSParams.OrderedKernels = getOrderedIndirectLDSAccessingKernels(
FuncLDSAccessInfo.KernelsWithIndirectLDSAccess);
NKLDSParams.OrdereLDSGlobals = getOrderedNonKernelAllLDSGlobals(
FuncLDSAccessInfo.AllNonKernelLDSAccess);
buildNonKernelLDSBaseTable(NKLDSParams);
buildNonKernelLDSOffsetTable(NKLDSParams);
for (auto &K : FuncLDSAccessInfo.NonKernelToLDSAccessMap) {
Function *Func = K.first;
DenseSet<GlobalVariable *> &LDSGlobals = K.second;
SetVector<GlobalVariable *> OrderedLDSGlobals = sortByName(
std::vector<GlobalVariable *>(LDSGlobals.begin(), LDSGlobals.end()));
lowerNonKernelLDSAccesses(Func, OrderedLDSGlobals, NKLDSParams);
}
for (Function *Func : FuncLDSAccessInfo.NonKernelsWithLDSArgument) {
auto &K = FuncLDSAccessInfo.NonKernelToLDSAccessMap;
if (K.contains(Func))
continue;
SetVector<llvm::GlobalVariable *> Vec;
lowerNonKernelLDSAccesses(Func, Vec, NKLDSParams);
}
Changed = true;
}
if (!Changed)
return Changed;
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();
}
}
if (AsanInstrumentLDS) {
SmallVector<InterestingMemoryOperand, 16> OperandsToInstrument;
for (Instruction *Inst : AsanInfo.Instructions) {
SmallVector<InterestingMemoryOperand, 1> InterestingOperands;
getInterestingMemoryOperands(M, Inst, InterestingOperands);
llvm::append_range(OperandsToInstrument, InterestingOperands);
}
for (auto &Operand : OperandsToInstrument) {
Value *Addr = Operand.getPtr();
instrumentAddress(M, IRB, Operand.getInsn(), Operand.getInsn(), Addr,
Operand.Alignment.valueOrOne(), Operand.TypeStoreSize,
Operand.IsWrite, nullptr, false, false, AsanInfo.Scale,
AsanInfo.Offset);
Changed = true;
}
}
return Changed;
}
class AMDGPUSwLowerLDSLegacy : public ModulePass {
public:
const AMDGPUTargetMachine *AMDGPUTM;
static char ID;
AMDGPUSwLowerLDSLegacy(const AMDGPUTargetMachine *TM)
: ModulePass(ID), AMDGPUTM(TM) {}
bool runOnModule(Module &M) override;
void getAnalysisUsage(AnalysisUsage &AU) const override {
AU.addPreserved<DominatorTreeWrapperPass>();
}
};
} // namespace
char AMDGPUSwLowerLDSLegacy::ID = 0;
char &llvm::AMDGPUSwLowerLDSLegacyPassID = AMDGPUSwLowerLDSLegacy::ID;
INITIALIZE_PASS_BEGIN(AMDGPUSwLowerLDSLegacy, "amdgpu-sw-lower-lds",
"AMDGPU Software lowering of LDS", false, false)
INITIALIZE_PASS_DEPENDENCY(TargetPassConfig)
INITIALIZE_PASS_END(AMDGPUSwLowerLDSLegacy, "amdgpu-sw-lower-lds",
"AMDGPU Software lowering of LDS", false, false)
bool AMDGPUSwLowerLDSLegacy::runOnModule(Module &M) {
// AddressSanitizer pass adds "nosanitize_address" module flag if it has
// instrumented the IR. Return early if the flag is not present.
if (!M.getModuleFlag("nosanitize_address"))
return false;
DominatorTreeWrapperPass *const DTW =
getAnalysisIfAvailable<DominatorTreeWrapperPass>();
auto DTCallback = [&DTW](Function &F) -> DominatorTree * {
return DTW ? &DTW->getDomTree() : nullptr;
};
if (!AMDGPUTM) {
auto &TPC = getAnalysis<TargetPassConfig>();
AMDGPUTM = &TPC.getTM<AMDGPUTargetMachine>();
}
AMDGPUSwLowerLDS SwLowerLDSImpl(M, *AMDGPUTM, DTCallback);
bool IsChanged = SwLowerLDSImpl.run();
return IsChanged;
}
ModulePass *
llvm::createAMDGPUSwLowerLDSLegacyPass(const AMDGPUTargetMachine *TM) {
return new AMDGPUSwLowerLDSLegacy(TM);
}
PreservedAnalyses AMDGPUSwLowerLDSPass::run(Module &M,
ModuleAnalysisManager &AM) {
// AddressSanitizer pass adds "nosanitize_address" module flag if it has
// instrumented the IR. Return early if the flag is not present.
if (!M.getModuleFlag("nosanitize_address"))
return PreservedAnalyses::all();
auto &FAM = AM.getResult<FunctionAnalysisManagerModuleProxy>(M).getManager();
auto DTCallback = [&FAM](Function &F) -> DominatorTree * {
return &FAM.getResult<DominatorTreeAnalysis>(F);
};
AMDGPUSwLowerLDS SwLowerLDSImpl(M, TM, DTCallback);
bool IsChanged = SwLowerLDSImpl.run();
if (!IsChanged)
return PreservedAnalyses::all();
PreservedAnalyses PA;
PA.preserve<DominatorTreeAnalysis>();
return PA;
}