Address comments.

Author: kerbowa

llvm/include/llvm/IR/Argument.h

@@ -178,6 +178,8 @@ class Argument final : public Value {
   /// Check if an argument has a given attribute.
   bool hasAttribute(Attribute::AttrKind Kind) const;
 
+  bool hasAttribute(StringRef Kind) const;
+
   Attribute getAttribute(Attribute::AttrKind Kind) const;
 
   /// Method for support type inquiry through isa, cast, and dyn_cast.

llvm/include/llvm/IR/Function.h

@@ -433,6 +433,9 @@ class LLVM_ABI Function : public GlobalObject, public ilist_node<Function> {
   /// check if an attributes is in the list of attributes.
   bool hasParamAttribute(unsigned ArgNo, Attribute::AttrKind Kind) const;
 
+  /// check if an attributes is in the list of attributes.
+  bool hasParamAttribute(unsigned ArgNo, StringRef Kind) const;
+
   /// gets the attribute from the list of attributes.
   Attribute getAttributeAtIndex(unsigned i, Attribute::AttrKind Kind) const;
 

llvm/lib/IR/Function.cpp

@@ -351,6 +351,10 @@ bool Argument::hasAttribute(Attribute::AttrKind Kind) const {
   return getParent()->hasParamAttribute(getArgNo(), Kind);
 }
 
+bool Argument::hasAttribute(StringRef Kind) const {
+  return getParent()->hasParamAttribute(getArgNo(), Kind);
+}
+
 Attribute Argument::getAttribute(Attribute::AttrKind Kind) const {
   return getParent()->getParamAttribute(getArgNo(), Kind);
 }
@@ -738,6 +742,10 @@ bool Function::hasParamAttribute(unsigned ArgNo,
   return AttributeSets.hasParamAttr(ArgNo, Kind);
 }
 
+bool Function::hasParamAttribute(unsigned ArgNo, StringRef Kind) const {
+  return AttributeSets.hasParamAttr(ArgNo, Kind);
+}
+
 Attribute Function::getAttributeAtIndex(unsigned i,
                                         Attribute::AttrKind Kind) const {
   return AttributeSets.getAttributeAtIndex(i, Kind);

llvm/lib/Target/AMDGPU/AMDGPUHSAMetadataStreamer.cpp

@@ -261,9 +261,7 @@ void MetadataStreamerMsgPackV4::emitKernelArgs(const MachineFunction &MF,
   unsigned Offset = 0;
   auto Args = HSAMetadataDoc->getArrayNode();
   for (auto &Arg : Func.args()) {
-    if (Func.getAttributes().hasAttributeAtIndex(AttributeList::FirstArgIndex +
-                                                     Arg.getArgNo(),
-                                                 "amdgpu-hidden-argument"))
+    if (Arg.hasAttribute("amdgpu-hidden-argument"))
       continue;
 
     emitKernelArg(Arg, Offset, Args);

llvm/lib/Target/AMDGPU/AMDGPULowerKernelArguments.cpp

@@ -212,11 +212,7 @@ class PreloadKernelArgInfo {
 
     // Allocate loads in order of offset. We need to be sure that the implicit
     // argument can actually be preloaded.
-    std::sort(ImplicitArgLoads.begin(), ImplicitArgLoads.end(),
-              [](const std::pair<LoadInst *, unsigned> &A,
-                 const std::pair<LoadInst *, unsigned> &B) {
-                return A.second < B.second;
-              });
+    std::sort(ImplicitArgLoads.begin(), ImplicitArgLoads.end(), less_second());
 
     uint64_t LastExplicitArgOffset = ImplicitArgsBaseOffset;
     // If we fail to preload any implicit argument we know we don't have SGPRs

llvm/lib/Target/AMDGPU/AMDGPUSubtarget.cpp

@@ -314,9 +314,7 @@ uint64_t AMDGPUSubtarget::getExplicitKernArgSize(const Function &F,
   MaxAlign = Align(1);
 
   for (const Argument &Arg : F.args()) {
-    if (F.getAttributes().hasAttributeAtIndex(AttributeList::FirstArgIndex +
-                                                  Arg.getArgNo(),
-                                              "amdgpu-hidden-argument"))
+    if (Arg.hasAttribute("amdgpu-hidden-argument"))
       continue;
 
     const bool IsByRef = Arg.hasByRefAttr();

llvm/lib/Target/AMDGPU/SIISelLowering.cpp

@@ -2539,10 +2539,7 @@ void SITargetLowering::allocatePreloadKernArgSGPRs(
           alignTo(ArgLoc.getLocVT().getFixedSizeInBits(), 32) / 32;
 
       // Add padding SPGR to fix alignment for hidden arguments.
-      if (!AlignedForImplictArgs &&
-          F.getAttributes().hasAttributeAtIndex(AttributeList::FirstArgIndex +
-                                                    Arg.getArgNo(),
-                                                "amdgpu-hidden-argument")) {
+      if (!AlignedForImplictArgs && Arg.hasAttribute("amdgpu-work-group-id")) {
         unsigned OffsetBefore = LastExplicitArgOffset;
         LastExplicitArgOffset = alignTo(
             LastExplicitArgOffset, Subtarget->getAlignmentForImplicitArgPtr());

llvm/test/CodeGen/AMDGPU/preload-implicit-kernargs-IR-lowering.ll

@@ -1,6 +1,6 @@
 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
-; RUN: opt -mtriple=amdgcn-amd-amdhsa -mcpu=gfx940 -amdgpu-attributor -amdgpu-lower-kernel-arguments -S < %s | FileCheck -check-prefix=NO-PRELOAD %s
-; RUN: opt -mtriple=amdgcn-amd-amdhsa -mcpu=gfx940 -amdgpu-attributor -amdgpu-lower-kernel-arguments -amdgpu-kernarg-preload-count=16 -S < %s | FileCheck -check-prefix=PRELOAD %s
+; RUN: opt -mtriple=amdgcn-amd-amdhsa -mcpu=gfx940 -passes='amdgpu-attributor,function(amdgpu-lower-kernel-arguments)' -S < %s | FileCheck -check-prefix=NO-PRELOAD %s
+; RUN: opt -mtriple=amdgcn-amd-amdhsa -mcpu=gfx940 -passes='amdgpu-attributor,function(amdgpu-lower-kernel-arguments)' -amdgpu-kernarg-preload-count=16 -S < %s | FileCheck -check-prefix=PRELOAD %s
 
 define amdgpu_kernel void @preload_block_count_x(ptr addrspace(1) %out) {
 ; NO-PRELOAD-LABEL: define amdgpu_kernel void @preload_block_count_x(

llvm/test/CodeGen/AMDGPU/preload-implicit-kernargs.ll

@@ -594,4 +594,36 @@ define amdgpu_kernel void @preloadremainder_xyz(ptr addrspace(1) inreg %out) #0
   ret void
 }
 
+define amdgpu_kernel void @no_free_sgprs_preloadremainder_z(ptr addrspace(1) inreg %out) {
+; GFX940-LABEL: no_free_sgprs_preloadremainder_z:
+; GFX940:         s_trap 2 ; Kernarg preload header. Trap with incompatible firmware that doesn't support preloading kernel arguments.
+; GFX940-NEXT:    .fill 63, 4, 0xbf800000 ; s_nop 0
+; GFX940-NEXT:  ; %bb.0:
+; GFX940-NEXT:    s_load_dword s0, s[2:3], 0x1c
+; GFX940-NEXT:    v_mov_b32_e32 v0, 0
+; GFX940-NEXT:    s_waitcnt lgkmcnt(0)
+; GFX940-NEXT:    s_lshr_b32 s0, s0, 16
+; GFX940-NEXT:    v_mov_b32_e32 v1, s0
+; GFX940-NEXT:    global_store_dword v0, v1, s[6:7] sc0 sc1
+; GFX940-NEXT:    s_endpgm
+;
+; GFX90a-LABEL: no_free_sgprs_preloadremainder_z:
+; GFX90a:         s_trap 2 ; Kernarg preload header. Trap with incompatible firmware that doesn't support preloading kernel arguments.
+; GFX90a-NEXT:    .fill 63, 4, 0xbf800000 ; s_nop 0
+; GFX90a-NEXT:  ; %bb.0:
+; GFX90a-NEXT:    s_load_dword s0, s[6:7], 0x1c
+; GFX90a-NEXT:    v_mov_b32_e32 v0, 0
+; GFX90a-NEXT:    s_waitcnt lgkmcnt(0)
+; GFX90a-NEXT:    s_lshr_b32 s0, s0, 16
+; GFX90a-NEXT:    v_mov_b32_e32 v1, s0
+; GFX90a-NEXT:    global_store_dword v0, v1, s[10:11]
+; GFX90a-NEXT:    s_endpgm
+  %imp_arg_ptr = call ptr addrspace(4) @llvm.amdgcn.implicitarg.ptr()
+  %gep = getelementptr i8, ptr addrspace(4) %imp_arg_ptr, i32 22
+  %load = load i16, ptr addrspace(4) %gep
+  %conv = zext i16 %load to i32
+  store i32 %conv, ptr addrspace(1) %out
+  ret void
+}
+
 attributes #0 = { "amdgpu-no-agpr" "amdgpu-no-completion-action" "amdgpu-no-default-queue" "amdgpu-no-dispatch-id" "amdgpu-no-dispatch-ptr" "amdgpu-no-heap-ptr" "amdgpu-no-hostcall-ptr" "amdgpu-no-lds-kernel-id" "amdgpu-no-multigrid-sync-arg" "amdgpu-no-queue-ptr" "amdgpu-no-workgroup-id-x" "amdgpu-no-workgroup-id-y" "amdgpu-no-workgroup-id-z" "amdgpu-no-workitem-id-x" "amdgpu-no-workitem-id-y" "amdgpu-no-workitem-id-z" "uniform-work-group-size"="false" }