Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

set kIsHeavy member variables #1012

Merged
merged 4 commits into from
Oct 4, 2023
Merged

set kIsHeavy member variables #1012

Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
aa4db41
set kIsHeavy member variables
FabianSchuetze Jul 23, 2023
190af21
correct kIsHeavy value for Tanh
FabianSchuetze Jul 23, 2023
f71f653
set kIsHeavy=false for HardSwish
FabianSchuetze Aug 30, 2023
64e1b9b
Merge branch 'main' of https://github.com/NVIDIA/cutlass into add_kis…
hwu36 Oct 4, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
95 changes: 89 additions & 6 deletions include/cutlass/epilogue/thread/activation.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,7 @@

#include "cutlass/cutlass.h"
#include "cutlass/numeric_types.h"
#include "cutlass/numeric_conversion.h"
#include "cutlass/constants.h"
#include "cutlass/complex.h"
#include "cutlass/array.h"
Expand All @@ -54,7 +55,7 @@ namespace thread {
// Identity operator
template <typename T>
struct Identity {
static const bool kIsHeavy=false;
static const bool kIsHeavy = false;

CUTLASS_HOST_DEVICE
T operator()(T value) const {
Expand Down Expand Up @@ -128,7 +129,8 @@ struct Scale<Activation<T>> {
/// Always put threshold in the right hand side of max to propagate NaN.
template <typename T>
struct ReLu {
static const bool kIsHeavy=false;
static const bool kIsHeavy = false;

CUTLASS_HOST_DEVICE
T operator()(T const & threshold, T value) const {
maximum<T> mx;
Expand All @@ -149,7 +151,8 @@ using ReLU = ReLu<T>;

template <typename T, int N>
struct ReLu<Array<T, N>> {
static const bool kIsHeavy=false;
static const bool kIsHeavy = false;

CUTLASS_HOST_DEVICE
Array<T, N> operator()(T const & threshold, Array<T, N> const &frag) const {
maximum<Array<T, N>> mx;
Expand Down Expand Up @@ -207,6 +210,9 @@ struct Clamp<Array<T,N>> {
// Leaky Relu operator
template <typename T>
struct LeakyReLU {

static const bool kIsHeavy = false;

struct Arguments {
T leaky_alpha = T(0);
};
Expand All @@ -225,6 +231,9 @@ struct LeakyReLU {

template <typename T, int N>
struct LeakyReLU<Array<T, N> > {

static const bool kIsHeavy = false;

using Arguments = typename LeakyReLU<T>::Arguments;

CUTLASS_HOST_DEVICE
Expand All @@ -249,6 +258,8 @@ struct LeakyReLU<Array<T, N> > {
// Tanh operator
template <typename T>
struct Tanh {
static const bool kIsHeavy = true;

CUTLASS_HOST_DEVICE
T operator()(T const &value) const {
return fast_tanh(value);
Expand All @@ -257,6 +268,8 @@ struct Tanh {

template <typename T, int N>
struct Tanh<Array<T, N> > {
static const bool kIsHeavy = true;

CUTLASS_HOST_DEVICE
Array<T, N> operator()(Array<T, N> const &value) const {
Array<T, N> y;
Expand All @@ -274,6 +287,7 @@ struct Tanh<Array<T, N> > {
template <int N>
struct Tanh<Array<half_t, N>> {
using T = half_t;
static const bool kIsHeavy = true;

CUTLASS_HOST_DEVICE
Array<T, N> operator()(Array<T, N> const& z) const {
Expand All @@ -285,6 +299,8 @@ struct Tanh<Array<half_t, N>> {
// Sigmoid operator
template <typename T>
struct Sigmoid {
static const bool kIsHeavy = true;

CUTLASS_HOST_DEVICE
T operator()(T const &value) const {
return T(1) / (T(1) + fast_exp(-value));
Expand All @@ -293,6 +309,8 @@ struct Sigmoid {

template <typename T, int N>
struct Sigmoid<Array<T, N> > {
static const bool kIsHeavy = true;

CUTLASS_HOST_DEVICE
Array<T, N> operator()(Array<T, N> const &value) const {
Array<T, N> y;
Expand All @@ -310,6 +328,7 @@ struct Sigmoid<Array<T, N> > {
template <int N>
struct Sigmoid<Array<half_t, N>> {
using T = half_t;
static const bool kIsHeavy = true;

CUTLASS_HOST_DEVICE
Array<T, N> operator()(Array<T, N> const& z) const {
Expand Down Expand Up @@ -338,6 +357,8 @@ struct Sigmoid<Array<half_t, N>> {
// Reference: https://pytorch.org/docs/stable/generated/torch.nn.SiLU.html
template <typename T>
struct SiLu {
static const bool kIsHeavy = true;

CUTLASS_HOST_DEVICE
T operator()(T const &value) const {
Sigmoid<T> sigmoid;
Expand All @@ -347,6 +368,8 @@ struct SiLu {

template <typename T, int N>
struct SiLu<Array<T, N>> {
static const bool kIsHeavy = true;

CUTLASS_HOST_DEVICE
Array<T, N> operator()(Array<T, N> const &value) const {
Sigmoid<Array<T, N>> sigmoid_op;
Expand All @@ -362,6 +385,8 @@ struct SiLu<Array<T, N>> {
// Reference: https://pytorch.org/docs/stable/generated/torch.nn.Hardswish.html
template <typename T>
struct HardSwish {
static const bool kIsHeavy = false;

CUTLASS_HOST_DEVICE
T operator()(T const &x) const {
minimum<T> mn;
Expand All @@ -374,6 +399,7 @@ struct HardSwish {
template <>
struct HardSwish<float> {
using T = float;
static const bool kIsHeavy = false;

CUTLASS_HOST_DEVICE
T operator()(T const &x) const {
Expand All @@ -386,6 +412,8 @@ struct HardSwish<float> {

template <typename T, int N>
struct HardSwish<Array<T, N> > {
static const bool kIsHeavy = false;

CUTLASS_HOST_DEVICE
Array<T, N> operator()(Array<T, N> const &value) const {
Array<T, N> y;
Expand All @@ -403,6 +431,7 @@ struct HardSwish<Array<T, N> > {
template <int N>
struct HardSwish<Array<half_t, N> > {
using T = half_t;
static const bool kIsHeavy = false;

CUTLASS_HOST_DEVICE
Array<T, N> operator()(Array<T, N> const &value) const {
Expand All @@ -427,6 +456,8 @@ struct HardSwish<Array<half_t, N> > {
// GELU operator
template <typename T>
struct GELU {
static const bool kIsHeavy = true;

CUTLASS_HOST_DEVICE
T operator()(T const &value) const {
return T(cutlass::constants::half<T>() * value *
Expand All @@ -436,6 +467,8 @@ struct GELU {

template <>
struct GELU<float> {
static const bool kIsHeavy = true;

CUTLASS_HOST_DEVICE
float operator()(float const &value) const {
return cutlass::constants::half<float>() * value *
Expand All @@ -445,6 +478,8 @@ struct GELU<float> {

template <>
struct GELU<double> {
static const bool kIsHeavy = true;

CUTLASS_HOST_DEVICE
double operator()(double const &value) const {
return cutlass::constants::half<double>() * value *
Expand All @@ -454,6 +489,8 @@ struct GELU<double> {

template <typename T, int N>
struct GELU<Array<T, N> > {
static const bool kIsHeavy = true;

CUTLASS_HOST_DEVICE
Array<T, N> operator()(Array<T, N> const &value) const {
Array<T, N> y;
Expand All @@ -474,7 +511,8 @@ using ScaledGELU = Scale<GELU<T>>;
// GELU operator implemented using the Taylor series approximation
template <typename T>
struct GELU_taylor {
static const bool kIsHeavy=true;
static const bool kIsHeavy = true;

CUTLASS_HOST_DEVICE
T operator()(T const &z) const {

Expand All @@ -488,7 +526,8 @@ struct GELU_taylor {

template <int N>
struct GELU_taylor<Array<half_t, N> > {
static const bool kIsHeavy=true;
static const bool kIsHeavy = true;

CUTLASS_HOST_DEVICE
Array<half_t, N> operator()(Array<half_t, N> const &z) const {

Expand All @@ -514,7 +553,8 @@ struct GELU_taylor<Array<half_t, N> > {

template <typename T, int N>
struct GELU_taylor<Array<T, N> > {
static const bool kIsHeavy=true;
static const bool kIsHeavy = true;

CUTLASS_HOST_DEVICE
Array<T, N> operator()(Array<T, N> const &value) const {
Array<T, N> y;
Expand All @@ -536,6 +576,8 @@ using ScaledGELU_taylor = Scale<GELU_taylor<T>>;
/// z is computed from the forward pass.
template <typename T>
struct dGELU {
static const bool kIsHeavy = true;

CUTLASS_HOST_DEVICE
T operator()(T const &d_t, T const &z) const {

Expand All @@ -554,6 +596,8 @@ struct dGELU {

template <typename T, int N>
struct dGELU<Array<T, N> > {
static const bool kIsHeavy = true;

CUTLASS_HOST_DEVICE
Array<T, N> operator()(Array<T, N> const &d_t, Array<T, N> const &z) const {
Array<T, N> y;
Expand All @@ -568,6 +612,45 @@ struct dGELU<Array<T, N> > {
}
};

template <typename T>
struct dReLU {
CUTLASS_HOST_DEVICE
T operator()(T const& d_t, bool d_relu) const {
return d_relu ? d_t : T(0);
}

CUTLASS_HOST_DEVICE
T operator()(T const& d_t, uint1b_t d_relu) const {
return operator()(d_t, static_cast<bool>(d_relu));
}
};

template <typename T, int N>
struct dReLU<Array<T, N>> {
CUTLASS_HOST_DEVICE
Array<T, N> operator()(Array<T, N> const& d_t, bool const (&d_relu)[N]) const {
Array<T, N> y;
dReLU<T> relu_op;

CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < N; ++i) {
y[i] = relu_op(d_t[i], d_relu[i]);
}

return y;
}

CUTLASS_HOST_DEVICE
Array<T, N> operator()(Array<T, N> const& d_t, Array<uint1b_t, N> const& d_relu) const {
UnpackPredicates<N> unpack_op;

bool preds[N];
unpack_op(preds, d_relu);

return operator()(d_t, preds);
}
};

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

} // namespace thread
Expand Down