llvm/lib/Target/NVPTX/NVVMIntrRange.cpp

   1 //===- NVVMIntrRange.cpp - Set range attributes for NVVM intrinsics -------===//
   2 //
   3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
   4 // See https://llvm.org/LICENSE.txt for license information.
   5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
   6 //
   7 //===----------------------------------------------------------------------===//
   8 //
   9 // This pass adds appropriate range attributes for calls to NVVM
  10 // intrinsics that return a limited range of values.
  11 //
  12 //===----------------------------------------------------------------------===//
  13
  14 #include "NVPTX.h"
  15 #include "NVPTXUtilities.h"
  16 #include "llvm/IR/InstIterator.h"
  17 #include "llvm/IR/Instructions.h"
  18 #include "llvm/IR/IntrinsicInst.h"
  19 #include "llvm/IR/Intrinsics.h"
  20 #include "llvm/IR/IntrinsicsNVPTX.h"
  21 #include "llvm/IR/PassManager.h"
  22 #include "llvm/Support/CommandLine.h"
  23 #include <cstdint>
  24
  25 using namespace llvm;
  26
  27 #define DEBUG_TYPE "nvvm-intr-range"
  28
  29 namespace llvm { void initializeNVVMIntrRangePass(PassRegistry &); }
  30
  31 namespace {
  32 class NVVMIntrRange : public FunctionPass {
  33 public:
  34   static char ID;
  35   NVVMIntrRange() : FunctionPass(ID) {
  36
  37     initializeNVVMIntrRangePass(*PassRegistry::getPassRegistry());
  38   }
  39
  40   bool runOnFunction(Function &) override;
  41 };
  42 } // namespace
  43
  44 FunctionPass *llvm::createNVVMIntrRangePass() { return new NVVMIntrRange(); }
  45
  46 char NVVMIntrRange::ID = 0;
  47 INITIALIZE_PASS(NVVMIntrRange, "nvvm-intr-range",
  48                 "Add !range metadata to NVVM intrinsics.", false, false)
  49
  50 // Adds the passed-in [Low,High) range information as metadata to the
  51 // passed-in call instruction.
  52 static bool addRangeAttr(uint64_t Low, uint64_t High, IntrinsicInst *II) {
  53   if (II->getMetadata(LLVMContext::MD_range))
  54     return false;
  55
  56   const uint64_t BitWidth = II->getType()->getIntegerBitWidth();
  57   ConstantRange Range(APInt(BitWidth, Low), APInt(BitWidth, High));
  58
  59   if (auto CurrentRange = II->getRange())
  60     Range = Range.intersectWith(CurrentRange.value());
  61
  62   II->addRangeRetAttr(Range);
  63   return true;
  64 }
  65
  66 static bool runNVVMIntrRange(Function &F) {
  67   struct {
  68     unsigned x, y, z;
  69   } MaxBlockSize, MaxGridSize;
  70
  71   const unsigned MetadataNTID = getReqNTID(F).value_or(
  72       getMaxNTID(F).value_or(std::numeric_limits<unsigned>::max()));
  73
  74   MaxBlockSize.x = std::min(1024u, MetadataNTID);
  75   MaxBlockSize.y = std::min(1024u, MetadataNTID);
  76   MaxBlockSize.z = std::min(64u, MetadataNTID);
  77
  78   MaxGridSize.x = 0x7fffffff;
  79   MaxGridSize.y = 0xffff;
  80   MaxGridSize.z = 0xffff;
  81
  82   // Go through the calls in this function.
  83   bool Changed = false;
  84   for (Instruction &I : instructions(F)) {
  85     IntrinsicInst *II = dyn_cast<IntrinsicInst>(&I);
  86     if (!II)
  87       continue;
  88
  89     switch (II->getIntrinsicID()) {
  90     // Index within block
  91     case Intrinsic::nvvm_read_ptx_sreg_tid_x:
  92       Changed |= addRangeAttr(0, MaxBlockSize.x, II);
  93       break;
  94     case Intrinsic::nvvm_read_ptx_sreg_tid_y:
  95       Changed |= addRangeAttr(0, MaxBlockSize.y, II);
  96       break;
  97     case Intrinsic::nvvm_read_ptx_sreg_tid_z:
  98       Changed |= addRangeAttr(0, MaxBlockSize.z, II);
  99       break;
 100
 101     // Block size
 102     case Intrinsic::nvvm_read_ptx_sreg_ntid_x:
 103       Changed |= addRangeAttr(1, MaxBlockSize.x + 1, II);
 104       break;
 105     case Intrinsic::nvvm_read_ptx_sreg_ntid_y:
 106       Changed |= addRangeAttr(1, MaxBlockSize.y + 1, II);
 107       break;
 108     case Intrinsic::nvvm_read_ptx_sreg_ntid_z:
 109       Changed |= addRangeAttr(1, MaxBlockSize.z + 1, II);
 110       break;
 111
 112     // Index within grid
 113     case Intrinsic::nvvm_read_ptx_sreg_ctaid_x:
 114       Changed |= addRangeAttr(0, MaxGridSize.x, II);
 115       break;
 116     case Intrinsic::nvvm_read_ptx_sreg_ctaid_y:
 117       Changed |= addRangeAttr(0, MaxGridSize.y, II);
 118       break;
 119     case Intrinsic::nvvm_read_ptx_sreg_ctaid_z:
 120       Changed |= addRangeAttr(0, MaxGridSize.z, II);
 121       break;
 122
 123     // Grid size
 124     case Intrinsic::nvvm_read_ptx_sreg_nctaid_x:
 125       Changed |= addRangeAttr(1, MaxGridSize.x + 1, II);
 126       break;
 127     case Intrinsic::nvvm_read_ptx_sreg_nctaid_y:
 128       Changed |= addRangeAttr(1, MaxGridSize.y + 1, II);
 129       break;
 130     case Intrinsic::nvvm_read_ptx_sreg_nctaid_z:
 131       Changed |= addRangeAttr(1, MaxGridSize.z + 1, II);
 132       break;
 133
 134     // warp size is constant 32.
 135     case Intrinsic::nvvm_read_ptx_sreg_warpsize:
 136       Changed |= addRangeAttr(32, 32 + 1, II);
 137       break;
 138
 139     // Lane ID is [0..warpsize)
 140     case Intrinsic::nvvm_read_ptx_sreg_laneid:
 141       Changed |= addRangeAttr(0, 32, II);
 142       break;
 143
 144     default:
 145       break;
 146     }
 147   }
 148
 149   return Changed;
 150 }
 151
 152 bool NVVMIntrRange::runOnFunction(Function &F) { return runNVVMIntrRange(F); }
 153
 154 PreservedAnalyses NVVMIntrRangePass::run(Function &F,
 155                                          FunctionAnalysisManager &AM) {
 156   return runNVVMIntrRange(F) ? PreservedAnalyses::none()
 157                              : PreservedAnalyses::all();
 158 }