/***************************************************************************************************
 * Copyright (c) 2017 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 * SPDX-License-Identifier: BSD-3-Clause
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 *
 * 1. Redistributions of source code must retain the above copyright notice, this
 * list of conditions and the following disclaimer.
 *
 * 2. 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.
 *
 * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 Template for a multistage threadblock-scoped Implicit GEMM Convolution kernel.
*/

#pragma once

#include "cutlass/aligned_buffer.h"
#include "cutlass/arch/memory.h"
#include "cutlass/array.h"
#include "cutlass/cutlass.h"
#include "cutlass/gemm/gemm.h"
#include "cutlass/matrix_shape.h"
#include "cutlass/numeric_types.h"
#include "cutlass/arch/cache_operation.h"
#include "cutlass/conv/threadblock/depthwise_mma_base.h"

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

namespace cutlass {
namespace conv {
namespace threadblock {

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

/// Structure to compute the matrix product targeting CUDA cores and SIMT math
/// instructions.
template <
    /// Size of the Gemm problem - concept: gemm::GemmShape<>
    typename Shape_,
    /// Iterates over tiles of A operand in global memory
    //  (concept: ReadableTileIterator | ForwardTileIterator |
    //  MaskedTileIterator)
    typename IteratorA_,
    /// Iterates over tiles of A operand in shared memory
    /// (concept: WriteableTileIterator | RandomAccessTileIterator)
    typename SmemIteratorA_,
    /// Cache operation for operand A
    cutlass::arch::CacheOperation::Kind CacheOpA,
    /// Iterates over tiles of B operand in global memory
    //  (concept: ReadableTileIterator | ForwardTileIterator |
    //  MaskedTileIterator)
    typename IteratorB_,
    /// Iterates over tiles of B operand in shared memory
    /// (concept: WriteableTileIterator | RandomAccessTileIterator)
    typename SmemIteratorB_,
    /// Cache operation for operand B
    cutlass::arch::CacheOperation::Kind CacheOpB,
    /// Policy describing tuning details (concept: MmaPolicy)
    typename Policy_,
    /// Number of stages,
    int Stages,
    /// Epilogue stores the data into global memory
    typename Epilogue_,
    /// iterator implementation variants
    conv::IteratorAlgorithm IteratorAlgorithm_ = conv::IteratorAlgorithm::kOptimized,
    /// Used for partial specialization
    typename Enable = bool>
class DepthwiseFpropDirectConvMultipleStage :
   public DepthwiseDirectConvMmaBase<Shape_, Policy_, Stages> {
public:
  ///< Base class
  using Base = DepthwiseDirectConvMmaBase<Shape_, Policy_, Stages>;
  ///< Size of the Gemm problem - concept: gemm::GemmShape<>
  using Shape = Shape_;
  ///< Iterates over tiles of A operand in global memory
  using IteratorA = IteratorA_;
  ///< Iterates over tiles of B operand in global memory
  using IteratorB = IteratorB_;
  ///< Policy describing tuning details
  using Policy = Policy_;

  using Epilogue = Epilogue_;

  using SmemIteratorA = SmemIteratorA_;
  using SmemIteratorB = SmemIteratorB_;

  static cutlass::arch::CacheOperation::Kind const kCacheOpA = CacheOpA;
  static cutlass::arch::CacheOperation::Kind const kCacheOpB = CacheOpB;

  static conv::IteratorAlgorithm const kItertorAlgorithm = IteratorAlgorithm_;

  //
  // Dependent types
  //

  /// Fragment of accumulator tile

  using ElementC = typename Policy::Operator::ElementC;
  using FragmentC = typename Policy::Operator::FragmentC;

  /// Warp-level Mma
  using Operator = typename Policy::Operator;
  
  /// Internal structure exposed for introspection.
  struct Detail {

    /// Number of cp.async instructions to load one stage of operand A
    static int const AsyncCopyIterationsPerStageA =
        IteratorA::ThreadMap::Iterations::kCount;

    /// Number of cp.async instructions to load one stage of operand B
    static int const AsyncCopyIterationsPerStageB =
        IteratorB::ThreadMap::Iterations::kCount;

    /// Number of stages
    static int const kStages = Stages;

    /// Number of cp.async instructions to load on group of operand B
    static int const kAccessesPerGroupB = 
        (AsyncCopyIterationsPerStageB + Base::kWarpGemmIterations - 1) / Base::kWarpGemmIterations;
  };

 private:

  using WarpLoadedFragmentA = typename Operator::FragmentA;
  using WarpLoadedFragmentB = typename Operator::FragmentB;
  using WarpTransformedFragmentA = typename Operator::TransformedFragmentA;
  using WarpTransformedFragmentB = typename Operator::TransformedFragmentB;

 private:

  //
  // Data members
  //

  /// Iterator to write threadblock-scoped tile of A operand to shared memory
  SmemIteratorA smem_iterator_A_;

  /// Iterator to write threadblock-scoped tile of B operand to shared memory
  SmemIteratorB smem_iterator_B_;

public:

  /// Construct from tensor references
  CUTLASS_DEVICE
  DepthwiseFpropDirectConvMultipleStage(
      ///< Shared storage needed for internal use by threadblock-scoped GEMM
      typename Base::SharedStorage &shared_storage,
      ///< ID within the threadblock
      int thread_idx,
      ///< ID of warp
      int warp_idx,
      ///< ID of each thread within a warp
      int lane_idx
    ):
      Base(shared_storage, thread_idx, warp_idx, lane_idx),
      smem_iterator_A_(shared_storage.operand_A_ref(), thread_idx),
      smem_iterator_B_(shared_storage.operand_B_ref(), thread_idx)
  {
    // Compute warp location within threadblock tile by mapping the warp_id to
    // three coordinates:
    //   _m: the warp's position within the threadblock along the M dimension
    //   _n: the warp's position within the threadblock along the N dimension
    //   _k: the warp's position within the threadblock along the K dimension

    int warp_idx_mn = warp_idx % (Base::WarpCount::kM * Base::WarpCount::kN);
    int warp_idx_k = warp_idx / (Base::WarpCount::kM * Base::WarpCount::kN);

    int warp_idx_m = warp_idx_mn % Base::WarpCount::kM;
    int warp_idx_n = warp_idx_mn / Base::WarpCount::kM;

    // Add per-warp offsets in units of warp-level tiles
    this->warp_tile_iterator_A_.add_tile_offset(
        {warp_idx_m, Base::kWarpGemmIterations * warp_idx_k});
    this->warp_tile_iterator_B_.add_tile_offset(
        {Base::kWarpGemmIterations * warp_idx_k, warp_idx_n});
  }

  CUTLASS_DEVICE
  void copy_tiles_and_advance(IteratorA &iterator_A,
                              IteratorB &iterator_B,
                              int group_start_A = 0,
                              int group_start_B = 0) {
    if (kItertorAlgorithm == conv::IteratorAlgorithm::kFixedStrideDilation) {
      // Number of iterators is a static value.
      iterator_A.set_iteration_index(group_start_A * IteratorA::kAccessesPerVector);
      this->smem_iterator_A_.set_iteration_index(group_start_A);

      // Async Copy for operand A
      CUTLASS_PRAGMA_UNROLL
      for (int j = 0; j < Detail::AsyncCopyIterationsPerStageA; ++j) {
        typename IteratorA::AccessType *dst_ptr =
            reinterpret_cast<typename IteratorA::AccessType *>(this->smem_iterator_A_.get());

        int const kSrcBytes = sizeof_bits<typename IteratorA::Element>::value *
                              IteratorA::ThreadMap::kElementsPerAccess /
                              IteratorA::kAccessesPerVector / 8;

        CUTLASS_PRAGMA_UNROLL
        for (int v = 0; v < IteratorA::kAccessesPerVector; ++v) {
          cutlass::arch::cp_async_zfill<kSrcBytes, kCacheOpA>(
              dst_ptr + v, iterator_A.get(), iterator_A.valid());

          ++iterator_A;
        }
        ++this->smem_iterator_A_;
      }
    } else {
      // Number of iterators is a runtime value.
      iterator_A.set_iteration_index(group_start_A * IteratorA::kAccessesPerVector);
      this->smem_iterator_A_.set_iteration_index(group_start_A);

      // Async Copy for operand A
      CUTLASS_PRAGMA_UNROLL
      for (int j = 0; j < iterator_A.get_iteration_num(); ++j) {
        typename IteratorA::AccessType *dst_ptr =
            reinterpret_cast<typename IteratorA::AccessType *>(this->smem_iterator_A_.get());

        int const kSrcBytes = sizeof_bits<typename IteratorA::Element>::value *
                              IteratorA::ThreadMap::kElementsPerAccess /
                              IteratorA::kAccessesPerVector / 8;

        CUTLASS_PRAGMA_UNROLL
        for (int v = 0; v < IteratorA::kAccessesPerVector; ++v) {
          cutlass::arch::cp_async_zfill<kSrcBytes, kCacheOpA>(
              dst_ptr + v, iterator_A.get(), iterator_A.valid());

          ++iterator_A;
        }
        ++this->smem_iterator_A_;
      }
    }
  }

  /// Perform a threadblock-scoped matrix multiply-accumulate
  CUTLASS_DEVICE
  void operator()(
      ///< problem size of GEMM
      int gemm_k_iterations,
      ///< destination accumulator tile
      FragmentC &accum,
      ///< iterator over A operand in global memory
      IteratorA &iterator_A,
      ///< Params of global memory iterator
      typename IteratorA::Params const &iterator_a_params,
      ///< iterator over B operand in global memory
      IteratorB &iterator_B,
      ///< Params of global memory iterator
      typename IteratorB::Params const &iterator_b_params,
      ///< initial value of accumulator
      FragmentC const &src_accum,
      /// Epilogue
      Epilogue &epilogue, 
      ///< Output operator
      typename Epilogue::OutputOp const &output_op, 
      ///< Tile iterator for destination 
      typename Epilogue::OutputTileIterator &destination_iterator,
      ///< Threadblock tile coordinate in GEMM (in units of threadblock tiles)
      typename Epilogue::OutputTileIterator &source_iterator,

      int split_k_slices = 1
      ) {

    //
    // Prologue
    //

    // Issue several complete stages
    CUTLASS_PRAGMA_UNROLL
    for (int stage = 0; stage < Base::kStages - 1; ++stage, --gemm_k_iterations) {

      if (stage == 0) {
        iterator_B.set_iteration_index(0);
        this->smem_iterator_B_.set_iteration_index(0);

        // Async Copy for operand B
        CUTLASS_PRAGMA_UNROLL
        for (int j = 0; j < Detail::AsyncCopyIterationsPerStageB; ++j) {
          typename IteratorB::AccessType *dst_ptr =
              reinterpret_cast<typename IteratorB::AccessType *>(this->smem_iterator_B_.get());

          CUTLASS_PRAGMA_UNROLL
          for (int v = 0; v < IteratorB::kAccessesPerVector; ++v) {
            int const kSrcBytes = sizeof_bits<typename IteratorB::Element>::value *
                                  IteratorB::ThreadMap::kElementsPerAccess /
                                  IteratorB::kAccessesPerVector / 8;

            cutlass::arch::cp_async_zfill<kSrcBytes, kCacheOpB>(
                dst_ptr + v, iterator_B.get(), iterator_B.valid());
            
            ++iterator_B;
          }

          ++this->smem_iterator_B_;
        }
      }

      if(kItertorAlgorithm == conv::IteratorAlgorithm::kFixedStrideDilation){
        // Number of iterators is compilation static.
        iterator_A.set_iteration_index(0);
        this->smem_iterator_A_.set_iteration_index(0);

        // Async Copy for operand A
        CUTLASS_PRAGMA_UNROLL
        for (int j = 0; j < Detail::AsyncCopyIterationsPerStageA; ++j) {
          typename IteratorA::AccessType *dst_ptr =
              reinterpret_cast<typename IteratorA::AccessType *>(this->smem_iterator_A_.get());

          CUTLASS_PRAGMA_UNROLL
          for (int v = 0; v < IteratorA::kAccessesPerVector; ++v) {
            int const kSrcBytes = sizeof_bits<typename IteratorA::Element>::value *
                                  IteratorA::ThreadMap::kElementsPerAccess /
                                  IteratorA::kAccessesPerVector / 8;

            cutlass::arch::cp_async_zfill<kSrcBytes, kCacheOpA>(
                dst_ptr + v, iterator_A.get(), iterator_A.valid());
          
            ++iterator_A;
          }

          ++this->smem_iterator_A_;
        }

      } else {
        // Number of iterators is a runtime value.
        iterator_A.set_iteration_index(0);
        this->smem_iterator_A_.set_iteration_num(iterator_A.get_iteration_num());
        this->smem_iterator_A_.set_iteration_index(0);


        // Async Copy for operand A
        CUTLASS_PRAGMA_UNROLL
        for (int j = 0; j < iterator_A.get_iteration_num(); ++j) {
          typename IteratorA::AccessType *dst_ptr =
              reinterpret_cast<typename IteratorA::AccessType *>(this->smem_iterator_A_.get());

          CUTLASS_PRAGMA_UNROLL
          for (int v = 0; v < IteratorA::kAccessesPerVector; ++v) {
            int const kSrcBytes = sizeof_bits<typename IteratorA::Element>::value *
                                  IteratorA::ThreadMap::kElementsPerAccess /
                                  IteratorA::kAccessesPerVector / 8;

            cutlass::arch::cp_async_zfill<kSrcBytes, kCacheOpA>(
                dst_ptr + v, iterator_A.get(), iterator_A.valid());

            ++iterator_A;
          }

          ++this->smem_iterator_A_;
        }
      }

      // Move to the next stage
      iterator_A.advance();

      this->smem_iterator_A_.add_tile_offset({1, 0});

      // Inserts a fence to group cp.async instructions into stages.
      cutlass::arch::cp_async_fence();
    }

    /////////////////////////////////////////////////////////////////////////////
    // Waits until kStages-2 stages have committed. 
    cutlass::arch::cp_async_wait<Base::kStages - 2>();
    __syncthreads();

    // Pair of fragments used to overlap shared memory loads and math
    // instructions
    WarpLoadedFragmentA warp_loaded_frag_A[2];
    WarpLoadedFragmentB warp_loaded_frag_B[2];
    WarpTransformedFragmentA warp_transformed_frag_A[2];
    WarpTransformedFragmentB warp_transformed_frag_B[2];

    Operator warp_mma;

    this->warp_tile_iterator_A_.set_kgroup_index(0);
    this->warp_tile_iterator_B_.set_kgroup_index(0);

    this->warp_tile_iterator_A_.setup_initial_status(iterator_a_params);


    this->warp_tile_iterator_A_.load(warp_loaded_frag_A[0]);
    this->warp_tile_iterator_B_.load(warp_loaded_frag_B[0]);

    ++this->warp_tile_iterator_A_;
    ++this->warp_tile_iterator_B_;

    int smem_write_stage_idx = Base::kStages - 1;
    int smem_read_stage_idx = 0;

    warp_mma.transform(warp_transformed_frag_A[0], warp_transformed_frag_B[0],
                       warp_loaded_frag_A[0], warp_loaded_frag_B[0]);

    //
    // Mainloop
    //

    unsigned int iterations = 0;
    constexpr int inner_loop_iterations = round_up(Base::kWarpGemmIterations, 2);

    CUTLASS_GEMM_LOOP
    for (; gemm_k_iterations > (-Base::kStages + 1);) {   // Each iteration is a cta tile.

      accum.clear();
    
      //
      // Loop over GEMM K dimension
      //

      // Computes a warp-level GEMM on data held in shared memory
      // Each "warp_mma_k" refers to a warp-level matrix multiply-accumulate

      CUTLASS_PRAGMA_UNROLL
      for (int warp_mma_k = 0; warp_mma_k < inner_loop_iterations; ++warp_mma_k) {
        if (Base::kWarpGemmIterations % 2 == 0 || warp_mma_k + 1 != Base::kWarpGemmIterations) {
          // Load warp-level tiles from shared memory, wrapping to k offset if
          // this is the last group as the case may be.

          this->warp_tile_iterator_A_.set_kgroup_index((warp_mma_k + 1) % Shape::kK);
          this->warp_tile_iterator_B_.set_kgroup_index((warp_mma_k + 1) % Shape::kK);

          this->warp_tile_iterator_A_.load(warp_loaded_frag_A[(warp_mma_k + 1) % 2]);
          this->warp_tile_iterator_B_.load(warp_loaded_frag_B[(warp_mma_k + 1) % 2]);

          ++this->warp_tile_iterator_A_;
          ++this->warp_tile_iterator_B_;
        }

        if (warp_mma_k > 0)
          warp_mma.transform(warp_transformed_frag_A[warp_mma_k % 2],
                              warp_transformed_frag_B[warp_mma_k % 2],
                              warp_loaded_frag_A[warp_mma_k % 2],
                              warp_loaded_frag_B[warp_mma_k % 2]);

        // Issue global->shared copies for the next stage
        int group_start_iteration_A, group_start_iteration_B;

        if (warp_mma_k == 0) {
          group_start_iteration_A = 0;
          group_start_iteration_B = 0;
          copy_tiles_and_advance(
              iterator_A, iterator_B, group_start_iteration_A, group_start_iteration_B);
        }

        if (warp_mma_k < Base::kWarpGemmIterations) {
          warp_mma(
            accum, 
            warp_transformed_frag_A[warp_mma_k % 2],
            warp_transformed_frag_B[warp_mma_k % 2], 
            accum
          );
        }

        if (warp_mma_k + 1 == inner_loop_iterations)
          warp_mma.transform(warp_transformed_frag_A[(warp_mma_k + 1) % 2],
                              warp_transformed_frag_B[(warp_mma_k + 1) % 2],
                              warp_loaded_frag_A[(warp_mma_k + 1) % 2],
                              warp_loaded_frag_B[(warp_mma_k + 1) % 2]);

        if (warp_mma_k + 2 == inner_loop_iterations) {
          // Inserts a fence to group cp.async instructions into stages.
          cutlass::arch::cp_async_fence();

          // Waits until kStages-2 stages of cp.async have committed
          arch::cp_async_wait<Base::kStages - 2>();
          __syncthreads();

          // Move to the next cta
          iterator_A.advance();

          this->smem_iterator_A_.add_tile_offset({1, 0});

          // Add negative offsets to return iterators to the 'start' of the
          // circular buffer in shared memory
          if (smem_write_stage_idx == (Base::kStages - 1)) {
            this->smem_iterator_A_.add_tile_offset({-Base::kStages, 0});
   
            smem_write_stage_idx = 0;
          } else {
            ++smem_write_stage_idx;
          }

          if (smem_read_stage_idx == (Base::kStages - 1)) {
            this->warp_tile_iterator_A_.advance(- (Base::kStages-1) * iterator_A.get_load_size());
            smem_read_stage_idx = 0;
          } else {
            this->warp_tile_iterator_A_.advance(iterator_A.get_load_size());
            ++smem_read_stage_idx;
          }

          if (kItertorAlgorithm == conv::IteratorAlgorithm::kFixedStrideDilation) {
            this->warp_tile_iterator_A_.setup_initial_status(iterator_a_params);
          }

          // goback to start position. B has no multiple stage
          this->warp_tile_iterator_B_.add_tile_offset({-Policy::kPartitionsK * Shape::kK, 0});

          --gemm_k_iterations;
        }
      }

      //
      // Epilogue
      //
      int32_t smem_base_offset = iterator_B.get_load_size() + (iterations % Base::kStages) * iterator_A.get_load_size();

      destination_iterator.set_tile_index(iterations * split_k_slices);
      
      source_iterator.set_tile_index(iterations * split_k_slices);
    
      epilogue(output_op, destination_iterator, accum, source_iterator, smem_base_offset);

      ++iterations;
    }

    // Insert fence and wait for all outstanding cp.async operations to commit.
    cutlass::arch::cp_async_fence();
    cutlass::arch::cp_async_wait<0>();
    __syncthreads();

  }

};

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

}  // namespace threadblock
}  // namespace gemm
}  // namespace cutlass

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