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

[Manuel Carrasco via llvm-commits]

================
@@ -296,6 +297,203 @@ struct AMDGPUOutgoingArgHandler : public AMDGPUOutgoingValueHandler {
     assignValueToAddress(ValVReg, Addr, MemTy, MPO, VA);
   }
 };
+
+// Return the virtual register and type for a preloaded physical SGPR live-in.
+static std::pair<Register, LLT> getPreloadedLiveIn(MachineIRBuilder &B,
+                                                   MCRegister PhysReg,
+                                                   const SIRegisterInfo &TRI) {
+  MachineFunction &MF = B.getMF();
+  const TargetInstrInfo *TII = MF.getSubtarget().getInstrInfo();
+  const TargetRegisterClass *RC = TRI.getPhysRegBaseClass(PhysReg);
+  assert(RC && "expected a register class for preload SGPR");
+  LLT RegTy = LLT::scalar(TRI.getRegSizeInBits(*RC).getFixedValue());
+  return {getFunctionLiveInPhysReg(MF, *TII, PhysReg, *RC, B.getDL(), RegTy),
+          RegTy};
+}
+
+// Convert raw SGPR preload bits to the type expected by the formal argument.
+static Register adjustPreloadedArgType(MachineIRBuilder &B, Register Src,
+                                       LLT SrcTy, LLT DstTy) {
+  if (SrcTy == DstTy)
+    return Src;
+
+  LLT IntDstTy = (DstTy.isPointer() || DstTy.isVector())
+                     ? LLT::scalar(DstTy.getSizeInBits())
+                     : DstTy;
+  if (SrcTy.getSizeInBits() != IntDstTy.getSizeInBits())
+    Src = B.buildAnyExtOrTrunc(IntDstTy, Src).getReg(0);
+  else if (SrcTy != IntDstTy)
+    Src = B.buildBitcast(IntDstTy, Src).getReg(0);
+
+  if (DstTy.isPointer())
+    return B.buildIntToPtr(DstTy, Src).getReg(0);
+
+  if (DstTy.isVector())
+    return B.buildBitcast(DstTy, Src).getReg(0);
+
+  return Src;
+}
+
+// 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,
+                                  const SIRegisterInfo &TRI) {
+  assert(!PreloadRegs.empty());
+
+  // Kernarg preloads are assigned in 32-bit SGPR units.
+  const LLT S32 = LLT::scalar(32);
+
+  Register Value;
+  LLT ValueTy;
+  const bool IsPackedSubDword = DstTy.getSizeInBits() < 32 && Alignment < 4;
+
+  if (PreloadRegs.size() == 1) {
+    auto [LiveIn, LiveInTy] = getPreloadedLiveIn(B, PreloadRegs[0], TRI);
+    if (IsPackedSubDword) {
+      // Extract sub-dword preloads from their containing 32-bit SGPR word.
+      Register Raw = B.buildCopy(S32, LiveIn).getReg(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 {
+      // A single preload entry may be a 32-bit SGPR or a wider SGPR tuple.
+      ValueTy = LiveInTy;
+      Value = B.buildCopy(ValueTy, LiveIn).getReg(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) {
+      Register LiveIn = getPreloadedLiveIn(B, Reg, TRI).first;
+      Regs.push_back(B.buildCopy(S32, LiveIn).getReg(0));
+    }
+
+    ValueTy = LLT::scalar(PreloadRegs.size() * 32);
+    Value = B.buildMergeLikeInstr(ValueTy, Regs).getReg(0);
+  }
+
+  Register Adjusted = adjustPreloadedArgType(B, Value, ValueTy, DstTy);
+  B.buildCopy(DstReg, Adjusted);
+}
+
+// CCValAssign records an MVT, so round odd-sized split argument parts to the
+// simple memory type used for kernarg layout.
+static MVT getLocVTForSplitArg(const TargetLowering &TLI, const DataLayout &DL,
+                               Type *Ty, LLVMContext &Ctx) {
+  EVT LocVT = TLI.getValueType(DL, Ty, true);
+
+  if (LocVT.isVector() && LocVT.getVectorNumElements() == 1)
+    LocVT = LocVT.getScalarType();
+
+  if (LocVT.isVector() && !LocVT.isPow2VectorType())
+    LocVT = LocVT.getPow2VectorType(Ctx);
+  else if (!LocVT.isSimple() && !LocVT.isVector())
+    LocVT = LocVT.getRoundIntegerType(Ctx);
+
+  assert(LocVT.isSimple());
+  return LocVT.getSimpleVT();
+}
+
+// Assign SGPRs for the contiguous sequence of kernel argument parts marked for
+// kernarg preload.
+static void allocatePreloadKernArgSGPRs(
+    CCState &CCInfo, ArrayRef<CallLowering::ArgInfo> SplitArgs,
+    ArrayRef<CCValAssign> ArgLocs, MachineFunction &MF,
+    const GCNSubtarget &Subtarget, const SIRegisterInfo &TRI,
+    SIMachineFunctionInfo &Info) {
+  assert(SplitArgs.size() == ArgLocs.size());
+
+  const Function &F = MF.getFunction();
+  const DataLayout &DL = F.getDataLayout();
+  unsigned LastExplicitArgOffset = Subtarget.getExplicitKernelArgOffset();
+  GCNUserSGPRUsageInfo &SGPRInfo = Info.getUserSGPRInfo();
+  bool InPreloadSequence = true;
+  unsigned InIdx = 0;
+  bool AlignedForImplicitArgs = false;
+  unsigned ImplicitArgOffset = 0;
+
+  for (const Argument &Arg : F.args()) {
+    // Preload assignment follows the original argument order. Each original
+    // argument may already have been split into one or more lowered parts.
+    const bool IsByRef = Arg.hasByRefAttr();
+    Type *ArgTy = IsByRef ? Arg.getParamByRefType() : Arg.getType();
+    if (DL.getTypeAllocSize(ArgTy) == 0)
+      continue;
+
+    // Hardware preloads a contiguous prefix of the kernarg segment. Once a
+    // non-preloaded argument is reached, no later argument can be preloaded.
+    if (!InPreloadSequence || !Arg.hasInRegAttr())
+      break;
+
+    unsigned ArgIdx = Arg.getArgNo();
+    if (InIdx < SplitArgs.size() && SplitArgs[InIdx].OrigArgIndex != ArgIdx)
----------------
mgcarrasco wrote:

Sorry, just in case, I meant that there was no test. I didn't mean to say that the line was not required. 

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