[llvm] [AMDGPU] Split struct kernel arguments (PR #133786)

[Shilei Tian via llvm-commits]

================
@@ -42,8 +51,494 @@ static cl::opt<bool>
                          cl::desc("Enable preload kernel arguments to SGPRs"),
                          cl::init(true));
 
+static cl::opt<bool> EnableKernargLayoutChange(
+    "amdgpu-kernarg-layout-change",
+    cl::desc("Allow changing kernel argument segment layout when splitting "
+             "byref structs (remove unused fields, reorder for packing). "
+             "When disabled (default), all struct fields are preserved in "
+             "their original order."),
+    cl::init(false));
+
 namespace {
 
+//===----------------------------------------------------------------------===//
+// Kernel Argument Splitting Logic
+//
+// The following functions handle splitting of byref struct kernel arguments
+// into scalar arguments. This enables preloading of struct fields that would
+// otherwise not be preloadable due to the byref attribute.
+//===----------------------------------------------------------------------===//
+
+// Attribute name for tracking original argument index and offset
+static constexpr StringRef OriginalArgAttr = "amdgpu-original-arg";
+
+// Prefix for backup declaration of original kernel (used for metadata
+// generation)
+static constexpr StringRef OriginalKernelPrefix = "__amdgpu_orig_kernel_";
+
+// Attribute to store the name of the backup declaration
+static constexpr StringRef OriginalKernelAttr = "amdgpu-original-kernel";
+
+static bool parseOriginalArgAttribute(StringRef S, unsigned &RootIdx,
+                                      uint64_t &BaseOff) {
+  auto Parts = S.split(':');
+  if (Parts.second.empty())
+    return false;
+  if (Parts.first.getAsInteger(10, RootIdx))
+    return false;
+  if (Parts.second.getAsInteger(10, BaseOff))
+    return false;
+  return true;
+}
+
+/// Traverses all users of an argument to check if it's suitable for
+/// splitting. A suitable argument is only used by a chain of
+/// GEPs that terminate in LoadInsts.
+static bool
+areArgUsersValidForSplit(Argument &Arg, SmallVectorImpl<LoadInst *> &Loads,
+                         SmallVectorImpl<GetElementPtrInst *> &GEPs) {
+  SmallVector<User *, 16> Worklist(Arg.user_begin(), Arg.user_end());
+  SetVector<User *> Visited;
+
+  while (!Worklist.empty()) {
+    User *U = Worklist.pop_back_val();
+    if (!Visited.insert(U))
+      continue;
+
+    if (auto *LI = dyn_cast<LoadInst>(U)) {
+      Loads.push_back(LI);
+    } else if (auto *GEP = dyn_cast<GetElementPtrInst>(U)) {
+      GEPs.push_back(GEP);
+      for (User *GEPUser : GEP->users()) {
+        Worklist.push_back(GEPUser);
+      }
+    } else
+      return false;
+  }
+
+  const DataLayout &DL = Arg.getParent()->getParent()->getDataLayout();
+  for (const LoadInst *LI : Loads) {
+    APInt Offset(DL.getPointerSizeInBits(), 0);
+    const Value *Base =
+        LI->getPointerOperand()->stripAndAccumulateConstantOffsets(
+            DL, Offset, /*AllowNonInbounds=*/false);
+    if (Base != &Arg)
+      return false;
+  }
+
+  return true;
+}
+
+/// Information about a struct field to be flattened into a scalar argument.
+struct FieldInfo {
+  Type *Ty;
+  uint64_t Offset;
+  LoadInst *Load; // nullptr if field is unused
----------------
shiltian wrote:

```suggestion
  Type *Ty = nullptr;
  uint64_t Offset = 0;
  LoadInst *Load = nullptr; // nullptr if field is unused
```

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