[llvm] [AMDGPU] Infer amdgpu-no-flat-scratch-init attribute in AMDGPUAttributor (PR #94647)

[Jun Wang via llvm-commits]

================
@@ -677,6 +693,33 @@ struct AAAMDAttributesFunction : public AAAMDAttributes {
     return !A.checkForAllCallLikeInstructions(DoesNotRetrieve, *this,
                                               UsedAssumedInformation);
   }
+
+  // Returns true if FlatScratchInit is needed, i.e., no-flat-scratch-init is
+  // not to be set.
+  bool needFlatScratchInit(Attributor &A) {
+    // This is called on each callee; false means callee shouldn't have
+    // no-flat-scratch-init.
+    auto CheckForNoFlatScratchInit = [&](Instruction &I) {
+      const auto &CB = cast<CallBase>(I);
+      const Value *CalleeOp = CB.getCalledOperand();
+      const Function *Callee = dyn_cast<Function>(CalleeOp);
+      if (!Callee) // indirect call
+        return CB.isInlineAsm();
----------------
jwanggit86 wrote:

If we were to use `callForAllCallees()`, based on my understanding, the code would look something like this:
```diff 
   bool needFlatScratchInit(Attributor &A) {
     // This is called on each callee; false means callee shouldn't have
     // no-flat-scratch-init.
     auto CheckForNoFlatScratchInit = [&](Instruction &I) {
       const auto &CB = cast<CallBase>(I);
-      const Value *CalleeOp = CB.getCalledOperand();
-      const Function *Callee = dyn_cast<Function>(CalleeOp);
-      if (!Callee) // indirect call
-        return CB.isInlineAsm();
-
-      if (Callee->isIntrinsic())
-        return Callee->getIntrinsicID() != Intrinsic::amdgcn_addrspacecast_nonnull;
-
-      const auto *CalleeInfo = A.getAAFor<AAAMDAttributes>(
-          *this, IRPosition::function(*Callee), DepClassTy::REQUIRED);
-      return CalleeInfo && CalleeInfo->isAssumed(FLAT_SCRATCH_INIT);
+      auto PredOnCallees = [&](ArrayRef<const Function *> Callees) {
+        bool Ret = true;
+        for (const Function *Callee : Callees) {
+          if (!Callee) { // indirect call
+            // non-asm indirect call is already handled in updateImpl()
+            assert(CB.isInlineAsm());
+            continue;
+          }
+
+          if (Callee->isIntrinsic()) {
+            Ret &= (Callee->getIntrinsicID() != Intrinsic::amdgcn_addrspacecast_nonnull);
+            if (!Ret)
+              return false;
+            continue;
+          }
+
+          const auto *CalleeInfo = A.getAAFor<AAAMDAttributes>(
+              *this, IRPosition::function(*Callee), DepClassTy::REQUIRED);
+          Ret &= (CalleeInfo && CalleeInfo->isAssumed(FLAT_SCRATCH_INIT));
+          if (!Ret)
+            return false;
+        }
+        return Ret;
+      };
+
+      return A.checkForAllCallees(PredOnCallees, *this, CB);
     };

     bool UsedAssumedInformation = false;
    // If any callee is false (i.e. need FlatScratchInit),
    // checkForAllCallLikeInstructions returns false, in which case this
    // function returns true.
    return !A.checkForAllCallLikeInstructions(CheckForNoFlatScratchInit, *this,
                                              UsedAssumedInformation);
  }
```
Here, in the lambda `CheckForNoFlatScratchInit`, a new lambda is created (`PredOnCallees`), and it simply calls `checkForAllCallees()` passing the new lambda. 
@arsenm @jdoerfert Pls let me know if this is what you had in mind.


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