[llvm] [AMDGPU] Introduce "amdgpu-sw-lower-lds" pass to lower LDS accesses. (PR #87265)
Jon Chesterfield via llvm-commits
llvm-commits at lists.llvm.org
Tue Aug 13 04:39:10 PDT 2024
================
@@ -0,0 +1,1335 @@
+//===-- 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.
+//
+// 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
+// with in 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 "AMDGPUTargetMachine.h"
+#include "Utils/AMDGPUMemoryUtils.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/DenseSet.h"
+#include "llvm/ADT/SetOperations.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 infor 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 *getTranslatedGlobalMemoryGEPOfLDSPointer(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>(V.begin(), V.end())};
+}
+
+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 std::move(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)
+ 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->hasFnAttribute(Attribute::SanitizeAddress))
+ 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);
+ auto *MinC = ConstantAsMetadata::get(ConstantInt::get(IntTy, Address));
+ auto *MaxC = ConstantAsMetadata::get(ConstantInt::get(IntTy, Address + 1));
+ GV->setMetadata(LLVMContext::MD_absolute_symbol,
+ MDNode::get(Ctx, {MinC, MaxC}));
+}
+
+static void addLDSSizeAttribute(Function *Func, uint32_t Offset,
+ bool IsDynLDS) {
+ if (Offset != 0) {
+ std::string Buffer;
+ raw_string_ostream SS{Buffer};
+ SS << format("%u", Offset);
+ if (IsDynLDS)
+ SS << format(",%u", 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::getDeclaration(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);
+ return;
+}
+
+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);
+ return;
+}
+
+void AMDGPUSwLowerLDS::populateSwLDSAttributeAndMetadata(Function *Func) {
+ auto &LDSParams = FuncLDSAccessInfo.KernelToLDSParametersMap[Func];
+ bool IsDynLDSUsed = LDSParams.SwDynLDS ? true : false;
+ 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 (std::find(UniqueLDSGlobals.begin(), UniqueLDSGlobals.end(), GV) !=
+ UniqueLDSGlobals.end())
+ continue;
+ else
+ 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);
+ return;
+}
+
+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 (std::find(UniqueLDSGlobals.begin(), UniqueLDSGlobals.end(), GV) !=
+ UniqueLDSGlobals.end())
+ continue;
+ else
+ 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);
+ return;
+}
+
+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 (std::find(UniqueLDSGlobals.begin(), UniqueLDSGlobals.end(), GV) !=
+ UniqueLDSGlobals.end())
+ continue;
+ else
+ UniqueLDSGlobals.insert(GV);
+ auto &Indices = LDSParams.LDSToReplacementIndicesMap[GV];
+ assert(Indices.size() == 3);
+ uint32_t Idx0 = Indices[0];
+ uint32_t Idx1 = Indices[1];
+ uint32_t Idx2 = Indices[2];
+ Constant *GEPIdx[] = {ConstantInt::get(Int32Ty, Idx0),
+ ConstantInt::get(Int32Ty, Idx1),
+ ConstantInt::get(Int32Ty, Idx2)};
+ Constant *GEP = ConstantExpr::getGetElementPtr(
+ SwLDSMetadataStructType, SwLDSMetadata, GEPIdx, true);
+ Value *Offset = IRB.CreateLoad(Int32Ty, GEP);
+ Value *BasePlusOffset =
+ IRB.CreateInBoundsGEP(IRB.getInt8Ty(), SwLDS, {Offset});
+ LLVM_DEBUG(dbgs() << "Sw LDS Lowering, Replacing LDS " << GV->getName());
+ 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) {
----------------
JonChesterfield wrote:
This is weird - we already know what instructions access LDS, we walked the uses of addrspace(3) variables at the start of the pass, why is there this heuristic search going on as well?
https://github.com/llvm/llvm-project/pull/87265
More information about the llvm-commits
mailing list