@@ -180,6 +180,7 @@ struct GELU<Array<T, N> > {
|
||||
// GELU operator implemented using the Taylor series approximation
|
||||
template <typename T>
|
||||
struct GELU_taylor {
|
||||
static const bool kIsHeavy=true;
|
||||
CUTLASS_HOST_DEVICE
|
||||
T operator()(T const &z) const {
|
||||
|
||||
@@ -193,6 +194,7 @@ struct GELU_taylor {
|
||||
|
||||
template <typename T, int N>
|
||||
struct GELU_taylor<Array<T, N> > {
|
||||
static const bool kIsHeavy=true;
|
||||
CUTLASS_HOST_DEVICE
|
||||
Array<T, N> operator()(Array<T, N> const &rhs) const {
|
||||
Array<T, N> y;
|
||||
@@ -250,4 +252,3 @@ struct dGELU<Array<T, N> > {
|
||||
} // namespace cutlass
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
|
||||
@@ -65,6 +65,8 @@ public:
|
||||
|
||||
static FloatRoundStyle const kRound = Round;
|
||||
|
||||
static bool const kIsHeavy = false;
|
||||
|
||||
/// Host-constructable parameters structure
|
||||
struct Params {
|
||||
|
||||
|
||||
@@ -49,7 +49,9 @@ namespace thread {
|
||||
///
|
||||
template <
|
||||
typename ElementOutput_, ///< Data type used to load and store tensors
|
||||
int Count, ///< Number of elements computed per operation
|
||||
int Count, ///< Number of elements computed per operation.
|
||||
///< Usually it is 128/sizeof_bits<ElementOutput_>,
|
||||
///< but we use 64 or 32 sometimes when there are not enough data to store
|
||||
typename ElementAccumulator_ = ElementOutput_, ///< Accumulator data type
|
||||
typename ElementCompute_ = ElementOutput_, ///< Data type used to compute linear combination
|
||||
ScaleType::Kind Scale = ScaleType::Default, ///< Control Alpha and Beta scaling
|
||||
@@ -146,6 +148,8 @@ public:
|
||||
|
||||
if (Scale == ScaleType::OnlyAlphaScaling) return false;
|
||||
|
||||
if (Scale == ScaleType::Nothing) return false;
|
||||
|
||||
return beta_ != ElementCompute(0);
|
||||
}
|
||||
|
||||
@@ -167,11 +171,17 @@ public:
|
||||
NumericArrayConverter<ElementCompute, ElementOutput, kCount, Round> source_converter;
|
||||
NumericArrayConverter<ElementCompute, ElementAccumulator, kCount, Round> accumulator_converter;
|
||||
|
||||
ComputeFragment converted_source = source_converter(source);
|
||||
// Convert to destination numeric type
|
||||
NumericArrayConverter<ElementOutput, ElementCompute, kCount, Round> destination_converter;
|
||||
|
||||
ComputeFragment converted_accumulator = accumulator_converter(accumulator);
|
||||
|
||||
// Perform binary operations
|
||||
if (Scale == ScaleType::Nothing)
|
||||
return destination_converter(converted_accumulator);
|
||||
|
||||
ComputeFragment converted_source = source_converter(source);
|
||||
|
||||
// Perform binary operations
|
||||
ComputeFragment intermediate;
|
||||
|
||||
multiplies<ComputeFragment> mul_add_source;
|
||||
@@ -180,13 +190,10 @@ public:
|
||||
if (Scale == ScaleType::NoBetaScaling)
|
||||
intermediate = converted_source;
|
||||
else
|
||||
intermediate = mul_add_source(beta_, converted_source); // X = beta * C + uniform
|
||||
intermediate = mul_add_source(beta_, converted_source); // X = beta * C + uniform
|
||||
|
||||
intermediate = mul_add_accumulator(alpha_, converted_accumulator, intermediate); // D = alpha * Accum + X
|
||||
|
||||
// Convert to destination numeric type
|
||||
NumericArrayConverter<ElementOutput, ElementCompute, kCount, Round> destination_converter;
|
||||
|
||||
return destination_converter(intermediate);
|
||||
}
|
||||
|
||||
@@ -198,17 +205,20 @@ public:
|
||||
// Convert source to interal compute numeric type
|
||||
NumericArrayConverter<ElementCompute, ElementAccumulator, kCount, Round> accumulator_converter;
|
||||
|
||||
// Convert to destination numeric type
|
||||
NumericArrayConverter<ElementOutput, ElementCompute, kCount, Round> destination_converter;
|
||||
|
||||
ComputeFragment converted_accumulator = accumulator_converter(accumulator);
|
||||
|
||||
if (Scale == ScaleType::Nothing)
|
||||
return destination_converter(converted_accumulator);
|
||||
|
||||
// Perform binary operations
|
||||
ComputeFragment intermediate;
|
||||
multiplies<ComputeFragment> mul_accumulator;
|
||||
|
||||
intermediate = mul_accumulator(alpha_, converted_accumulator); // D = alpha * Accum
|
||||
|
||||
// Convert to destination numeric type
|
||||
NumericArrayConverter<ElementOutput, ElementCompute, kCount, Round> destination_converter;
|
||||
|
||||
return destination_converter(intermediate);
|
||||
}
|
||||
};
|
||||
|
||||
@@ -0,0 +1,251 @@
|
||||
/***************************************************************************************************
|
||||
* 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 Functor performing linear combination operations used by epilogues.
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cutlass/numeric_types.h"
|
||||
#include "cutlass/array.h"
|
||||
#include "cutlass/functional.h"
|
||||
#include "cutlass/numeric_conversion.h"
|
||||
|
||||
#include "cutlass/epilogue/thread/activation.h"
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace cutlass {
|
||||
namespace epilogue {
|
||||
namespace thread {
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// This base class is meant to define the concept required of the
|
||||
/// EpilogueWithBroadcast::OutputOp
|
||||
template <
|
||||
typename ElementC_,
|
||||
typename ElementAccumulator_,
|
||||
typename ElementCompute_,
|
||||
typename ElementZ_,
|
||||
typename ElementT_,
|
||||
int ElementsPerAccess,
|
||||
typename ElementwiseOp_ = Identity<ElementCompute_>,
|
||||
typename BinaryOp_ = plus<ElementCompute_>
|
||||
>
|
||||
class LinearCombinationBiasElementwise {
|
||||
public:
|
||||
|
||||
using ElementOutput = ElementC_;
|
||||
using ElementC = ElementC_;
|
||||
using ElementAccumulator = ElementAccumulator_;
|
||||
using ElementCompute = ElementCompute_;
|
||||
using ElementZ = ElementZ_;
|
||||
using ElementT = ElementT_;
|
||||
static int const kElementsPerAccess = ElementsPerAccess;
|
||||
static int const kCount = kElementsPerAccess;
|
||||
|
||||
using ElementwiseOp = ElementwiseOp_;
|
||||
using BinaryOp = BinaryOp_;
|
||||
|
||||
using FragmentAccumulator = Array<ElementAccumulator, kElementsPerAccess>;
|
||||
using FragmentCompute = Array<ElementCompute, kElementsPerAccess>;
|
||||
using FragmentC = Array<ElementOutput, kElementsPerAccess>;
|
||||
using FragmentZ = Array<ElementZ, kElementsPerAccess>;
|
||||
using FragmentT = Array<ElementT, kElementsPerAccess>;
|
||||
|
||||
using FragmentOutput = FragmentZ;
|
||||
|
||||
static bool const kIsHeavy = ElementwiseOp::kIsHeavy;
|
||||
|
||||
/// If true, the 'Z' tensor is stored
|
||||
static bool const kStoreZ = true;
|
||||
|
||||
/// If true, the 'T' tensor is stored
|
||||
static bool const kStoreT = true;
|
||||
|
||||
/// 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 alpha
|
||||
): alpha(alpha), beta(0), 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) {
|
||||
|
||||
}
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
Params(
|
||||
ElementCompute const *alpha_ptr
|
||||
): alpha(0), beta(0), alpha_ptr(alpha_ptr), beta_ptr(nullptr) {
|
||||
|
||||
}
|
||||
};
|
||||
|
||||
private:
|
||||
|
||||
//
|
||||
// Data members
|
||||
//
|
||||
|
||||
ElementCompute alpha_;
|
||||
ElementCompute beta_;
|
||||
bool skip_elementwise_;
|
||||
|
||||
public:
|
||||
|
||||
//
|
||||
// Methods
|
||||
//
|
||||
|
||||
/// Constructor from Params
|
||||
CUTLASS_HOST_DEVICE
|
||||
LinearCombinationBiasElementwise(Params const ¶ms) {
|
||||
|
||||
alpha_ = (params.alpha_ptr ? *params.alpha_ptr : params.alpha);
|
||||
beta_ = (params.beta_ptr ? *params.beta_ptr : params.beta);
|
||||
skip_elementwise_ = false;
|
||||
}
|
||||
|
||||
/// 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, int k_partition_count) {
|
||||
if (k_partition) {
|
||||
beta_ = ElementCompute(1);
|
||||
}
|
||||
|
||||
if (k_partition != k_partition_count - 1) {
|
||||
skip_elementwise_ = true;
|
||||
}
|
||||
}
|
||||
|
||||
/// Applies the operation when is_source_needed() is true
|
||||
CUTLASS_HOST_DEVICE
|
||||
void operator()(
|
||||
FragmentZ &frag_Z,
|
||||
FragmentT &frag_T,
|
||||
FragmentAccumulator const &AB,
|
||||
FragmentC const &frag_C,
|
||||
FragmentCompute const &V) const {
|
||||
|
||||
ElementwiseOp elementwise_op;
|
||||
BinaryOp binary_op;
|
||||
|
||||
FragmentCompute tmp_Accum = NumericArrayConverter<ElementCompute, ElementAccumulator, kElementsPerAccess>()(AB);
|
||||
FragmentCompute tmp_C = NumericArrayConverter<ElementCompute, ElementC, kElementsPerAccess>()(frag_C);
|
||||
FragmentCompute result_Z;
|
||||
FragmentCompute result_T;
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int i = 0; i < kElementsPerAccess; ++i) {
|
||||
ElementCompute z = binary_op(alpha_ * tmp_Accum[i] + beta_ * tmp_C[i], V[i]);
|
||||
result_Z[i] = z;
|
||||
result_T[i] = skip_elementwise_ ? z : elementwise_op(z);
|
||||
}
|
||||
|
||||
NumericArrayConverter<ElementZ, ElementCompute, kElementsPerAccess> convert_z;
|
||||
frag_Z = convert_z(result_Z);
|
||||
|
||||
NumericArrayConverter<ElementT, ElementCompute, kElementsPerAccess> convert_t;
|
||||
frag_T = convert_t(result_T);
|
||||
}
|
||||
|
||||
/// Applies the operation when is_source_needed() is false
|
||||
CUTLASS_HOST_DEVICE
|
||||
void operator()(
|
||||
FragmentZ &frag_Z,
|
||||
FragmentT &frag_T,
|
||||
FragmentAccumulator const &AB,
|
||||
FragmentCompute const &V) const {
|
||||
|
||||
ElementwiseOp elementwise_op;
|
||||
BinaryOp binary_op;
|
||||
|
||||
FragmentCompute tmp_Accum = NumericArrayConverter<ElementCompute, ElementAccumulator, kElementsPerAccess>()(AB);
|
||||
FragmentCompute result_Z;
|
||||
FragmentCompute result_T;
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int i = 0; i < kElementsPerAccess; ++i) {
|
||||
ElementCompute z = binary_op(alpha_ * tmp_Accum[i], V[i]);
|
||||
result_Z[i] = z;
|
||||
result_T[i] = skip_elementwise_ ? z : elementwise_op(z);
|
||||
}
|
||||
|
||||
NumericArrayConverter<ElementZ, ElementCompute, kElementsPerAccess> convert_z;
|
||||
frag_Z = convert_z(result_Z);
|
||||
|
||||
NumericArrayConverter<ElementT, ElementCompute, kElementsPerAccess> convert_t;
|
||||
frag_T = convert_t(result_T);
|
||||
}
|
||||
};
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace thread
|
||||
} // namespace epilogue
|
||||
} // namespace cutlass
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
@@ -28,6 +28,8 @@
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <cuda_fp16.h>
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cutlass/numeric_types.h"
|
||||
#include "cutlass/array.h"
|
||||
@@ -41,6 +43,146 @@ namespace cutlass {
|
||||
namespace epilogue {
|
||||
namespace thread {
|
||||
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace detail {
|
||||
|
||||
template <typename Element, int ElementsPerAccess>
|
||||
struct ArrayMaximum {
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
Array<Element, ElementsPerAccess> operator()(
|
||||
Array<Element, ElementsPerAccess> const &lhs,
|
||||
Array<Element, ElementsPerAccess> const &rhs) const {
|
||||
|
||||
Array<Element, ElementsPerAccess> result;
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int i = 0; i < ElementsPerAccess; ++i) {
|
||||
result[i] = fmax(lhs[i], rhs[i]);
|
||||
}
|
||||
|
||||
return result;
|
||||
}
|
||||
};
|
||||
|
||||
template <int ElementsPerAccess>
|
||||
struct ArrayMaximum<half_t, ElementsPerAccess> {
|
||||
|
||||
CUTLASS_DEVICE
|
||||
Array<half_t, ElementsPerAccess> operator()(
|
||||
Array<half_t, ElementsPerAccess> const &lhs,
|
||||
Array<half_t, ElementsPerAccess> const &rhs) const {
|
||||
|
||||
Array<half_t, ElementsPerAccess> result;
|
||||
|
||||
#if __CUDA_ARCH__ >= 800
|
||||
int const kVectorCount = ElementsPerAccess / 2;
|
||||
|
||||
|
||||
__half2 const *lhs_ptr = reinterpret_cast<__half2 const *>(lhs.raw_data());
|
||||
__half2 const *rhs_ptr = reinterpret_cast<__half2 const *>(rhs.raw_data());
|
||||
__half2 *res_ptr = reinterpret_cast<__half2 *>(result.raw_data());
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int i = 0; i < kVectorCount; ++i) {
|
||||
res_ptr[i] = __hmax2(lhs_ptr[i], rhs_ptr[i]);
|
||||
}
|
||||
|
||||
#else
|
||||
__half const *lhs_ptr = reinterpret_cast<__half const *>(lhs.raw_data());
|
||||
__half const *rhs_ptr = reinterpret_cast<__half const *>(rhs.raw_data());
|
||||
__half *res_ptr = reinterpret_cast<__half *>(result.raw_data());
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int i = 0; i < ElementsPerAccess; ++i) {
|
||||
res_ptr[i] = ((lhs_ptr[i] < rhs_ptr[i]) ? rhs_ptr[i] : lhs_ptr[i]);
|
||||
}
|
||||
|
||||
#endif
|
||||
|
||||
return result;
|
||||
}
|
||||
|
||||
CUTLASS_DEVICE
|
||||
Array<half_t, ElementsPerAccess> operator()(
|
||||
Array<half_t, ElementsPerAccess> const &lhs,
|
||||
half_t const &rhs) const {
|
||||
|
||||
Array<half_t, ElementsPerAccess> result;
|
||||
|
||||
#if __CUDA_ARCH__ >= 800
|
||||
int const kVectorCount = ElementsPerAccess / 2;
|
||||
|
||||
|
||||
__half rhs_raw = reinterpret_cast<__half const &>(rhs);
|
||||
__half2 rhs_pair = __half2half2(rhs_raw);
|
||||
|
||||
__half2 const *lhs_ptr = reinterpret_cast<__half2 const *>(lhs.raw_data());
|
||||
__half2 *res_ptr = reinterpret_cast<__half2 *>(result.raw_data());
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int i = 0; i < kVectorCount; ++i) {
|
||||
res_ptr[i] = __hmax2(lhs_ptr[i], rhs_pair);
|
||||
}
|
||||
|
||||
#else
|
||||
|
||||
__half const *lhs_ptr = reinterpret_cast<__half const *>(lhs.raw_data());
|
||||
__half const rhs_raw = reinterpret_cast<__half const &>(rhs);
|
||||
__half *res_ptr = reinterpret_cast<__half *>(result.raw_data());
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int i = 0; i < ElementsPerAccess; ++i) {
|
||||
res_ptr[i] = ((lhs_ptr[i] < rhs_raw) ? rhs_raw : lhs_ptr[i]);
|
||||
}
|
||||
|
||||
#endif
|
||||
|
||||
return result;
|
||||
}
|
||||
};
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <typename Element, int ElementsPerAccess>
|
||||
struct ReluConditional {
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
void operator()(
|
||||
bool conditional[],
|
||||
Array<Element, ElementsPerAccess> const &fragment,
|
||||
Element threshold) const {
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int i = 0; i < ElementsPerAccess; ++i) {
|
||||
conditional[i] = !(fragment[i] < threshold);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
template <int ElementsPerAccess>
|
||||
struct ReluConditional<half_t, ElementsPerAccess> {
|
||||
|
||||
CUTLASS_DEVICE
|
||||
void operator()(
|
||||
bool conditional[],
|
||||
Array<half_t, ElementsPerAccess> const &fragment,
|
||||
half_t threshold) const {
|
||||
|
||||
__half y = reinterpret_cast<__half const &>(threshold);
|
||||
__half const *x = reinterpret_cast<__half const *>(fragment.raw_data());
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int i = 0; i < ElementsPerAccess; ++i) {
|
||||
conditional[i] = !__hlt(x[i], y);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace detail
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// This is a partial specialization for fused Bias and ReLU. It supports the option of packing
|
||||
@@ -94,8 +236,11 @@ public:
|
||||
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
|
||||
ElementCompute threshold; ///< ReLu threshold
|
||||
ElementZ threshold; ///< ReLu threshold
|
||||
|
||||
//
|
||||
// Methods
|
||||
//
|
||||
//
|
||||
// Methods
|
||||
//
|
||||
@@ -112,16 +257,19 @@ public:
|
||||
Params(
|
||||
ElementCompute alpha,
|
||||
ElementCompute beta,
|
||||
ElementCompute threshold = ElementCompute()
|
||||
ElementCompute threshold_ = ElementCompute()
|
||||
):
|
||||
alpha(alpha), beta(beta), alpha_ptr(nullptr), beta_ptr(nullptr), threshold(threshold) {
|
||||
alpha(alpha), beta(beta), alpha_ptr(nullptr), beta_ptr(nullptr) {
|
||||
|
||||
NumericConverter<ElementZ, ElementCompute> convert_threshold;
|
||||
|
||||
threshold = convert_threshold(threshold_);
|
||||
}
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
Params(
|
||||
ElementCompute alpha
|
||||
): alpha(alpha), beta(0), alpha_ptr(nullptr), beta_ptr(nullptr), threshold(threshold) {
|
||||
): alpha(alpha), beta(0), alpha_ptr(nullptr), beta_ptr(nullptr), threshold(ElementZ()) {
|
||||
|
||||
}
|
||||
|
||||
@@ -129,17 +277,20 @@ public:
|
||||
Params(
|
||||
ElementCompute const *alpha_ptr,
|
||||
ElementCompute const *beta_ptr,
|
||||
ElementCompute threshold = ElementCompute()
|
||||
): alpha(0), beta(0), alpha_ptr(alpha_ptr), beta_ptr(beta_ptr), threshold(threshold) {
|
||||
ElementCompute threshold_ = ElementCompute()
|
||||
): alpha(0), beta(0), alpha_ptr(alpha_ptr), beta_ptr(beta_ptr) {
|
||||
|
||||
NumericConverter<ElementZ, ElementCompute> convert_threshold;
|
||||
|
||||
threshold = convert_threshold(threshold_);
|
||||
}
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
Params(
|
||||
ElementCompute const *alpha_ptr
|
||||
): alpha(0), beta(0), alpha_ptr(alpha_ptr), beta_ptr(nullptr), threshold(threshold) {
|
||||
|
||||
): alpha(0), beta(0), alpha_ptr(alpha_ptr), beta_ptr(nullptr), threshold(ElementZ()) {
|
||||
}
|
||||
|
||||
};
|
||||
|
||||
private:
|
||||
@@ -150,7 +301,7 @@ private:
|
||||
|
||||
ElementCompute alpha_;
|
||||
ElementCompute beta_;
|
||||
ElementCompute threshold_;
|
||||
ElementZ threshold_;
|
||||
|
||||
public:
|
||||
|
||||
@@ -179,6 +330,12 @@ public:
|
||||
if (k_partition) {
|
||||
beta_ = ElementCompute(1);
|
||||
}
|
||||
|
||||
if (k_partition != k_partition_count - 1) {
|
||||
// set to NaN to make ReLU no-op for all except last k partitions
|
||||
int64_t allones = -1;
|
||||
threshold_ = reinterpret_cast<ElementZ const &>(allones);
|
||||
}
|
||||
}
|
||||
|
||||
/// Applies the operation when is_source_needed() is true
|
||||
@@ -201,18 +358,27 @@ public:
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int i = 0; i < kElementsPerAccess; ++i) {
|
||||
ElementCompute z = binary_op(alpha_ * tmp_Accum[i] + beta_ * tmp_C[i], V[i]);
|
||||
|
||||
bool condition = !(z < threshold_);
|
||||
z = fmax(z, threshold_);
|
||||
ElementCompute z = alpha_ * tmp_Accum[i];
|
||||
z += beta_ * tmp_C[i];
|
||||
|
||||
z = binary_op(z, V[i]);
|
||||
result_Z[i] = z;
|
||||
conditions[i] = condition;
|
||||
}
|
||||
|
||||
NumericArrayConverter<ElementZ, ElementCompute, kElementsPerAccess> convert_z;
|
||||
frag_Z = convert_z(result_Z);
|
||||
|
||||
//
|
||||
// Compute condition
|
||||
//
|
||||
|
||||
detail::ReluConditional<ElementZ, kElementsPerAccess> relu_conditional;
|
||||
relu_conditional(conditions, frag_Z, threshold_);
|
||||
|
||||
detail::ArrayMaximum<ElementZ, kElementsPerAccess> maximum_op;
|
||||
frag_Z = maximum_op(frag_Z, threshold_);
|
||||
|
||||
if (kStoreT) {
|
||||
PackPredicates<kElementsPerAccess> pack_predicates;
|
||||
frag_T = pack_predicates(conditions);
|
||||
@@ -238,17 +404,29 @@ public:
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int i = 0; i < kElementsPerAccess; ++i) {
|
||||
ElementCompute z = binary_op(alpha_ * tmp_Accum[i], V[i]);
|
||||
|
||||
bool condition = !(z < threshold_);
|
||||
z = fmax(z, threshold_);
|
||||
|
||||
result_Z[i] = z;
|
||||
conditions[i] = condition;
|
||||
}
|
||||
|
||||
NumericArrayConverter<ElementZ, ElementCompute, kElementsPerAccess> convert_z;
|
||||
frag_Z = convert_z(result_Z);
|
||||
|
||||
//
|
||||
// Compute condition
|
||||
//
|
||||
|
||||
detail::ReluConditional<ElementZ, kElementsPerAccess> relu_conditional;
|
||||
relu_conditional(conditions, frag_Z, threshold_);
|
||||
|
||||
detail::ArrayMaximum<ElementZ, kElementsPerAccess> maximum_op;
|
||||
frag_Z = maximum_op(frag_Z, threshold_);
|
||||
|
||||
//
|
||||
// Compute conditions
|
||||
//
|
||||
|
||||
//
|
||||
// Store
|
||||
//
|
||||
if (kStoreT) {
|
||||
PackPredicates<kElementsPerAccess> pack_predicates;
|
||||
frag_T = pack_predicates(conditions);
|
||||
|
||||
@@ -43,6 +43,17 @@ namespace thread {
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace detail {
|
||||
|
||||
/// Single source of truth for whether to unroll for `LinearCombinationClamp()`
|
||||
constexpr bool LinearCombinationClampIsHeavy() {
|
||||
return false;
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Applies a linear combination operator to an array of elements then clamps the output before
|
||||
/// converting to the output element type.
|
||||
///
|
||||
@@ -51,6 +62,8 @@ namespace thread {
|
||||
template <
|
||||
typename ElementOutput_, ///< Data type used to load and store tensors
|
||||
int Count, ///< Number of elements computed per operation
|
||||
///< Usually it is 128/sizeof_bits<ElementOutput_>,
|
||||
///< but we use 64 or 32 sometimes when there are not enough data to store
|
||||
typename ElementAccumulator_ = ElementOutput_, ///< Accumulator data type
|
||||
typename ElementCompute_ = ElementOutput_, ///< Data type used to compute linear combination
|
||||
ScaleType::Kind Scale = ScaleType::Default, ///< Control Alpha and Beta scaling
|
||||
@@ -71,6 +84,8 @@ public:
|
||||
|
||||
static FloatRoundStyle const kRound = Round;
|
||||
|
||||
static bool const kIsHeavy = detail::LinearCombinationClampIsHeavy();
|
||||
|
||||
/// Host-constructable parameters structure
|
||||
struct Params {
|
||||
|
||||
@@ -282,6 +297,8 @@ public:
|
||||
|
||||
static FloatRoundStyle const kRound = Round;
|
||||
|
||||
static bool const kIsHeavy = detail::LinearCombinationClampIsHeavy();
|
||||
|
||||
/// Host-constructable parameters structure
|
||||
struct Params {
|
||||
|
||||
@@ -396,10 +413,9 @@ public:
|
||||
// Convert floats back to INT
|
||||
FragmentAccumulator scaled_accumulator;
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int i = 0; i < kCount; ++i) {
|
||||
scaled_accumulator[i] = __float2int_rn(intermediate[i]);
|
||||
}
|
||||
NumericArrayConverter<int, ElementCompute, kCount, Round> compute_converter;
|
||||
|
||||
scaled_accumulator = compute_converter(intermediate);
|
||||
|
||||
// Convert to destination numeric type
|
||||
NumericArrayConverter<ElementOutput, int, kCount, Round> destination_converter;
|
||||
@@ -427,10 +443,9 @@ public:
|
||||
// Convert floats back to INT
|
||||
FragmentAccumulator scaled_accumulator;
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int i = 0; i < kCount; ++i) {
|
||||
scaled_accumulator[i] = __float2int_rn(intermediate[i]);
|
||||
}
|
||||
NumericArrayConverter<int, ElementCompute, kCount, Round> compute_converter;
|
||||
|
||||
scaled_accumulator = compute_converter(intermediate);
|
||||
|
||||
// Convert to destination numeric type
|
||||
NumericArrayConverter<ElementOutput, int, kCount, Round> destination_converter;
|
||||
@@ -487,6 +502,8 @@ class FastLinearCombinationClamp {
|
||||
|
||||
static FloatRoundStyle const kRound = Round;
|
||||
|
||||
static bool const kIsHeavy = false;
|
||||
|
||||
/// Host-constructable parameters structure
|
||||
struct Params {
|
||||
/// scales accumulators
|
||||
|
||||
@@ -0,0 +1,244 @@
|
||||
/***************************************************************************************************
|
||||
* 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 Functor performing linear combination followed by dGelu operation
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <cutlass/half.h>
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cutlass/numeric_types.h"
|
||||
#include "cutlass/array.h"
|
||||
#include "cutlass/constants.h"
|
||||
#include "cutlass/fast_math.h"
|
||||
#include "cutlass/functional.h"
|
||||
#include "cutlass/numeric_conversion.h"
|
||||
#include "cutlass/epilogue/thread/activation.h"
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace cutlass {
|
||||
namespace epilogue {
|
||||
namespace thread {
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Applies a linear combination operator to an array of elements.
|
||||
///
|
||||
/// D = alpha * accumulator + beta * source + uniform
|
||||
///
|
||||
template <
|
||||
typename ElementCompute_, ///< Data type returned by this functor
|
||||
typename ElementAccumulator_, ///< Data type of accumulators
|
||||
typename ElementSource_, ///< Data type of source tensor
|
||||
typename ElementTensor_, ///< Data type of additional tensor
|
||||
int Count, ///< Number of elements computed per operation
|
||||
///< Usually it is 128/sizeof_bits<ElementOutput_>,
|
||||
///< but we use 64 or 32 sometimes when there are not enough data to store
|
||||
FloatRoundStyle Round = FloatRoundStyle::round_to_nearest
|
||||
>
|
||||
class LinearCombinationDGelu {
|
||||
public:
|
||||
|
||||
using ElementOutput = ElementSource_;
|
||||
using ElementCompute = ElementCompute_;
|
||||
using ElementAccumulator = ElementAccumulator_;
|
||||
using ElementSource = ElementSource_;
|
||||
using ElementTensor = ElementTensor_;
|
||||
|
||||
static bool const kIsHeavy = true;
|
||||
|
||||
static int const kCount = Count;
|
||||
|
||||
using FragmentCompute = Array<ElementCompute, kCount>;
|
||||
using FragmentAccumulator = Array<ElementAccumulator, kCount>;
|
||||
using FragmentSource = Array<ElementSource, kCount>;
|
||||
using FragmentTensor = Array<ElementTensor, kCount>;
|
||||
|
||||
static FloatRoundStyle const kRound = Round;
|
||||
|
||||
/// Host-constructable parameters structure
|
||||
struct Params {
|
||||
|
||||
ElementCompute alpha; ///< scales accumulators
|
||||
ElementCompute beta; ///< scales source tensor
|
||||
ElementCompute threshold; ///< minimum value that is output
|
||||
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)),
|
||||
threshold(ElementCompute(0)),
|
||||
alpha_ptr(nullptr),
|
||||
beta_ptr(nullptr) { }
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
Params(
|
||||
ElementCompute alpha,
|
||||
ElementCompute beta,
|
||||
ElementCompute threshold = ElementCompute(0)
|
||||
): alpha(alpha), beta(beta), threshold(threshold), alpha_ptr(nullptr), beta_ptr(nullptr) {
|
||||
|
||||
}
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
Params(
|
||||
ElementCompute const *alpha_ptr,
|
||||
ElementCompute const *beta_ptr,
|
||||
ElementCompute threshold = ElementCompute(0)
|
||||
): alpha(0), beta(0), threshold(threshold), alpha_ptr(alpha_ptr), beta_ptr(beta_ptr) {
|
||||
|
||||
}
|
||||
};
|
||||
|
||||
private:
|
||||
|
||||
//
|
||||
// Data members
|
||||
//
|
||||
|
||||
ElementCompute alpha_;
|
||||
ElementCompute beta_;
|
||||
ElementCompute threshold_;
|
||||
bool participates_in_reduction_;
|
||||
|
||||
public:
|
||||
|
||||
/// Constructs the function object, possibly loading from pointers in host memory
|
||||
CUTLASS_HOST_DEVICE
|
||||
LinearCombinationDGelu(Params const ¶ms) {
|
||||
|
||||
alpha_ = (params.alpha_ptr ? *params.alpha_ptr : params.alpha);
|
||||
beta_ = (params.beta_ptr ? *params.beta_ptr : params.beta);
|
||||
threshold_ = params.threshold;
|
||||
participates_in_reduction_ = true;
|
||||
}
|
||||
|
||||
/// Returns true if source is needed
|
||||
CUTLASS_HOST_DEVICE
|
||||
bool is_source_needed() const {
|
||||
return beta_ != ElementCompute(0);
|
||||
}
|
||||
|
||||
/// Returns true if the threadblock computes the reduction
|
||||
CUTLASS_HOST_DEVICE
|
||||
bool participates_in_reduction() const {
|
||||
return participates_in_reduction_;
|
||||
}
|
||||
|
||||
/// Functionally required for serial reduction in the epilogue
|
||||
CUTLASS_HOST_DEVICE
|
||||
void set_k_partition(int k_partition, int k_partition_count) {
|
||||
if (k_partition) {
|
||||
beta_ = ElementCompute(1);
|
||||
}
|
||||
|
||||
if (k_partition != k_partition_count - 1) {
|
||||
// set to NaN to make ReLU no-op for all except last k partitions
|
||||
int64_t allones = -1;
|
||||
threshold_ = reinterpret_cast<ElementCompute const &>(allones);
|
||||
// Avoid computing the reduction if this isn't the final Split-K slice
|
||||
participates_in_reduction_ = false;
|
||||
}
|
||||
}
|
||||
|
||||
/// Computes linear scaling: D = alpha * accumulator + beta * source
|
||||
CUTLASS_HOST_DEVICE
|
||||
FragmentCompute operator()(
|
||||
FragmentAccumulator const &accumulator,
|
||||
FragmentSource const &source,
|
||||
FragmentTensor const &tensor) const {
|
||||
|
||||
// Convert source to interal compute numeric type
|
||||
NumericArrayConverter<ElementCompute, ElementSource, kCount, Round> source_converter;
|
||||
NumericArrayConverter<ElementCompute, ElementAccumulator, kCount, Round> accumulator_converter;
|
||||
|
||||
FragmentCompute converted_source = source_converter(source);
|
||||
FragmentCompute converted_accumulator = accumulator_converter(accumulator);
|
||||
|
||||
// Perform binary operations
|
||||
FragmentCompute intermediate;
|
||||
|
||||
multiplies<FragmentCompute> mul_add_source;
|
||||
multiply_add<FragmentCompute> mul_add_accumulator;
|
||||
|
||||
intermediate = mul_add_source(beta_, converted_source); // X = beta * C + uniform
|
||||
intermediate = mul_add_accumulator(alpha_, converted_accumulator, intermediate); // D = alpha * Accum + X
|
||||
|
||||
dGELU<ElementCompute> gelu_op;
|
||||
|
||||
// dGelu
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int i = 0; i < kCount; ++i) {
|
||||
intermediate[i] = gelu_op(intermediate[i], ElementCompute(tensor[i]));
|
||||
}
|
||||
|
||||
return intermediate;
|
||||
}
|
||||
|
||||
/// Computes linear scaling: D = alpha * accumulator
|
||||
CUTLASS_HOST_DEVICE
|
||||
FragmentCompute operator()(
|
||||
FragmentAccumulator const &accumulator,
|
||||
FragmentTensor const &tensor) const {
|
||||
|
||||
// Convert source to interal compute numeric type
|
||||
NumericArrayConverter<ElementCompute, ElementAccumulator, kCount, Round> accumulator_converter;
|
||||
|
||||
FragmentCompute converted_accumulator = accumulator_converter(accumulator);
|
||||
|
||||
// Perform binary operations
|
||||
FragmentCompute intermediate;
|
||||
|
||||
multiplies<FragmentCompute> mul_accumulator;
|
||||
|
||||
intermediate = mul_accumulator(alpha_, converted_accumulator); // D = alpha * Accum
|
||||
|
||||
dGELU<ElementCompute> gelu_op;
|
||||
|
||||
// dGelu with conversion
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int i = 0; i < kCount; ++i) {
|
||||
intermediate[i] = gelu_op(intermediate[i], ElementCompute(tensor[i]));
|
||||
}
|
||||
|
||||
return intermediate;
|
||||
}
|
||||
};
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace thread
|
||||
} // namespace epilogue
|
||||
} // namespace cutlass
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
@@ -0,0 +1,446 @@
|
||||
/***************************************************************************************************
|
||||
* 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 Functor performing linear combination with a maximum operation used by epilogues.
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <cutlass/half.h>
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cutlass/numeric_types.h"
|
||||
#include "cutlass/array.h"
|
||||
#include "cutlass/functional.h"
|
||||
#include "cutlass/numeric_conversion.h"
|
||||
#include "cutlass/epilogue/thread/activation.h"
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace cutlass {
|
||||
namespace epilogue {
|
||||
namespace thread {
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Applies a linear combination operator to an array of elements.
|
||||
///
|
||||
/// D = alpha * accumulator + beta * source + uniform
|
||||
///
|
||||
template <
|
||||
typename ElementCompute_, ///< Data type returned by this functor
|
||||
typename ElementAccumulator_, ///< Data type of accumulators
|
||||
typename ElementSource_, ///< Data type of source tensor
|
||||
typename ElementTensor_, ///< Data type of additional tensor
|
||||
int Count, ///< Number of elements computed per operation
|
||||
///< Usually it is 128/sizeof_bits<ElementOutput_>,
|
||||
///< but we use 64 or 32 sometimes when there are not enough data to store
|
||||
FloatRoundStyle Round = FloatRoundStyle::round_to_nearest
|
||||
>
|
||||
class LinearCombinationDRelu {
|
||||
public:
|
||||
|
||||
using ElementOutput = ElementSource_;
|
||||
using ElementCompute = ElementCompute_;
|
||||
using ElementAccumulator = ElementAccumulator_;
|
||||
using ElementSource = ElementSource_;
|
||||
using ElementTensor = ElementTensor_;
|
||||
|
||||
static int const kCount = Count;
|
||||
|
||||
using FragmentCompute = Array<ElementCompute, kCount>;
|
||||
using FragmentAccumulator = Array<ElementAccumulator, kCount>;
|
||||
using FragmentSource = Array<ElementSource, kCount>;
|
||||
using FragmentTensor = Array<ElementTensor, kCount>;
|
||||
|
||||
static FloatRoundStyle const kRound = Round;
|
||||
|
||||
/// Host-constructable parameters structure
|
||||
struct Params {
|
||||
|
||||
ElementCompute alpha; ///< scales accumulators
|
||||
ElementCompute beta; ///< scales source tensor
|
||||
ElementCompute threshold; ///< minimum value that is output
|
||||
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)),
|
||||
threshold(ElementCompute(0)),
|
||||
alpha_ptr(nullptr),
|
||||
beta_ptr(nullptr) { }
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
Params(
|
||||
ElementCompute alpha,
|
||||
ElementCompute beta,
|
||||
ElementCompute threshold = ElementCompute(0)
|
||||
): alpha(alpha), beta(beta), threshold(threshold), alpha_ptr(nullptr), beta_ptr(nullptr) {
|
||||
|
||||
}
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
Params(
|
||||
ElementCompute const *alpha_ptr,
|
||||
ElementCompute const *beta_ptr,
|
||||
ElementCompute threshold = ElementCompute(0)
|
||||
): alpha(0), beta(0), threshold(threshold), alpha_ptr(alpha_ptr), beta_ptr(beta_ptr) {
|
||||
|
||||
}
|
||||
};
|
||||
|
||||
private:
|
||||
|
||||
//
|
||||
// Data members
|
||||
//
|
||||
|
||||
ElementCompute alpha_;
|
||||
ElementCompute beta_;
|
||||
ElementTensor threshold_;
|
||||
bool participates_in_reduction_;
|
||||
|
||||
public:
|
||||
|
||||
/// Constructs the function object, possibly loading from pointers in host memory
|
||||
CUTLASS_HOST_DEVICE
|
||||
LinearCombinationDRelu(Params const ¶ms) {
|
||||
|
||||
alpha_ = (params.alpha_ptr ? *params.alpha_ptr : params.alpha);
|
||||
beta_ = (params.beta_ptr ? *params.beta_ptr : params.beta);
|
||||
threshold_ = ElementTensor(params.threshold);
|
||||
participates_in_reduction_ = true;
|
||||
}
|
||||
|
||||
/// Returns true if source is needed
|
||||
CUTLASS_HOST_DEVICE
|
||||
bool is_source_needed() const {
|
||||
return beta_ != ElementCompute(0);
|
||||
}
|
||||
|
||||
/// Returns true if the threadblock computes the reduction
|
||||
CUTLASS_HOST_DEVICE
|
||||
bool participates_in_reduction() const {
|
||||
return participates_in_reduction_;
|
||||
}
|
||||
|
||||
/// Functionally required for serial reduction in the epilogue
|
||||
CUTLASS_DEVICE
|
||||
void set_k_partition(int k_partition, int k_partition_count) {
|
||||
if (k_partition) {
|
||||
beta_ = ElementCompute(1);
|
||||
}
|
||||
|
||||
if (k_partition != k_partition_count - 1) {
|
||||
// set to NaN to make ReLU no-op for all except last k partitions
|
||||
int64_t allones = -1;
|
||||
threshold_ = reinterpret_cast<ElementTensor const &>(allones);
|
||||
participates_in_reduction_ = false;
|
||||
}
|
||||
}
|
||||
|
||||
/// Computes linear scaling: D = alpha * accumulator + beta * source
|
||||
CUTLASS_HOST_DEVICE
|
||||
FragmentCompute operator()(
|
||||
FragmentAccumulator const &accumulator,
|
||||
FragmentSource const &source,
|
||||
FragmentTensor const &tensor) const {
|
||||
|
||||
// Convert source to interal compute numeric type
|
||||
NumericArrayConverter<ElementCompute, ElementSource, kCount, Round> source_converter;
|
||||
NumericArrayConverter<ElementCompute, ElementAccumulator, kCount, Round> accumulator_converter;
|
||||
|
||||
FragmentCompute converted_source = source_converter(source);
|
||||
FragmentCompute converted_accumulator = accumulator_converter(accumulator);
|
||||
|
||||
// Perform binary operations
|
||||
FragmentCompute intermediate;
|
||||
|
||||
multiplies<FragmentCompute> mul_add_source;
|
||||
multiply_add<FragmentCompute> mul_add_accumulator;
|
||||
|
||||
intermediate = mul_add_source(beta_, converted_source); // X = beta * C
|
||||
intermediate = mul_add_accumulator(alpha_, converted_accumulator, intermediate); // D = alpha * Accum + X
|
||||
|
||||
// dReLU = (cond ? dy : 0)
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int i = 0; i < kCount; ++i) {
|
||||
ElementTensor cond = tensor[i];
|
||||
if (cond <= threshold_) {
|
||||
intermediate[i] = ElementCompute();
|
||||
}
|
||||
}
|
||||
|
||||
return intermediate;
|
||||
}
|
||||
|
||||
/// Computes linear scaling: D = alpha * accumulator
|
||||
CUTLASS_HOST_DEVICE
|
||||
FragmentCompute operator()(
|
||||
FragmentAccumulator const &accumulator,
|
||||
FragmentTensor const &tensor) const {
|
||||
|
||||
// Convert source to interal compute numeric type
|
||||
NumericArrayConverter<ElementCompute, ElementAccumulator, kCount, Round> accumulator_converter;
|
||||
|
||||
FragmentCompute converted_accumulator = accumulator_converter(accumulator);
|
||||
|
||||
// Perform binary operations
|
||||
FragmentCompute intermediate;
|
||||
|
||||
multiplies<FragmentCompute> mul_accumulator;
|
||||
|
||||
intermediate = mul_accumulator(alpha_, converted_accumulator); // D = alpha * Accum
|
||||
|
||||
// dReLU = (cond ? dy : 0)
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int i = 0; i < kCount; ++i) {
|
||||
ElementTensor cond = tensor[i];
|
||||
if (cond <= threshold_) {
|
||||
intermediate[i] = ElementCompute();
|
||||
}
|
||||
}
|
||||
|
||||
return intermediate;
|
||||
}
|
||||
};
|
||||
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Applies a linear combination operator to an array of elements.
|
||||
///
|
||||
/// D = alpha * accumulator + beta * source + uniform
|
||||
///
|
||||
template <
|
||||
typename ElementCompute_, ///< Data type returned by this functor
|
||||
typename ElementAccumulator_, ///< Data type of accumulators
|
||||
typename ElementSource_, ///< Data type of source tensor
|
||||
int Count, ///< Number of elements computed per operation
|
||||
FloatRoundStyle Round = FloatRoundStyle::round_to_nearest
|
||||
>
|
||||
class LinearCombinationDReluConditionalBits {
|
||||
public:
|
||||
|
||||
using ElementOutput = ElementSource_;
|
||||
using ElementCompute = ElementCompute_;
|
||||
using ElementAccumulator = ElementAccumulator_;
|
||||
using ElementSource = ElementSource_;
|
||||
using ElementTensor = uint1b_t;
|
||||
|
||||
static bool const kIsHeavy = false;
|
||||
|
||||
static int const kCount = Count;
|
||||
|
||||
using FragmentCompute = Array<ElementCompute, kCount>;
|
||||
using FragmentAccumulator = Array<ElementAccumulator, kCount>;
|
||||
using FragmentSource = Array<ElementSource, kCount>;
|
||||
using FragmentTensor = Array<ElementTensor, 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_;
|
||||
FragmentTensor predicate_mask_;
|
||||
bool participates_in_reduction_;
|
||||
|
||||
public:
|
||||
|
||||
/// Constructs the function object, possibly loading from pointers in host memory
|
||||
CUTLASS_HOST_DEVICE
|
||||
LinearCombinationDReluConditionalBits(Params const ¶ms) {
|
||||
|
||||
alpha_ = (params.alpha_ptr ? *params.alpha_ptr : params.alpha);
|
||||
beta_ = (params.beta_ptr ? *params.beta_ptr : params.beta);
|
||||
participates_in_reduction_ = true;
|
||||
predicate_mask_.clear();
|
||||
}
|
||||
|
||||
/// Returns true if source is needed
|
||||
CUTLASS_HOST_DEVICE
|
||||
bool is_source_needed() const {
|
||||
return beta_ != ElementCompute(0);
|
||||
}
|
||||
|
||||
/// Returns true if the threadblock computes the reduction
|
||||
CUTLASS_HOST_DEVICE
|
||||
bool participates_in_reduction() const {
|
||||
return participates_in_reduction_;
|
||||
}
|
||||
|
||||
/// Functionally required for serial reduction in the epilogue
|
||||
CUTLASS_HOST_DEVICE
|
||||
void set_k_partition(int k_partition, int k_partition_count) {
|
||||
predicate_mask_.clear();
|
||||
|
||||
if (k_partition) {
|
||||
beta_ = ElementCompute(1);
|
||||
}
|
||||
|
||||
if (k_partition != k_partition_count - 1) {
|
||||
// Avoid computing the reduction if this isn't the final Split-K slice
|
||||
participates_in_reduction_ = false;
|
||||
|
||||
bit_not<FragmentTensor> not_op;
|
||||
predicate_mask_ = not_op(predicate_mask_);
|
||||
}
|
||||
}
|
||||
|
||||
/// Computes linear scaling: D = alpha * accumulator + beta * source
|
||||
CUTLASS_DEVICE
|
||||
FragmentCompute operator()(
|
||||
FragmentAccumulator const &accumulator,
|
||||
FragmentSource const &source,
|
||||
FragmentTensor const &tensor) const {
|
||||
|
||||
// Convert source to interal compute numeric type
|
||||
NumericArrayConverter<ElementCompute, ElementSource, kCount, Round> source_converter;
|
||||
NumericArrayConverter<ElementCompute, ElementAccumulator, kCount, Round> accumulator_converter;
|
||||
|
||||
FragmentCompute converted_source = source_converter(source);
|
||||
FragmentCompute converted_accumulator = accumulator_converter(accumulator);
|
||||
|
||||
// Perform binary operations
|
||||
FragmentCompute intermediate;
|
||||
|
||||
multiplies<FragmentCompute> mul_add_source;
|
||||
multiply_add<FragmentCompute> mul_add_accumulator;
|
||||
|
||||
intermediate = mul_add_source(beta_, converted_source); // X = beta * C + uniform
|
||||
intermediate = mul_add_accumulator(alpha_, converted_accumulator, intermediate); // D = alpha * Accum + X
|
||||
|
||||
bit_or<FragmentTensor> or_op;
|
||||
|
||||
FragmentTensor predicates = or_op(tensor, predicate_mask_);
|
||||
|
||||
// Obtain from packed bits
|
||||
bool conditions[kCount];
|
||||
UnpackPredicates<kCount> unpack_predicates;
|
||||
|
||||
unpack_predicates(conditions, predicates);
|
||||
|
||||
// dReLU = (cond ? dy : 0)
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int i = 0; i < kCount; ++i) {
|
||||
if (!conditions[i]) {
|
||||
intermediate[i] = ElementCompute();
|
||||
}
|
||||
}
|
||||
|
||||
return intermediate;
|
||||
}
|
||||
|
||||
/// Computes linear scaling: D = alpha * accumulator
|
||||
CUTLASS_HOST_DEVICE
|
||||
FragmentCompute operator()(
|
||||
FragmentAccumulator const &accumulator,
|
||||
FragmentTensor const &tensor) const {
|
||||
|
||||
// Convert source to interal compute numeric type
|
||||
NumericArrayConverter<ElementCompute, ElementAccumulator, kCount, Round> accumulator_converter;
|
||||
|
||||
FragmentCompute converted_accumulator = accumulator_converter(accumulator);
|
||||
|
||||
// Perform binary operations
|
||||
FragmentCompute intermediate;
|
||||
|
||||
multiplies<FragmentCompute> mul_accumulator;
|
||||
|
||||
intermediate = mul_accumulator(alpha_, converted_accumulator); // D = alpha * Accum
|
||||
|
||||
bit_or<FragmentTensor> or_op;
|
||||
|
||||
FragmentTensor predicates = or_op(tensor, predicate_mask_);
|
||||
|
||||
// Obtain from packed bits
|
||||
bool conditions[kCount];
|
||||
UnpackPredicates<kCount> unpack_predicates;
|
||||
|
||||
unpack_predicates(conditions, predicates);
|
||||
|
||||
// dReLU = (cond ? dy : 0)
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int i = 0; i < kCount; ++i) {
|
||||
if (!conditions[i]) {
|
||||
intermediate[i] = ElementCompute();
|
||||
}
|
||||
}
|
||||
|
||||
return intermediate;
|
||||
}
|
||||
};
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace thread
|
||||
} // namespace epilogue
|
||||
} // namespace cutlass
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
@@ -51,6 +51,8 @@ namespace thread {
|
||||
template <
|
||||
typename ElementOutput_, ///< Data type used to load and store tensors
|
||||
int Count, ///< Number of elements computed per operation
|
||||
///< Usually it is 128/sizeof_bits<ElementOutput_>,
|
||||
///< but we use 64 or 32 sometimes when there are not enough data to store
|
||||
typename ElementAccumulator_ = ElementOutput_, ///< Accumulator data type
|
||||
typename ElementCompute_ = ElementOutput_, ///< Data type used to compute linear combination
|
||||
FloatRoundStyle Round = FloatRoundStyle::round_to_nearest
|
||||
@@ -62,6 +64,8 @@ public:
|
||||
using ElementAccumulator = ElementAccumulator_;
|
||||
using ElementCompute = ElementCompute_;
|
||||
|
||||
static bool const kIsHeavy = true;
|
||||
|
||||
static int const kCount = Count;
|
||||
|
||||
using FragmentOutput = Array<ElementOutput, kCount>;
|
||||
@@ -134,10 +138,11 @@ public:
|
||||
/// Functionally required for serial reduction in the epilogue
|
||||
CUTLASS_HOST_DEVICE
|
||||
void set_k_partition(int k_partition, int k_partition_count) {
|
||||
CUTLASS_UNUSED(k_partition_count);
|
||||
if (k_partition) {
|
||||
beta_ = ElementCompute(1);
|
||||
}
|
||||
|
||||
CUTLASS_UNUSED(k_partition_count);
|
||||
}
|
||||
|
||||
/// Computes: D = gelu( alpha * accumulator + beta * source )
|
||||
|
||||
@@ -52,6 +52,8 @@ namespace thread {
|
||||
template <
|
||||
typename ElementOutput_, ///< Data type used to load and store tensors
|
||||
int Count, ///< Number of elements computed per operation
|
||||
///< Usually it is 128/sizeof_bits<ElementOutput_>,
|
||||
///< but we use 64 or 32 sometimes when there are not enough data to store
|
||||
typename ElementAccumulator_ = ElementOutput_, ///< Accumulator data type
|
||||
typename ElementCompute_ = ElementOutput_, ///< Data type used to compute linear combination
|
||||
FloatRoundStyle Round = FloatRoundStyle::round_to_nearest
|
||||
|
||||
@@ -45,6 +45,17 @@ namespace thread {
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace detail {
|
||||
|
||||
/// Single source of truth for whether to unroll for `LinearCombinationClamp()`
|
||||
constexpr bool LinearCombinationReluIsHeavy() {
|
||||
return false;
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Applies a linear combination operator to an array of elements.
|
||||
///
|
||||
/// D = alpha * accumulator + beta * source + uniform
|
||||
@@ -52,6 +63,8 @@ namespace thread {
|
||||
template <
|
||||
typename ElementOutput_, ///< Data type used to load and store tensors
|
||||
int Count, ///< Number of elements computed per operation
|
||||
///< Usually it is 128/sizeof_bits<ElementOutput_>,
|
||||
///< but we use 64 or 32 sometimes when there are not enough data to store
|
||||
typename ElementAccumulator_ = ElementOutput_, ///< Accumulator data type
|
||||
typename ElementCompute_ = ElementOutput_, ///< Data type used to compute linear combination
|
||||
ScaleType::Kind Scale = ScaleType::Default, ///< Control Alpha and Beta scaling
|
||||
@@ -72,6 +85,8 @@ public:
|
||||
|
||||
static FloatRoundStyle const kRound = Round;
|
||||
|
||||
static bool const kIsHeavy = detail::LinearCombinationReluIsHeavy();
|
||||
|
||||
/// Host-constructable parameters structure
|
||||
struct Params {
|
||||
|
||||
@@ -244,6 +259,8 @@ public:
|
||||
using ElementAccumulator = int;
|
||||
using ElementCompute = float;
|
||||
|
||||
static bool const kIsHeavy = detail::LinearCombinationReluIsHeavy();
|
||||
|
||||
static int const kCount = Count;
|
||||
|
||||
using FragmentOutput = Array<ElementOutput, kCount>;
|
||||
@@ -357,10 +374,10 @@ public:
|
||||
ReLu<ComputeFragment> relu;
|
||||
|
||||
if (Scale == ScaleType::NoBetaScaling)
|
||||
intermediate = converted_source;
|
||||
intermediate = converted_source;
|
||||
else
|
||||
intermediate = mul_add_source(beta_, converted_source); // X = beta * C + uniform
|
||||
|
||||
intermediate = mul_add_source(beta_, converted_source); // X = beta * C + uniform
|
||||
|
||||
intermediate = mul_add_accumulator(alpha_, converted_accumulator, intermediate); // D = alpha * Accum + X
|
||||
|
||||
// Compute threshold optionally
|
||||
@@ -378,10 +395,9 @@ public:
|
||||
// Convert floats back to INT
|
||||
FragmentAccumulator scaled_accumulator;
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int i = 0; i < kCount; ++i) {
|
||||
scaled_accumulator[i] = __float2int_rn(intermediate[i]);
|
||||
}
|
||||
NumericArrayConverter<int, ElementCompute, kCount, Round> compute_converter;
|
||||
|
||||
scaled_accumulator = compute_converter(intermediate);
|
||||
|
||||
// Convert to destination numeric type
|
||||
NumericArrayConverter<ElementOutput, int, kCount, Round>
|
||||
@@ -416,14 +432,6 @@ public:
|
||||
// Compute threshold optionally
|
||||
intermediate = relu(threshold_, intermediate);
|
||||
|
||||
// Convert floats back to INT
|
||||
FragmentAccumulator scaled_accumulator;
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int i = 0; i < kCount; ++i) {
|
||||
scaled_accumulator[i] = __float2int_rn(intermediate[i]);
|
||||
}
|
||||
|
||||
if (platform::is_same<ElementOutput, int32_t>::value ||
|
||||
platform::is_same<ElementOutput, uint32_t>::value ||
|
||||
platform::is_same<ElementOutput, int16_t>::value ||
|
||||
@@ -436,10 +444,9 @@ public:
|
||||
// Convert floats back to INT
|
||||
FragmentAccumulator scaled_accumulator;
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int i = 0; i < kCount; ++i) {
|
||||
scaled_accumulator[i] = __float2int_rn(intermediate[i]);
|
||||
}
|
||||
NumericArrayConverter<int, ElementCompute, kCount, Round> compute_converter;
|
||||
|
||||
scaled_accumulator = compute_converter(intermediate);
|
||||
|
||||
// Convert to destination numeric type
|
||||
NumericArrayConverter<ElementOutput, int, kCount, Round>
|
||||
|
||||
@@ -51,6 +51,8 @@ namespace thread {
|
||||
template <
|
||||
typename ElementOutput_, ///< Data type used to load and store tensors
|
||||
int Count, ///< Number of elements computed per operation
|
||||
///< Usually it is 128/sizeof_bits<ElementOutput_>,
|
||||
///< but we use 64 or 32 sometimes when there are not enough data to store
|
||||
typename ElementAccumulator_ = ElementOutput_, ///< Accumulator data type
|
||||
typename ElementCompute_ = ElementOutput_, ///< Data type used to compute linear combination
|
||||
FloatRoundStyle Round = FloatRoundStyle::round_to_nearest
|
||||
|
||||
@@ -0,0 +1,228 @@
|
||||
/***************************************************************************************************
|
||||
* 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 Functor performing linear combination with elementwise
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <cutlass/half.h>
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cutlass/numeric_types.h"
|
||||
#include "cutlass/array.h"
|
||||
#include "cutlass/constants.h"
|
||||
#include "cutlass/fast_math.h"
|
||||
#include "cutlass/functional.h"
|
||||
#include "cutlass/numeric_conversion.h"
|
||||
#include "cutlass/epilogue/thread/activation.h"
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace cutlass {
|
||||
namespace epilogue {
|
||||
namespace thread {
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Applies a linear combination operator to an array of elements.
|
||||
///
|
||||
/// D = alpha * accumulator + beta * source + uniform
|
||||
///
|
||||
template <
|
||||
typename ElementCompute_, ///< Data type returned by this functor
|
||||
typename ElementAccumulator_, ///< Data type of accumulators
|
||||
typename ElementSource_, ///< Data type of source tensor
|
||||
typename ElementTensor_, ///< Data type of additional tensor
|
||||
int Count, ///< Number of elements computed per operation
|
||||
///< Usually it is 128/sizeof_bits<ElementOutput_>,
|
||||
///< but we use 64 or 32 sometimes when there are not enough data to store
|
||||
FloatRoundStyle Round = FloatRoundStyle::round_to_nearest
|
||||
>
|
||||
class LinearCombinationWithElementwise {
|
||||
public:
|
||||
|
||||
using ElementOutput = ElementSource_;
|
||||
using ElementCompute = ElementCompute_;
|
||||
using ElementAccumulator = ElementAccumulator_;
|
||||
using ElementSource = ElementSource_;
|
||||
using ElementTensor = ElementTensor_;
|
||||
|
||||
static bool const kIsHeavy = true;
|
||||
|
||||
static int const kCount = Count;
|
||||
|
||||
using FragmentCompute = Array<ElementCompute, kCount>;
|
||||
using FragmentAccumulator = Array<ElementAccumulator, kCount>;
|
||||
using FragmentSource = Array<ElementSource, kCount>;
|
||||
using FragmentTensor = Array<ElementTensor, kCount>;
|
||||
|
||||
static FloatRoundStyle const kRound = Round;
|
||||
|
||||
/// Host-constructable parameters structure
|
||||
struct Params {
|
||||
|
||||
ElementCompute alpha; ///< scales accumulators
|
||||
ElementCompute beta; ///< scales source tensor
|
||||
ElementCompute threshold; ///< minimum value that is output
|
||||
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)),
|
||||
threshold(ElementCompute(0)),
|
||||
alpha_ptr(nullptr),
|
||||
beta_ptr(nullptr) { }
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
Params(
|
||||
ElementCompute alpha,
|
||||
ElementCompute beta,
|
||||
ElementCompute threshold = ElementCompute(0)
|
||||
): alpha(alpha), beta(beta), threshold(threshold), alpha_ptr(nullptr), beta_ptr(nullptr) {
|
||||
|
||||
}
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
Params(
|
||||
ElementCompute const *alpha_ptr,
|
||||
ElementCompute const *beta_ptr,
|
||||
ElementCompute threshold = ElementCompute(0)
|
||||
): alpha(0), beta(0), threshold(threshold), alpha_ptr(alpha_ptr), beta_ptr(beta_ptr) {
|
||||
|
||||
}
|
||||
};
|
||||
|
||||
private:
|
||||
|
||||
//
|
||||
// Data members
|
||||
//
|
||||
|
||||
ElementCompute alpha_;
|
||||
ElementCompute beta_;
|
||||
ElementCompute threshold_;
|
||||
bool participates_in_reduction_;
|
||||
|
||||
public:
|
||||
|
||||
/// Constructs the function object, possibly loading from pointers in host memory
|
||||
CUTLASS_HOST_DEVICE
|
||||
LinearCombinationWithElementwise(Params const ¶ms) {
|
||||
|
||||
alpha_ = (params.alpha_ptr ? *params.alpha_ptr : params.alpha);
|
||||
beta_ = (params.beta_ptr ? *params.beta_ptr : params.beta);
|
||||
threshold_ = params.threshold;
|
||||
participates_in_reduction_ = true;
|
||||
}
|
||||
|
||||
/// Returns true if source is needed
|
||||
CUTLASS_HOST_DEVICE
|
||||
bool is_source_needed() const {
|
||||
return beta_ != ElementCompute(0);
|
||||
}
|
||||
|
||||
/// Returns true if the threadblock computes the reduction
|
||||
CUTLASS_HOST_DEVICE
|
||||
bool participates_in_reduction() const {
|
||||
return participates_in_reduction_;
|
||||
}
|
||||
|
||||
/// Functionally required for serial reduction in the epilogue
|
||||
CUTLASS_HOST_DEVICE
|
||||
void set_k_partition(int k_partition, int k_partition_count) {
|
||||
if (k_partition) {
|
||||
beta_ = ElementCompute(1);
|
||||
}
|
||||
|
||||
if (k_partition != k_partition_count - 1) {
|
||||
// set to NaN to make ReLU no-op for all except last k partitions
|
||||
int64_t allones = -1;
|
||||
threshold_ = reinterpret_cast<ElementCompute const &>(allones);
|
||||
// Avoid computing the reduction if this isn't the final Split-K slice
|
||||
participates_in_reduction_ = false;
|
||||
}
|
||||
}
|
||||
|
||||
/// Computes linear scaling: D = alpha * accumulator + beta * source
|
||||
CUTLASS_HOST_DEVICE
|
||||
FragmentCompute operator()(
|
||||
FragmentAccumulator const &accumulator,
|
||||
FragmentSource const &source,
|
||||
FragmentTensor const &tensor) const {
|
||||
|
||||
// Convert source to interal compute numeric type
|
||||
NumericArrayConverter<ElementCompute, ElementSource, kCount, Round> source_converter;
|
||||
NumericArrayConverter<ElementCompute, ElementAccumulator, kCount, Round> accumulator_converter;
|
||||
|
||||
FragmentCompute converted_source = source_converter(source);
|
||||
FragmentCompute converted_accumulator = accumulator_converter(accumulator);
|
||||
|
||||
// Perform binary operations
|
||||
FragmentCompute intermediate;
|
||||
|
||||
multiplies<FragmentCompute> mul_add_source;
|
||||
multiply_add<FragmentCompute> mul_add_accumulator;
|
||||
|
||||
intermediate = mul_add_source(beta_, converted_source); // X = beta * C + uniform
|
||||
intermediate = mul_add_accumulator(alpha_, converted_accumulator, intermediate); // D = alpha * Accum + X
|
||||
|
||||
return intermediate;
|
||||
}
|
||||
|
||||
/// Computes linear scaling: D = alpha * accumulator
|
||||
CUTLASS_HOST_DEVICE
|
||||
FragmentCompute operator()(
|
||||
FragmentAccumulator const &accumulator,
|
||||
FragmentTensor const &tensor) const {
|
||||
|
||||
// Convert source to interal compute numeric type
|
||||
NumericArrayConverter<ElementCompute, ElementAccumulator, kCount, Round> accumulator_converter;
|
||||
|
||||
FragmentCompute converted_accumulator = accumulator_converter(accumulator);
|
||||
|
||||
// Perform binary operations
|
||||
FragmentCompute intermediate;
|
||||
|
||||
multiplies<FragmentCompute> mul_accumulator;
|
||||
|
||||
intermediate = mul_accumulator(alpha_, converted_accumulator); // D = alpha * Accum
|
||||
|
||||
return intermediate;
|
||||
}
|
||||
};
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace thread
|
||||
} // namespace epilogue
|
||||
} // namespace cutlass
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
@@ -41,9 +41,10 @@ namespace thread {
|
||||
/// Specifies internal data type for computation
|
||||
struct ScaleType {
|
||||
enum Kind {
|
||||
Default, // alpha x C + beta x D
|
||||
NoBetaScaling, // alpha x C + D
|
||||
OnlyAlphaScaling // alpha x C
|
||||
Default, // alpha x C + beta x D
|
||||
NoBetaScaling, // alpha x C + D
|
||||
OnlyAlphaScaling, // alpha x C
|
||||
Nothing // C
|
||||
};
|
||||
};
|
||||
|
||||
|
||||
Reference in New Issue
Block a user