llvm/lib/Target/AMDGPU/AMDGPULowerKernelAttributes.cpp

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