text: /***************************************************************************************************
   * Copyright (c) 2017-2021, 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 PredicatedTileIteratorPredicates.
  
    PredicatedTileIteratorPredicates enables both upper and lower bounds for predicates.
  
  */
  
  #pragma once
  
  #include "cutlass/cutlass.h"
  #include "cutlass/numeric_types.h"
  #include "cutlass/array.h"
  #include "cutlass/layout/matrix.h"
  #include "cutlass/layout/tensor.h"
  #include "cutlass/matrix_shape.h"
  #include "cutlass/tensor_ref.h"
  #include "cutlass/transform/pitch_linear_thread_map.h"
  #include "cutlass/epilogue/threadblock/output_tile_thread_map.h"
  #include "cutlass/arch/arch.h"
  #include "cutlass/arch/memory.h"
  #include "cutlass/epilogue/threadblock/predicated_tile_iterator_params.h"
  
  ////////////////////////////////////////////////////////////////////////////////
  
  namespace cutlass {
  
  ////////////////////////////////////////////////////////////////////////////////
  
  namespace epilogue {
  namespace threadblock {
  
  ////////////////////////////////////////////////////////////////////////////////
  
  /// Tile iterator predicates used to bound computations in epilogue.
  ///
  /// Satisfies: ReadableTileIterator | PredicatedTileIterator | ForwardTileIterator
  ///
  template <
    typename ThreadMap_,       ///< Thread map (conept: OutputTileThreadMap)
    typename Element_          ///< Element data type
  >
  class PredicatedTileIteratorPredicates {
  public:
    using ThreadMap = ThreadMap_;
    using Shape = typename ThreadMap::Shape;
  
    using Element = Element_;
  
    using Layout = layout::RowMajor;
    using TensorRef = TensorRef<Element, Layout>;
    using ConstTensorRef = typename TensorRef::ConstTensorRef;
  
    using Index = typename Layout::Index;
    using LongIndex = typename Layout::LongIndex;
    using TensorCoord = MatrixCoord;
  
    static int const kElementsPerAccess = ThreadMap::kElementsPerAccess;
    static int const kThreads = ThreadMap::kThreads;
    static int const kIterations = ThreadMap::Count::kTile;
  
    static_assert( ThreadMap::Iterations::kRow > 0,"ThreadMap::Iterations::kRow must be > 0");
    static_assert( ThreadMap::Iterations::kGroup > 0,"ThreadMap::Iterations::kGroup must be > 0");
    static_assert( ThreadMap::Iterations::kCluster > 0,"ThreadMap::Iterations::kCluster must be > 0");
    static_assert( ThreadMap::Iterations::kColumn > 0,"ThreadMap::Iterations::kColumn must be > 0");
  
    /// Fragment object
    using Fragment = Array<
      Element, 
      ThreadMap::Iterations::kColumn * 
      ThreadMap::Iterations::kRow * 
      ThreadMap::Iterations::kGroup * 
      ThreadMap::Iterations::kCluster * ThreadMap::kElementsPerAccess>;
  
    /// Memory access size
    using AccessType = AlignedArray<Element, ThreadMap::kElementsPerAccess>;
  
    //
    // Parameters struct
    //
  
    /// Uses a non-template class
    struct Params : PredicatedTileIteratorParams {
  
      CUTLASS_HOST_DEVICE
      Params() { }
  
      CUTLASS_HOST_DEVICE
      Params(Layout const &layout): 
        PredicatedTileIteratorParams(
          layout.stride(0) * int(sizeof(AccessType)) / kElementsPerAccess,
          make_OutputTileThreadMapDesc<ThreadMap>()
        ) 
      {
          
      }
    };
  
    /// Mask object
    struct Mask {
  
      static int const kCount = ThreadMap::Iterations::kColumn;
  
      /// Predicate state
      bool predicates[kCount];
  
      //
      // Mask
      //
      CUTLASS_HOST_DEVICE
      Mask() {
        enable();
      }
  
      ///< Efficiently disables all accesses guarded by mask
      CUTLASS_HOST_DEVICE void clear() {
        CUTLASS_PRAGMA_UNROLL
        for (int i = 0; i < kCount; ++i) {
          predicates[i] = false;
        }
      }
  
      ///< CUTLASS_HOST_DEVICE enables all accesses guarded by mask
      CUTLASS_DEVICE void enable() {
        CUTLASS_PRAGMA_UNROLL
        for (int i = 0; i < kCount; ++i) {
          predicates[i] = true;
        }
      }
    };
  
  private:
  
    //
    // Data members
    //
  
    /// Parameters structure containing reference and precomputed state.
    PredicatedTileIteratorParams params_;
  
    /// Array of boolean values to contain steady-state predicates
    Mask mask_;
  
    /// Extent of the matrix tile in rows
    Index lower_extent_row_;
    Index upper_extent_row_;
  
    /// A thread's starting row position (assuming steady-state predicates have been computed)
    Index thread_start_row_;
  
    /// Internal state counter
    int state_[3];
   
    //
    // Static asserts about internal strides
    //
  
    static_assert(sizeof(lower_extent_row_) == 4, "Expected 32b extents");
    static_assert(sizeof(upper_extent_row_) == 4, "Expected 32b extents");
    static_assert(sizeof(thread_start_row_) == 4, "Expected 32b extents");
    static_assert(sizeof(PredicatedTileIteratorParams::stride) == 8, "Expected 64b strides");
  
  private:
  
    //
    // Methods
    //
  
  public:
  
    //
    // Methods
    //
  
    /// Constructor
    CUTLASS_DEVICE
    PredicatedTileIteratorPredicates(
      PredicatedTileIteratorParams const & params,
      TensorCoord lower_extent,
      TensorCoord upper_extent,
      int thread_idx,
      TensorCoord threadblock_offset = TensorCoord()
    ): 
      params_(params)
    {
  
      TensorCoord thread_offset = ThreadMap::initial_offset(thread_idx) + threadblock_offset;
  
      lower_extent_row_ = lower_extent.row();
      upper_extent_row_ = upper_extent.row();
      thread_start_row_ = thread_offset.row();
  
      // Initialize predicates
      CUTLASS_PRAGMA_UNROLL
      for (int c = 0; c < ThreadMap::Iterations::kColumn; ++c) {
  
        mask_.predicates[c] = ((thread_offset.column() 
          + ThreadMap::Delta::kColumn * c) < upper_extent.column()) &&
          ((thread_offset.column() + ThreadMap::Delta::kColumn * c) >= lower_extent.column());
      }
  
      // Initialize internal state counter
      state_[0] = state_[1] = state_[2] = 0;
    }
  
    /// Advances to the next position to load or store
    CUTLASS_HOST_DEVICE
    PredicatedTileIteratorPredicates &operator++() {
  
      ++state_[0];
      thread_start_row_ += ThreadMap::Shape::kRow;
  
      if (state_[0] == ThreadMap::Count::kRow) {
  
        state_[0] = 0;
        ++state_[1];
  
        thread_start_row_ += (ThreadMap::Shape::kGroup - 1) *
          ThreadMap::Shape::kRow * ThreadMap::Count::kRow;
  
        if (state_[1] == ThreadMap::Count::kGroup) {
  
          state_[1] = 0;
          ++state_[2];
  
          thread_start_row_ += ThreadMap::Count::kGroup *
            ThreadMap::Shape::kGroup * ThreadMap::Count::kRow * ThreadMap::Shape::kRow;
  
          if (state_[2] == ThreadMap::Count::kCluster) {
            state_[2] = 0;
          }
        }
      }
  
      return *this;
    }
  
    ///< Efficiently disables all accesses guarded by mask
    CUTLASS_DEVICE void clear_mask() {
      mask_.clear();
    }
  
    ///< Efficiently enables all accesses guarded by mask
    CUTLASS_DEVICE void enable_mask() {
      mask_.enable();
    }
  
    ///< Gets the mask
    CUTLASS_DEVICE void get_mask(Mask &mask) {
      mask = mask_;
    }
  
    ///< Sets the mask
    CUTLASS_DEVICE void set_mask(Mask const &mask) {
      mask_ = mask;
    }
  
    ///< Gets lower_extent_row_
    CUTLASS_DEVICE Index get_lower_extent_row() {
      return lower_extent_row_;
    }
  
    ///< Gets upper_extent_row_
    CUTLASS_DEVICE Index get_upper_extent_row() {
      return upper_extent_row_;
    }
  
    ///< Gets thread_start_row_
    CUTLASS_DEVICE Index get_thread_start_row() {
      return thread_start_row_;
    }
  };
  
  ///////////////////////////////////////////////////////////////////////////////
  
  } // namespace threadblock
  } // namespace epilogue
  } // namespace cutlass
  
  ////////////////////////////////////////////////////////////////////////////////
  