///////////////////////////////////////////////////////////////////////////////
//                                                                           //
// DxilPIXAddTidToAmplificationShaderPayload.cpp                             //
// Copyright (C) Microsoft Corporation. All rights reserved.                 //
// This file is distributed under the University of Illinois Open Source     //
// License. See LICENSE.TXT for details.                                     //
//                                                                           //
///////////////////////////////////////////////////////////////////////////////

#include "dxc/DXIL/DxilOperations.h"
#include "dxc/DXIL/DxilUtil.h"

#include "dxc/DXIL/DxilInstructions.h"
#include "dxc/DXIL/DxilModule.h"
#include "dxc/DxilPIXPasses/DxilPIXPasses.h"

#include "llvm/IR/InstIterator.h"
#include "llvm/IR/PassManager.h"
#include "llvm/Support/FormattedStream.h"
#include "llvm/Transforms/Utils/Local.h"

#include "PixPassHelpers.h"

using namespace llvm;
using namespace hlsl;
using namespace PIXPassHelpers;

class DxilPIXAddTidToAmplificationShaderPayload : public ModulePass {
  uint32_t m_DispatchArgumentY = 1;
  uint32_t m_DispatchArgumentZ = 1;

public:
  static char ID; // Pass identification, replacement for typeid
  DxilPIXAddTidToAmplificationShaderPayload() : ModulePass(ID) {}
  StringRef getPassName() const override {
    return "DXIL Add flat thread id to payload from AS to MS";
  }
  bool runOnModule(Module &M) override;
  void applyOptions(PassOptions O) override;
};

void DxilPIXAddTidToAmplificationShaderPayload::applyOptions(PassOptions O) {
  GetPassOptionUInt32(O, "dispatchArgY", &m_DispatchArgumentY, 1);
  GetPassOptionUInt32(O, "dispatchArgZ", &m_DispatchArgumentZ, 1);
}

void AddValueToExpandedPayload(OP *HlslOP, llvm::IRBuilder<> &B,
                               ExpandedStruct &expanded,
                               AllocaInst *NewStructAlloca,
                               unsigned int expandedValueIndex, Value *value) {
  Constant *Zero32Arg = HlslOP->GetU32Const(0);
  SmallVector<Value *, 2> IndexToAppendedValue;
  IndexToAppendedValue.push_back(Zero32Arg);
  IndexToAppendedValue.push_back(HlslOP->GetU32Const(expandedValueIndex));
  auto *PointerToEmbeddedNewValue = B.CreateInBoundsGEP(
      expanded.ExpandedPayloadStructType, NewStructAlloca, IndexToAppendedValue,
      "PointerToEmbeddedNewValue" + std::to_string(expandedValueIndex));
  B.CreateStore(value, PointerToEmbeddedNewValue);
}

bool DxilPIXAddTidToAmplificationShaderPayload::runOnModule(Module &M) {

  DxilModule &DM = M.GetOrCreateDxilModule();
  LLVMContext &Ctx = M.getContext();
  OP *HlslOP = DM.GetOP();

  Type *OriginalPayloadStructPointerType = nullptr;
  Type *OriginalPayloadStructType = nullptr;
  ExpandedStruct expanded;
  llvm::Function *entryFunction = PIXPassHelpers::GetEntryFunction(DM);
  for (inst_iterator I = inst_begin(entryFunction), E = inst_end(entryFunction);
       I != E; ++I) {
    if (auto *Instr = llvm::cast<Instruction>(&*I)) {
      if (hlsl::OP::IsDxilOpFuncCallInst(Instr,
                                         hlsl::OP::OpCode::DispatchMesh)) {
        DxilInst_DispatchMesh DispatchMesh(Instr);
        OriginalPayloadStructPointerType =
            DispatchMesh.get_payload()->getType();
        OriginalPayloadStructType =
            OriginalPayloadStructPointerType->getPointerElementType();
        expanded = ExpandStructType(Ctx, OriginalPayloadStructType);
      }
    }
  }

  AllocaInst *OldStructAlloca = nullptr;
  AllocaInst *NewStructAlloca = nullptr;
  std::vector<AllocaInst *> allocasOfPayloadType;
  for (inst_iterator I = inst_begin(entryFunction), E = inst_end(entryFunction);
       I != E; ++I) {
    auto *Inst = &*I;
    if (llvm::isa<AllocaInst>(Inst)) {
      auto *Alloca = llvm::cast<AllocaInst>(Inst);
      if (Alloca->getType() == OriginalPayloadStructPointerType) {
        allocasOfPayloadType.push_back(Alloca);
      }
    }
  }
  for (auto &Alloca : allocasOfPayloadType) {
    OldStructAlloca = Alloca;
    llvm::IRBuilder<> B(Alloca->getContext());
    NewStructAlloca = B.CreateAlloca(expanded.ExpandedPayloadStructType,
                                     HlslOP->GetU32Const(1), "NewPayload");
    NewStructAlloca->setAlignment(Alloca->getAlignment());
    NewStructAlloca->insertAfter(Alloca);

    ReplaceAllUsesOfInstructionWithNewValueAndDeleteInstruction(
        Alloca, NewStructAlloca, expanded.ExpandedPayloadStructType);
  }

  auto F = HlslOP->GetOpFunc(DXIL::OpCode::DispatchMesh,
                             expanded.ExpandedPayloadStructPtrType);
  for (auto FI = F->user_begin(); FI != F->user_end();) {
    auto *FunctionUser = *FI++;
    auto *UserInstruction = llvm::cast<Instruction>(FunctionUser);
    DxilInst_DispatchMesh DispatchMesh(UserInstruction);

    llvm::IRBuilder<> B(UserInstruction);

    auto ThreadIdFunc =
        HlslOP->GetOpFunc(DXIL::OpCode::ThreadId, Type::getInt32Ty(Ctx));
    Constant *Opcode = HlslOP->GetU32Const((unsigned)DXIL::OpCode::ThreadId);
    Constant *Zero32Arg = HlslOP->GetU32Const(0);
    Constant *One32Arg = HlslOP->GetU32Const(1);
    Constant *Two32Arg = HlslOP->GetU32Const(2);

    auto ThreadIdX =
        B.CreateCall(ThreadIdFunc, {Opcode, Zero32Arg}, "ThreadIdX");
    auto ThreadIdY =
        B.CreateCall(ThreadIdFunc, {Opcode, One32Arg}, "ThreadIdY");
    auto ThreadIdZ =
        B.CreateCall(ThreadIdFunc, {Opcode, Two32Arg}, "ThreadIdZ");

    auto *XxY =
        B.CreateMul(ThreadIdX, HlslOP->GetU32Const(m_DispatchArgumentY));
    auto *XplusY = B.CreateAdd(ThreadIdY, XxY);
    auto *XYxZ = B.CreateMul(XplusY, HlslOP->GetU32Const(m_DispatchArgumentZ));
    auto *XYZ = B.CreateAdd(ThreadIdZ, XYxZ);

    AddValueToExpandedPayload(HlslOP, B, expanded, NewStructAlloca,
                              OriginalPayloadStructType->getStructNumElements(),
                              XYZ);
    AddValueToExpandedPayload(
        HlslOP, B, expanded, NewStructAlloca,
        OriginalPayloadStructType->getStructNumElements() + 1,
        DispatchMesh.get_threadGroupCountY());
    AddValueToExpandedPayload(
        HlslOP, B, expanded, NewStructAlloca,
        OriginalPayloadStructType->getStructNumElements() + 2,
        DispatchMesh.get_threadGroupCountZ());
  }

  DM.ReEmitDxilResources();

  return true;
}

char DxilPIXAddTidToAmplificationShaderPayload::ID = 0;

ModulePass *llvm::createDxilPIXAddTidToAmplificationShaderPayloadPass() {
  return new DxilPIXAddTidToAmplificationShaderPayload();
}

INITIALIZE_PASS(DxilPIXAddTidToAmplificationShaderPayload,
                "hlsl-dxil-PIX-add-tid-to-as-payload",
                "HLSL DXIL Add flat thread id to payload from AS to MS", false,
                false)