/***************************************************************************************************
 * Copyright (c) 2017 - 2024 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 Templates implementing warp-level matrix multiply-accumulate operations.
*/

#pragma once

#include "cutlass/cutlass.h"
#include "cutlass/array.h"
#include "cutlass/numeric_types.h"
#include "cutlass/matrix_shape.h"
#include "cutlass/gemm/gemm.h"

#include "cutlass/array_planar_complex.h"

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

namespace cutlass {
namespace gemm {
namespace warp {

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

template <typename TileIterator_>
class TileIteratorPlanarComplex {
public:

  /// Underlying iterator over real-valued tiles
  using TileIterator = TileIterator_;

  /// Underlying element type
  using Element = typename TileIterator::Element;

  /// Underlying layout type
  using Layout = typename TileIterator::Layout;

  /// TensorRef type for loading element from a tensor
  using TensorRef = typename TileIterator::TensorRef;

  /// Index type
  using Index = typename TensorRef::Index;

  /// Long Index type
  using LongIndex = typename TensorRef::LongIndex;

  /// Coordinate for an element in the tensor
  using TensorCoord = typename TensorRef::TensorCoord;

  /// Planar complex fragment
  using Fragment = ArrayPlanarComplex<Element, TileIterator::Fragment::kElements>;

public:

  /// Underlying tile iterator
  TileIterator tile_iterator_;

  /// Offset (in units of bytes) to the imaginary part of the planar complex matrix
  LongIndex imaginary_offset_;

public:
    /// Default ctor constructs null iterator
  CUTLASS_HOST_DEVICE
  TileIteratorPlanarComplex(): imaginary_offset_(0) { }

  /// Constructor from TensorRef
  CUTLASS_DEVICE
  TileIteratorPlanarComplex(
    TensorRef const &ref, 
    int lane_id,
    LongIndex imaginary_offset
  ):
    tile_iterator_(ref, lane_id),
    imaginary_offset_((imaginary_offset * sizeof_bits<Element>::value) / 8) { }


  /// Adds a pointer offset to internal pointer(s) to advance through memory
  CUTLASS_DEVICE
  TileIteratorPlanarComplex &add_pointer_offset(LongIndex offset) {

    tile_iterator_.add_pointer_offset(offset);

    return *this;
  }

  /// Advances an iterator along logical dimensions of matrix in units of whole tiles
  CUTLASS_HOST_DEVICE
  TileIteratorPlanarComplex &add_tile_offset(TensorCoord const &tile_offset) {

    tile_iterator_.add_tile_offset(tile_offset);

    return *this;
  }

  /// Advances the iterator along the advance dimension
  CUTLASS_DEVICE
  TileIteratorPlanarComplex & operator++() {
    ++tile_iterator_;
    return *this;
  }

  //
  // WIP
  //

  /// Advances the iterator along the opposite of the advance dimension
  CUTLASS_HOST_DEVICE
  TileIteratorPlanarComplex & operator--() {
    --tile_iterator_;
    return *this;
  }

  ///< advances in units of whole tiles along the logical coordinate space of the tensor
  CUTLASS_DEVICE
  TileIteratorPlanarComplex & operator+=(TensorCoord const &tile_offset) {
    tile_iterator_.add_tile_offset(tile_offset);
    return *this;
  }

  ///< advances in units of whole tiles along the logical coordinate space of the tensor
  CUTLASS_DEVICE
  TileIteratorPlanarComplex & operator-=(TensorCoord const &tile_offset) {
    tile_iterator_.add_tile_offset(-tile_offset);
    return *this;
  }

  /// Loads a fragment from memory at the location pointed to by the iterator.
  CUTLASS_HOST_DEVICE
  void load(Fragment &frag) const {

    tile_iterator_.load_with_byte_offset(frag.real, 0);
    tile_iterator_.load_with_byte_offset(frag.imag, imaginary_offset_);
  }

  /// Loads a fragment from memory with additional logical offset
  CUTLASS_DEVICE
  void load_with_byte_offset(
      /// fragment to load from the tensor
      Fragment &frag,
      /// loads a tile with a linear offset in units of bytes
      Index byte_offset) const {

    tile_iterator_.load_with_byte_offset(frag.real, byte_offset);
    tile_iterator_.load_with_byte_offset(frag.imag, byte_offset + imaginary_offset_);
  }

  /// Loads a fragment from memory with additional logical offset
  CUTLASS_DEVICE
  void load_with_pointer_offset(
      /// fragment to load from the tensor
      Fragment &frag,
      /// loads a tile with a linear offset
      Index pointer_offset) const {

    Index byte_offset = (pointer_offset * sizeof_bits<Element>::value)/8;

    tile_iterator_.load_with_byte_offset(frag.real, byte_offset);
    tile_iterator_.load_with_byte_offset(frag.imag, byte_offset + imaginary_offset_);
  }

  /// Loads a fragment from memory with logical offset in units of whole tiles.
  CUTLASS_DEVICE
  void load(
      /// fragment to load from the tensor
      Fragment &frag,
      /// loads a tile with a logical offset in units of whole tiles
      TensorCoord const &tile_offset) const {

    tile_iterator_.load_with_byte_offset(frag.real, tile_offset, 0);
    tile_iterator_.load_with_byte_offset(frag.imag, tile_offset, imaginary_offset_);
  }

  /// Loads a fragment from memory with logical offset in units of whole tiles.
  CUTLASS_DEVICE
  void load(
      /// fragment to load from the tensor
      Fragment &frag,
      /// loads a tile with a logical offset in units of whole tiles
      TensorCoord const &tile_offset,
      /// loads a tile with a logical offset AND a pointer offset
      Index pointer_offset) const {

    Index byte_offset = (pointer_offset * sizeof_bits<Element>::value)/8;

    tile_iterator_.load_with_byte_offset(frag.real, tile_offset, byte_offset);
    tile_iterator_.load_with_byte_offset(frag.real, tile_offset, byte_offset + imaginary_offset_);
  }

  /// Loads a fragment from memory with logical offset in units of whole tiles.
  CUTLASS_DEVICE
  void load_with_byte_offset(
      /// fragment to load from the tensor
      Fragment &frag,
      /// loads a tile with a logical offset in units of whole tiles
      TensorCoord const &tile_offset,
      /// loads a tile with a logical offset AND a pointer offset
      Index byte_offset) const {

    tile_iterator_.load_with_byte_offset(frag.real, tile_offset, byte_offset);
    tile_iterator_.load_with_byte_offset(frag.imag, tile_offset, byte_offset + imaginary_offset_);
  }

  /// Notify the iterator which k-group it is currently pointing to.
  ///
  /// This does not advance the iterator. Rather, it overrides its internal
  /// tracking with constant-valued k-group index to enable the compiler to
  /// fold constants and achieve more efficient code.
  ///
  /// This is used by some nontrivial permuted layouts.
  CUTLASS_DEVICE
  void set_kgroup_index(int k_group) {
    tile_iterator_.set_kgroup_index(k_group);
  }
};

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

} // namespace warp
} // namespace gemm
} // namespace cutlass

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