[NewPM][NVPTX] Port NVPTX opt passes

[aeubanks]

Summary

There are only two used in the IR optimization pipeline.
Port these and add them to the default pipeline.

Similar to https://reviews.llvm.org/D93863.

I added -mtriple to some tests since under the new PM, the passes are
only available when the TargetMachine is specified.

[rnk]

lgtm with nit

llvm/lib/Target/NVPTX/NVVMReflect.cpp

77

static?

[aeubanks]

rebase
make function static

[tra]

LGTM

llvm/lib/Target/NVPTX/NVPTX.h

 //===-- NVPTX.h - Top-level interface for NVPTX representation --*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file contains the entry points for global functions defined in
 // the LLVM NVPTX back-end.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_LIB_TARGET_NVPTX_NVPTX_H
 #define LLVM_LIB_TARGET_NVPTX_NVPTX_H
 
+#include "llvm/IR/PassManager.h"
 #include "llvm/Pass.h"
 #include "llvm/Support/CodeGen.h"
 
 namespace llvm {
 class NVPTXTargetMachine;
 class FunctionPass;
 class MachineFunctionPass;
 

 MachineFunctionPass *createNVPTXPrologEpilogPass();
 MachineFunctionPass *createNVPTXReplaceImageHandlesPass();
 FunctionPass *createNVPTXImageOptimizerPass();
 FunctionPass *createNVPTXLowerArgsPass(const NVPTXTargetMachine *TM);
 FunctionPass *createNVPTXLowerAllocaPass();
 MachineFunctionPass *createNVPTXPeephole();
 MachineFunctionPass *createNVPTXProxyRegErasurePass();
 
+struct NVVMIntrRangePass : PassInfoMixin<NVVMIntrRangePass> {
+  NVVMIntrRangePass();
+  NVVMIntrRangePass(unsigned SmVersion) : SmVersion(SmVersion) {}
+  PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM);
+
+private:
+  unsigned SmVersion;
+};
+
+struct NVVMReflectPass : PassInfoMixin<NVVMReflectPass> {
+  NVVMReflectPass();
+  NVVMReflectPass(unsigned SmVersion) : SmVersion(SmVersion) {}
+  PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM);
+
+private:
+  unsigned SmVersion;
+};
+
 namespace NVPTX {
 enum DrvInterface {
   NVCL,
   CUDA
 };
 
 // A field inside TSFlags needs a shift and a mask. The usage is
 // always as follows :

llvm/lib/Target/NVPTX/NVPTXTargetMachine.h

bool addPassesToEmitMC(PassManagerBase &, MCContext *&, raw_pwrite_stream &,
                          bool = true) override {
     return true;
   }
   TargetLoweringObjectFile *getObjFileLowering() const override {
     return TLOF.get();
   }
 
   void adjustPassManager(PassManagerBuilder &) override;
+  void registerPassBuilderCallbacks(PassBuilder &PB,
+                                    bool DebugPassManager) override;
 
   TargetTransformInfo getTargetTransformInfo(const Function &F) override;
 
   bool isMachineVerifierClean() const override {
     return false;
   }
 }; // NVPTXTargetMachine.
 

llvm/lib/Target/NVPTX/NVPTXTargetMachine.cpp

 #include "TargetInfo/NVPTXTargetInfo.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/Triple.h"
 #include "llvm/Analysis/TargetTransformInfo.h"
 #include "llvm/CodeGen/Passes.h"
 #include "llvm/CodeGen/TargetPassConfig.h"
 #include "llvm/IR/LegacyPassManager.h"
 #include "llvm/Pass.h"
+#include "llvm/Passes/PassBuilder.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/TargetRegistry.h"
 #include "llvm/Target/TargetMachine.h"
 #include "llvm/Target/TargetOptions.h"
 #include "llvm/Transforms/IPO/PassManagerBuilder.h"
 #include "llvm/Transforms/Scalar.h"
 #include "llvm/Transforms/Scalar/GVN.h"
 #include "llvm/Transforms/Vectorize.h"
void NVPTXTargetMachine::adjustPassManager(PassManagerBuilder &Builder) {
   Builder.addExtension(
     PassManagerBuilder::EP_EarlyAsPossible,
     [&](const PassManagerBuilder &, legacy::PassManagerBase &PM) {
       PM.add(createNVVMReflectPass(Subtarget.getSmVersion()));
       PM.add(createNVVMIntrRangePass(Subtarget.getSmVersion()));
     });
 }
 
+void NVPTXTargetMachine::registerPassBuilderCallbacks(PassBuilder &PB,
+                                                      bool DebugPassManager) {
+  PB.registerPipelineParsingCallback(
+      [](StringRef PassName, FunctionPassManager &PM,
+         ArrayRef<PassBuilder::PipelineElement>) {
+        if (PassName == "nvvm-reflect") {
+          PM.addPass(NVVMReflectPass());
+          return true;
+        }
+        if (PassName == "nvvm-intr-range") {
+          PM.addPass(NVVMIntrRangePass());
+          return true;
+        }
+        return false;
+      });
+
+  PB.registerPipelineStartEPCallback(
+      [this, DebugPassManager](ModulePassManager &PM,
+                               PassBuilder::OptimizationLevel Level) {
+        FunctionPassManager FPM(DebugPassManager);
+        FPM.addPass(NVVMReflectPass(Subtarget.getSmVersion()));
+        FPM.addPass(NVVMIntrRangePass(Subtarget.getSmVersion()));
+        PM.addPass(createModuleToFunctionPassAdaptor(std::move(FPM)));
+      });
+}
+
 TargetTransformInfo
 NVPTXTargetMachine::getTargetTransformInfo(const Function &F) {
   return TargetTransformInfo(NVPTXTTIImpl(this, F));
 }
 
 void NVPTXPassConfig::addEarlyCSEOrGVNPass() {
   if (getOptLevel() == CodeGenOpt::Aggressive)
     addPass(createGVNPass());

llvm/lib/Target/NVPTX/NVVMIntrRange.cpp

 //===----------------------------------------------------------------------===//
 
 #include "NVPTX.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/InstIterator.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/Intrinsics.h"
 #include "llvm/IR/IntrinsicsNVPTX.h"
+#include "llvm/IR/PassManager.h"
 #include "llvm/Support/CommandLine.h"
 
 using namespace llvm;
 
 #define DEBUG_TYPE "nvvm-intr-range"
 
 namespace llvm { void initializeNVVMIntrRangePass(PassRegistry &); }
 
 // Add !range metadata based on limits of given SM variant.
 static cl::opt<unsigned> NVVMIntrRangeSM("nvvm-intr-range-sm", cl::init(20),
                                          cl::Hidden, cl::desc("SM variant"));
 
 namespace {
 class NVVMIntrRange : public FunctionPass {
  private:
-   struct {
-     unsigned x, y, z;
-   } MaxBlockSize, MaxGridSize;
+   unsigned SmVersion;
 
  public:
    static char ID;
    NVVMIntrRange() : NVVMIntrRange(NVVMIntrRangeSM) {}
-   NVVMIntrRange(unsigned int SmVersion) : FunctionPass(ID) {
-     MaxBlockSize.x = 1024;
-     MaxBlockSize.y = 1024;
-     MaxBlockSize.z = 64;
-
-     MaxGridSize.x = SmVersion >= 30 ? 0x7fffffff : 0xffff;
-     MaxGridSize.y = 0xffff;
-     MaxGridSize.z = 0xffff;
+   NVVMIntrRange(unsigned int SmVersion)
+       : FunctionPass(ID), SmVersion(SmVersion) {
 
      initializeNVVMIntrRangePass(*PassRegistry::getPassRegistry());
    }
 
    bool runOnFunction(Function &) override;
 };
 }
 
static bool addRangeMetadata(uint64_t Low, uint64_t High, CallInst *C) {
   IntegerType *Int32Ty = Type::getInt32Ty(Context);
   Metadata *LowAndHigh[] = {
       ConstantAsMetadata::get(ConstantInt::get(Int32Ty, Low)),
       ConstantAsMetadata::get(ConstantInt::get(Int32Ty, High))};
   C->setMetadata(LLVMContext::MD_range, MDNode::get(Context, LowAndHigh));
   return true;
 }
 
-bool NVVMIntrRange::runOnFunction(Function &F) {
+static bool runNVVMIntrRange(Function &F, unsigned SmVersion) {
+  struct {
+    unsigned x, y, z;
+  } MaxBlockSize, MaxGridSize;
+  MaxBlockSize.x = 1024;
+  MaxBlockSize.y = 1024;
+  MaxBlockSize.z = 64;
+
+  MaxGridSize.x = SmVersion >= 30 ? 0x7fffffff : 0xffff;
+  MaxGridSize.y = 0xffff;
+  MaxGridSize.z = 0xffff;
+
   // Go through the calls in this function.
   bool Changed = false;
   for (Instruction &I : instructions(F)) {
     CallInst *Call = dyn_cast<CallInst>(&I);
     if (!Call)
       continue;
 
     if (Function *Callee = Call->getCalledFunction()) {
if (Function *Callee = Call->getCalledFunction()) {
       default:
         break;
       }
     }
   }
 
   return Changed;
 }
+
+bool NVVMIntrRange::runOnFunction(Function &F) {
+  return runNVVMIntrRange(F, SmVersion);
+}
+
+NVVMIntrRangePass::NVVMIntrRangePass() : NVVMIntrRangePass(NVVMIntrRangeSM) {}
+
+PreservedAnalyses NVVMIntrRangePass::run(Function &F,
+                                         FunctionAnalysisManager &AM) {
+  return runNVVMIntrRange(F, SmVersion) ? PreservedAnalyses::none()
+                                        : PreservedAnalyses::all();
+}

llvm/lib/Target/NVPTX/NVVMReflect.cpp

 #include "llvm/IR/Constants.h"
 #include "llvm/IR/DerivedTypes.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/InstIterator.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/Intrinsics.h"
 #include "llvm/IR/IntrinsicsNVPTX.h"
 #include "llvm/IR/Module.h"
+#include "llvm/IR/PassManager.h"
 #include "llvm/IR/Type.h"
 #include "llvm/Pass.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_os_ostream.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Transforms/Scalar.h"
 #include <sstream>

 NVVMReflectEnabled("nvvm-reflect-enable", cl::init(true), cl::Hidden,
                    cl::desc("NVVM reflection, enabled by default"));
 
 char NVVMReflect::ID = 0;
 INITIALIZE_PASS(NVVMReflect, "nvvm-reflect",
                 "Replace occurrences of __nvvm_reflect() calls with 0/1", false,
                 false)
 
-bool NVVMReflect::runOnFunction(Function &F) {
+bool runNVVMReflect(Function &F, unsigned SmVersion) {
   if (!NVVMReflectEnabled)
     return false;
 
   if (F.getName() == NVVM_REFLECT_FUNCTION) {
     assert(F.isDeclaration() && "_reflect function should not have a body");
     assert(F.getReturnType()->isIntegerTy() &&
            "_reflect's return type should be integer");
     return false;
for (Instruction &I : instructions(F)) {
     ToRemove.push_back(Call);
   }
 
   for (Instruction *I : ToRemove)
     I->eraseFromParent();
 
   return ToRemove.size() > 0;
 }
+
+bool NVVMReflect::runOnFunction(Function &F) {
+  return runNVVMReflect(F, SmVersion);
+}
+
+NVVMReflectPass::NVVMReflectPass() : NVVMReflectPass(0) {}
+
+PreservedAnalyses NVVMReflectPass::run(Function &F,
+                                       FunctionAnalysisManager &AM) {
+  return runNVVMReflect(F, SmVersion) ? PreservedAnalyses::none()
+                                      : PreservedAnalyses::all();
+}

llvm/test/CodeGen/NVPTX/intrinsic-old.ll

 ; RUN: llc < %s -march=nvptx -mcpu=sm_20 | FileCheck -allow-deprecated-dag-overlap %s
 ; RUN: llc < %s -march=nvptx64 -mcpu=sm_20 | FileCheck -allow-deprecated-dag-overlap %s
 ; RUN: opt < %s -S -mtriple=nvptx-nvidia-cuda -nvvm-intr-range \
 ; RUN:   | FileCheck -allow-deprecated-dag-overlap --check-prefix=RANGE --check-prefix=RANGE_20 %s
+; RUN: opt < %s -S -mtriple=nvptx-nvidia-cuda -passes=nvvm-intr-range \
+; RUN:   | FileCheck -allow-deprecated-dag-overlap --check-prefix=RANGE --check-prefix=RANGE_20 %s
 ; RUN: opt < %s -S -mtriple=nvptx-nvidia-cuda \
 ; RUN:    -nvvm-intr-range -nvvm-intr-range-sm=30 \
 ; RUN:   | FileCheck -allow-deprecated-dag-overlap --check-prefix=RANGE --check-prefix=RANGE_30 %s
+; RUN: opt < %s -S -mtriple=nvptx-nvidia-cuda \
+; RUN:    -passes=nvvm-intr-range -nvvm-intr-range-sm=30 \
+; RUN:   | FileCheck -allow-deprecated-dag-overlap --check-prefix=RANGE --check-prefix=RANGE_30 %s
 
 define ptx_device i32 @test_tid_x() {
 ; CHECK: mov.u32 %r{{[0-9]+}}, %tid.x;
 ; RANGE: call i32 @llvm.nvvm.read.ptx.sreg.tid.x(), !range ![[BLK_IDX_XY:[0-9]+]]
 ; CHECK: ret;
   %x = call i32 @llvm.nvvm.read.ptx.sreg.tid.x()
   ret i32 %x
 }

llvm/test/CodeGen/NVPTX/nvvm-reflect-module-flag.ll

-; RUN: opt < %s -S -nvvm-reflect | FileCheck %s
+; RUN: opt < %s -S -mtriple=nvptx-nvidia-cuda -nvvm-reflect | FileCheck %s
+; RUN: opt < %s -S -mtriple=nvptx-nvidia-cuda -passes=nvvm-reflect | FileCheck %s
 
 declare i32 @__nvvm_reflect(i8*)
 @str = private unnamed_addr addrspace(1) constant [11 x i8] c"__CUDA_FTZ\00"
 
 define i32 @foo() {
   %call = call i32 @__nvvm_reflect(i8* addrspacecast (i8 addrspace(1)* getelementptr inbounds ([11 x i8], [11 x i8] addrspace(1)* @str, i32 0, i32 0) to i8*))
   ; CHECK: ret i32 42
   ret i32 %call
 }
 
 !llvm.module.flags = !{!0}
 !0 = !{i32 4, !"nvvm-reflect-ftz", i32 42}

llvm/test/CodeGen/NVPTX/nvvm-reflect.ll

 ; We run nvvm-reflect (and then optimize) this module twice, once with metadata
 ; that enables FTZ, and again with metadata that disables it.
 
 ; RUN: cat %s > %t.noftz
 ; RUN: echo '!0 = !{i32 4, !"nvvm-reflect-ftz", i32 0}' >> %t.noftz
-; RUN: opt %t.noftz -S -nvvm-reflect -O2 \
+; RUN: opt %t.noftz -S -mtriple=nvptx-nvidia-cuda -nvvm-reflect -O2 \
 ; RUN:   | FileCheck %s --check-prefix=USE_FTZ_0 --check-prefix=CHECK
 
 ; RUN: cat %s > %t.ftz
 ; RUN: echo '!0 = !{i32 4, !"nvvm-reflect-ftz", i32 1}' >> %t.ftz
-; RUN: opt %t.ftz -S -nvvm-reflect -O2 \
+; RUN: opt %t.ftz -S -mtriple=nvptx-nvidia-cuda -nvvm-reflect -O2 \
 ; RUN:   | FileCheck %s --check-prefix=USE_FTZ_1 --check-prefix=CHECK
 
 @str = private unnamed_addr addrspace(4) constant [11 x i8] c"__CUDA_FTZ\00"
 
 declare i32 @__nvvm_reflect(i8*)
 declare i8* @llvm.nvvm.ptr.constant.to.gen.p0i8.p4i8(i8 addrspace(4)*)
 
 ; CHECK-LABEL: @foo

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-lower-kernel-attributes"};
+      "nvvm-reflect",
+      "nvvm-intr-range",
+      "amdgpu-simplifylib",
+      "amdgpu-usenative",
+      "amdgpu-promote-alloca",
+      "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-"};