diff --git a/llvm/lib/Target/AMDGPU/AMDGPU.h b/llvm/lib/Target/AMDGPU/AMDGPU.h --- a/llvm/lib/Target/AMDGPU/AMDGPU.h +++ b/llvm/lib/Target/AMDGPU/AMDGPU.h @@ -115,6 +115,11 @@ void initializeAMDGPULowerKernelAttributesPass(PassRegistry &); extern char &AMDGPULowerKernelAttributesID; +struct AMDGPULowerKernelAttributesPass + : PassInfoMixin { + PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM); +}; + void initializeAMDGPUPropagateAttributesEarlyPass(PassRegistry &); extern char &AMDGPUPropagateAttributesEarlyID; diff --git a/llvm/lib/Target/AMDGPU/AMDGPULowerKernelAttributes.cpp b/llvm/lib/Target/AMDGPU/AMDGPULowerKernelAttributes.cpp --- a/llvm/lib/Target/AMDGPU/AMDGPULowerKernelAttributes.cpp +++ b/llvm/lib/Target/AMDGPU/AMDGPULowerKernelAttributes.cpp @@ -19,7 +19,9 @@ #include "llvm/CodeGen/TargetPassConfig.h" #include "llvm/IR/Constants.h" #include "llvm/IR/Function.h" +#include "llvm/IR/InstIterator.h" #include "llvm/IR/Instructions.h" +#include "llvm/IR/PassManager.h" #include "llvm/IR/PatternMatch.h" #include "llvm/Pass.h" @@ -41,16 +43,11 @@ }; class AMDGPULowerKernelAttributes : public ModulePass { - Module *Mod = nullptr; - public: static char ID; AMDGPULowerKernelAttributes() : ModulePass(ID) {} - bool processUse(CallInst *CI); - - bool doInitialization(Module &M) override; bool runOnModule(Module &M) override; StringRef getPassName() const override { @@ -64,12 +61,7 @@ } // end anonymous namespace -bool AMDGPULowerKernelAttributes::doInitialization(Module &M) { - Mod = &M; - return false; -} - -bool AMDGPULowerKernelAttributes::processUse(CallInst *CI) { +static bool processUse(CallInst *CI) { Function *F = CI->getParent()->getParent(); auto MD = F->getMetadata("reqd_work_group_size"); @@ -89,7 +81,7 @@ Value *GridSizeY = nullptr; Value *GridSizeZ = nullptr; - const DataLayout &DL = Mod->getDataLayout(); + const DataLayout &DL = F->getParent()->getDataLayout(); // We expect to see several GEP users, casted to the appropriate type and // loaded. @@ -239,7 +231,7 @@ StringRef DispatchPtrName = Intrinsic::getName(Intrinsic::amdgcn_dispatch_ptr); - Function *DispatchPtr = Mod->getFunction(DispatchPtrName); + Function *DispatchPtr = M.getFunction(DispatchPtrName); if (!DispatchPtr) // Dispatch ptr not used. return false; @@ -267,3 +259,22 @@ ModulePass *llvm::createAMDGPULowerKernelAttributesPass() { return new AMDGPULowerKernelAttributes(); } + +PreservedAnalyses +AMDGPULowerKernelAttributesPass::run(Function &F, FunctionAnalysisManager &AM) { + StringRef DispatchPtrName = + Intrinsic::getName(Intrinsic::amdgcn_dispatch_ptr); + + Function *DispatchPtr = F.getParent()->getFunction(DispatchPtrName); + if (!DispatchPtr) // Dispatch ptr not used. + return PreservedAnalyses::all(); + + for (Instruction &I : instructions(F)) { + if (CallInst *CI = dyn_cast(&I)) { + if (CI->getCalledFunction() == DispatchPtr) + processUse(CI); + } + } + + return PreservedAnalyses::all(); +} diff --git a/llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp b/llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp --- a/llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp +++ b/llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp @@ -508,6 +508,10 @@ PM.addPass(AMDGPUPromoteAllocaToVectorPass(*this)); return true; } + if (PassName == "amdgpu-lower-kernel-attributes") { + PM.addPass(AMDGPULowerKernelAttributesPass()); + return true; + } return false; }); @@ -530,6 +534,10 @@ // but before SROA to increase SROA opportunities. FPM.addPass(InferAddressSpacesPass()); + // This should run after inlining to have any chance of doing + // anything, and before other cleanup optimizations. + FPM.addPass(AMDGPULowerKernelAttributesPass()); + if (Level != PassBuilder::OptimizationLevel::O0) { // Promote alloca to vector before SROA and loop unroll. If we // manage to eliminate allocas before unroll we may choose to unroll diff --git a/llvm/test/CodeGen/AMDGPU/reqd-work-group-size.ll b/llvm/test/CodeGen/AMDGPU/reqd-work-group-size.ll --- a/llvm/test/CodeGen/AMDGPU/reqd-work-group-size.ll +++ b/llvm/test/CodeGen/AMDGPU/reqd-work-group-size.ll @@ -1,4 +1,5 @@ ; RUN: opt -mtriple=amdgcn-amd-amdhsa -S -amdgpu-lower-kernel-attributes -instcombine %s | FileCheck -enable-var-scope %s +; RUN: opt -mtriple=amdgcn-amd-amdhsa -S -passes=amdgpu-lower-kernel-attributes,instcombine %s | FileCheck -enable-var-scope %s target datalayout = "n32" diff --git a/llvm/tools/opt/opt.cpp b/llvm/tools/opt/opt.cpp --- a/llvm/tools/opt/opt.cpp +++ b/llvm/tools/opt/opt.cpp @@ -464,7 +464,7 @@ static bool shouldPinPassToLegacyPM(StringRef Pass) { std::vector PassNameExactToIgnore = { "amdgpu-simplifylib", "amdgpu-usenative", "amdgpu-promote-alloca", - "amdgpu-promote-alloca-to-vector"}; + "amdgpu-promote-alloca-to-vector", "amdgpu-lower-kernel-attributes"}; for (const auto &P : PassNameExactToIgnore) if (Pass == P) return false;