lib/Target/AMDGPU/AMDGPULowerKernelAttributes.cpp - llvm-project/llvm - Git at Google

 //===-- AMDGPULowerKernelAttributes.cpp ------------------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 /// \file This pass does attempts to make use of reqd_work_group_size metadata
 /// to eliminate loads from the dispatch packet and to constant fold OpenCL
 /// get_local_size-like functions.
 //
 //===----------------------------------------------------------------------===//

 #include "AMDGPU.h"
 #include "llvm/Analysis/ValueTracking.h"
 #include "llvm/CodeGen/Passes.h"
 #include "llvm/CodeGen/TargetPassConfig.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/InstIterator.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/IntrinsicsAMDGPU.h"
 #include "llvm/IR/PatternMatch.h"
 #include "llvm/Pass.h"

 #define DEBUG_TYPE "amdgpu-lower-kernel-attributes"

 using namespace llvm;

 namespace {

 // Field offsets in hsa_kernel_dispatch_packet_t.
 enum DispatchPackedOffsets {
   WORKGROUP_SIZE_X = 4,
   WORKGROUP_SIZE_Y = 6,
   WORKGROUP_SIZE_Z = 8,

   GRID_SIZE_X = 12,
   GRID_SIZE_Y = 16,
   GRID_SIZE_Z = 20
 };

 class AMDGPULowerKernelAttributes : public ModulePass {
 public:
   static char ID;

   AMDGPULowerKernelAttributes() : ModulePass(ID) {}

   bool runOnModule(Module &M) override;

   StringRef getPassName() const override {
     return "AMDGPU Kernel Attributes";
   }

   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.setPreservesAll();
  }
 };

 } // end anonymous namespace

 static bool processUse(CallInst *CI) {
   Function *F = CI->getParent()->getParent();

   auto MD = F->getMetadata("reqd_work_group_size");
   const bool HasReqdWorkGroupSize = MD && MD->getNumOperands() == 3;

   const bool HasUniformWorkGroupSize =
     F->getFnAttribute("uniform-work-group-size").getValueAsString() == "true";

   if (!HasReqdWorkGroupSize && !HasUniformWorkGroupSize)
     return false;

   Value *WorkGroupSizeX = nullptr;
   Value *WorkGroupSizeY = nullptr;
   Value *WorkGroupSizeZ = nullptr;

   Value *GridSizeX = nullptr;
   Value *GridSizeY = nullptr;
   Value *GridSizeZ = nullptr;

   const DataLayout &DL = F->getParent()->getDataLayout();

   // We expect to see several GEP users, casted to the appropriate type and
   // loaded.
   for (User *U : CI->users()) {
     if (!U->hasOneUse())
       continue;

     int64_t Offset = 0;
     if (GetPointerBaseWithConstantOffset(U, Offset, DL) != CI)
       continue;

     auto *BCI = dyn_cast<BitCastInst>(*U->user_begin());
     if (!BCI || !BCI->hasOneUse())
       continue;

     auto *Load = dyn_cast<LoadInst>(*BCI->user_begin());
     if (!Load || !Load->isSimple())
       continue;

     unsigned LoadSize = DL.getTypeStoreSize(Load->getType());

     // TODO: Handle merged loads.
     switch (Offset) {
     case WORKGROUP_SIZE_X:
       if (LoadSize == 2)
         WorkGroupSizeX = Load;
       break;
     case WORKGROUP_SIZE_Y:
       if (LoadSize == 2)
         WorkGroupSizeY = Load;
       break;
     case WORKGROUP_SIZE_Z:
       if (LoadSize == 2)
         WorkGroupSizeZ = Load;
       break;
     case GRID_SIZE_X:
       if (LoadSize == 4)
         GridSizeX = Load;
       break;
     case GRID_SIZE_Y:
       if (LoadSize == 4)
         GridSizeY = Load;
       break;
     case GRID_SIZE_Z:
       if (LoadSize == 4)
         GridSizeZ = Load;
       break;
     default:
       break;
     }
   }

   // Pattern match the code used to handle partial workgroup dispatches in the
   // library implementation of get_local_size, so the entire function can be
   // constant folded with a known group size.
   //
   // uint r = grid_size - group_id * group_size;
   // get_local_size = (r < group_size) ? r : group_size;
   //
   // If we have uniform-work-group-size (which is the default in OpenCL 1.2),
   // the grid_size is required to be a multiple of group_size). In this case:
   //
   // grid_size - (group_id * group_size) < group_size
   // ->
   // grid_size < group_size + (group_id * group_size)
   //
   // (grid_size / group_size) < 1 + group_id
   //
   // grid_size / group_size is at least 1, so we can conclude the select
   // condition is false (except for group_id == 0, where the select result is
   // the same).

   bool MadeChange = false;
   Value *WorkGroupSizes[3] = { WorkGroupSizeX, WorkGroupSizeY, WorkGroupSizeZ };
   Value *GridSizes[3] = { GridSizeX, GridSizeY, GridSizeZ };

   for (int I = 0; HasUniformWorkGroupSize && I < 3; ++I) {
     Value *GroupSize = WorkGroupSizes[I];
     Value *GridSize = GridSizes[I];
     if (!GroupSize || !GridSize)
       continue;

     for (User *U : GroupSize->users()) {
       auto *ZextGroupSize = dyn_cast<ZExtInst>(U);
       if (!ZextGroupSize)
         continue;

       for (User *ZextUser : ZextGroupSize->users()) {
         auto *SI = dyn_cast<SelectInst>(ZextUser);
         if (!SI)
           continue;

         using namespace llvm::PatternMatch;
         auto GroupIDIntrin = I == 0 ?
           m_Intrinsic<Intrinsic::amdgcn_workgroup_id_x>() :
             (I == 1 ? m_Intrinsic<Intrinsic::amdgcn_workgroup_id_y>() :
                       m_Intrinsic<Intrinsic::amdgcn_workgroup_id_z>());

         auto SubExpr = m_Sub(m_Specific(GridSize),
                              m_Mul(GroupIDIntrin, m_Specific(ZextGroupSize)));

         ICmpInst::Predicate Pred;
         if (match(SI,
                   m_Select(m_ICmp(Pred, SubExpr, m_Specific(ZextGroupSize)),
                            SubExpr,
                            m_Specific(ZextGroupSize))) &&
             Pred == ICmpInst::ICMP_ULT) {
           if (HasReqdWorkGroupSize) {
             ConstantInt *KnownSize
               = mdconst::extract<ConstantInt>(MD->getOperand(I));
             SI->replaceAllUsesWith(ConstantExpr::getIntegerCast(KnownSize,
                                                                 SI->getType(),
                                                                 false));
           } else {
             SI->replaceAllUsesWith(ZextGroupSize);
           }

           MadeChange = true;
         }
       }
     }
   }

   if (!HasReqdWorkGroupSize)
     return MadeChange;

   // Eliminate any other loads we can from the dispatch packet.
   for (int I = 0; I < 3; ++I) {
     Value *GroupSize = WorkGroupSizes[I];
     if (!GroupSize)
       continue;

     ConstantInt *KnownSize = mdconst::extract<ConstantInt>(MD->getOperand(I));
     GroupSize->replaceAllUsesWith(
       ConstantExpr::getIntegerCast(KnownSize,
                                    GroupSize->getType(),
                                    false));
     MadeChange = true;
   }

   return MadeChange;
 }

 // TODO: Move makeLIDRangeMetadata usage into here. Seem to not get
 // TargetPassConfig for subtarget.
 bool AMDGPULowerKernelAttributes::runOnModule(Module &M) {
   StringRef DispatchPtrName
     = Intrinsic::getName(Intrinsic::amdgcn_dispatch_ptr);

   Function *DispatchPtr = M.getFunction(DispatchPtrName);
   if (!DispatchPtr) // Dispatch ptr not used.
     return false;

   bool MadeChange = false;

   SmallPtrSet<Instruction *, 4> HandledUses;
   for (auto *U : DispatchPtr->users()) {
     CallInst *CI = cast<CallInst>(U);
     if (HandledUses.insert(CI).second) {
       if (processUse(CI))
         MadeChange = true;
     }
   }

   return MadeChange;
 }

 INITIALIZE_PASS_BEGIN(AMDGPULowerKernelAttributes, DEBUG_TYPE,
                       "AMDGPU IR optimizations", false, false)
 INITIALIZE_PASS_END(AMDGPULowerKernelAttributes, DEBUG_TYPE, "AMDGPU IR optimizations",
                     false, false)

 char AMDGPULowerKernelAttributes::ID = 0;

 ModulePass *llvm::createAMDGPULowerKernelAttributesPass() {
   return new AMDGPULowerKernelAttributes();
 }

 PreservedAnalyses
 AMDGPULowerKernelAttributesPass::run(Function &F, FunctionAnalysisManager &AM) {
   StringRef DispatchPtrName =
       Intrinsic::getName(Intrinsic::amdgcn_dispatch_ptr);

   Function *DispatchPtr = F.getParent()->getFunction(DispatchPtrName);
   if (!DispatchPtr) // Dispatch ptr not used.
     return PreservedAnalyses::all();

   for (Instruction &I : instructions(F)) {
     if (CallInst *CI = dyn_cast<CallInst>(&I)) {
       if (CI->getCalledFunction() == DispatchPtr)
         processUse(CI);
     }
   }

   return PreservedAnalyses::all();
 }
	//===-- AMDGPULowerKernelAttributes.cpp ------------------------------------------===//
	//
	// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
	// See https://llvm.org/LICENSE.txt for license information.
	// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
	//
	//===----------------------------------------------------------------------===//
	//
	/// \file This pass does attempts to make use of reqd_work_group_size metadata
	/// to eliminate loads from the dispatch packet and to constant fold OpenCL
	/// get_local_size-like functions.
	//
	//===----------------------------------------------------------------------===//

	#include "AMDGPU.h"
	#include "llvm/Analysis/ValueTracking.h"
	#include "llvm/CodeGen/Passes.h"
	#include "llvm/CodeGen/TargetPassConfig.h"
	#include "llvm/IR/Constants.h"
	#include "llvm/IR/Function.h"
	#include "llvm/IR/InstIterator.h"
	#include "llvm/IR/Instructions.h"
	#include "llvm/IR/IntrinsicsAMDGPU.h"
	#include "llvm/IR/PatternMatch.h"
	#include "llvm/Pass.h"

	#define DEBUG_TYPE "amdgpu-lower-kernel-attributes"

	using namespace llvm;

	namespace {

	// Field offsets in hsa_kernel_dispatch_packet_t.
	enum DispatchPackedOffsets {
	WORKGROUP_SIZE_X = 4,
	WORKGROUP_SIZE_Y = 6,
	WORKGROUP_SIZE_Z = 8,

	GRID_SIZE_X = 12,
	GRID_SIZE_Y = 16,
	GRID_SIZE_Z = 20
	};

	class AMDGPULowerKernelAttributes : public ModulePass {
	public:
	static char ID;

	AMDGPULowerKernelAttributes() : ModulePass(ID) {}

	bool runOnModule(Module &M) override;

	StringRef getPassName() const override {
	return "AMDGPU Kernel Attributes";
	}

	void getAnalysisUsage(AnalysisUsage &AU) const override {
	AU.setPreservesAll();
	}
	};

	} // end anonymous namespace

	static bool processUse(CallInst *CI) {
	Function *F = CI->getParent()->getParent();

	auto MD = F->getMetadata("reqd_work_group_size");
	const bool HasReqdWorkGroupSize = MD && MD->getNumOperands() == 3;

	const bool HasUniformWorkGroupSize =
	F->getFnAttribute("uniform-work-group-size").getValueAsString() == "true";

	if (!HasReqdWorkGroupSize && !HasUniformWorkGroupSize)
	return false;

	Value *WorkGroupSizeX = nullptr;
	Value *WorkGroupSizeY = nullptr;
	Value *WorkGroupSizeZ = nullptr;

	Value *GridSizeX = nullptr;
	Value *GridSizeY = nullptr;
	Value *GridSizeZ = nullptr;

	const DataLayout &DL = F->getParent()->getDataLayout();

	// We expect to see several GEP users, casted to the appropriate type and
	// loaded.
	for (User *U : CI->users()) {
	if (!U->hasOneUse())
	continue;

	int64_t Offset = 0;
	if (GetPointerBaseWithConstantOffset(U, Offset, DL) != CI)
	continue;

	auto BCI = dyn_cast<BitCastInst>(U->user_begin());
	if (!BCI \|\| !BCI->hasOneUse())
	continue;

	auto Load = dyn_cast<LoadInst>(BCI->user_begin());
	if (!Load \|\| !Load->isSimple())
	continue;

	unsigned LoadSize = DL.getTypeStoreSize(Load->getType());

	// TODO: Handle merged loads.
	switch (Offset) {
	case WORKGROUP_SIZE_X:
	if (LoadSize == 2)
	WorkGroupSizeX = Load;
	break;
	case WORKGROUP_SIZE_Y:
	if (LoadSize == 2)
	WorkGroupSizeY = Load;
	break;
	case WORKGROUP_SIZE_Z:
	if (LoadSize == 2)
	WorkGroupSizeZ = Load;
	break;
	case GRID_SIZE_X:
	if (LoadSize == 4)
	GridSizeX = Load;
	break;
	case GRID_SIZE_Y:
	if (LoadSize == 4)
	GridSizeY = Load;
	break;
	case GRID_SIZE_Z:
	if (LoadSize == 4)
	GridSizeZ = Load;
	break;
	default:
	break;
	}
	}

	// Pattern match the code used to handle partial workgroup dispatches in the
	// library implementation of get_local_size, so the entire function can be
	// constant folded with a known group size.
	//
	// uint r = grid_size - group_id * group_size;
	// get_local_size = (r < group_size) ? r : group_size;
	//
	// If we have uniform-work-group-size (which is the default in OpenCL 1.2),
	// the grid_size is required to be a multiple of group_size). In this case:
	//
	// grid_size - (group_id * group_size) < group_size
	// ->
	// grid_size < group_size + (group_id * group_size)
	//
	// (grid_size / group_size) < 1 + group_id
	//
	// grid_size / group_size is at least 1, so we can conclude the select
	// condition is false (except for group_id == 0, where the select result is
	// the same).

	bool MadeChange = false;
	Value *WorkGroupSizes[3] = { WorkGroupSizeX, WorkGroupSizeY, WorkGroupSizeZ };
	Value *GridSizes[3] = { GridSizeX, GridSizeY, GridSizeZ };

	for (int I = 0; HasUniformWorkGroupSize && I < 3; ++I) {
	Value *GroupSize = WorkGroupSizes[I];
	Value *GridSize = GridSizes[I];
	if (!GroupSize \|\| !GridSize)
	continue;

	for (User *U : GroupSize->users()) {
	auto *ZextGroupSize = dyn_cast<ZExtInst>(U);
	if (!ZextGroupSize)
	continue;

	for (User *ZextUser : ZextGroupSize->users()) {
	auto *SI = dyn_cast<SelectInst>(ZextUser);
	if (!SI)
	continue;

	using namespace llvm::PatternMatch;
	auto GroupIDIntrin = I == 0 ?
	m_Intrinsic<Intrinsic::amdgcn_workgroup_id_x>() :
	(I == 1 ? m_Intrinsic<Intrinsic::amdgcn_workgroup_id_y>() :
	m_Intrinsic<Intrinsic::amdgcn_workgroup_id_z>());

	auto SubExpr = m_Sub(m_Specific(GridSize),
	m_Mul(GroupIDIntrin, m_Specific(ZextGroupSize)));

	ICmpInst::Predicate Pred;
	if (match(SI,
	m_Select(m_ICmp(Pred, SubExpr, m_Specific(ZextGroupSize)),
	SubExpr,
	m_Specific(ZextGroupSize))) &&
	Pred == ICmpInst::ICMP_ULT) {
	if (HasReqdWorkGroupSize) {
	ConstantInt *KnownSize
	= mdconst::extract<ConstantInt>(MD->getOperand(I));
	SI->replaceAllUsesWith(ConstantExpr::getIntegerCast(KnownSize,
	SI->getType(),
	false));
	} else {
	SI->replaceAllUsesWith(ZextGroupSize);
	}

	MadeChange = true;
	}
	}
	}
	}

	if (!HasReqdWorkGroupSize)
	return MadeChange;

	// Eliminate any other loads we can from the dispatch packet.
	for (int I = 0; I < 3; ++I) {
	Value *GroupSize = WorkGroupSizes[I];
	if (!GroupSize)
	continue;

	ConstantInt *KnownSize = mdconst::extract<ConstantInt>(MD->getOperand(I));
	GroupSize->replaceAllUsesWith(
	ConstantExpr::getIntegerCast(KnownSize,
	GroupSize->getType(),
	false));
	MadeChange = true;
	}

	return MadeChange;
	}

	// TODO: Move makeLIDRangeMetadata usage into here. Seem to not get
	// TargetPassConfig for subtarget.
	bool AMDGPULowerKernelAttributes::runOnModule(Module &M) {
	StringRef DispatchPtrName
	= Intrinsic::getName(Intrinsic::amdgcn_dispatch_ptr);

	Function *DispatchPtr = M.getFunction(DispatchPtrName);
	if (!DispatchPtr) // Dispatch ptr not used.
	return false;

	bool MadeChange = false;

	SmallPtrSet<Instruction *, 4> HandledUses;
	for (auto *U : DispatchPtr->users()) {
	CallInst *CI = cast<CallInst>(U);
	if (HandledUses.insert(CI).second) {
	if (processUse(CI))
	MadeChange = true;
	}
	}

	return MadeChange;
	}

	INITIALIZE_PASS_BEGIN(AMDGPULowerKernelAttributes, DEBUG_TYPE,
	"AMDGPU IR optimizations", false, false)
	INITIALIZE_PASS_END(AMDGPULowerKernelAttributes, DEBUG_TYPE, "AMDGPU IR optimizations",
	false, false)

	char AMDGPULowerKernelAttributes::ID = 0;

	ModulePass *llvm::createAMDGPULowerKernelAttributesPass() {
	return new AMDGPULowerKernelAttributes();
	}

	PreservedAnalyses
	AMDGPULowerKernelAttributesPass::run(Function &F, FunctionAnalysisManager &AM) {
	StringRef DispatchPtrName =
	Intrinsic::getName(Intrinsic::amdgcn_dispatch_ptr);

	Function *DispatchPtr = F.getParent()->getFunction(DispatchPtrName);
	if (!DispatchPtr) // Dispatch ptr not used.
	return PreservedAnalyses::all();

	for (Instruction &I : instructions(F)) {
	if (CallInst *CI = dyn_cast<CallInst>(&I)) {
	if (CI->getCalledFunction() == DispatchPtr)
	processUse(CI);
	}
	}

	return PreservedAnalyses::all();
	}