[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
Wed Aug 14 05:24:55 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) {
+ 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
+ continue;
+ }
+ }
+}
+
+Value *
+AMDGPUSwLowerLDS::getTranslatedGlobalMemoryGEPOfLDSPointer(Value *LoadMallocPtr,
+ Value *LDSPtr) {
+ assert(LDSPtr && "Invalid LDS pointer operand");
+ Value *PtrToInt = IRB.CreatePtrToInt(LDSPtr, IRB.getInt32Ty());
+ Value *GEP =
+ IRB.CreateInBoundsGEP(IRB.getInt8Ty(), LoadMallocPtr, {PtrToInt});
+ return GEP;
+}
+
+void AMDGPUSwLowerLDS::translateLDSMemoryOperationsToGlobalMemory(
+ Function *Func, Value *LoadMallocPtr,
+ SetVector<Instruction *> &LDSInstructions) {
+ LLVM_DEBUG(dbgs() << "Translating LDS memory operations to global memory : "
+ << Func->getName());
----------------
JonChesterfield wrote:
Alright, so this is straight broken.
What we do is rewrite some instructions to point at the newly heap allocated memory but not all of them. There is no magic aliasing which keeps the malloc block up to date with the LDS one.
So when a function accesses LDS and some of the uses get pointed to the heap and others don't, that's a straight miscompile.
https://github.com/llvm/llvm-project/pull/87265
More information about the llvm-commits
mailing list