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

CUTLASS 3.6.0 (#1850)

Browse Source 
* v3.6

* update changelog

* update readme

* fix typo

* fixing typos

* hopper gemm with weight prefetch

---------

Co-authored-by: yuzhai <yuzhai@nvidia.com>
Co-authored-by: Haicheng Wu <haichengw@nvidia.com>
This commit is contained in:
Yujia Zhai
2024-10-09 12:33:27 -07:00
committed by GitHub
parent 0837a2a00a
commit cc3c29a81a
354 changed files with 105943 additions and 8203 deletions
Download Patch File Download Diff File Expand all files Collapse all files

215
include/cutlass/array.h
View File

@@ -37,6 +37,7 @@
#include "cutlass/cutlass.h"
#include "cutlass/functional.h"
#include "cutlass/numeric_types.h"
#include "cutlass/platform/platform.h"
namespace cutlass {
////////////////////////////////////////////////////////////////////////////////////////////////////
@@ -49,6 +50,23 @@ template <
>
struct Array;
namespace detail {
template<class T>
struct is_Array : platform::false_type {};
template <
typename T,
int N,
bool RegisterSized
>
struct is_Array<Array<T, N, RegisterSized> > : platform::true_type {};
template<typename T>
constexpr bool is_Array_v = is_Array<T>::value;
} // namespace detail
////////////////////////////////////////////////////////////////////////////////////////////////////
/// Defines the size of an Array<> in bits
@@ -803,14 +821,14 @@ struct reciprocal_approximate_ftz<Array<T, N>> {
}
};
template <typename T, int N>
struct maximum<Array<T, N>, false> {
template <typename T, int N, bool PropagateNaN>
struct maximum<Array<T, N>, PropagateNaN> {
CUTLASS_HOST_DEVICE
Array<T, N> operator()(Array<T, N> const &lhs, Array<T, N> const &rhs) const {
Array<T, N> result;
maximum<T, false> scalar_op;
maximum<T, PropagateNaN> scalar_op;
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < N; ++i) {
@@ -824,7 +842,7 @@ struct maximum<Array<T, N>, false> {
Array<T, N> operator()(Array<T, N> const &lhs, T const &scalar) const {
Array<T, N> result;
maximum<T, false> scalar_op;
maximum<T, PropagateNaN> scalar_op;
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < N; ++i) {
@@ -838,7 +856,7 @@ struct maximum<Array<T, N>, false> {
Array<T, N> operator()(T const &scalar, Array<T, N> const &rhs) const {
Array<T, N> result;
maximum<T, false> scalar_op;
maximum<T, PropagateNaN> scalar_op;
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < N; ++i) {
@@ -849,54 +867,8 @@ struct maximum<Array<T, N>, false> {
}
};
template <typename T, int N>
struct maximum<Array<T, N>, true> {
CUTLASS_HOST_DEVICE
Array<T, N> operator()(Array<T, N> const &lhs, Array<T, N> const &rhs) const {
Array<T, N> result;
maximum<T, true> scalar_op;
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < N; ++i) {
result[i] = scalar_op(lhs[i], rhs[i]);
}
return result;
}
CUTLASS_HOST_DEVICE
Array<T, N> operator()(Array<T, N> const &lhs, T const &scalar) const {
Array<T, N> result;
maximum<T, true> scalar_op;
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < N; ++i) {
result[i] = scalar_op(lhs[i], scalar);
}
return result;
}
CUTLASS_HOST_DEVICE
Array<T, N> operator()(T const &scalar, Array<T, N> const &rhs) const {
Array<T, N> result;
maximum<T, true> scalar_op;
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < N; ++i) {
result[i] = scalar_op(scalar, rhs[i]);
}
return result;
}
};
template <typename T, int N>
struct minimum<Array<T, N>, false> {
template <typename T, int N, bool PropagateNaN>
struct minimum<Array<T, N>, PropagateNaN> {
CUTLASS_HOST_DEVICE
static T scalar_op(T const &lhs, T const &rhs) {
@@ -907,7 +879,7 @@ struct minimum<Array<T, N>, false> {
Array<T, N> operator()(Array<T, N> const &lhs, Array<T, N> const &rhs) const {
Array<T, N> result;
minimum<T, false> scalar_op;
minimum<T, PropagateNaN> scalar_op;
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < N; ++i) {
@@ -921,7 +893,7 @@ struct minimum<Array<T, N>, false> {
Array<T, N> operator()(Array<T, N> const &lhs, T const &scalar) const {
Array<T, N> result;
minimum<T, false> scalar_op;
minimum<T, PropagateNaN> scalar_op;
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < N; ++i) {
@@ -935,58 +907,7 @@ struct minimum<Array<T, N>, false> {
Array<T, N> operator()(T const &scalar, Array<T, N> const &rhs) const {
Array<T, N> result;
minimum<T, false> scalar_op;
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < N; ++i) {
result[i] = scalar_op(scalar, rhs[i]);
}
return result;
}
};
template <typename T, int N>
struct minimum<Array<T, N>, true> {
CUTLASS_HOST_DEVICE
static T scalar_op(T const &lhs, T const &rhs) {
return (rhs < lhs ? rhs : lhs);
}
CUTLASS_HOST_DEVICE
Array<T, N> operator()(Array<T, N> const &lhs, Array<T, N> const &rhs) const {
Array<T, N> result;
minimum<T, true> scalar_op;
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < N; ++i) {
result[i] = scalar_op(lhs[i], rhs[i]);
}
return result;
}
CUTLASS_HOST_DEVICE
Array<T, N> operator()(Array<T, N> const &lhs, T const &scalar) const {
Array<T, N> result;
minimum<T, true> scalar_op;
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < N; ++i) {
result[i] = scalar_op(lhs[i], scalar);
}
return result;
}
CUTLASS_HOST_DEVICE
Array<T, N> operator()(T const &scalar, Array<T, N> const &rhs) const {
Array<T, N> result;
minimum<T, true> scalar_op;
minimum<T, PropagateNaN> scalar_op;
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < N; ++i) {
@@ -2030,8 +1951,8 @@ struct multiply_add_relu0<Array<half_t, N>, Array<half_t, N>, Array<half_t, N>>
}
};
template <int N>
struct minimum<Array<half_t, N>, false> {
template <int N, bool PropagateNaN>
struct minimum<Array<half_t, N>, PropagateNaN> {
CUTLASS_HOST_DEVICE
Array<half_t, N> operator()(Array<half_t, N> const & lhs, Array<half_t, N> const &rhs) const {
Array<half_t, N> result;
@@ -2043,25 +1964,27 @@ struct minimum<Array<half_t, N>, false> {
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < N / 2; ++i) {
result_ptr[i] = __hmin2(lhs_ptr[i], rhs_ptr[i]);
result_ptr[i] = PropagateNaN ? __hmin2_nan(lhs_ptr[i], rhs_ptr[i])
: __hmin2(lhs_ptr[i], rhs_ptr[i]);
}
if constexpr (N % 2) {
__half const *a_residual_ptr = reinterpret_cast<__half const *>(&lhs);
__half const *b_residual_ptr = reinterpret_cast<__half const *>(&rhs);
__half d_residual = __hmin(
a_residual_ptr[N - 1],
b_residual_ptr[N - 1]);
__half d_residual = PropagateNaN ? __hmin_nan(a_residual_ptr[N - 1], b_residual_ptr[N - 1])
: __hmin(a_residual_ptr[N - 1], b_residual_ptr[N - 1]);
result[N - 1] = reinterpret_cast<half_t const &>(d_residual);
}
#else
minimum<half_t,PropagateNaN> mn;
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < N; ++i) {
result[i] = (rhs[i] < lhs[i] ? rhs[i] : lhs[i]);
result[i] = mn(lhs[i],rhs[i]);
}
#endif
@@ -2079,24 +2002,26 @@ struct minimum<Array<half_t, N>, false> {
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < N / 2; ++i) {
result_ptr[i] = __hmin2(lhs_pair, rhs_ptr[i]);
result_ptr[i] = PropagateNaN ? __hmin2_nan(lhs_pair, rhs_ptr[i])
: __hmin2(lhs_pair, rhs_ptr[i]);
}
if constexpr (N % 2) {
__half const *b_residual_ptr = reinterpret_cast<__half const *>(&rhs);
__half d_residual = __hmin(
reinterpret_cast<__half const &>(lhs),
b_residual_ptr[N - 1]);
__half d_residual = PropagateNaN ? __hmin_nan(reinterpret_cast<__half const &>(lhs), b_residual_ptr[N - 1])
: __hmin(reinterpret_cast<__half const &>(lhs), b_residual_ptr[N - 1]);
result[N - 1] = reinterpret_cast<half_t const &>(d_residual);
}
#else
minimum<half_t,PropagateNaN> mn;
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < N; ++i) {
result[i] = (rhs[i] < lhs ? rhs[i] : lhs);
result[i] = mn(lhs, rhs[i]);
}
#endif
@@ -2114,24 +2039,26 @@ struct minimum<Array<half_t, N>, false> {
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < N / 2; ++i) {
result_ptr[i] = __hmin2(lhs_ptr[i], rhs_pair);
result_ptr[i] = PropagateNaN ? __hmin2_nan(lhs_ptr[i], rhs_pair)
: __hmin2(lhs_ptr[i], rhs_pair);
}
if constexpr (N % 2) {
__half const *a_residual_ptr = reinterpret_cast<__half const *>(&lhs);
__half d_residual = __hmin(
a_residual_ptr[N - 1],
reinterpret_cast<__half const &>(rhs));
__half d_residual = PropagateNaN ? __hmin_nan(a_residual_ptr[N - 1], reinterpret_cast<__half const &>(rhs))
: __hmin(a_residual_ptr[N - 1], reinterpret_cast<__half const &>(rhs));
result[N - 1] = reinterpret_cast<half_t const &>(d_residual);
}
#else
minimum<half_t, PropagateNaN> mn;
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < N; ++i) {
result[i] = (rhs < lhs[i] ? rhs : lhs[i]);
result[i] = mn(lhs[i], rhs);
}
#endif
@@ -2139,8 +2066,8 @@ struct minimum<Array<half_t, N>, false> {
}
};
template <int N>
struct maximum<Array<half_t, N>, false> {
template <int N, bool PropagateNaN>
struct maximum<Array<half_t, N>, PropagateNaN> {
CUTLASS_HOST_DEVICE
Array<half_t, N> operator()(Array<half_t, N> const & lhs, Array<half_t, N> const &rhs) const {
Array<half_t, N> result;
@@ -2152,25 +2079,27 @@ struct maximum<Array<half_t, N>, false> {
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < N / 2; ++i) {
result_ptr[i] = __hmax2(lhs_ptr[i], rhs_ptr[i]);
result_ptr[i] = PropagateNaN ? __hmax2_nan(lhs_ptr[i], rhs_ptr[i])
: __hmax2(lhs_ptr[i], rhs_ptr[i]);
}
if constexpr (N % 2) {
__half const *a_residual_ptr = reinterpret_cast<__half const *>(&lhs);
__half const *b_residual_ptr = reinterpret_cast<__half const *>(&rhs);
__half d_residual = __hmax(
a_residual_ptr[N - 1],
b_residual_ptr[N - 1]);
__half d_residual = PropagateNaN ? __hmax(a_residual_ptr[N - 1], b_residual_ptr[N - 1])
: __hmax_nan(a_residual_ptr[N - 1], b_residual_ptr[N - 1]);
result[N - 1] = reinterpret_cast<half_t const &>(d_residual);
}
#else
maximum<half_t,PropagateNaN> mx;
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < N; ++i) {
result[i] = (lhs[i] < rhs[i] ? rhs[i] : lhs[i]);
result[i] = mx(lhs[i], rhs[i]);
}
#endif
@@ -2188,24 +2117,26 @@ struct maximum<Array<half_t, N>, false> {
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < N / 2; ++i) {
result_ptr[i] = __hmax2(lhs_pair, rhs_ptr[i]);
result_ptr[i] = PropagateNaN ? __hmax2_nan(lhs_pair, rhs_ptr[i])
: __hmax2(lhs_pair, rhs_ptr[i]);
}
if constexpr (N % 2) {
__half const *b_residual_ptr = reinterpret_cast<__half const *>(&rhs);
__half d_residual = __hmax(
reinterpret_cast<__half const &>(lhs),
b_residual_ptr[N - 1]);
__half d_residual = PropagateNaN ? __hmax_nan(reinterpret_cast<__half const &>(lhs), b_residual_ptr[N - 1])
: __hmax(reinterpret_cast<__half const &>(lhs), b_residual_ptr[N - 1]);
result[N - 1] = reinterpret_cast<half_t const &>(d_residual);
}
#else
maximum<half_t,PropagateNaN> mx;
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < N; ++i) {
result[i] = (lhs < rhs[i] ? rhs[i] : lhs);
result[i] = mx(lhs, rhs[i]);
}
#endif
@@ -2223,24 +2154,26 @@ struct maximum<Array<half_t, N>, false> {
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < N / 2; ++i) {
result_ptr[i] = __hmax2(lhs_ptr[i], rhs_pair);
result_ptr[i] = PropagateNaN ? __hmax2_nan(lhs_ptr[i], rhs_pair)
: __hmax2(lhs_ptr[i], rhs_pair);
}
if constexpr (N % 2) {
__half const *a_residual_ptr = reinterpret_cast<__half const *>(&lhs);
__half d_residual = __hmax(
a_residual_ptr[N - 1],
reinterpret_cast<__half const &>(rhs));
__half d_residual = PropagateNaN ? __hmax_nan(a_residual_ptr[N - 1], reinterpret_cast<__half const &>(rhs))
: __hmax(a_residual_ptr[N - 1], reinterpret_cast<__half const &>(rhs));
result[N - 1] = reinterpret_cast<half_t const &>(d_residual);
}
#else
maximum<half_t,PropagateNaN> mx;
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < N; ++i) {
result[i] = (lhs[i] < rhs ? rhs : lhs[i]);
result[i] = mx(lhs[i], rhs);
}
#endif