leavelet/cutlass
1
0
Fork 0
Code Issues Pull Requests Actions 26 Packages Projects Releases Wiki Activity

CUTLASS 2.10 updates (#622)

Browse Source 
Co-authored-by: Aniket Shivam <ashivam@nvidia.com>
This commit is contained in:
ANIKET SHIVAM
2022-09-12 21:26:30 -04:00
committed by GitHub
co-authored by Aniket Shivam
parent beae168f90
commit e773429f7e
96 changed files with 8365 additions and 1667 deletions
Download Patch File Download Diff File Expand all files Collapse all files
include/cutlass/epilogue/thread/activation.h
+217
View File
@@ -59,6 +59,38 @@ struct Identity {
}
};
/////////////////////////////////////////////////////////////////////////////////////////////////
template <typename T>
struct LinearCombinationGenericParams {
T alpha; ///< scales accumulators
T beta; ///< scales source tensor
T const *alpha_ptr; ///< pointer to accumulator scalar - if not null, loads it from memory
T const *beta_ptr; ///< pointer to source scalar - if not null, loads it from memory
//
// Methods
//
CUTLASS_HOST_DEVICE
LinearCombinationGenericParams():
alpha(T(1)),
beta(T(0)),
alpha_ptr(nullptr),
beta_ptr(nullptr) { }
CUTLASS_HOST_DEVICE
LinearCombinationGenericParams(
T alpha,
T beta = T(0)
): alpha(alpha), beta(beta), alpha_ptr(nullptr), beta_ptr(nullptr) { }
CUTLASS_HOST_DEVICE
LinearCombinationGenericParams(
T const *alpha_ptr,
T const *beta_ptr = nullptr
): alpha(0), beta(0), alpha_ptr(alpha_ptr), beta_ptr(beta_ptr) { }
};
/////////////////////////////////////////////////////////////////////////////////////////////////
/// ReLu operator - propagates NaNs
@@ -79,6 +111,14 @@ struct ReLu {
return mx(value, T(0));
}
/// Host-constructable parameters structure
using Params = LinearCombinationGenericParams<T>;
CUTLASS_HOST_DEVICE
T operator()(T value, Params const &params_) const {
return this->operator()(value);
}
};
template <typename T, int N>
@@ -96,20 +136,74 @@ struct ReLu<Array<T, N>> {
maximum<Array<T, N> > mx;
return mx(frag, T(0));
}
/// Host-constructable parameters structure
using Params = LinearCombinationGenericParams<T>;
CUTLASS_HOST_DEVICE
Array<T, N> operator()(Array<T, N> const &frag, Params const &params_) const {
return this->operator()(frag);
}
};
// Leaky Relu operator
template <typename T>
struct LeakyReLU {
struct Params: LinearCombinationGenericParams<T> {
T leaky_alpha; ///< leaky_alpha
// Methods
using LinearCombinationGenericParams<T>::LinearCombinationGenericParams;
CUTLASS_HOST_DEVICE
Params():
LinearCombinationGenericParams<T>(),
leaky_alpha(T(1)) {}
CUTLASS_HOST_DEVICE
Params(
T alpha,
T beta,
T leaky_alpha = T(1)
): LinearCombinationGenericParams<T>(alpha, beta), leaky_alpha(leaky_alpha) {}
};
CUTLASS_HOST_DEVICE
T operator()(T const &value, T const & alpha_recip) const {
T res = value > T(0) ? value : value * alpha_recip;
return res;
}
CUTLASS_HOST_DEVICE
T operator()(T const &value, Params const &params_) const {
this->operator()(value, params_.leaky_alpha);
}
};
template <typename T, int N>
struct LeakyReLU<Array<T, N> > {
struct Params: LinearCombinationGenericParams<T> {
T leaky_alpha; ///< leaky_alpha
using LinearCombinationGenericParams<T>::LinearCombinationGenericParams;
// Methods
CUTLASS_HOST_DEVICE
Params():
LinearCombinationGenericParams<T>(),
leaky_alpha(T(1)) {}
CUTLASS_HOST_DEVICE
Params(
T alpha,
T beta,
T leaky_alpha = T(1)
): LinearCombinationGenericParams<T>(alpha, beta), leaky_alpha(leaky_alpha) {}
};
CUTLASS_HOST_DEVICE
Array<T, N> operator()(Array<T, N> const &rhs, T const & alpha_recip) const {
Array<T, N> y;
@@ -122,6 +216,11 @@ struct LeakyReLU<Array<T, N> > {
return y;
}
CUTLASS_HOST_DEVICE
Array<T, N> operator()(Array<T, N> const &rhs, Params const &params_) const {
return this->operator()(rhs, params_.leaky_alpha);
}
};
// Tanh operator
@@ -131,6 +230,13 @@ struct Tanh {
T operator()(T const &scalar) const {
return fast_tanh(scalar);
}
using Params = LinearCombinationGenericParams<T>;
CUTLASS_HOST_DEVICE
T operator()(T const &scalar, Params const &params_) const {
return this->operator()(scalar);
}
};
template <typename T, int N>
@@ -147,6 +253,13 @@ struct Tanh<Array<T, N> > {
return y;
}
using Params = LinearCombinationGenericParams<T>;
CUTLASS_HOST_DEVICE
Array<T, N> operator()(Array<T, N> const &rhs, Params const &params_) const {
return this->operator()(rhs);
}
};
template <int N>
@@ -159,6 +272,13 @@ struct Tanh<Array<half_t, N>> {
return tanh(z);
}
using Params = LinearCombinationGenericParams<T>;
CUTLASS_HOST_DEVICE
Array<T, N> operator()(Array<T, N> const &rhs, Params const &params_) const {
return this->operator()(rhs);
}
};
// Sigmoid operator
@@ -168,6 +288,13 @@ struct Sigmoid {
T operator()(T const &scalar) const {
return T(1) / (T(1) + fast_exp(-scalar));
}
using Params = LinearCombinationGenericParams<T>;
CUTLASS_HOST_DEVICE
T operator()(T const &scalar, Params const &params_) const {
return this->operator()(scalar);
}
};
template <typename T, int N>
@@ -184,6 +311,13 @@ struct Sigmoid<Array<T, N> > {
return y;
}
using Params = LinearCombinationGenericParams<T>;
CUTLASS_HOST_DEVICE
Array<T, N> operator()(Array<T, N> const &rhs, Params const &params_) const {
return this->operator()(rhs);
}
};
template <int N>
@@ -208,6 +342,12 @@ struct Sigmoid<Array<half_t, N>> {
fast_exp(neg(z))));
#endif
}
using Params = LinearCombinationGenericParams<T>;
Array<T, N> operator()(Array<T, N> const &z, Params const &params_) const {
return this->operator()(z);
}
};
// SiLu (swish) operator introduced by Elfwing et al. in the following paper
@@ -222,6 +362,13 @@ struct SiLu {
Sigmoid<T> sigmoid;
return scalar * sigmoid(scalar);
}
using Params = LinearCombinationGenericParams<T>;
CUTLASS_HOST_DEVICE
T operator()(T const &scalar, Params const &params_) const {
return this->operator()(scalar);
}
};
template <typename T, int N>
@@ -232,6 +379,13 @@ struct SiLu<Array<T, N>> {
multiplies<Array<T, N>> mul;
return mul(rhs, sigmoid_op(rhs));
}
using Params = LinearCombinationGenericParams<T>;
CUTLASS_HOST_DEVICE
Array<T, N> operator()(Array<T, N> const &rhs, Params const &params_) const {
return this->operator()(rhs);
}
};
// Hardswish operator introduced by Howard et al. in the following paper
@@ -248,6 +402,13 @@ struct HardSwish {
T relu6 = mn(mx(x + T(3), T(0)), T(6));
return x * relu6 / T(6);
}
using Params = LinearCombinationGenericParams<T>;
CUTLASS_HOST_DEVICE
T operator()(T const &x, Params const &params_) const {
return this->operator()(x);
}
};
template <>
@@ -261,6 +422,13 @@ struct HardSwish<float> {
T relu6 = mn(mx(x + T(3), T(0)), T(6));
return x * relu6 * 0.16666667f;
}
using Params = LinearCombinationGenericParams<T>;
CUTLASS_HOST_DEVICE
T operator()(T const &x, Params const &params_) const {
return this->operator()(x);
}
};
template <typename T, int N>
@@ -277,6 +445,13 @@ struct HardSwish<Array<T, N> > {
return y;
}
using Params = LinearCombinationGenericParams<T>;
CUTLASS_HOST_DEVICE
Array<T, N> operator()(Array<T, N> const &x, Params const &params_) const {
return this->operator()(x);
}
};
template <int N>
@@ -292,6 +467,13 @@ struct HardSwish<Array<half_t, N> > {
return mul(mul(mn(mx(add(rhs, T(3)), T(0)), T(6)), rhs), T(0.16666667f));
}
using Params = LinearCombinationGenericParams<T>;
CUTLASS_HOST_DEVICE
Array<T, N> operator()(Array<T, N> const &x, Params const &params_) const {
return this->operator()(x);
}
};
//
@@ -311,6 +493,13 @@ struct GELU {
return T(cutlass::constants::half<T>() * scalar *
(cutlass::constants::one<T>() + (T)erff((float)(scalar / cutlass::constants::root_two<T>()))));
}
using Params = LinearCombinationGenericParams<T>;
CUTLASS_HOST_DEVICE
T operator()(T const &scalar, Params const &params_) const {
return this->operator()(scalar);
}
};
template <>
@@ -320,6 +509,13 @@ struct GELU<float> {
return cutlass::constants::half<float>() * scalar *
(cutlass::constants::one<float>() + erff( scalar / cutlass::constants::root_two<float>() ));
}
using Params = LinearCombinationGenericParams<float>;
CUTLASS_HOST_DEVICE
float operator()(float const &scalar, Params const &params_) const {
return this->operator()(scalar);
}
};
template <>
@@ -329,6 +525,13 @@ struct GELU<double> {
return cutlass::constants::half<double>() * scalar *
(cutlass::constants::one<double>() + erf( scalar / cutlass::constants::root_two<double>() ));
}
using Params = LinearCombinationGenericParams<double>;
CUTLASS_HOST_DEVICE
double operator()(double const &scalar, Params const &params_) const {
return this->operator()(scalar);
}
};
template <typename T, int N>
@@ -345,6 +548,13 @@ struct GELU<Array<T, N> > {
return y;
}
using Params = LinearCombinationGenericParams<T>;
CUTLASS_HOST_DEVICE
Array<T, N> operator()(Array<T, N> const &rhs, Params const &params_) const {
return this->operator()(rhs);
}
};
// GELU operator implemented using the Taylor series approximation
@@ -360,6 +570,9 @@ struct GELU_taylor {
return T(cutlass::constants::half<T>() * z *
(cutlass::constants::one<T>() + fast_tanh(k0 * z * (cutlass::constants::one<T>() + k1 * z * z))));
}
using Params = LinearCombinationGenericParams<T>;
};
template <int N>
@@ -386,6 +599,8 @@ struct GELU_taylor<Array<half_t, N> > {
return y;
}
using Params = LinearCombinationGenericParams<half_t>;
};
template <typename T, int N>
@@ -403,6 +618,8 @@ struct GELU_taylor<Array<T, N> > {
return y;
}
using Params = LinearCombinationGenericParams<T>;
};
/// Computes backwards pass for GELU operator assuming d_t is the layer gradient and