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

[Keshav Vinayak Jha via llvm-commits]

================
@@ -451,46 +606,55 @@ void AMDGPUCallLowering::lowerParameterPtr(Register DstReg, MachineIRBuilder &B,
   B.buildPtrAdd(DstReg, KernArgSegmentVReg, OffsetReg);
 }
 
-void AMDGPUCallLowering::lowerParameter(MachineIRBuilder &B, ArgInfo &OrigArg,
-                                        uint64_t Offset,
-                                        Align Alignment) const {
+bool AMDGPUCallLowering::lowerParameter(MachineIRBuilder &B, ArgInfo &Arg,
+                                        uint64_t Offset, Align Alignment,
+                                        unsigned InputArgIndex) const {
   MachineFunction &MF = B.getMF();
   const Function &F = MF.getFunction();
   const DataLayout &DL = F.getDataLayout();
+  const Argument *IRArg = dyn_cast_if_present<Argument>(Arg.OrigValue);
+  const bool IsHiddenArg =
+      IRArg && IRArg->hasAttribute("amdgpu-hidden-argument");
+  const SIMachineFunctionInfo *Info = MF.getInfo<SIMachineFunctionInfo>();
   const SITargetLowering &TLI = *getTLI<SITargetLowering>();
   MachinePointerInfo PtrInfo = TLI.getKernargSegmentPtrInfo(MF);
 
   LLT PtrTy = LLT::pointer(AMDGPUAS::CONSTANT_ADDRESS, 64);
+  LLT ArgTy = getLLTForType(*Arg.Ty, DL);
+  if (Arg.Flags[0].isPointer()) {
+    // Compensate for losing pointeriness in splitValueTypes.
+    LLT PtrTy = LLT::pointer(Arg.Flags[0].getPointerAddrSpace(),
+                             ArgTy.getScalarSizeInBits());
+    ArgTy =
+        ArgTy.isVector() ? LLT::vector(ArgTy.getElementCount(), PtrTy) : PtrTy;
+  }
 
-  SmallVector<ArgInfo, 32> SplitArgs;
-  SmallVector<TypeSize> FieldOffsets;
-  splitToValueTypes(OrigArg, SplitArgs, DL, F.getCallingConv(), &FieldOffsets);
+  assert(Arg.Regs.size() == 1);
 
-  unsigned Idx = 0;
-  for (ArgInfo &SplitArg : SplitArgs) {
-    Register PtrReg = B.getMRI()->createGenericVirtualRegister(PtrTy);
-    lowerParameterPtr(PtrReg, B, Offset + FieldOffsets[Idx]);
-
-    LLT ArgTy = getLLTForType(*SplitArg.Ty, DL);
-    if (SplitArg.Flags[0].isPointer()) {
-      // Compensate for losing pointeriness in splitValueTypes.
-      LLT PtrTy = LLT::pointer(SplitArg.Flags[0].getPointerAddrSpace(),
-                               ArgTy.getScalarSizeInBits());
-      ArgTy = ArgTy.isVector() ? LLT::vector(ArgTy.getElementCount(), PtrTy)
-                               : PtrTy;
-    }
+  auto PreloadArg = Info->getArgInfo().PreloadKernArgs.find(InputArgIndex);
+  if (PreloadArg != Info->getArgInfo().PreloadKernArgs.end()) {
+    lowerPreloadedKernArg(B, Arg.Regs[0], ArgTy, Offset, Alignment,
+                          PreloadArg->getSecond().Regs);
+    return true;
+  }
 
-    MachineMemOperand *MMO = MF.getMachineMemOperand(
-        PtrInfo,
-        MachineMemOperand::MOLoad | MachineMemOperand::MODereferenceable |
-            MachineMemOperand::MOInvariant,
-        ArgTy, commonAlignment(Alignment, FieldOffsets[Idx]));
+  if (IsHiddenArg) {
+    F.getContext().diagnose(DiagnosticInfoUnsupported(
+        F, "hidden argument in kernel signature was not preloaded"));
+    return false;
+  }
----------------
keshavvinayak01 wrote:

Removing `IsHiddenArg` and moved the direct attribute check next to its use. The check remains after the preload lookup because successfully preloaded hidden arguments must continue through the supported path.

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