text: /***************************************************************************************************
   * Copyright (c) 2017-2020, 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 Functor performing linear scaling operations used by epilogues. Values are clamped before
           converting to the output element type.
  */
  
  #pragma once
  
  #include "cutlass/cutlass.h"
  #include "cutlass/numeric_types.h"
  #include "cutlass/array.h"
  #include "cutlass/functional.h"
  #include "cutlass/numeric_conversion.h"
  
  /////////////////////////////////////////////////////////////////////////////////////////////////
  
  namespace cutlass {
  namespace epilogue {
  namespace thread {
  
  /////////////////////////////////////////////////////////////////////////////////////////////////
  
  /// Applies a linear combination operator to an array of elements then clamps the output before
  /// converting to the output element type.
  ///
  /// D = alpha * accumulator + beta * source + uniform
  ///
  template <
    typename ElementOutput_,                             ///< Data type used to load and store tensors
    int Count,                                           ///< Number of elements computed per operation
    typename ElementAccumulator_ = ElementOutput_,       ///< Accumulator data type
    typename ElementCompute_ = ElementOutput_,           ///< Data type used to compute linear combination
    FloatRoundStyle Round = FloatRoundStyle::round_to_nearest
  >
  class LinearCombinationClamp {
  public:
  
    using ElementOutput = ElementOutput_;
    using ElementAccumulator = ElementAccumulator_;
    using ElementCompute = ElementCompute_;
  
    static int const kCount = Count;
  
    using FragmentOutput = Array<ElementOutput, kCount>;
    using FragmentAccumulator = Array<ElementAccumulator, kCount>;
    using ComputeFragment = Array<ElementCompute, kCount>;
  
    static FloatRoundStyle const kRound = Round;
  
    /// Host-constructable parameters structure
    struct Params {
  
      ElementCompute alpha;                  ///< scales accumulators
      ElementCompute beta;                   ///< scales source tensor
      ElementCompute const *alpha_ptr;       ///< pointer to accumulator scalar - if not null, loads it from memory
      ElementCompute const *beta_ptr;        ///< pointer to source scalar - if not null, loads it from memory
  
      //
      // Methods
      //
  
      CUTLASS_HOST_DEVICE
      Params(): 
        alpha(ElementCompute(1)), 
        beta(ElementCompute(0)), 
        alpha_ptr(nullptr), 
        beta_ptr(nullptr) { }
  
      CUTLASS_HOST_DEVICE
      Params(
        ElementCompute alpha,
        ElementCompute beta
      ): alpha(alpha), beta(beta), alpha_ptr(nullptr), beta_ptr(nullptr) {
  
      }
  
      CUTLASS_HOST_DEVICE
      Params(
        ElementCompute const *alpha_ptr,
        ElementCompute const *beta_ptr
      ): alpha(0), beta(0), alpha_ptr(alpha_ptr), beta_ptr(beta_ptr) {
  
      }
    };
  
  private:
  
    //
    // Data members
    //
  
    ElementCompute alpha_;
    ElementCompute beta_;
  
  public:
  
    /// Constructs the function object, possibly loading from pointers in host memory
    CUTLASS_HOST_DEVICE
    LinearCombinationClamp(Params const &params) {
  
      alpha_ = (params.alpha_ptr ? *params.alpha_ptr : params.alpha);
      beta_ = (params.beta_ptr ? *params.beta_ptr : params.beta);
    }
  
    /// Returns true if source is needed
    CUTLASS_HOST_DEVICE
    bool is_source_needed() const {
      return beta_ != ElementCompute(0);
    }
  
    /// Functionally required for serial reduction in the epilogue
    CUTLASS_HOST_DEVICE
    void set_k_partition(int k_partition) {
      if (k_partition) {
        beta_ = ElementCompute(1);
      }
    }
  
    /// Computes linear scaling: D = alpha * accumulator + beta * source
    CUTLASS_HOST_DEVICE
    FragmentOutput operator()(
      FragmentAccumulator const &accumulator, 
      FragmentOutput const &source,
      ElementCompute uniform = ElementCompute(0)) const {
  
      // Convert source to interal compute numeric type
      NumericArrayConverter<ElementCompute, ElementOutput, kCount, Round> source_converter;
      NumericArrayConverter<ElementCompute, ElementAccumulator, kCount, Round> accumulator_converter;
  
      ComputeFragment converted_source = source_converter(source);
      ComputeFragment converted_accumulator = accumulator_converter(accumulator);
  
      // Perform binary operations
  
      ComputeFragment intermediate;
  
      multiplies<ComputeFragment> mul_add_source;
      multiply_add<ComputeFragment> mul_add_accumulator;
      
      minimum<ComputeFragment> min_accumulator;
      maximum<ComputeFragment> max_accumulator;
  
      intermediate = mul_add_source(beta_, converted_source);                             // X =  beta * C + uniform
      intermediate = mul_add_accumulator(alpha_, converted_accumulator, intermediate);    // D = alpha * Accum + X
  
      /// Clamping constant value
      ElementCompute const kClamp =
          ElementCompute((1U << (sizeof_bits<ElementOutput>::value - 1)) - 1);
  
      intermediate = max_accumulator(intermediate, -kClamp - ElementCompute(1));
      intermediate = min_accumulator(intermediate, kClamp);
  
      // Convert to destination numeric type
      NumericArrayConverter<ElementOutput, ElementCompute, kCount, Round> destination_converter;
  
      return destination_converter(intermediate);
    }
  
    /// Computes linear scaling: D = alpha * accumulator 
    CUTLASS_HOST_DEVICE
    FragmentOutput operator()(
      FragmentAccumulator const &accumulator) const {
  
      // Convert source to interal compute numeric type
      NumericArrayConverter<ElementCompute, ElementAccumulator, kCount, Round> accumulator_converter;
  
      ComputeFragment converted_accumulator = accumulator_converter(accumulator);
  
      // Perform binary operations
  
      ComputeFragment intermediate;
  
      multiplies<ComputeFragment> mul_accumulator;
      
      minimum<ComputeFragment> min_accumulator;
      maximum<ComputeFragment> max_accumulator;
  
      intermediate = mul_accumulator(alpha_, converted_accumulator);    // D = alpha * Accum
  
      /// Clamping constant value
      ElementCompute const kClamp =
          ElementCompute((1U << (sizeof_bits<ElementOutput>::value - 1)) - 1);
  
      intermediate = max_accumulator(intermediate, -kClamp - ElementCompute(1));
      intermediate = min_accumulator(intermediate, kClamp);
  
      // Convert to destination numeric type
      NumericArrayConverter<ElementOutput, ElementCompute, kCount, Round> destination_converter;
  
      return destination_converter(intermediate);
    }
  };
  
  /////////////////////////////////////////////////////////////////////////////////////////////////
  
  // Conditional guards to enable partial specialization for packed integers
  #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 720) && ((__CUDACC_VER_MAJOR__ > 10) || ((__CUDACC_VER_MAJOR__ >= 10) && (__CUDACC_VER_MINOR__ >= 2)))
  
  /// Applies a linear combination operator to an array of elements then clamps the output before
  /// converting to the output element type.
  ///
  /// D = alpha * accumulator + beta * source + uniform
  ///
  template <
    typename ElementOutput_,                             ///< Data type used to load and store tensors
    int Count,                                           ///< Number of elements computed per operation
    FloatRoundStyle Round
  >
  class LinearCombinationClamp<ElementOutput_, Count, int, float, Round> {
  public:
  
    using ElementOutput = ElementOutput_;
    using ElementAccumulator = int;
    using ElementCompute = float;
  
    static_assert(
        platform::is_same<ElementOutput, int32_t>::value ||
            platform::is_same<ElementOutput, uint32_t>::value ||
            platform::is_same<ElementOutput, int16_t>::value ||
            platform::is_same<ElementOutput, uint16_t>::value ||
            platform::is_same<ElementOutput, int8_t>::value ||
            platform::is_same<ElementOutput, uint8_t>::value ||
            platform::is_same<ElementOutput, cutlass::int4b_t>::value ||
            platform::is_same<ElementOutput, cutlass::uint4b_t>::value ||
            platform::is_same<ElementOutput, cutlass::uint1b_t>::value,
        "This elementwise op expects the output to be int.");
  
    static int const kCount = Count;
  
    using FragmentOutput = Array<ElementOutput, kCount>;
    using FragmentAccumulator = Array<ElementAccumulator, kCount>;
    using ComputeFragment = Array<ElementCompute, kCount>;
  
    static FloatRoundStyle const kRound = Round;
  
    /// Host-constructable parameters structure
    struct Params {
  
      ElementCompute alpha;                  ///< scales accumulators
      ElementCompute beta;                   ///< scales source tensor
      ElementCompute const *alpha_ptr;       ///< pointer to accumulator scalar - if not null, loads it from memory
      ElementCompute const *beta_ptr;        ///< pointer to source scalar - if not null, loads it from memory
  
      //
      // Methods
      //
  
      CUTLASS_HOST_DEVICE
      Params(): 
        alpha(ElementCompute(1)), 
        beta(ElementCompute(0)), 
        alpha_ptr(nullptr), 
        beta_ptr(nullptr) { }
  
      CUTLASS_HOST_DEVICE
      Params(
        ElementCompute alpha,
        ElementCompute beta
      ): alpha(alpha), beta(beta), alpha_ptr(nullptr), beta_ptr(nullptr) {
  
      }
  
      CUTLASS_HOST_DEVICE
      Params(
        ElementCompute const *alpha_ptr,
        ElementCompute const *beta_ptr
      ): alpha(0), beta(0), alpha_ptr(alpha_ptr), beta_ptr(beta_ptr) {
  
      }
    };
  
  private:
  
    //
    // Data members
    //
  
    ElementCompute alpha_;
    ElementCompute beta_;
  
  public:
  
    /// Constructs the function object, possibly loading from pointers in host memory
    CUTLASS_HOST_DEVICE
    LinearCombinationClamp(Params const &params) {
  
      alpha_ = (params.alpha_ptr ? *params.alpha_ptr : params.alpha);
      beta_ = (params.beta_ptr ? *params.beta_ptr : params.beta);
    }
  
    /// Returns true if source is needed
    CUTLASS_HOST_DEVICE
    bool is_source_needed() const {
      return beta_ != ElementCompute(0);
    }
  
    /// Functionally required for serial reduction in the epilogue
    CUTLASS_HOST_DEVICE
    void set_k_partition(int k_partition) {
      if (k_partition) {
        beta_ = ElementCompute(1);
      }
    }
    
    /// Computes linear scaling: D = alpha * accumulator + beta * source
    CUTLASS_HOST_DEVICE
    FragmentOutput operator()(
      FragmentAccumulator const &accumulator, 
      FragmentOutput const &source,
      ElementCompute uniform = ElementCompute(0)) const {
  
      // Convert source to interal compute numeric type
      NumericArrayConverter<ElementCompute, ElementOutput, kCount, Round> source_converter;
      NumericArrayConverter<ElementCompute, ElementAccumulator, kCount, Round> accumulator_converter;
  
      ComputeFragment converted_source = source_converter(source);
      ComputeFragment converted_accumulator = accumulator_converter(accumulator);
  
      // Compute linear scaling in floating point
      ComputeFragment intermediate;
  
      multiplies<ComputeFragment> mul_add_source;
      multiply_add<ComputeFragment> mul_add_accumulator;
      
      // Float min-max
      intermediate = mul_add_source(beta_, converted_source);                             // X =  beta * C + uniform
      intermediate = mul_add_accumulator(alpha_, converted_accumulator, intermediate);    // D = alpha * Accum + X
  
      // Convert floats back to INT
      FragmentAccumulator scaled_accumulator;
  
      CUTLASS_PRAGMA_UNROLL
      for (int i = 0; i < kCount; ++i) {
        scaled_accumulator[i] = static_cast<int>(intermediate[i]);
      }
  
      // Convert to destination numeric type
      NumericArrayConverter<ElementOutput, int, kCount, Round> destination_converter;
  
      return destination_converter(scaled_accumulator);
    }
  
    /// Computes linear scaling: D = alpha * accumulator
    CUTLASS_HOST_DEVICE
    FragmentOutput operator()(FragmentAccumulator const &accumulator) const {
  
      // Convert source to interal compute numeric type
      NumericArrayConverter<ElementCompute, ElementAccumulator, kCount, Round> accumulator_converter;
  
      ComputeFragment converted_accumulator = accumulator_converter(accumulator);
  
      // Compute linear scaling in floating point
      ComputeFragment intermediate;
  
      multiplies<ComputeFragment> mul_add_accumulator;
      
      // Float min-max
      intermediate = mul_add_accumulator(alpha_, converted_accumulator);    // D = alpha * Accum
  
      // Convert floats back to INT
      FragmentAccumulator scaled_accumulator;
  
      CUTLASS_PRAGMA_UNROLL
      for (int i = 0; i < kCount; ++i) {
        scaled_accumulator[i] = static_cast<int>(intermediate[i]);
      }
  
      // Convert to destination numeric type
      NumericArrayConverter<ElementOutput, int, kCount, Round> destination_converter;
  
      return destination_converter(scaled_accumulator);
    }
  };
  
  #endif // Conditional guards to enable partial specialization for packed integers
  
  ////////////////////////////////////////////////////////////////////////////////
  
  /// Applies a linear combination operator to an array of elements then clamps
  /// the output before converting to the output element type.
  ///
  /// D = alpha * accumulator + beta * source + uniform
  ///
  /// Note: The below method only when problem_size_K <= 256 for signed int8 gemm
  /// or problem_size_K <= 128 for unsigned int8 gemm. The default approach is
  /// above.
  /// TODO: Add logic to fallback to the default approach
  template <
      /// Data type used to load and store< tensors
      typename ElementOutput_,
      /// Number of elements computed per operation
      int Count,
      /// Rounding mode
      FloatRoundStyle Round = FloatRoundStyle::round_to_nearest>
  class FastLinearCombinationClamp {
   public:
    using ElementOutput = ElementOutput_;
    using ElementAccumulator = int;
    using ElementCompute = float;
  
    static_assert(
        platform::is_same<ElementOutput, int32_t>::value ||
            platform::is_same<ElementOutput, uint32_t>::value ||
            platform::is_same<ElementOutput, int16_t>::value ||
            platform::is_same<ElementOutput, uint16_t>::value ||
            platform::is_same<ElementOutput, int8_t>::value ||
            platform::is_same<ElementOutput, uint8_t>::value ||
            platform::is_same<ElementOutput, cutlass::int4b_t>::value ||
            platform::is_same<ElementOutput, cutlass::uint4b_t>::value ||
            platform::is_same<ElementOutput, cutlass::uint1b_t>::value,
        "This elementwise op expects the output to be int.");
  
    static int const kCount = Count;
  
    using FragmentOutput = Array<ElementOutput, kCount>;
    using FragmentAccumulator = Array<ElementAccumulator, kCount>;
    using ComputeFragment = Array<ElementCompute, kCount>;
  
    static FloatRoundStyle const kRound = Round;
  
    /// Host-constructable parameters structure
    struct Params {
      /// scales accumulators
      ElementCompute alpha;
      /// scales source tensor
      ElementCompute beta;
      /// pointer to accumulator scalar - if not null, loads it from memory
      ElementCompute const *alpha_ptr;
      /// pointer to source scalar - if not null, loads it from memory
      ElementCompute const *beta_ptr;
  
      //
      // Methods
      //
  
      CUTLASS_HOST_DEVICE
      Params()
          : alpha(ElementCompute(1)),
            beta(ElementCompute(0)),
            alpha_ptr(nullptr),
            beta_ptr(nullptr) {}
  
      CUTLASS_HOST_DEVICE
      Params(ElementCompute alpha, ElementCompute beta)
          : alpha(alpha), beta(beta), alpha_ptr(nullptr), beta_ptr(nullptr) {}
  
      CUTLASS_HOST_DEVICE
      Params(ElementCompute const *alpha_ptr, ElementCompute const *beta_ptr)
          : alpha(0), beta(0), alpha_ptr(alpha_ptr), beta_ptr(beta_ptr) {}
    };
  
   private:
    //
    // Data members
    //
  
    ElementCompute alpha_;
    ElementCompute beta_;
  
   public:
    /// Constructs the function object, possibly loading from pointers in host
    /// memory
    CUTLASS_HOST_DEVICE
    FastLinearCombinationClamp(Params const &params) {
      alpha_ = (params.alpha_ptr ? *params.alpha_ptr : params.alpha);
      beta_ = (params.beta_ptr ? *params.beta_ptr : params.beta);
    }
  
    /// Returns true if source is needed
    CUTLASS_HOST_DEVICE
    bool is_source_needed() const { return beta_ != ElementCompute(0); }
  
    /// Functionally required for serial reduction in the epilogue
    CUTLASS_HOST_DEVICE
    void set_k_partition(int k_partition) {
      if (k_partition) {
        beta_ = ElementCompute(1);
      }
    }
    
    /// Computes linear scaling: D = alpha * accumulator + beta * source
    CUTLASS_HOST_DEVICE
    FragmentOutput operator()(FragmentAccumulator const &accumulator,
                              FragmentOutput const &source,
                              ElementCompute uniform = ElementCompute(0)) const {
      // Convert source to interal compute numeric type
      FastNumericArrayConverter<ElementCompute, ElementOutput, kCount, Round>
          source_converter;
      FastNumericArrayConverter<ElementCompute, ElementAccumulator, kCount, Round>
          accumulator_converter;
  
      ComputeFragment converted_source = source_converter(source);
      ComputeFragment converted_accumulator = accumulator_converter(accumulator);
  
      // Compute linear scaling in floating point
      ComputeFragment intermediate;
  
      multiplies<ComputeFragment> mul_add_source;
      multiply_add<ComputeFragment> mul_add_accumulator;
  
      minimum<ComputeFragment> min_accumulator;
      maximum<ComputeFragment> max_accumulator;
  
      // Float min-max
      intermediate =
          mul_add_source(beta_, converted_source);  // X =  beta * C + uniform
      intermediate = mul_add_accumulator(alpha_, converted_accumulator,
                                         intermediate);  // D = alpha * Accum + X
  
      /// Clamping constant value
      ElementCompute const kClamp =
          ElementCompute(1 << (sizeof_bits<ElementOutput>::value - 1));
  
      intermediate = max_accumulator(intermediate, -kClamp);
      intermediate = min_accumulator(intermediate, kClamp - ElementCompute(1));
  
      // Convert to destination numeric type
      FastNumericArrayConverter<ElementOutput, ElementCompute, kCount, Round>
          destination_converter;
  
      return destination_converter(intermediate);
    }
  
    /// Computes linear scaling: D = alpha * accumulator + beta * source
    CUTLASS_HOST_DEVICE
    FragmentOutput operator()(FragmentAccumulator const &accumulator) const {
  
      // Convert source to interal compute numeric type
      FastNumericArrayConverter<ElementCompute, ElementAccumulator, kCount, Round>
          accumulator_converter;
  
      ComputeFragment converted_accumulator = accumulator_converter(accumulator);
  
      // Compute linear scaling in floating point
      ComputeFragment intermediate;
  
      multiplies<ComputeFragment> mul_accumulator;
  
      minimum<ComputeFragment> min_accumulator;
      maximum<ComputeFragment> max_accumulator;
  
      // Float min-max
      intermediate = mul_accumulator(alpha_, converted_accumulator);
  
      /// Clamping constant value
      ElementCompute const kClamp =
          ElementCompute(1 << (sizeof_bits<ElementOutput>::value - 1));
  
      intermediate = max_accumulator(intermediate, -kClamp);
      intermediate = min_accumulator(intermediate, kClamp - ElementCompute(1));
  
      // Convert to destination numeric type
      FastNumericArrayConverter<ElementOutput, ElementCompute, kCount, Round>
          destination_converter;
  
      return destination_converter(intermediate);
    }
  };
  
  ////////////////////////////////////////////////////////////////////////////////
  
  } // namespace thread
  } // namespace epilogue
  } // namespace cutlass
  