(1) Use getCalledFunction instead of getCalledOperand (2) other

minor code change based on reviews (3) fix test files.

Author: jwanggit86

llvm/lib/Target/AMDGPU/AMDGPUAttributor.cpp

@@ -434,17 +434,11 @@ struct AAAMDAttributesFunction : public AAAMDAttributes {
       return;
     }
 
-    bool HasAllocaOrASCast = false;
-    for (BasicBlock &BB : *F) {
-      for (Instruction &I : BB) {
-        if (isa<AllocaInst>(I) || isa<AddrSpaceCastInst>(I)) {
-          HasAllocaOrASCast = true;
-          removeAssumedBits(FLAT_SCRATCH_INIT);
-          break;
-        }
+    for (Instruction &I : instructions(F)) {
+      if (isa<AllocaInst>(I) || isa<AddrSpaceCastInst>(I)) {
+        removeAssumedBits(FLAT_SCRATCH_INIT);
+        return;
       }
-      if (HasAllocaOrASCast)
-        break;
     }
   }
 
@@ -701,13 +695,12 @@ struct AAAMDAttributesFunction : public AAAMDAttributes {
     // 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);
+      const Function *Callee = CB.getCalledFunction();
       if (!Callee) // indirect call
         return CB.isInlineAsm();
 
       if (Callee->isIntrinsic())
-        return true;
+        return Callee->getIntrinsicID() != Intrinsic::amdgcn_addrspacecast_nonnull;
 
       const auto *CalleeInfo = A.getAAFor<AAAMDAttributes>(
           *this, IRPosition::function(*Callee), DepClassTy::REQUIRED);
@@ -716,7 +709,8 @@ struct AAAMDAttributesFunction : public AAAMDAttributes {
 
     bool UsedAssumedInformation = false;
     // If any callee is false (i.e. need FlatScratchInit),
-    // checkForAllCallLikeInstructions returns false
+    // checkForAllCallLikeInstructions returns false, in which case this
+    // function returns true.
     return !A.checkForAllCallLikeInstructions(CheckForNoFlatScratchInit, *this,
                                               UsedAssumedInformation);
   }

llvm/lib/Target/AMDGPU/AMDGPUSubtarget.cpp

@@ -1054,8 +1054,6 @@ GCNUserSGPRUsageInfo::GCNUserSGPRUsageInfo(const Function &F,
   const CallingConv::ID CC = F.getCallingConv();
   const bool IsKernel =
       CC == CallingConv::AMDGPU_KERNEL || CC == CallingConv::SPIR_KERNEL;
-  const bool NoFlatScratchInit =
-      F.hasFnAttribute("amdgpu-no-flat-scratch-init");
 
   if (IsKernel && (!F.arg_empty() || ST.getImplicitArgNumBytes(F) != 0))
     KernargSegmentPtr = true;
@@ -1078,6 +1076,8 @@ GCNUserSGPRUsageInfo::GCNUserSGPRUsageInfo(const Function &F,
       DispatchID = true;
   }
 
+  const bool NoFlatScratchInit =
+      F.hasFnAttribute("amdgpu-no-flat-scratch-init");
   // TODO: This could be refined a lot. The attribute is a poor way of
   // detecting calls or stack objects that may require it before argument
   // lowering.