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