Use IndexMap for PrelaodKernArgs.

Author: kerbowa

llvm/lib/Target/AMDGPU/AMDGPUArgumentUsageInfo.cpp

@@ -182,16 +182,12 @@ AMDGPUFunctionArgInfo AMDGPUFunctionArgInfo::fixedABILayout() {
 SmallVector<const KernArgPreloadDescriptor *, 4>
 AMDGPUFunctionArgInfo::getPreloadDescriptorsForArgIdx(unsigned ArgIdx) const {
   SmallVector<const KernArgPreloadDescriptor *, 4> Results;
-  for (const auto &KV : PreloadKernArgs) {
-    if (KV.second.OrigArgIdx == ArgIdx)
-      Results.push_back(&KV.second);
+  for (unsigned PartIdx = 0; PartIdx < PreloadKernArgs.size(); ++PartIdx) {
+    const auto &Desc = PreloadKernArgs[PartIdx];
+    if (Desc.OrigArgIdx == ArgIdx)
+      Results.push_back(&Desc);
   }
 
-  stable_sort(Results, [](const KernArgPreloadDescriptor *A,
-                          const KernArgPreloadDescriptor *B) {
-    return A->PartIdx < B->PartIdx;
-  });
-
   return Results;
 }
 
@@ -203,11 +199,9 @@ AMDGPUFunctionArgInfo::getHiddenArgPreloadDescriptor(HiddenArg HA) const {
   if (HiddenArgIt == PreloadHiddenArgsIndexMap.end())
     return nullptr;
 
-  auto KernArgIt = PreloadKernArgs.find(HiddenArgIt->second);
-  if (KernArgIt == PreloadKernArgs.end())
-    return nullptr;
-
-  return &KernArgIt->second;
+  const KernArgPreloadDescriptor &Desc = PreloadKernArgs[HiddenArgIt->second];
+  assert(Desc.IsValid && "Hidden argument preload descriptor not valid.");
+  return &Desc;
 }
 
 const AMDGPUFunctionArgInfo &

llvm/lib/Target/AMDGPU/AMDGPUArgumentUsageInfo.h

@@ -11,6 +11,7 @@
 
 #include "MCTargetDesc/AMDGPUMCTargetDesc.h"
 #include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/IndexedMap.h"
 #include "llvm/Analysis/ValueTracking.h"
 #include "llvm/CodeGen/Register.h"
 #include "llvm/IR/LLVMContext.h"
@@ -167,6 +168,8 @@ struct KernArgPreloadDescriptor {
   // The registers that the argument is preloaded into. The argument may be
   // split across multiple registers.
   SmallVector<MCRegister, 2> Regs;
+
+  bool IsValid = false;
 };
 
 } // namespace KernArgPreload
@@ -231,9 +234,15 @@ struct AMDGPUFunctionArgInfo {
   ArgDescriptor WorkItemIDY;
   ArgDescriptor WorkItemIDZ;
 
+  struct PreloadArgIndexFunctor {
+    using argument_type = unsigned;
+    unsigned operator()(unsigned Idx) const { return Idx; }
+  };
+
   // Map the index of preloaded kernel arguments to its descriptor.
-  SmallDenseMap<int, KernArgPreload::KernArgPreloadDescriptor>
-      PreloadKernArgs{};
+  IndexedMap<KernArgPreload::KernArgPreloadDescriptor, PreloadArgIndexFunctor>
+      PreloadKernArgs;
+
   // Map hidden argument to the index of it's descriptor.
   SmallDenseMap<KernArgPreload::HiddenArg, int> PreloadHiddenArgsIndexMap{};
   // The first user SGPR allocated for kernarg preloading.

llvm/lib/Target/AMDGPU/SIISelLowering.cpp

@@ -2745,17 +2745,20 @@ void SITargetLowering::allocatePreloadKernArgSGPRs(
       // Arg is preloaded into the previous SGPR.
       if (ArgLoc.getLocVT().getStoreSize() < 4 && Alignment < 4) {
         assert(InIdx >= 1 && "No previous SGPR");
-        auto [It, Inserted] =
-            Info.getArgInfo().PreloadKernArgs.try_emplace(InIdx);
-        assert(Inserted && "Preload kernel argument allocated twice.");
-        KernArgPreloadDescriptor &PreloadDesc = It->second;
-
-        const KernArgPreloadDescriptor &PrevDesc =
-            Info.getArgInfo().PreloadKernArgs[InIdx - 1];
+        auto &PreloadKernArgs = Info.getArgInfo().PreloadKernArgs;
+        PreloadKernArgs.grow(InIdx);
+        KernArgPreloadDescriptor &PreloadDesc = PreloadKernArgs[InIdx];
+        assert(!PreloadDesc.IsValid &&
+               "Preload kernel argument allocated twice.");
+
+        const KernArgPreloadDescriptor &PrevDesc = PreloadKernArgs[InIdx - 1];
+        assert(PrevDesc.IsValid &&
+               "Previous preload kernel argument not allocated.");
         PreloadDesc.Regs.push_back(PrevDesc.Regs[0]);
 
         PreloadDesc.OrigArgIdx = Arg.getArgNo();
         PreloadDesc.PartIdx = InIdx;
+        PreloadDesc.IsValid = true;
         if (Arg.hasAttribute("amdgpu-hidden-argument"))
           mapHiddenArgToPreloadIndex(Info.getArgInfo(), ArgOffset,
                                      ImplicitArgOffset, InIdx);
@@ -3183,7 +3186,9 @@ SDValue SITargetLowering::LowerFormalArguments(
       }
 
       SDValue NewArg;
-      if (Arg.isOrigArg() && Info->getArgInfo().PreloadKernArgs.count(i)) {
+      auto &PreloadKernArgs = Info->getArgInfo().PreloadKernArgs;
+      if (Arg.isOrigArg() && PreloadKernArgs.inBounds(i) &&
+          PreloadKernArgs[i].IsValid) {
         if (MemVT.getStoreSize() < 4 && Alignment < 4) {
           // In this case the argument is packed into the previous preload SGPR.
           int64_t AlignDownOffset = alignDown(Offset, 4);
@@ -3193,8 +3198,7 @@ SDValue SITargetLowering::LowerFormalArguments(
           const SIMachineFunctionInfo *Info =
               MF.getInfo<SIMachineFunctionInfo>();
           MachineRegisterInfo &MRI = DAG.getMachineFunction().getRegInfo();
-          Register Reg =
-              Info->getArgInfo().PreloadKernArgs.find(i)->getSecond().Regs[0];
+          Register Reg = Info->getArgInfo().PreloadKernArgs[i].Regs[0];
 
           assert(Reg);
           Register VReg = MRI.getLiveInVirtReg(Reg);
@@ -3214,7 +3218,7 @@ SDValue SITargetLowering::LowerFormalArguments(
               MF.getInfo<SIMachineFunctionInfo>();
           MachineRegisterInfo &MRI = DAG.getMachineFunction().getRegInfo();
           const SmallVectorImpl<MCRegister> &PreloadRegs =
-              Info->getArgInfo().PreloadKernArgs.find(i)->getSecond().Regs;
+              Info->getArgInfo().PreloadKernArgs[i].Regs;
 
           SDValue Copy;
           if (PreloadRegs.size() == 1) {

llvm/lib/Target/AMDGPU/SIMachineFunctionInfo.cpp

@@ -278,11 +278,13 @@ SmallVectorImpl<MCRegister> *SIMachineFunctionInfo::addPreloadedKernArg(
     const SIRegisterInfo &TRI, const TargetRegisterClass *RC,
     unsigned AllocSizeDWord, unsigned PartIdx, unsigned ArgIdx,
     unsigned PaddingSGPRs) {
-  auto [It, Inserted] = ArgInfo.PreloadKernArgs.try_emplace(PartIdx);
-  assert(Inserted && "Preload kernel argument allocated twice.");
-  KernArgPreload::KernArgPreloadDescriptor &PreloadDesc = It->second;
+  ArgInfo.PreloadKernArgs.grow(PartIdx);
+  KernArgPreload::KernArgPreloadDescriptor &PreloadDesc =
+      ArgInfo.PreloadKernArgs[PartIdx];
+  assert(!PreloadDesc.IsValid && "Preload kernel argument allocated twice.");
   PreloadDesc.PartIdx = PartIdx;
   PreloadDesc.OrigArgIdx = ArgIdx;
+  PreloadDesc.IsValid = true;
 
   NumUserSGPRs += PaddingSGPRs;
   // If the available register tuples are aligned with the kernarg to be