/***************************************************************************************************
 * Copyright (c) 2017-2018, NVIDIA CORPORATION.  All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without modification, are permitted
 * provided that the following conditions are met:
 *     * Redistributions of source code must retain the above copyright notice, this list of
 *       conditions and the following disclaimer.
 *     * Redistributions in binary form must reproduce the above copyright notice, this list of
 *       conditions and the following disclaimer in the documentation and/or other materials
 *       provided with the distribution.
 *     * Neither the name of the NVIDIA CORPORATION nor the names of its contributors may be used
 *       to endorse or promote products derived from this software without specific prior written
 *       permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR
 * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND
 * FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE
 * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
 * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS;
 * OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
 * STRICT LIABILITY, OR TOR (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 *
 **************************************************************************************************/
/*! \file
    \brief Implements a software-pipelined efficient GEMM.
*/
#pragma once

#if !defined(__CUDACC_RTC__)
#include <cuda.h>
#endif

#include <cutlass/coord.h>
#include <cutlass/util/platform.h>

namespace cutlass {
namespace gemm {

////////////////////////////////////////////////////////////////////////////////////////////////////

template <typename Gemm_>
__global__ /*__launch_bounds__(Gemm_::kThreads)*/ void gemm_kernel(typename Gemm_::Params params) {
  // Declare shared memory.
  __shared__ typename Gemm_::SharedStorage shared_storage;

  // Construct the GEMM object.
  Gemm_ gemm(params, shared_storage);
  // Run GEMM.
  gemm.multiply_add();
}

////////////////////////////////////////////////////////////////////////////////////////////////////

template <typename Scalar_, typename Index_ = int>
struct GemmDesc {
  /// The dimensions of the GEMM.
  Index_ m, n, k;
  /// The alpha/beta scaling values.
  Scalar_ alpha, beta;
  /// The source matrix A.
  void const* d_a;
  /// The stride for A.
  Index_ lda;
  /// The source matrix B.
  void const* d_b;
  /// The stride for B.
  Index_ ldb;
  /// The source matrix C.
  void const* d_c;
  /// The stride for C.
  Index_ ldc;
  /// The destination matrix D.
  void* d_d;
  /// The stride for D.
  Index_ ldd;
};

////////////////////////////////////////////////////////////////////////////////////////////////////

template <typename GemmTraits_>
struct Gemm {
  /// This class.
  typedef Gemm<GemmTraits_> This_;
  /// The traits.
  typedef GemmTraits_ Traits;
  /// The shared storage.
  typedef typename Traits::SharedStorage SharedStorage;

  /// The scalar for A.
  typedef typename Traits::ScalarA ScalarA;
  /// The scalar for B.
  typedef typename Traits::ScalarB ScalarB;
  /// The scalar in the epilogue.
  typedef typename Traits::Epilogue::Scalar ScalarEpilogue;
  /// The scalar for C.
  typedef typename Traits::Epilogue::ScalarC ScalarC;
  /// The scalar for D.
  typedef typename Traits::Epilogue::ScalarD ScalarD;
  /// The index.
  typedef typename Traits::Index Index;

  /// The number of threads.
  static int const kThreads = Traits::GemmConfig::kThreads;

  /// The params.
  struct Params : public Traits::Params {
    CUTLASS_HOST_DEVICE int initialize(Index m,
                                       Index n,
                                       Index k,
                                       ScalarEpilogue alpha,
                                       ScalarA const* d_a,
                                       Index lda,
                                       ScalarB const* d_b,
                                       Index ldb,
                                       ScalarEpilogue beta,
                                       ScalarC const* d_c,
                                       Index ldc,
                                       ScalarD* d_d,
                                       Index ldd) {
      GemmDesc<ScalarEpilogue, Index> desc;
      desc.m = m;
      desc.n = n;
      desc.k = k;
      desc.alpha = alpha;
      desc.beta = beta;
      desc.d_a = reinterpret_cast<void const*>(d_a);
      desc.lda = lda;
      desc.d_b = reinterpret_cast<void const*>(d_b);
      desc.ldb = ldb;
      desc.d_c = reinterpret_cast<void const*>(d_c);
      desc.ldc = ldc;
      desc.d_d = reinterpret_cast<void*>(d_d);
      desc.ldd = ldd;
      return Traits::Params::initialize(desc);
    }
  };

#if !defined(__CUDACC_RTC__)
  /// Launch the kernel.
  static __host__ cudaError_t launch(Params const& params,
                                     cudaStream_t stream = cudaStreamDefault) {
    // Setup the grid.
    dim3 grid;
    grid.x = (params.m + Traits::OutputTile::kW - 1) / Traits::OutputTile::kW;
    grid.y = (params.n + Traits::OutputTile::kH - 1) / Traits::OutputTile::kH;

    // The number of threads.
    dim3 block;
    block.x = kThreads;

    // Launch the kernel.
    void const* params_ = reinterpret_cast<void const*>(&params);

    return cudaLaunchKernel(reinterpret_cast<void*>(&gemm_kernel<This_>),
                            grid,
                            block,
                            const_cast<void**>(&params_),
                            0,
                            stream);
  }

  /// Launch the kernel.
  static __host__ cudaError_t launch(CUfunction kernel,
                                     Params const& params,
                                     CUstream stream = CU_STREAM_LEGACY) {
    // Setup the grid.
    dim3 grid;
    grid.x = (params.m + Traits::OutputTile::kW - 1) / Traits::OutputTile::kW;
    grid.y = (params.n + Traits::OutputTile::kH - 1) / Traits::OutputTile::kH;

    // The number of threads.
    dim3 block;
    block.x = kThreads;

    // Launch the kernel.
    void* params_[] = {const_cast<void*>(reinterpret_cast<void const*>(&params))};

    // return cudaLaunchKernel(reinterpret_cast<void*>(&gemm_kernel<This_>), grid, block,
    //  const_cast<void**>(&params_), 0, stream);
    CUresult result = cuLaunchKernel(
        kernel, grid.x, grid.y, grid.z, block.x, block.y, block.z, 0, stream, params_, 0);

    if (result != CUDA_SUCCESS) {
      return cudaErrorLaunchFailure;
    }
    return cudaSuccess;
  }

#endif

  /// Ctor.
  CUTLASS_DEVICE Gemm(Params const& params_, SharedStorage& shared_storage_)
      : params(params_), shared_storage(shared_storage_) {}

  /// Consume a single iteration of the loop.
  template <bool kIsLastIteration>
  CUTLASS_DEVICE void consume_tile(typename Traits::GlobalLoadStream& global_stream,
                                   typename Traits::SharedLoadStream& shared_load_stream,
                                   typename Traits::MultiplyAdd::Accumulators& accumulators,
                                   Index outer_k) {
    // If that's the last "load iteration" update the predicates.
    if (!kIsLastIteration) {
      global_stream.move_to_residue<false>(outer_k);
    }

    // Load data for the next iteration of the main loop.
    if (!kIsLastIteration) {
      global_stream.copy();
    }

    // The unrolling steps for the main loop.
    int const kUnrollingSteps =
        Traits::MultiplyAdd::AccumulatorsPerWarp::kD / Traits::MultiplyAdd::InstructionShape::kD;

    CUTLASS_PRAGMA_UNROLL
    for (int step = 0; step < kUnrollingSteps - 1; ++step) {
      // Trigger the copy from shared memory for the next A/B values.
      shared_load_stream.copy(step + 1);
      // Make sure the values are available for the current iteration to do the multiply-add.
      shared_load_stream.commit(step);

      // Do the math on the fragments of the current iteration.
      typename Traits::MultiplyAdd multiply_add;
      multiply_add.multiply_add(shared_load_stream.fragment_a(step),
                                shared_load_stream.fragment_b(step),
                                accumulators,
                                accumulators);
    }

    // Make sure the data from shared memory has been entirely consumed.
    Traits::shared_load_fence(true);

    // Commit the data in shared memory for A/B.
    if (!kIsLastIteration) {
      global_stream.commit();
    }

    // Make sure the data is in shared memory.
    Traits::shared_store_fence(true);

    // Trigger the loads for the next iteration (if needed).
    if (!kIsLastIteration) {
      // Move to the next stage for the load (if it makes sense).
      shared_load_stream.inc_stage();
      // Trigger the copy from shared memory for the next loop iteration.
      shared_load_stream.copy(0);
    }

    // Make sure the values are available for the current iteration to do the multiply-add.
    shared_load_stream.commit(kUnrollingSteps - 1);

    // Do the math on the fragments of the current iteration.
    typename Traits::MultiplyAdd multiply_add;
    multiply_add.multiply_add(shared_load_stream.fragment_a(kUnrollingSteps - 1),
                              shared_load_stream.fragment_b(kUnrollingSteps - 1),
                              accumulators,
                              accumulators);
  }

  /// Do the GEMM.
  CUTLASS_DEVICE void multiply_add() {
    // Swizzle the IDs of the block (to enable better cache behavior).
    typename Traits::BlockSwizzle block_swizzle;
    dim3 block = block_swizzle.swizzle();

    // Scale the id.
    block.x *= Traits::OutputTile::kW;
    block.y *= Traits::OutputTile::kH;

    // We may want to use shared memory to clear the registers.
    typedef typename Traits::ClearAccumulators ClearAccumulators;

    // The streams to read A/B from global memory to shared memory.
    typename Traits::GlobalLoadStream global_stream(params, shared_storage, block);

    // Create the accumulator clear.
    ClearAccumulators clear(shared_storage.main_loop.clear);

    // By how much we unroll the main loop.
    Index const kUnroll = static_cast<Index>(Traits::OutputTile::kD);

    // If we do not have enough steps in the main loop, trigger the residue code.
    global_stream.move_to_residue<true>(params.k);

    // Fetch the fragments for A and B from global memory.
    global_stream.copy();

    // Copy the elements to shared memory (after transformation if needed).
    global_stream.commit();

    // Make sure the data is in shared memory.
    Traits::shared_store_fence(false);

    // Rollback to the beginning of the GEMM-K dimension. It may have no impact.
    global_stream.rollback();

    // The unrolling steps for the main loop.
    int const kUnrollingSteps =
        Traits::MultiplyAdd::AccumulatorsPerWarp::kD / Traits::MultiplyAdd::InstructionShape::kD;

    // Make sure we have at least 2 unrolling steps or our pipeling is not going to work.
    static_assert(kUnrollingSteps >= 2, "The pipelining assumes at least two steps");

    // The stream of data from shared memory to fragments.
    typename Traits::SharedLoadStream shared_load_stream(params, shared_storage);

    // Trigger the copy from shared memory for the 1st stream.
    shared_load_stream.copy(0);

    // Allocate the accumulators.
    typename Traits::MultiplyAdd::Accumulators accumulators;
    // Clear the accumulators.
    clear.clear(accumulators);

    // The loop index.
    Index outer_k = params.k - kUnroll;

    // Enter the main loop and iterate.
    for (; outer_k > 0; outer_k -= kUnroll) {
      consume_tile<false>(global_stream, shared_load_stream, accumulators, outer_k);
    }

    // Residual loop.
    for (; outer_k > -kUnroll; outer_k -= kUnroll) {
      consume_tile<true>(global_stream, shared_load_stream, accumulators, outer_k);
    }

    // Epilogue.
    typedef typename Traits::Epilogue Epilogue;
    Epilogue epilogue(params.epilogue, shared_storage.epilogue, params.m, params.n);
    epilogue.epilogue(cutlass::make_Coord(0, block.y, block.x), accumulators);
  }

  /// The params.
  Params const& params;
  /// The shared storage.
  SharedStorage& shared_storage;
};

////////////////////////////////////////////////////////////////////////////////////////////////////

}  // namespace gemm
}  // namespace cutlass