//===-- 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();
}