/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 * All rights reserved.
 *
 * This source code is licensed under the BSD-style license found in the
 * LICENSE file in the root directory of this source tree.
 */

#define FBGEMM_EXPORTS
#include "fbgemm/FbgemmI64.h"

#if defined(__x86_64__) || defined(__i386__) || \
    (defined(_MSC_VER) && (defined(_M_X64) || defined(_M_IX86)))
#include <immintrin.h>
#endif
#include <cmath>
#include <iostream>
#include <vector>

#include "./GenerateKernel.h"
#include "./RefImplementations.h"
#include "fbgemm/PackingTraits-inl.h"

using namespace std;

namespace fbgemm {

/**
 * Generate AVX2 instructions for computing block in the rank-k update of 32-bit
 * Accmulation kernel.
 */
template <>
template <inst_set_t instSet>
void CodeGenBase<int64_t, int64_t, int64_t, int64_t>::genComputeBlock(
    x86::Emitter* a,
    x86::Gp buffer_A,
    x86::Gp buffer_B,
    x86::Gp B_pf,
    int rowRegs,
    int colRegs,
    int lda) {
  using VecRegT = typename simd_info<instSet>::vec_reg_t;
  constexpr int vectorLen = simd_info<instSet>::WIDTH_BITS / 64;

  // used for matrix B
  VecRegT BReg(31);

  // temporary register
  VecRegT res1(30);

  for (int j = 0; j < colRegs; ++j) {
    // load B
    a->vmovaps(
        BReg,
        x86::Mem(
            buffer_B,
            j * vectorLen * sizeof(int64_t),
            simd_info<instSet>::WIDTH_BYTES));
    // load A, broadcast and fmas
    for (int i = 0; i < rowRegs; ++i) {
      a->vpmullq(
          res1,
          BReg,
          x86::qword_ptr(buffer_A, (i * lda) * sizeof(int64_t))._1to8());
      a->vpaddq(VecRegT(i * colRegs + j), res1, VecRegT(i * colRegs + j));
    }
    // TODO: need to tune
    a->prefetcht0(x86::dword_ptr(B_pf, j * vectorLen * sizeof(int64_t)));
  }
}

/**
 * Generate AVX2 instructions for storing the C registers back to the memory in
 * 32-bit Accumulation kernel.
 */
template <>
template <inst_set_t instSet>
void CodeGenBase<int64_t, int64_t, int64_t, int64_t>::storeCRegs(
    x86::Emitter* a,
    int rowRegs,
    int colRegs,
    x86::Gp C_Offset,
    x86::Gp ldcReg,
    bool accum) {
  using VecT = typename simd_info<instSet>::vec_reg_t;
  static constexpr int vectorLen = simd_info<instSet>::WIDTH_BITS / 64;

  for (int i = 0; i < rowRegs; ++i) {
    if (i != 0) {
      a->add(C_Offset, ldcReg);
    } else {
      a->xor_(C_Offset.r32(), C_Offset.r32());
    }
    for (int j = 0; j < colRegs; ++j) {
      if (accum) {
        a->vpaddq(
            VecT(i * colRegs + j),
            VecT(i * colRegs + j),
            x86::dword_ptr(
                a->zcx(), C_Offset, 0, j * vectorLen * sizeof(int64_t)));
      }
      a->vmovups(
          x86::dword_ptr(
              a->zcx(), C_Offset, 0, j * vectorLen * sizeof(int64_t)),
          VecT(i * colRegs + j));
    }
  }
}

/**
 * Get or Create the avx512 instructions for int64_t GEMM macro-kernel.
 */
template <>
template <inst_set_t instSet>
CodeGenBase<int64_t, int64_t, int64_t, int64_t>::jit_micro_kernel_fp
CodeGenBase<int64_t, int64_t, int64_t, int64_t>::getOrCreate(
    bool accum,
    int32_t mc,
    int32_t nc,
    int32_t /* unused */) {
  static constexpr int vectorLen = simd_info<instSet>::WIDTH_BITS / 64;

  tuple<bool, int, int, int, int, int, int> kernelSig;
  int kBlock;
  int nBlock;
  int mRegBlockSize;
  int nRegBlockSize;

  if (blocking_params) {
    kBlock = blocking_params->KCB;
    nBlock = blocking_params->NCB;
    mRegBlockSize = blocking_params->MR;
    nRegBlockSize = blocking_params->NR;
  } else {
    kBlock = PackingTraits<int64_t, int64_t, instSet>::KCB;
    nBlock = PackingTraits<int64_t, int64_t, instSet>::NCB;
    mRegBlockSize = PackingTraits<int64_t, int64_t, instSet>::MR;
    nRegBlockSize = PackingTraits<int64_t, int64_t, instSet>::NR;
  }

  kernelSig =
      make_tuple(accum, mc, nc, nBlock, kBlock, mRegBlockSize, nRegBlockSize);

  return codeCache_.getOrCreate(kernelSig, [&]() -> jit_micro_kernel_fp {
    asmjit::CodeHolder code;
    code.init(runtime().environment());
    x86::Assembler assembler(&code);
    x86::Emitter* a = assembler.as<x86::Emitter>();
#ifdef FBGEMM_LOG_CODE
    // generated code logging
    FILE* codeLogfile = fopen(
        getCodeLoggingFile<instSet>(
            accum, mc, nc, nBlock, kBlock, mRegBlockSize, nRegBlockSize)
            .c_str(),
        "w");
    asmjit::FileLogger* codeLogger = new asmjit::FileLogger(codeLogfile);
    if (codeLogger) {
      code.setLogger(codeLogger);
    }
#endif

    const int maxMRegs = mRegBlockSize;
    (void)maxMRegs; // Suppress unused variable warning
    const int maxNRegs = nRegBlockSize / vectorLen;
    assert(
        maxMRegs * maxNRegs <= 30 &&
        "MR*(NR*64/512) \
        must be <= 29 (available registers constraint)");

    const int mRegBlocks = mc / mRegBlockSize;
    const int mRegBlocksRem = mc % mRegBlockSize;

    // arguments to the function created
    x86::Gp buffer_A = a->zdi();
    x86::Gp buffer_B = a->zsi();
    x86::Gp B_pf = a->zdx();
    x86::Gp CBase = a->zcx();
    x86::Gp kSize = a->gpz(8);
    x86::Gp ldcReg = a->gpz(9);

    asmjit::FuncDetail func;
    func.init(
        asmjit::FuncSignatureT<
            void,
            int64_t*,
            int64_t*,
            int64_t*,
            int64_t*,
            int,
            int>(asmjit::CallConvId::kHost),
        a->environment());

    asmjit::FuncFrame frame;
    frame.init(func);

    frame.setDirtyRegs(
        asmjit::RegGroup::kVec,
        asmjit::Support::bitMask(0, 1, 2, 3, 4, 5, 6, 7) |
            asmjit::Support::bitMask(8, 9, 10, 11, 12, 13, 14, 15) |
            asmjit::Support::bitMask(16, 17, 18, 19, 20, 21, 22, 23) |
            asmjit::Support::bitMask(24, 25, 26, 27, 28, 29, 30, 31));
    frame.setDirtyRegs(
        asmjit::RegGroup::kGp,
        asmjit::Support::bitMask(8, 9, 10, 11, 12, 13, 14, 15));

    asmjit::FuncArgsAssignment args(&func);
    args.assignAll(buffer_A, buffer_B, B_pf, CBase, kSize, ldcReg);

    args.updateFuncFrame(frame);
    frame.finalize();

    a->emitProlog(frame);
    a->emitArgsAssignment(frame, args);

    asmjit::Label LoopMBlocks = a->newLabel();
    asmjit::Label LoopNBlocks = a->newLabel();
    asmjit::Label Loopk = a->newLabel();

    x86::Gp buffer_B_saved = a->gpz(10);
    x86::Gp C_Offset = a->gpz(11);
    x86::Gp B_pf_saved = a->gpz(12);
    x86::Gp iIdx = a->gpz(13);
    x86::Gp jIdx = a->gpz(14);
    x86::Gp kIdx = a->gpz(15);

    a->imul(ldcReg, ldcReg, static_cast<asmjit::Imm>(sizeof(int64_t)));
    a->imul(kSize, kSize, static_cast<asmjit::Imm>(sizeof(int64_t)));

    // save B_buffer address
    a->mov(buffer_B_saved, buffer_B);
    a->mov(B_pf_saved, B_pf);

    int currColRegs = nc / vectorLen;
    int colRegs = std::min(currColRegs, maxNRegs);
    if (mRegBlocks > 0) {
      // move 0 to iteration variables
      a->xor_(iIdx.r32(), iIdx.r32());

      a->bind(LoopMBlocks);
      a->inc(iIdx);
      a->xor_(jIdx.r32(), jIdx.r32());

      a->bind(LoopNBlocks);
      a->inc(jIdx);

      int rowRegs = mRegBlockSize;

      // init C registers
      initCRegs(a, rowRegs, colRegs);

      // init k loop index
      a->xor_(kIdx.r32(), kIdx.r32());
      a->bind(Loopk);

      // k is incremented by 1
      a->add(kIdx, static_cast<asmjit::Imm>(sizeof(int64_t)));

      genComputeBlock<instSet>(
          a, buffer_A, buffer_B, B_pf, rowRegs, colRegs, kBlock);

      // update buffer_A address for next k iteration
      a->add(buffer_A, static_cast<asmjit::Imm>(sizeof(int64_t)));

      // update buffer_B address for next k iteration
      a->add(buffer_B, static_cast<asmjit::Imm>(nBlock * sizeof(int64_t)));
      a->add(B_pf, static_cast<asmjit::Imm>(nBlock * sizeof(int64_t)));

      a->cmp(kIdx, kSize);
      a->jl(Loopk);

      // store C matrix
      storeCRegs<instSet>(a, rowRegs, colRegs, C_Offset, ldcReg, accum);

      // reset A
      a->sub(buffer_A, kSize);

      // B for next block
      a->mov(buffer_B, buffer_B_saved);
      // using C_Offset as temp reg
      a->imul(
          C_Offset,
          jIdx,
          static_cast<asmjit::Imm>(nRegBlockSize * sizeof(int64_t)));
      a->add(buffer_B, C_Offset);
      a->mov(B_pf, B_pf_saved);
      a->add(B_pf, C_Offset);

      // increment C for next B block
      a->add(CBase, static_cast<asmjit::Imm>(nRegBlockSize * sizeof(int64_t)));

      int jLoopTrips = currColRegs / maxNRegs;
      // jLoopTrips should be at least 1
      jLoopTrips = jLoopTrips ? jLoopTrips : 1;
      a->cmp(jIdx, jLoopTrips);
      a->jl(LoopNBlocks);

      // increment A for next block
      a->add(
          buffer_A,
          static_cast<asmjit::Imm>(rowRegs * kBlock * sizeof(int64_t)));

      // increment C for next A block
      a->sub(
          CBase,
          static_cast<asmjit::Imm>(
              jLoopTrips * nRegBlockSize * sizeof(int64_t)));
      a->imul(C_Offset, ldcReg, static_cast<asmjit::Imm>(rowRegs));
      a->add(CBase, C_Offset);

      // reset B
      a->mov(buffer_B, buffer_B_saved);
      a->mov(B_pf, B_pf_saved);
      a->cmp(iIdx, mRegBlocks);
      a->jl(LoopMBlocks);
    }
    // generate code for remainder
    if (mRegBlocksRem > 0) {
      assert(false);
      asmjit::Label LoopNRem = a->newLabel();
      asmjit::Label LoopkRem = a->newLabel();
      int rowRegs = mRegBlocksRem;

      a->xor_(jIdx.r32(), jIdx.r32());
      a->bind(LoopNRem);
      a->inc(jIdx);

      // init C registers
      initCRegs(a, rowRegs, colRegs);

      // init k loop index
      a->xor_(kIdx.r32(), kIdx.r32());
      a->bind(LoopkRem);

      // k is incremented by 1
      a->add(kIdx, static_cast<asmjit::Imm>(sizeof(int64_t)));

      genComputeBlock<instSet>(
          a, buffer_A, buffer_B, B_pf, rowRegs, colRegs, kBlock);

      // update buffer_A address for next k iteration
      a->add(buffer_A, static_cast<asmjit::Imm>(sizeof(int64_t)));

      // update buffer_B address for next k iteration
      a->add(buffer_B, static_cast<asmjit::Imm>(nBlock * sizeof(int64_t)));
      a->add(B_pf, static_cast<asmjit::Imm>(nBlock * sizeof(int64_t)));

      a->cmp(kIdx, kSize);
      a->jl(LoopkRem);

      // reset A
      a->sub(buffer_A, kSize);

      // B for next block
      // using C_Offset as temp reg
      a->imul(
          C_Offset,
          jIdx,
          static_cast<asmjit::Imm>(nRegBlockSize * sizeof(int64_t)));
      a->mov(buffer_B, buffer_B_saved);
      a->add(buffer_B, C_Offset);
      a->mov(B_pf, B_pf_saved);
      a->add(B_pf, C_Offset);

      // store C matrix
      storeCRegs<instSet>(a, rowRegs, colRegs, C_Offset, ldcReg, accum);

      // increment C for next B block
      a->add(CBase, static_cast<asmjit::Imm>(nRegBlockSize * sizeof(int64_t)));

      int jLoopTrips = currColRegs / maxNRegs;
      // jLoopTrips should be at least 1
      jLoopTrips = jLoopTrips ? jLoopTrips : 1;
      a->cmp(jIdx, jLoopTrips);
      a->jl(LoopNRem);
    }

    a->emitEpilog(frame);

    jit_micro_kernel_fp fn;
    asmjit::Error err;
    {
      unique_lock<mutex> lock(rtMutex_);
      err = runtime().add(&fn, &code);
    }
    if (err) {
      cout << "Error: in fn add" << endl;
      return nullptr;
    }

#ifdef FBGEMM_LOG_CODE
    fclose(codeLogfile);
    delete codeLogger;
#endif

    return fn;
  });
}

/**
 * Instatiate the AVX512 instructions for int64_t GEMM macro-kernel.
 */
template CodeGenBase<int64_t, int64_t, int64_t, int64_t>::jit_micro_kernel_fp
CodeGenBase<int64_t, int64_t, int64_t, int64_t>::getOrCreate<
    inst_set_t::avx512>(bool accum, int32_t mc, int32_t nc, int32_t kc);

// Expected to have overflows
NO_SANITIZE("undefined")
void cblas_gemm_i64_i64acc(
    matrix_op_t transa,
    matrix_op_t transb,
    int M,
    int N,
    int K,
    const int64_t* A,
    int lda,
    const int64_t* B,
    int ldb,
    bool accumulate,
    int64_t* C,
    int ldc) {
  cpuinfo_initialize();
  if (!fbgemmHasAvx512Support()) {
    cblas_gemm_i64_i64acc_ref(
        transa, transb, M, N, K, A, lda, B, ldb, accumulate, C, ldc);
    return;
  }
  constexpr int MCB = PackingTraits<int64_t, int64_t, inst_set_t::avx512>::MCB;
  constexpr int NCB = PackingTraits<int64_t, int64_t, inst_set_t::avx512>::NCB;
  constexpr int KCB = PackingTraits<int64_t, int64_t, inst_set_t::avx512>::KCB;
  constexpr int MR = PackingTraits<int64_t, int64_t, inst_set_t::avx512>::MR;
  constexpr int NR = PackingTraits<int64_t, int64_t, inst_set_t::avx512>::NR;
  static_assert(MCB % MR == 0, "MR must divide MCB");
  static_assert(NCB % NR == 0, "NR must divide NCB");
  constexpr int VLEN =
      simd_info<inst_set_t::avx512>::WIDTH_BYTES / sizeof(int64_t);
  static_assert(NR % VLEN == 0, "VLEN must divide NR");

  using CodeGenType = CodeGenBase<int64_t, int64_t, int64_t, int64_t>;
  CodeGenType codeObj;
  CodeGenType::jit_micro_kernel_fp fn =
      codeObj.getOrCreate<inst_set_t::avx512>(true /* accum */, MCB, NCB, KCB);
  CodeGenType::jit_micro_kernel_fp fn_noacc;
  if (!accumulate) {
    fn_noacc = codeObj.getOrCreate<inst_set_t::avx512>(
        false /* accum */, MCB, NCB, KCB);
  }

  vector<int64_t> At, Bt;
  // TODO: handle transpose during packing
  if (transa == matrix_op_t::Transpose) {
    At.resize(M * K);
    for (int i = 0; i < M; ++i) {
      for (int k = 0; k < K; ++k) {
        At.at(i * K + k) = A[i + k * lda];
      }
    }
    A = At.data();
    lda = K;
  }
  if (transb == matrix_op_t::Transpose) {
    Bt.resize(K * N);
    for (int k = 0; k < K; ++k) {
      for (int j = 0; j < N; ++j) {
        Bt.at(k * N + j) = B[k + j * ldb];
      }
    }
    B = Bt.data();
    ldb = N;
  }

  alignas(64) array<int64_t, MCB * KCB> packA;
  alignas(64) array<int64_t, KCB * NCB> packB;
  alignas(64) array<int64_t, MCB * NCB> packC;

  for (int ic = 0; ic < M; ic += MCB) {
    for (int kc = 0; kc < K; kc += KCB) {
      // pack A
      for (int i = 0; i < std::min(MCB, M - ic); ++i) {
        memcpy(
            &packA[i * KCB],
            A + (ic + i) * lda + kc,
            std::min(K - kc, KCB) * sizeof(int64_t));
      }

      for (int jc = 0; jc < N; jc += NCB) {
        // pack B
        for (int i = 0; i < std::min(KCB, K - kc); ++i) {
          memcpy(
              &packB[i * NCB],
              B + (kc + i) * ldb + jc,
              std::min(NCB, N - jc) * sizeof(int64_t));
        }

        if (M - ic >= MCB && N - jc >= NCB) {
          if (kc == 0 && !accumulate) {
            fn_noacc(
                packA.data(),
                packB.data(),
                packB.data(),
                C + ic * ldc + jc,
                std::min(KCB, K - kc),
                ldc);
          } else {
            fn(packA.data(),
               packB.data(),
               packB.data(),
               C + ic * ldc + jc,
               std::min(KCB, K - kc),
               ldc);
          }
        } else {
          // remainder
          if (kc == 0 && !accumulate) {
            fn_noacc(
                packA.data(),
                packB.data(),
                packB.data(),
                packC.data(),
                std::min(KCB, K - kc),
                NCB);
          } else {
            for (int i = 0; i < std::min(MCB, M - ic); ++i) {
              memcpy(
                  &packC[i * NCB],
                  C + (ic + i) * ldc + jc,
                  std::min(NCB, N - jc) * sizeof(int64_t));
            }
            fn(packA.data(),
               packB.data(),
               packB.data(),
               packC.data(),
               std::min(KCB, K - kc),
               NCB);
          }
          for (int i = 0; i < std::min(MCB, M - ic); ++i) {
            memcpy(
                C + (ic + i) * ldc + jc,
                &packC[i * NCB],
                std::min(NCB, N - jc) * sizeof(int64_t));
          }
        }
      } // jc
    } // kc
  } // ic
}

} // namespace fbgemm