[AMDGPU] Make printf lowering faster when there are no printfs

[foad]

Summary

Printf lowering unconditionally visited every instruction in the module.
To make it faster in the common case where there are no printfs, look up
the printf function (if any) and iterate over its users instead.

[arsenm]

llvm/lib/Target/AMDGPU/AMDGPUPrintfRuntimeBinding.cpp

574–575

Braces

575

Perhaps CallSiteBase in cases printf ever ends up invoked maybe?

[foad]

llvm/lib/Target/AMDGPU/AMDGPUPrintfRuntimeBinding.cpp

574–575

I can add braces of course (just around multi-line constructs, yeah?) but whose rule is that? I don't see it in the coding standards and clang-format didn't add them.

575

I did consider that but the rest of the pass seems to be assuming CallInst, so it didn't seem worth generalizing here.

[arsenm]

LGTM

llvm/lib/Target/AMDGPU/AMDGPUPrintfRuntimeBinding.cpp

574–575

clang-format apparently isn't capable of this. If it were up to me, braces would be include in 100% of contexts. the style guide doesn't actually say anything about this last I checked

llvm/lib/Target/AMDGPU/AMDGPUPrintfRuntimeBinding.cpp

 #include "llvm/Analysis/InstructionSimplify.h"
 #include "llvm/Analysis/TargetLibraryInfo.h"
 #include "llvm/CodeGen/Passes.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/GlobalVariable.h"
 #include "llvm/IR/IRBuilder.h"
-#include "llvm/IR/InstVisitor.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/Type.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
 using namespace llvm;
 
 #define DEBUG_TYPE "printfToRuntime"
 #define DWORD_ALIGN 4
 
 namespace {
 class LLVM_LIBRARY_VISIBILITY AMDGPUPrintfRuntimeBinding final
-    : public ModulePass,
-      public InstVisitor<AMDGPUPrintfRuntimeBinding> {
+    : public ModulePass {
 
 public:
   static char ID;
 
   explicit AMDGPUPrintfRuntimeBinding();
 
-  void visitCallSite(CallSite CS) {
-    Function *F = CS.getCalledFunction();
-    if (F && F->hasName() && F->getName() == "printf")
-      Printfs.push_back(CS.getInstruction());
-  }
-
 private:
   bool runOnModule(Module &M) override;
   void getConversionSpecifiers(SmallVectorImpl<char> &OpConvSpecifiers,
                                StringRef fmt, size_t num_ops) const;
 
   bool shouldPrintAsStr(char Specifier, Type *OpType) const;
   bool
   lowerPrintfForGpu(Module &M,
                     function_ref<const TargetLibraryInfo &(Function &)> GetTLI);
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.addRequired<TargetLibraryInfoWrapperPass>();
     AU.addRequired<DominatorTreeWrapperPass>();
   }
 
   Value *simplify(Instruction *I, const TargetLibraryInfo *TLI) {
     return SimplifyInstruction(I, {*TD, TLI, DT});
   }
 
   const DataLayout *TD;
   const DominatorTree *DT;
-  SmallVector<Value *, 32> Printfs;
+  SmallVector<CallInst *, 32> Printfs;
 };
 } // namespace
 
 char AMDGPUPrintfRuntimeBinding::ID = 0;
 
 INITIALIZE_PASS_BEGIN(AMDGPUPrintfRuntimeBinding,
                       "amdgpu-printf-runtime-binding", "AMDGPU Printf lowering",
                       false, false)
bool AMDGPUPrintfRuntimeBinding::lowerPrintfForGpu(
     Module &M, function_ref<const TargetLibraryInfo &(Function &)> GetTLI) {
   LLVMContext &Ctx = M.getContext();
   IRBuilder<> Builder(Ctx);
   Type *I32Ty = Type::getInt32Ty(Ctx);
   unsigned UniqID = 0;
   // NB: This is important for this string size to be divizable by 4
   const char NonLiteralStr[4] = "???";
 
-  for (auto P : Printfs) {
-    auto CI = cast<CallInst>(P);
+  for (auto CI : Printfs) {
     unsigned NumOps = CI->getNumArgOperands();
 
     SmallString<16> OpConvSpecifiers;
     Value *Op = CI->getArgOperand(0);
 
     if (auto LI = dyn_cast<LoadInst>(Op)) {
       Op = LI->getPointerOperand();
       for (auto Use : Op->users()) {
if (ConstExpr) {
           LLVM_DEBUG(dbgs() << "inserting gep to the printf buffer:\n"
                             << *BufferIdx << '\n');
         }
       }
     }
   }
 
   // erase the printf calls
-  for (auto P : Printfs) {
-    auto CI = cast<CallInst>(P);
+  for (auto CI : Printfs)
     CI->eraseFromParent();
-  }
 
   Printfs.clear();
   return true;
 }
 
 bool AMDGPUPrintfRuntimeBinding::runOnModule(Module &M) {
   Triple TT(M.getTargetTriple());
   if (TT.getArch() == Triple::r600)
     return false;
 
-  visit(M);
+  auto PrintfFunction = M.getFunction("printf");
+  if (!PrintfFunction)
+    return false;
+
+  for (auto &U : PrintfFunction->uses())
+    if (auto *CI = dyn_cast<CallInst>(U.getUser()))
+      if (CI->isCallee(&U))
+        Printfs.push_back(CI);
 
   if (Printfs.empty())
     return false;
 
   TD = &M.getDataLayout();
   auto DTWP = getAnalysisIfAvailable<DominatorTreeWrapperPass>();
   DT = DTWP ? &DTWP->getDomTree() : nullptr;
   auto GetTLI = [this](Function &F) -> TargetLibraryInfo & {
     return this->getAnalysis<TargetLibraryInfoWrapperPass>().getTLI(F);
   };
 
   return lowerPrintfForGpu(M, GetTLI);
 }