/***************************************************************************************************
 * 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 Defines tile iterator traits for loading thread block-level tile from global memory.
*/
#pragma once

#include <cutlass/gemm/gemm_global_tile.h>

namespace cutlass {
namespace gemm {

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

template <typename Scalar_, typename Tile_, typename Threads_, int kAccessSize_>
struct WmmaGemmGlobalIteratorCdTraits : public GemmGlobalTileTraits<GemmOperand::kC,
                                                                    MatrixLayout::kColumnMajor,
                                                                    Scalar_,
                                                                    Tile_,
                                                                    Threads_,
                                                                    kAccessSize_> {
  /// The base class.
  typedef GemmGlobalTileTraits<GemmOperand::kC,
                               MatrixLayout::kColumnMajor,
                               Scalar_,
                               Tile_,
                               Threads_,
                               kAccessSize_>
      Base;

  /// Override the strides in each dimension between different loads/stores.
  typedef Shape<0, 0, Base::Delta::kW, Base::Delta::kC> Delta;

  /// Computes the thread offset in (H, W) based on thread ID
  struct ThreadOffset {
    CUTLASS_HOST_DEVICE
    Coord<4> operator()() const {
      int thread_offset_h = threadIdx.x / Base::Threads::kW;
      int thread_offset_w = threadIdx.x % Base::Threads::kW * Base::ThreadsDelta::kW;

      return make_Coord(0, thread_offset_h, thread_offset_w, 0);
    }
  };
};

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

template <typename TileTraits_, typename Index_ = int>
struct WmmaGemmGlobalIteratorCd : public TileIteratorBase<TileTraits_,
                                                          typename TileTraits_::Scalar,
                                                          IteratorAdvance::kH,
                                                          MemorySpace::kGlobal,
                                                          Index_> {
  /// This class.
  typedef WmmaGemmGlobalIteratorCd<TileTraits_, Index_> This_;
  /// The traits.
  typedef TileTraits_ Traits;
  /// The base class.
  typedef TileIteratorBase<Traits,
                           typename TileTraits_::Scalar,
                           IteratorAdvance::kH,
                           MemorySpace::kGlobal,
                           Index_>
      Base;
  /// Override the strides in each dimension between different loads/stores.
  typedef Shape<0, 0, Base::Delta::kW, Base::Delta::kC> ImmediateOffsetStrides;
  /// The layout.
  static MatrixLayout::Kind const kLayout = TileTraits_::kLayout;

  /// The scalar.
  typedef typename TileTraits_::Scalar Scalar;
  /// The pointer.
  typedef typename TileTraits_::Pointer Pointer;
  /// The threads.
  typedef typename TileTraits_::Threads Threads;
  /// The index.
  typedef Index_ Index;
  /// The thread offset functor.
  typedef typename TileTraits_::ThreadOffset ThreadOffset;

  /// The params.
  struct Params {
    /// The pointer.
    Pointer pointer;
    /// The stride in the H dimension to setup the thread in the block.
    Index stride_h;
    /// The strides to increment the pointer.
    Index inc_h, inc_advance;
    /// The column offset to compute the predicate for the columns.
    Index predicate_offset;
    /// The strides to increment the predicate offset.
    Index predicate_inc_h, predicate_inc_advance;

    /// Setup the params.
    CUTLASS_HOST_DEVICE int initialize(
        Pointer pointer, Index ld, Index n, Index epilogue_stride_w, Index epilogue_delta_w) {
      // The pointer.
      this->pointer = pointer;
      // Setup the base stride. One "group of threads" per column.
      stride_h = ld;
      // Each thread output 1 column per iteration. .
      inc_h = ld * TileTraits_::Threads::kH;
      inc_advance = inc_h + epilogue_stride_w;

      predicate_offset = n;
      predicate_inc_h = TileTraits_::Threads::kH;
      predicate_inc_advance = predicate_inc_h + epilogue_delta_w;

      // It worked.
      return 0;
    }
  };

  Params params;

  Coord<4> thread_offset;

  /// Ctor.
  CUTLASS_DEVICE WmmaGemmGlobalIteratorCd() {}

  /// Ctor.
  CUTLASS_DEVICE WmmaGemmGlobalIteratorCd(Params const& params,
                                          const Coord<3>& bounds,
                                          const Coord<3>& block,
                                          int const pointer_offset = 0,
                                          int const pred_offset = 0,
                                          ThreadOffset thread_offset_func = ThreadOffset())

      : params(params) {
    thread_offset = thread_offset_func();
    // Each warp works on a different column of the tile.
    int const h = thread_offset[1] + block[1];
    // Each lane writes a different element.
    int const w = thread_offset[2] + block[2];
    // Setup the pointer.
    this->params.pointer += ((h * params.stride_h + w) + pointer_offset);

    // Prepare the vector of predicates.
    for (int i = 0; i < Base::Iterations::kW; ++i) {
      predicates.set(i, w + i * Base::Delta::kW < bounds[2]);
    }
    this->params.predicate_offset -= (h + pred_offset);
  }

  /// The accessor.
  CUTLASS_DEVICE void get(typename Base::AccessType& value, int d, int h, int w, int c) const {
    int const imm =
        ComputeOffsetFromStrides<typename Base::ImmediateOffsetStrides>::get(0, 0, w, c);
    Load<Scalar, TileTraits_::kAccessSize, MemorySpace::kGlobal>::load(value, params.pointer, imm);
  }

  /// Increment the pointer in the C dimension.
  CUTLASS_DEVICE void inc_c() {}
  /// Increment the pointer in the W dimension.
  CUTLASS_DEVICE void inc_w() {}
  /// Increment the pointer in the H dimension.
  CUTLASS_DEVICE void inc_h() {
    params.pointer += params.inc_h;
    params.predicate_offset -= params.predicate_inc_h;
  }
  /// Increment the pointer in the D dimension.
  CUTLASS_DEVICE void inc_d() {}
  /// Increment the pointer to move to the next iteration.
  CUTLASS_DEVICE void inc_advance() {
    params.pointer += params.inc_advance;
    params.predicate_offset -= params.predicate_inc_advance;
  }

  /// The accessor.
  CUTLASS_DEVICE void set(typename Base::AccessType const& value, int d, int h, int w, int c) {
    int const imm =
        ComputeOffsetFromStrides<typename Base::ImmediateOffsetStrides>::get(d, h, w, 0);
    Store<Scalar, TileTraits_::kAccessSize, MemorySpace::kGlobal>::store(
        value, params.pointer, imm);
  }

  /// Test the predicate.
  CUTLASS_DEVICE bool valid(int d, int h, int w, int c) const {
    return predicates.at(w) && params.predicate_offset > 0;
  }

  /// The predicates for the row.
  cutlass::PredicateVector<Base::Iterations::kW> predicates;
};

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

}  // namespace gemm
}  // namespace cutlass