[AMDGPU] Split struct kernel arguments

AMDGPU backend has a pass which does transformations to allow
firmware to preload kernel arguments into sgpr's to avoid
loading them from kernel arg segment. This pass can improve
kernel latency but it cannot preload struct-type kernel
arguments.

This patch adds a pass to AMDGPU backend to split and flatten
struct-type kernel arguments so that later passes can
preload them into sgpr's.

Basically, the pass collects load or GEP/load instructions
with struct-type kenel args as operands and makes them
new arguments as the kernel. If all uses of a struct-type
kernel arg can be replaced, it will do the replacements
and create a new kernel with the new signature, and
translate all instructions of the old kernel to use
the new arguments in the new kernel. It adds a function
attribute to encode the mapping from the new kernel
argument index to the old kernel argument index and
offset. The streamer will generate kernel argument
metadata based on that and runtime will process
the kernel arguments based on the metadata.

The pass is disabled by default and can be enabled
by LLVM option `-amdgpu-enable-split-kernel-args`.

Author: yxsamliu

llvm/lib/Target/AMDGPU/AMDGPU.h

@@ -125,6 +125,15 @@ struct AMDGPUPromoteKernelArgumentsPass
   PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM);
 };
 
+ModulePass *createAMDGPUSplitKernelArgumentsPass();
+void initializeAMDGPUSplitKernelArgumentsPass(PassRegistry &);
+extern char &AMDGPUSplitKernelArgumentsID;
+
+struct AMDGPUSplitKernelArgumentsPass
+    : PassInfoMixin<AMDGPUSplitKernelArgumentsPass> {
+  PreservedAnalyses run(Module &M, ModuleAnalysisManager &AM);
+};
+
 ModulePass *createAMDGPULowerKernelAttributesPass();
 void initializeAMDGPULowerKernelAttributesPass(PassRegistry &);
 extern char &AMDGPULowerKernelAttributesID;

llvm/lib/Target/AMDGPU/AMDGPUHSAMetadataStreamer.cpp

@@ -357,17 +357,38 @@ void MetadataStreamerMsgPackV4::emitKernelArg(const Argument &Arg,
   Align ArgAlign;
   std::tie(ArgTy, ArgAlign) = getArgumentTypeAlign(Arg, DL);
 
+  // Assuming the argument is not split from struct-type argument by default,
+  // unless we find it in function attribute amdgpu-argument-mapping.
+  unsigned OriginalArgIndex = ~0U;
+  uint64_t OriginalArgOffset = 0;
+  Attribute Attr =
+      Func->getAttributes().getParamAttr(ArgNo, "amdgpu-original-arg");
+  if (Attr.isValid()) {
+    StringRef MappingStr = Attr.getValueAsString();
+    SmallVector<StringRef, 2> Elements;
+    MappingStr.split(Elements, ':');
+    if (Elements.size() == 2) {
+      if (Elements[0].getAsInteger(10, OriginalArgIndex))
+        report_fatal_error(
+            "Invalid original argument index in amdgpu-original-arg attribute");
+      if (Elements[1].getAsInteger(10, OriginalArgOffset))
+        report_fatal_error("Invalid original argument offset in "
+                           "amdgpu-original-arg attribute");
+    }
+  }
+
   emitKernelArg(DL, ArgTy, ArgAlign,
                 getValueKind(ArgTy, TypeQual, BaseTypeName), Offset, Args,
-                PointeeAlign, Name, TypeName, BaseTypeName, ActAccQual,
-                AccQual, TypeQual);
+                PointeeAlign, Name, TypeName, BaseTypeName, ActAccQual, AccQual,
+                TypeQual, OriginalArgIndex, OriginalArgOffset);
 }
 
 void MetadataStreamerMsgPackV4::emitKernelArg(
     const DataLayout &DL, Type *Ty, Align Alignment, StringRef ValueKind,
     unsigned &Offset, msgpack::ArrayDocNode Args, MaybeAlign PointeeAlign,
     StringRef Name, StringRef TypeName, StringRef BaseTypeName,
-    StringRef ActAccQual, StringRef AccQual, StringRef TypeQual) {
+    StringRef ActAccQual, StringRef AccQual, StringRef TypeQual,
+    unsigned OriginalArgIndex, uint64_t OriginalArgOffset) {
   auto Arg = Args.getDocument()->getMapNode();
 
   if (!Name.empty())
@@ -409,6 +430,12 @@ void MetadataStreamerMsgPackV4::emitKernelArg(
       Arg[".is_pipe"] = Arg.getDocument()->getNode(true);
   }
 
+  // Add original argument index and offset to the metadata
+  if (OriginalArgIndex != ~0U) {
+    Arg[".original_arg_index"] = Arg.getDocument()->getNode(OriginalArgIndex);
+    Arg[".original_arg_offset"] = Arg.getDocument()->getNode(OriginalArgOffset);
+  }
+
   Args.push_back(Arg);
 }
 

llvm/lib/Target/AMDGPU/AMDGPUHSAMetadataStreamer.h

@@ -116,7 +116,9 @@ class LLVM_EXTERNAL_VISIBILITY MetadataStreamerMsgPackV4
                      MaybeAlign PointeeAlign = std::nullopt,
                      StringRef Name = "", StringRef TypeName = "",
                      StringRef BaseTypeName = "", StringRef ActAccQual = "",
-                     StringRef AccQual = "", StringRef TypeQual = "");
+                     StringRef AccQual = "", StringRef TypeQual = "",
+                     unsigned OriginalArgIndex = ~0U,
+                     uint64_t OriginalArgOffset = 0);
 
   void emitHiddenKernelArgs(const MachineFunction &MF, unsigned &Offset,
                             msgpack::ArrayDocNode Args) override;

llvm/lib/Target/AMDGPU/AMDGPUPassRegistry.def

@@ -29,6 +29,7 @@ MODULE_PASS("amdgpu-printf-runtime-binding", AMDGPUPrintfRuntimeBindingPass())
 MODULE_PASS("amdgpu-remove-incompatible-functions", AMDGPURemoveIncompatibleFunctionsPass(*this))
 MODULE_PASS("amdgpu-sw-lower-lds", AMDGPUSwLowerLDSPass(*this))
 MODULE_PASS("amdgpu-unify-metadata", AMDGPUUnifyMetadataPass())
+MODULE_PASS("amdgpu-split-kernel-arguments", AMDGPUSplitKernelArgumentsPass())
 #undef MODULE_PASS
 
 #ifndef MODULE_PASS_WITH_PARAMS