[llvm] [AMDGPU][GlobalISel] Support lowering preloaded kernel arguments (PR #205049)

[Matt Arsenault via llvm-commits]

================
@@ -296,6 +297,147 @@ struct AMDGPUOutgoingArgHandler : public AMDGPUOutgoingValueHandler {
     assignValueToAddress(ValVReg, Addr, MemTy, MPO, VA);
   }
 };
+
+struct KernArgPart {
+  CallLowering::ArgInfo Arg;
+  uint64_t Offset;
+  Align Alignment;
+};
+
+// Materialize a preloaded kernarg from its assigned SGPRs into DstReg.
+static void lowerPreloadedKernArg(MachineIRBuilder &B, Register DstReg,
+                                  LLT DstTy, uint64_t Offset, Align Alignment,
+                                  ArrayRef<MCRegister> PreloadRegs) {
+  assert(!PreloadRegs.empty());
+
+  // Kernarg preloads are assigned in 32-bit SGPR units.
+  const LLT S32 = LLT::scalar(32);
+  auto BuildCopyFromPreload = [&B](LLT Ty, Register PhysReg) {
+    B.getMBB().addLiveIn(PhysReg.asMCReg());
+    return B.buildCopy(Ty, PhysReg).getReg(0);
+  };
+  auto getRawBitsTy = [](LLT Ty) { return LLT::scalar(Ty.getSizeInBits()); };
+  auto getPreloadStorageTy = [](LLT Ty) {
+    unsigned RoundedBits = alignTo(Ty.getSizeInBits(), 32);
+    if (!Ty.isPointer() && !Ty.isVector() && RoundedBits == Ty.getSizeInBits())
+      return Ty;
+    return LLT::scalar(RoundedBits);
+  };
+
+  Register Value;
+  LLT ValueTy;
+  const bool IsPackedSubDword = DstTy.getSizeInBits() < 32 && Alignment < 4;
+
+  if (PreloadRegs.size() == 1) {
+    if (IsPackedSubDword) {
+      // Extract sub-dword preloads from their containing 32-bit SGPR word.
+      Register Raw = BuildCopyFromPreload(S32, Register(PreloadRegs[0]));
+      uint64_t OffsetDiff = Offset - alignDown(Offset, 4);
+      Register ShiftAmt = B.buildConstant(S32, OffsetDiff * 8).getReg(0);
+      Value = B.buildLShr(S32, Raw, ShiftAmt).getReg(0);
+      ValueTy = S32;
+    } else {
+      // The physical register may cover more bits than the IR value uses.
+      ValueTy = getPreloadStorageTy(DstTy);
+      Value = BuildCopyFromPreload(ValueTy, Register(PreloadRegs[0]));
+    }
+  } else {
+    assert(!IsPackedSubDword && "packed sub-dword preload should use one SGPR");
+    // Reconstruct wider preloads from separate 32-bit SGPR entries.
+    SmallVector<Register, 4> Regs;
+    Regs.reserve(PreloadRegs.size());
+    for (MCRegister Reg : PreloadRegs) {
+      Regs.push_back(BuildCopyFromPreload(S32, Register(Reg)));
+    }
+
+    ValueTy = LLT::scalar(PreloadRegs.size() * 32);
+    Value = B.buildMergeLikeInstr(ValueTy, Regs).getReg(0);
+  }
+
+  LLT AdjustedTy =
+      (DstTy.isPointer() || DstTy.isVector()) ? getRawBitsTy(DstTy) : DstTy;
+  if (ValueTy != AdjustedTy)
+    Value = B.buildAnyExtOrTrunc(AdjustedTy, Value).getReg(0);
+
+  if (DstTy.isPointer())
+    Value = B.buildIntToPtr(DstTy, Value).getReg(0);
+  else if (DstTy.isVector())
+    Value = B.buildBitcast(DstTy, Value).getReg(0);
+
+  B.buildCopy(DstReg, Value);
+}
+
+// Assign SGPRs for the contiguous sequence of kernel argument parts marked for
+// kernarg preload.
+static void allocatePreloadKernArgSGPRs(CCState &CCInfo,
+                                        ArrayRef<KernArgPart> KernArgParts,
+                                        MachineFunction &MF,
+                                        const GCNSubtarget &Subtarget,
+                                        const SIRegisterInfo &TRI,
+                                        SIMachineFunctionInfo &Info) {
+  const DataLayout &DL = MF.getFunction().getDataLayout();
+  unsigned LastExplicitArgOffset = Subtarget.getExplicitKernelArgOffset();
+  GCNUserSGPRUsageInfo &SGPRInfo = Info.getUserSGPRInfo();
+  bool AlignedForImplicitArgs = false;
+  unsigned ImplicitArgOffset = 0;
+
+  for (auto [KernArgIdx, Part] : enumerate(KernArgParts)) {
+    const Argument &Arg = *cast<Argument>(Part.Arg.OrigValue);
+    // Hardware preloads a contiguous prefix of the kernarg segment. Once a
+    // non-preloaded argument is reached, no later argument can be preloaded.
+    if (!Arg.hasInRegAttr())
+      break;
+
+    unsigned ArgOffset = Part.Offset;
+    unsigned NumAllocSGPRs = alignTo(DL.getTypeAllocSize(Part.Arg.Ty), 4) / 4;
+
+    if (Arg.hasAttribute("amdgpu-hidden-argument")) {
+      // Hidden arguments are laid out after the explicit kernargs, aligned to
+      // the ABI-required implicit argument block boundary.
+      if (!AlignedForImplicitArgs) {
+        ImplicitArgOffset = alignTo(LastExplicitArgOffset,
+                                    Subtarget.getAlignmentForImplicitArgPtr()) -
+                            LastExplicitArgOffset;
+        AlignedForImplicitArgs = true;
+      }
+      ArgOffset += ImplicitArgOffset;
+    }
+
+    if (DL.getTypeStoreSize(Part.Arg.Ty) < 4 && Part.Alignment < 4) {
+      // Packed sub-dword values share the previous 32-bit SGPR. The extract
+      // happens later when materializing the argument value.
+      assert(KernArgIdx >= 1 && "no previous SGPR");
+      Info.getArgInfo().PreloadKernArgs[KernArgIdx].Regs.push_back(
+          Info.getArgInfo().PreloadKernArgs[KernArgIdx - 1].Regs[0]);
+      continue;
+    }
+
+    // Padding between the last preloaded byte and this argument still consumes
+    // preload SGPR slots because the hardware preload sequence is contiguous.
+    unsigned Padding = ArgOffset - LastExplicitArgOffset;
+    unsigned PaddingSGPRs = alignTo(Padding, 4) / 4;
+    // Kernarg preloads share the user-data SGPR budget with other ABI inputs.
+    if (PaddingSGPRs + NumAllocSGPRs > SGPRInfo.getNumFreeUserSGPRs())
+      break;
+
+    const TargetRegisterClass *RC =
+        TRI.getSGPRClassForBitWidth(NumAllocSGPRs * 32);
+    SmallVectorImpl<MCRegister> *PreloadRegs = Info.addPreloadedKernArg(
----------------
arsenm wrote:

addPreloadedKernArg really ought to return an ArrayRef, not a pointer to the underlying storage 

https://github.com/llvm/llvm-project/pull/205049