contrib/llvm-project/llvm/lib/Target/NVPTX/NVVMIntrRange.cpp

   1 //===- NVVMIntrRange.cpp - Set !range metadata 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 metadata for calls to NVVM
  10 // intrinsics that return a limited range of values.
  11 //
  12 //===----------------------------------------------------------------------===//
  13
  14 #include "NVPTX.h"
  15 #include "llvm/IR/Constants.h"
  16 #include "llvm/IR/InstIterator.h"
  17 #include "llvm/IR/Instructions.h"
  18 #include "llvm/IR/Intrinsics.h"
  19
  20 using namespace llvm;
  21
  22 #define DEBUG_TYPE "nvvm-intr-range"
  23
  24 namespace llvm { void initializeNVVMIntrRangePass(PassRegistry &); }
  25
  26 // Add !range metadata based on limits of given SM variant.
  27 static cl::opt<unsigned> NVVMIntrRangeSM("nvvm-intr-range-sm", cl::init(20),
  28                                          cl::Hidden, cl::desc("SM variant"));
  29
  30 namespace {
  31 class NVVMIntrRange : public FunctionPass {
  32  private:
  33    struct {
  34      unsigned x, y, z;
  35    } MaxBlockSize, MaxGridSize;
  36
  37  public:
  38    static char ID;
  39    NVVMIntrRange() : NVVMIntrRange(NVVMIntrRangeSM) {}
  40    NVVMIntrRange(unsigned int SmVersion) : FunctionPass(ID) {
  41      MaxBlockSize.x = 1024;
  42      MaxBlockSize.y = 1024;
  43      MaxBlockSize.z = 64;
  44
  45      MaxGridSize.x = SmVersion >= 30 ? 0x7fffffff : 0xffff;
  46      MaxGridSize.y = 0xffff;
  47      MaxGridSize.z = 0xffff;
  48
  49      initializeNVVMIntrRangePass(*PassRegistry::getPassRegistry());
  50    }
  51
  52    bool runOnFunction(Function &) override;
  53 };
  54 }
  55
  56 FunctionPass *llvm::createNVVMIntrRangePass(unsigned int SmVersion) {
  57   return new NVVMIntrRange(SmVersion);
  58 }
  59
  60 char NVVMIntrRange::ID = 0;
  61 INITIALIZE_PASS(NVVMIntrRange, "nvvm-intr-range",
  62                 "Add !range metadata to NVVM intrinsics.", false, false)
  63
  64 // Adds the passed-in [Low,High) range information as metadata to the
  65 // passed-in call instruction.
  66 static bool addRangeMetadata(uint64_t Low, uint64_t High, CallInst *C) {
  67   // This call already has range metadata, nothing to do.
  68   if (C->getMetadata(LLVMContext::MD_range))
  69     return false;
  70
  71   LLVMContext &Context = C->getParent()->getContext();
  72   IntegerType *Int32Ty = Type::getInt32Ty(Context);
  73   Metadata *LowAndHigh[] = {
  74       ConstantAsMetadata::get(ConstantInt::get(Int32Ty, Low)),
  75       ConstantAsMetadata::get(ConstantInt::get(Int32Ty, High))};
  76   C->setMetadata(LLVMContext::MD_range, MDNode::get(Context, LowAndHigh));
  77   return true;
  78 }
  79
  80 bool NVVMIntrRange::runOnFunction(Function &F) {
  81   // Go through the calls in this function.
  82   bool Changed = false;
  83   for (Instruction &I : instructions(F)) {
  84     CallInst *Call = dyn_cast<CallInst>(&I);
  85     if (!Call)
  86       continue;
  87
  88     if (Function *Callee = Call->getCalledFunction()) {
  89       switch (Callee->getIntrinsicID()) {
  90       // Index within block
  91       case Intrinsic::nvvm_read_ptx_sreg_tid_x:
  92         Changed |= addRangeMetadata(0, MaxBlockSize.x, Call);
  93         break;
  94       case Intrinsic::nvvm_read_ptx_sreg_tid_y:
  95         Changed |= addRangeMetadata(0, MaxBlockSize.y, Call);
  96         break;
  97       case Intrinsic::nvvm_read_ptx_sreg_tid_z:
  98         Changed |= addRangeMetadata(0, MaxBlockSize.z, Call);
  99         break;
 100
 101       // Block size
 102       case Intrinsic::nvvm_read_ptx_sreg_ntid_x:
 103         Changed |= addRangeMetadata(1, MaxBlockSize.x+1, Call);
 104         break;
 105       case Intrinsic::nvvm_read_ptx_sreg_ntid_y:
 106         Changed |= addRangeMetadata(1, MaxBlockSize.y+1, Call);
 107         break;
 108       case Intrinsic::nvvm_read_ptx_sreg_ntid_z:
 109         Changed |= addRangeMetadata(1, MaxBlockSize.z+1, Call);
 110         break;
 111
 112       // Index within grid
 113       case Intrinsic::nvvm_read_ptx_sreg_ctaid_x:
 114         Changed |= addRangeMetadata(0, MaxGridSize.x, Call);
 115         break;
 116       case Intrinsic::nvvm_read_ptx_sreg_ctaid_y:
 117         Changed |= addRangeMetadata(0, MaxGridSize.y, Call);
 118         break;
 119       case Intrinsic::nvvm_read_ptx_sreg_ctaid_z:
 120         Changed |= addRangeMetadata(0, MaxGridSize.z, Call);
 121         break;
 122
 123       // Grid size
 124       case Intrinsic::nvvm_read_ptx_sreg_nctaid_x:
 125         Changed |= addRangeMetadata(1, MaxGridSize.x+1, Call);
 126         break;
 127       case Intrinsic::nvvm_read_ptx_sreg_nctaid_y:
 128         Changed |= addRangeMetadata(1, MaxGridSize.y+1, Call);
 129         break;
 130       case Intrinsic::nvvm_read_ptx_sreg_nctaid_z:
 131         Changed |= addRangeMetadata(1, MaxGridSize.z+1, Call);
 132         break;
 133
 134       // warp size is constant 32.
 135       case Intrinsic::nvvm_read_ptx_sreg_warpsize:
 136         Changed |= addRangeMetadata(32, 32+1, Call);
 137         break;
 138
 139       // Lane ID is [0..warpsize)
 140       case Intrinsic::nvvm_read_ptx_sreg_laneid:
 141         Changed |= addRangeMetadata(0, 32, Call);
 142         break;
 143
 144       default:
 145         break;
 146       }
 147     }
 148   }
 149
 150   return Changed;
 151 }