[llvm] [WIP][AMDGPU][Attributor] Infer `inreg` attribute in `AMDGPUAttributor` (PR #101609)

Shilei Tian via llvm-commits llvm-commits at lists.llvm.org
Fri Aug 16 18:14:04 PDT 2024


https://github.com/shiltian updated https://github.com/llvm/llvm-project/pull/101609

>From 80c0a958620b546ceb58d2ee640f18b3944b3e1b Mon Sep 17 00:00:00 2001
From: Shilei Tian <i at tianshilei.me>
Date: Thu, 1 Aug 2024 23:39:05 -0400
Subject: [PATCH] [WIP][AMDGPU][Attributor] Infer `inreg` attribute in
 `AMDGPUAttributor`

---
 llvm/lib/Target/AMDGPU/AMDGPUAttributor.cpp | 83 ++++++++++++++++++++-
 llvm/test/CodeGen/AMDGPU/inreg-inference.ll | 66 ++++++++++++++++
 2 files changed, 148 insertions(+), 1 deletion(-)
 create mode 100644 llvm/test/CodeGen/AMDGPU/inreg-inference.ll

diff --git a/llvm/lib/Target/AMDGPU/AMDGPUAttributor.cpp b/llvm/lib/Target/AMDGPU/AMDGPUAttributor.cpp
index d65e0ae92308e6..075d44b814f24b 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUAttributor.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPUAttributor.cpp
@@ -1009,6 +1009,82 @@ struct AAAMDGPUNoAGPR
 
 const char AAAMDGPUNoAGPR::ID = 0;
 
+struct AAAMDGPUInreg
+    : public IRAttribute<Attribute::InReg,
+                         StateWrapper<BooleanState, AbstractAttribute>,
+                         AAAMDGPUInreg> {
+  AAAMDGPUInreg(const IRPosition &IRP, Attributor &A) : IRAttribute(IRP) {}
+
+  /// Create an abstract attribute view for the position \p IRP.
+  static AAAMDGPUInreg &createForPosition(const IRPosition &IRP, Attributor &A);
+
+  /// See AbstractAttribute::getName()
+  const std::string getName() const override { return "AAAMDGPUInreg"; }
+
+  const std::string getAsStr(Attributor *A) const override {
+    return getAssumed() ? "inreg" : "non-inreg";
+  }
+
+  void trackStatistics() const override {}
+
+  /// See AbstractAttribute::getIdAddr()
+  const char *getIdAddr() const override { return &ID; }
+
+  /// This function should return true if the type of the \p AA is AAAMDGPUInreg
+  static bool classof(const AbstractAttribute *AA) {
+    return (AA->getIdAddr() == &ID);
+  }
+
+  /// Unique ID (due to the unique address)
+  static const char ID;
+};
+
+const char AAAMDGPUInreg::ID = 0;
+
+namespace {
+
+struct AAAMDGPUInregArgument : public AAAMDGPUInreg {
+  AAAMDGPUInregArgument(const IRPosition &IRP, Attributor &A)
+      : AAAMDGPUInreg(IRP, A) {}
+
+  void initialize(Attributor &A) override {
+    if (getAssociatedArgument()->hasAttribute(Attribute::InReg))
+      indicateOptimisticFixpoint();
+  }
+
+  ChangeStatus updateImpl(Attributor &A) override {
+    const auto &InfoCache =
+        static_cast<AMDGPUInformationCache &>(A.getInfoCache());
+    unsigned ArgNo = getAssociatedArgument()->getArgNo();
+
+    auto Pred = [&](AbstractCallSite ACS) -> bool {
+      CallBase *CB = ACS.getInstruction();
+      auto TTI = InfoCache.TM.getTargetTransformInfo(*CB->getFunction());
+      Value *V = CB->getArgOperandUse(ArgNo);
+      return TTI.isAlwaysUniform(V);
+    };
+
+    bool UsedAssumedInformation = false;
+    if (!A.checkForAllCallSites(Pred, *this, /*RequireAllCallSites=*/true,
+                                UsedAssumedInformation))
+      return indicatePessimisticFixpoint();
+
+    return ChangeStatus::UNCHANGED;
+  }
+};
+
+} // namespace
+
+AAAMDGPUInreg &AAAMDGPUInreg::createForPosition(const IRPosition &IRP,
+                                                Attributor &A) {
+  switch (IRP.getPositionKind()) {
+  case IRPosition::IRP_ARGUMENT:
+    return *new (A.Allocator) AAAMDGPUInregArgument(IRP, A);
+  default:
+    llvm_unreachable("not a valid position for AAAMDGPUInreg");
+  }
+}
+
 static void addPreloadKernArgHint(Function &F, TargetMachine &TM) {
   const GCNSubtarget &ST = TM.getSubtarget<GCNSubtarget>(F);
   for (unsigned I = 0;
@@ -1041,7 +1117,7 @@ static bool runImpl(Module &M, AnalysisGetter &AG, TargetMachine &TM,
        &AAAMDWavesPerEU::ID, &AAAMDGPUNoAGPR::ID, &AACallEdges::ID,
        &AAPointerInfo::ID, &AAPotentialConstantValues::ID,
        &AAUnderlyingObjects::ID, &AAAddressSpace::ID, &AAIndirectCallInfo::ID,
-       &AAInstanceInfo::ID});
+       &AAInstanceInfo::ID, &AAAMDGPUInreg::ID});
 
   AttributorConfig AC(CGUpdater);
   AC.IsClosedWorldModule = Options.IsClosedWorld;
@@ -1091,6 +1167,11 @@ static bool runImpl(Module &M, AnalysisGetter &AG, TargetMachine &TM,
             IRPosition::value(*SI->getPointerOperand()));
       }
     }
+
+    if (F.getCallingConv() != CallingConv::AMDGPU_KERNEL) {
+      for (auto &Arg : F.args())
+        A.getOrCreateAAFor<AAAMDGPUInreg>(IRPosition::argument(Arg));
+    }
   }
 
   ChangeStatus Change = A.run();
diff --git a/llvm/test/CodeGen/AMDGPU/inreg-inference.ll b/llvm/test/CodeGen/AMDGPU/inreg-inference.ll
new file mode 100644
index 00000000000000..54c49c2629d2b2
--- /dev/null
+++ b/llvm/test/CodeGen/AMDGPU/inreg-inference.ll
@@ -0,0 +1,66 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --function-signature --check-globals
+; RUN: opt -S -mtriple=amdgcn-unknown-unknown -passes=amdgpu-attributor %s -o - | FileCheck %s
+
+ at g1 = protected addrspace(1) externally_initialized global i32 0, align 4
+ at g2 = protected addrspace(1) externally_initialized global i32 0, align 4
+
+;.
+; CHECK: @g1 = protected addrspace(1) externally_initialized global i32 0, align 4
+; CHECK: @g2 = protected addrspace(1) externally_initialized global i32 0, align 4
+;.
+define internal fastcc void @f(i32 %x.0.val, ptr nocapture noundef readonly %y) {
+; CHECK-LABEL: define {{[^@]+}}@f
+; CHECK-SAME: (i32 [[X_0_VAL:%.*]], ptr nocapture noundef readonly [[Y:%.*]]) #[[ATTR0:[0-9]+]] {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    store i32 [[X_0_VAL]], ptr addrspace(1) @g1, align 4
+; CHECK-NEXT:    [[LOAD:%.*]] = load i32, ptr [[Y]], align 4
+; CHECK-NEXT:    store i32 [[LOAD]], ptr addrspace(1) @g2, align 4
+; CHECK-NEXT:    ret void
+;
+entry:
+  store i32 %x.0.val, ptr addrspace(1) @g1, align 4
+  %load = load i32, ptr %y, align 4
+  store i32 %load, ptr addrspace(1) @g2, align 4
+  ret void
+}
+
+define protected amdgpu_kernel void @kernel(ptr addrspace(1) %x2, i32 %z) {
+; CHECK-LABEL: define {{[^@]+}}@kernel
+; CHECK-SAME: (ptr addrspace(1) [[X2:%.*]], i32 [[Z:%.*]]) #[[ATTR1:[0-9]+]] {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[X2_CAST:%.*]] = addrspacecast ptr addrspace(1) [[X2]] to ptr
+; CHECK-NEXT:    [[QUEUE_PTR:%.*]] = tail call ptr addrspace(4) @llvm.amdgcn.queue.ptr()
+; CHECK-NEXT:    [[QUEUE_PTR_CAST:%.*]] = addrspacecast ptr addrspace(4) [[QUEUE_PTR]] to ptr
+; CHECK-NEXT:    [[IMPLICITARG_PTR:%.*]] = tail call ptr addrspace(4) @llvm.amdgcn.implicitarg.ptr()
+; CHECK-NEXT:    [[IMPLICITARG_PTR_CAST:%.*]] = addrspacecast ptr addrspace(4) [[IMPLICITARG_PTR]] to ptr
+; CHECK-NEXT:    [[CMP:%.*]] = icmp sgt i32 [[Z]], 0
+; CHECK-NEXT:    [[COND:%.*]] = select i1 [[CMP]], ptr [[QUEUE_PTR_CAST]], ptr [[X2_CAST]]
+; CHECK-NEXT:    [[COND_VAL:%.*]] = load i32, ptr [[COND]], align 4
+; CHECK-NEXT:    tail call fastcc void @f(i32 [[COND_VAL]], ptr noundef [[IMPLICITARG_PTR_CAST]])
+; CHECK-NEXT:    [[DOTVAL:%.*]] = load i32, ptr addrspace(4) [[QUEUE_PTR]], align 4
+; CHECK-NEXT:    tail call fastcc void @f(i32 [[DOTVAL]], ptr noundef [[IMPLICITARG_PTR_CAST]])
+; CHECK-NEXT:    ret void
+;
+entry:
+  %x2.cast = addrspacecast ptr addrspace(1) %x2 to ptr
+  %queue.ptr = tail call ptr addrspace(4) @llvm.amdgcn.queue.ptr()
+  %queue.ptr.cast = addrspacecast ptr addrspace(4) %queue.ptr to ptr
+  %implicitarg.ptr = tail call ptr addrspace(4) @llvm.amdgcn.implicitarg.ptr()
+  %implicitarg.ptr.cast = addrspacecast ptr addrspace(4) %implicitarg.ptr to ptr
+  %cmp = icmp sgt i32 %z, 0
+  %cond = select i1 %cmp, ptr %queue.ptr.cast, ptr %x2.cast
+  %cond.val = load i32, ptr %cond, align 4
+  tail call fastcc void @f(i32 %cond.val, ptr noundef %implicitarg.ptr.cast)
+  %.val = load i32, ptr addrspace(4) %queue.ptr, align 4
+  tail call fastcc void @f(i32 %.val, ptr noundef %implicitarg.ptr.cast)
+  ret void
+}
+
+declare align 4 ptr addrspace(4) @llvm.amdgcn.queue.ptr()
+
+declare align 4 ptr addrspace(4) @llvm.amdgcn.implicitarg.ptr()
+;.
+; CHECK: attributes #[[ATTR0]] = { "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-implicitarg-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" "amdgpu-waves-per-eu"="4,10" "uniform-work-group-size"="false" }
+; CHECK: attributes #[[ATTR1]] = { "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-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" "amdgpu-waves-per-eu"="4,10" "uniform-work-group-size"="false" }
+; CHECK: attributes #[[ATTR2:[0-9]+]] = { nocallback nofree nosync nounwind speculatable willreturn memory(none) }
+;.



More information about the llvm-commits mailing list