lib/Target/AMDGPU/AMDGPULowerKernelAttributes.cpp

   1 //===-- AMDGPULowerKernelAttributes.cpp ------------------------------------------===//
   2 //
   3 //                     The LLVM Compiler Infrastructure
   4 //
   5 // This file is distributed under the University of Illinois Open Source
   6 // License. See LICENSE.TXT for details.
   7 //
   8 //===----------------------------------------------------------------------===//
   9 //
  10 /// \file This pass does attempts to make use of reqd_work_group_size metadata
  11 /// to eliminate loads from the dispatch packet and to constant fold OpenCL
  12 /// get_local_size-like functions.
  13 //
  14 //===----------------------------------------------------------------------===//
  15
  16 #include "AMDGPU.h"
  17 #include "AMDGPUTargetMachine.h"
  18 #include "llvm/Analysis/ValueTracking.h"
  19 #include "llvm/CodeGen/Passes.h"
  20 #include "llvm/CodeGen/TargetPassConfig.h"
  21 #include "llvm/IR/Constants.h"
  22 #include "llvm/IR/Function.h"
  23 #include "llvm/IR/Instructions.h"
  24 #include "llvm/IR/PatternMatch.h"
  25 #include "llvm/Pass.h"
  26
  27 #define DEBUG_TYPE "amdgpu-lower-kernel-attributes"
  28
  29 using namespace llvm;
  30
  31 namespace {
  32
  33 // Field offsets in hsa_kernel_dispatch_packet_t.
  34 enum DispatchPackedOffsets {
  35   WORKGROUP_SIZE_X = 4,
  36   WORKGROUP_SIZE_Y = 6,
  37   WORKGROUP_SIZE_Z = 8,
  38
  39   GRID_SIZE_X = 12,
  40   GRID_SIZE_Y = 16,
  41   GRID_SIZE_Z = 20
  42 };
  43
  44 class AMDGPULowerKernelAttributes : public ModulePass {
  45   Module *Mod = nullptr;
  46
  47 public:
  48   static char ID;
  49
  50   AMDGPULowerKernelAttributes() : ModulePass(ID) {}
  51
  52   bool processUse(CallInst *CI);
  53
  54   bool doInitialization(Module &M) override;
  55   bool runOnModule(Module &M) override;
  56
  57   StringRef getPassName() const override {
  58     return "AMDGPU Kernel Attributes";
  59   }
  60
  61   void getAnalysisUsage(AnalysisUsage &AU) const override {
  62     AU.setPreservesAll();
  63  }
  64 };
  65
  66 } // end anonymous namespace
  67
  68 bool AMDGPULowerKernelAttributes::doInitialization(Module &M) {
  69   Mod = &M;
  70   return false;
  71 }
  72
  73 bool AMDGPULowerKernelAttributes::processUse(CallInst *CI) {
  74   Function *F = CI->getParent()->getParent();
  75
  76   auto MD = F->getMetadata("reqd_work_group_size");
  77   const bool HasReqdWorkGroupSize = MD && MD->getNumOperands() == 3;
  78
  79   const bool HasUniformWorkGroupSize =
  80     F->getFnAttribute("uniform-work-group-size").getValueAsString() == "true";
  81
  82   if (!HasReqdWorkGroupSize && !HasUniformWorkGroupSize)
  83     return false;
  84
  85   Value *WorkGroupSizeX = nullptr;
  86   Value *WorkGroupSizeY = nullptr;
  87   Value *WorkGroupSizeZ = nullptr;
  88
  89   Value *GridSizeX = nullptr;
  90   Value *GridSizeY = nullptr;
  91   Value *GridSizeZ = nullptr;
  92
  93   const DataLayout &DL = Mod->getDataLayout();
  94
  95   // We expect to see several GEP users, casted to the appropriate type and
  96   // loaded.
  97   for (User *U : CI->users()) {
  98     if (!U->hasOneUse())
  99       continue;
 100
 101     int64_t Offset = 0;
 102     if (GetPointerBaseWithConstantOffset(U, Offset, DL) != CI)
 103       continue;
 104
 105     auto *BCI = dyn_cast<BitCastInst>(*U->user_begin());
 106     if (!BCI || !BCI->hasOneUse())
 107       continue;
 108
 109     auto *Load = dyn_cast<LoadInst>(*BCI->user_begin());
 110     if (!Load || !Load->isSimple())
 111       continue;
 112
 113     unsigned LoadSize = DL.getTypeStoreSize(Load->getType());
 114
 115     // TODO: Handle merged loads.
 116     switch (Offset) {
 117     case WORKGROUP_SIZE_X:
 118       if (LoadSize == 2)
 119         WorkGroupSizeX = Load;
 120       break;
 121     case WORKGROUP_SIZE_Y:
 122       if (LoadSize == 2)
 123         WorkGroupSizeY = Load;
 124       break;
 125     case WORKGROUP_SIZE_Z:
 126       if (LoadSize == 2)
 127         WorkGroupSizeZ = Load;
 128       break;
 129     case GRID_SIZE_X:
 130       if (LoadSize == 4)
 131         GridSizeX = Load;
 132       break;
 133     case GRID_SIZE_Y:
 134       if (LoadSize == 4)
 135         GridSizeY = Load;
 136       break;
 137     case GRID_SIZE_Z:
 138       if (LoadSize == 4)
 139         GridSizeZ = Load;
 140       break;
 141     default:
 142       break;
 143     }
 144   }
 145
 146   // Pattern match the code used to handle partial workgroup dispatches in the
 147   // library implementation of get_local_size, so the entire function can be
 148   // constant folded with a known group size.
 149   //
 150   // uint r = grid_size - group_id * group_size;
 151   // get_local_size = (r < group_size) ? r : group_size;
 152   //
 153   // If we have uniform-work-group-size (which is the default in OpenCL 1.2),
 154   // the grid_size is required to be a multiple of group_size). In this case:
 155   //
 156   // grid_size - (group_id * group_size) < group_size
 157   // ->
 158   // grid_size < group_size + (group_id * group_size)
 159   //
 160   // (grid_size / group_size) < 1 + group_id
 161   //
 162   // grid_size / group_size is at least 1, so we can conclude the select
 163   // condition is false (except for group_id == 0, where the select result is
 164   // the same).
 165
 166   bool MadeChange = false;
 167   Value *WorkGroupSizes[3] = { WorkGroupSizeX, WorkGroupSizeY, WorkGroupSizeZ };
 168   Value *GridSizes[3] = { GridSizeX, GridSizeY, GridSizeZ };
 169
 170   for (int I = 0; HasUniformWorkGroupSize && I < 3; ++I) {
 171     Value *GroupSize = WorkGroupSizes[I];
 172     Value *GridSize = GridSizes[I];
 173     if (!GroupSize || !GridSize)
 174       continue;
 175
 176     for (User *U : GroupSize->users()) {
 177       auto *ZextGroupSize = dyn_cast<ZExtInst>(U);
 178       if (!ZextGroupSize)
 179         continue;
 180
 181       for (User *ZextUser : ZextGroupSize->users()) {
 182         auto *SI = dyn_cast<SelectInst>(ZextUser);
 183         if (!SI)
 184           continue;
 185
 186         using namespace llvm::PatternMatch;
 187         auto GroupIDIntrin = I == 0 ?
 188           m_Intrinsic<Intrinsic::amdgcn_workgroup_id_x>() :
 189             (I == 1 ? m_Intrinsic<Intrinsic::amdgcn_workgroup_id_y>() :
 190                       m_Intrinsic<Intrinsic::amdgcn_workgroup_id_z>());
 191
 192         auto SubExpr = m_Sub(m_Specific(GridSize),
 193                              m_Mul(GroupIDIntrin, m_Specific(ZextGroupSize)));
 194
 195         ICmpInst::Predicate Pred;
 196         if (match(SI,
 197                   m_Select(m_ICmp(Pred, SubExpr, m_Specific(ZextGroupSize)),
 198                            SubExpr,
 199                            m_Specific(ZextGroupSize))) &&
 200             Pred == ICmpInst::ICMP_ULT) {
 201           if (HasReqdWorkGroupSize) {
 202             ConstantInt *KnownSize
 203               = mdconst::extract<ConstantInt>(MD->getOperand(I));
 204             SI->replaceAllUsesWith(ConstantExpr::getIntegerCast(KnownSize,
 205                                                                 SI->getType(),
 206                                                                 false));
 207           } else {
 208             SI->replaceAllUsesWith(ZextGroupSize);
 209           }
 210
 211           MadeChange = true;
 212         }
 213       }
 214     }
 215   }
 216
 217   if (!HasReqdWorkGroupSize)
 218     return MadeChange;
 219
 220   // Eliminate any other loads we can from the dispatch packet.
 221   for (int I = 0; I < 3; ++I) {
 222     Value *GroupSize = WorkGroupSizes[I];
 223     if (!GroupSize)
 224       continue;
 225
 226     ConstantInt *KnownSize = mdconst::extract<ConstantInt>(MD->getOperand(I));
 227     GroupSize->replaceAllUsesWith(
 228       ConstantExpr::getIntegerCast(KnownSize,
 229                                    GroupSize->getType(),
 230                                    false));
 231     MadeChange = true;
 232   }
 233
 234   return MadeChange;
 235 }
 236
 237 // TODO: Move makeLIDRangeMetadata usage into here. Seem to not get
 238 // TargetPassConfig for subtarget.
 239 bool AMDGPULowerKernelAttributes::runOnModule(Module &M) {
 240   StringRef DispatchPtrName
 241     = Intrinsic::getName(Intrinsic::amdgcn_dispatch_ptr);
 242
 243   Function *DispatchPtr = Mod->getFunction(DispatchPtrName);
 244   if (!DispatchPtr) // Dispatch ptr not used.
 245     return false;
 246
 247   bool MadeChange = false;
 248
 249   SmallPtrSet<Instruction *, 4> HandledUses;
 250   for (auto *U : DispatchPtr->users()) {
 251     CallInst *CI = cast<CallInst>(U);
 252     if (HandledUses.insert(CI).second) {
 253       if (processUse(CI))
 254         MadeChange = true;
 255     }
 256   }
 257
 258   return MadeChange;
 259 }
 260
 261 INITIALIZE_PASS_BEGIN(AMDGPULowerKernelAttributes, DEBUG_TYPE,
 262                       "AMDGPU IR optimizations", false, false)
 263 INITIALIZE_PASS_END(AMDGPULowerKernelAttributes, DEBUG_TYPE, "AMDGPU IR optimizations",
 264                     false, false)
 265
 266 char AMDGPULowerKernelAttributes::ID = 0;
 267
 268 ModulePass *llvm::createAMDGPULowerKernelAttributesPass() {
 269   return new AMDGPULowerKernelAttributes();
 270 }