[NewPM][AMDGPU] Port amdgpu-lower-kernel-attributes

[aeubanks]

Summary

And add it to the AMDGPU opt pipeline.

This is a function pass instead of a module pass (like the legacy pass)
because it's getting added to a CGSCCPassManager, and you can't put a
module pass in a CGSCCPassManager.

llvm/lib/Target/AMDGPU/AMDGPU.h

 FunctionPass *createAMDGPULowerKernelArgumentsPass();
 void initializeAMDGPULowerKernelArgumentsPass(PassRegistry &);
 extern char &AMDGPULowerKernelArgumentsID;
 
 ModulePass *createAMDGPULowerKernelAttributesPass();
 void initializeAMDGPULowerKernelAttributesPass(PassRegistry &);
 extern char &AMDGPULowerKernelAttributesID;
 
+struct AMDGPULowerKernelAttributesPass
+    : PassInfoMixin<AMDGPULowerKernelAttributesPass> {
+  PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM);
+};
+
 void initializeAMDGPUPropagateAttributesEarlyPass(PassRegistry &);
 extern char &AMDGPUPropagateAttributesEarlyID;
 
 void initializeAMDGPUPropagateAttributesLatePass(PassRegistry &);
 extern char &AMDGPUPropagateAttributesLateID;
 
 void initializeAMDGPURewriteOutArgumentsPass(PassRegistry &);
 extern char &AMDGPURewriteOutArgumentsID;

llvm/lib/Target/AMDGPU/AMDGPULowerKernelAttributes.cpp

 
 #include "AMDGPU.h"
 #include "AMDGPUTargetMachine.h"
 #include "llvm/Analysis/ValueTracking.h"
 #include "llvm/CodeGen/Passes.h"
 #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"
 
 #define DEBUG_TYPE "amdgpu-lower-kernel-attributes"
 
 using namespace llvm;
 
 namespace {
 
 // Field offsets in hsa_kernel_dispatch_packet_t.
 enum DispatchPackedOffsets {
   WORKGROUP_SIZE_X = 4,
   WORKGROUP_SIZE_Y = 6,
   WORKGROUP_SIZE_Z = 8,
 
   GRID_SIZE_X = 12,
   GRID_SIZE_Y = 16,
   GRID_SIZE_Z = 20
 };
 
 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 {
     return "AMDGPU Kernel Attributes";
   }
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.setPreservesAll();
  }
 };
 
 } // 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");
   const bool HasReqdWorkGroupSize = MD && MD->getNumOperands() == 3;
 
   const bool HasUniformWorkGroupSize =
     F->getFnAttribute("uniform-work-group-size").getValueAsString() == "true";
 
   if (!HasReqdWorkGroupSize && !HasUniformWorkGroupSize)
     return false;
 
   Value *WorkGroupSizeX = nullptr;
   Value *WorkGroupSizeY = nullptr;
   Value *WorkGroupSizeZ = nullptr;
 
   Value *GridSizeX = nullptr;
   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.
   for (User *U : CI->users()) {
     if (!U->hasOneUse())
       continue;
 
     int64_t Offset = 0;

 }
 
 // TODO: Move makeLIDRangeMetadata usage into here. Seem to not get
 // TargetPassConfig for subtarget.
 bool AMDGPULowerKernelAttributes::runOnModule(Module &M) {
   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;
 
   bool MadeChange = false;
 
   SmallPtrSet<Instruction *, 4> HandledUses;
   for (auto *U : DispatchPtr->users()) {
     CallInst *CI = cast<CallInst>(U);

 INITIALIZE_PASS_END(AMDGPULowerKernelAttributes, DEBUG_TYPE, "AMDGPU IR optimizations",
                     false, false)
 
 char AMDGPULowerKernelAttributes::ID = 0;
 
 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<CallInst>(&I)) {
+      if (CI->getCalledFunction() == DispatchPtr)
+        processUse(CI);
+    }
+  }
+
+  return PreservedAnalyses::all();
+}

llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp

PB.registerPipelineParsingCallback(
         if (PassName == "amdgpu-promote-alloca") {
           PM.addPass(AMDGPUPromoteAllocaPass(*this));
           return true;
         }
         if (PassName == "amdgpu-promote-alloca-to-vector") {
           PM.addPass(AMDGPUPromoteAllocaToVectorPass(*this));
           return true;
         }
+        if (PassName == "amdgpu-lower-kernel-attributes") {
+          PM.addPass(AMDGPULowerKernelAttributesPass());
+          return true;
+        }
         return false;
       });
 
   PB.registerPipelineStartEPCallback([DebugPassManager](
                                          ModulePassManager &PM,
                                          PassBuilder::OptimizationLevel Level) {
     FunctionPassManager FPM(DebugPassManager);
     FPM.addPass(AMDGPUUseNativeCallsPass());
     if (EnableLibCallSimplify && Level != PassBuilder::OptimizationLevel::O0)
       FPM.addPass(AMDGPUSimplifyLibCallsPass());
     PM.addPass(createModuleToFunctionPassAdaptor(std::move(FPM)));
   });
 
   PB.registerCGSCCOptimizerLateEPCallback(
       [this, DebugPassManager](CGSCCPassManager &PM,
                                PassBuilder::OptimizationLevel Level) {
           FunctionPassManager FPM(DebugPassManager);
 
           // Add infer address spaces pass to the opt pipeline after inlining
           // 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
             // less.
             FPM.addPass(AMDGPUPromoteAllocaToVectorPass(*this));
           }
 
           PM.addPass(createCGSCCToFunctionPassAdaptor(std::move(FPM)));

llvm/test/CodeGen/AMDGPU/reqd-work-group-size.ll

 ; 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"
 
 ; CHECK-LABEL: @invalid_reqd_work_group_size(
 ; CHECK: load i16,
 define amdgpu_kernel void @invalid_reqd_work_group_size(i16 addrspace(1)* %out) #0 !reqd_work_group_size !1 {
   %dispatch.ptr = tail call i8 addrspace(4)* @llvm.amdgcn.dispatch.ptr()
   %gep.group.size.x = getelementptr inbounds i8, i8 addrspace(4)* %dispatch.ptr, i64 4

llvm/tools/opt/opt.cpp

 
 // For use in NPM transition. Currently this contains most codegen-specific
 // passes. Remove passes from here when porting to the NPM.
 // TODO: use a codegen version of PassRegistry.def/PassBuilder::is*Pass() once
 // it exists.
 static bool shouldPinPassToLegacyPM(StringRef Pass) {
   std::vector<StringRef> 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;
 
   std::vector<StringRef> PassNamePrefix = {
       "x86-",  "xcore-", "wasm-",    "systemz-", "ppc-",   "nvvm-",   "nvptx-",
       "mips-", "lanai-", "hexagon-", "bpf-",     "avr-",   "thumb2-", "arm-",
       "si-",   "gcn-",   "amdgpu-",  "aarch64-", "amdgcn-"};