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

[Shilei Tian via llvm-commits]

================
@@ -1014,6 +1016,110 @@ struct AAAMDGPUNoAGPR
 
 const char AAAMDGPUNoAGPR::ID = 0;
 
+struct AAAMDGPUUniform
+    : public IRAttribute<Attribute::InReg,
+                         StateWrapper<BooleanState, AbstractAttribute>,
+                         AAAMDGPUUniform> {
+  AAAMDGPUUniform(const IRPosition &IRP, Attributor &A) : IRAttribute(IRP) {}
+
+  /// Create an abstract attribute view for the position \p IRP.
+  static AAAMDGPUUniform &createForPosition(const IRPosition &IRP,
+                                            Attributor &A);
+
+  /// See AbstractAttribute::getName()
+  const std::string getName() const override { return "AAAMDGPUUniform"; }
+
+  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
+  /// AAAMDGPUUniform
+  static bool classof(const AbstractAttribute *AA) {
+    return (AA->getIdAddr() == &ID);
+  }
+
+  /// Unique ID (due to the unique address)
+  static const char ID;
+};
+
+const char AAAMDGPUUniform::ID = 0;
+
+namespace {
+
+struct AAAMDGPUUniformArgument : public AAAMDGPUUniform {
+  AAAMDGPUUniformArgument(const IRPosition &IRP, Attributor &A)
+      : AAAMDGPUUniform(IRP, A) {}
+
+  void initialize(Attributor &A) override {
+    assert(
+        !AMDGPU::isEntryFunctionCC(getAssociatedFunction()->getCallingConv()));
+    if (getAssociatedArgument()->hasAttribute(Attribute::InReg))
+      indicateOptimisticFixpoint();
+  }
+
+  ChangeStatus updateImpl(Attributor &A) override {
+    unsigned ArgNo = getAssociatedArgument()->getArgNo();
+
+    auto isUniform = [&](AbstractCallSite ACS) -> bool {
+      CallBase *CB = ACS.getInstruction();
+      Value *V = CB->getArgOperandUse(ArgNo);
+      if (isa<Constant>(V))
+        return true;
+      if (auto *I = dyn_cast<Instruction>(V)) {
+        auto *UA = A.getInfoCache()
+                       .getAnalysisResultForFunction<UniformityInfoAnalysis>(
+                           *I->getFunction());
+        return UA && UA->isUniform(I);
+      }
+      if (auto *Arg = dyn_cast<Argument>(V)) {
+        auto *UA = A.getInfoCache()
+                       .getAnalysisResultForFunction<UniformityInfoAnalysis>(
+                           *Arg->getParent());
+        if (UA && UA->isUniform(Arg))
+          return true;
+        // We only rely on isArgPassedInSGPR when the function is terminal,
+        // assuming there is no call edge from a function to an entry function.
+        if (AMDGPU::isEntryFunctionCC(Arg->getParent()->getCallingConv()))
+          return AMDGPU::isArgPassedInSGPR(Arg);
----------------
shiltian wrote:

> UA should still be correct for an argument, so I still don't understand why you need the second check on the argument

UA might miss some cases where the argument can be uniform but it is unable to tell (probably due to the lack of sophisticated inter-procedural data flow analysis as what Attributor framework does), thus conservatively deems it as diverge. That is the whole point of this AA; otherwise we would just need to query UA on each function argument and add the attribute if UA says yes.

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