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

CUTLASS 3.2 (#1024)

Browse Source 
* CUTLASS 3.2
This commit is contained in:
ANIKET SHIVAM
2023-08-07 20:50:32 -04:00
committed by GitHub
parent a0d787b746
commit 4575443d44
392 changed files with 47559 additions and 7940 deletions
Download Patch File Download Diff File Expand all files Collapse all files
include/cutlass/array.h
+166 -49
View File
@@ -38,7 +38,6 @@
#include "cutlass/functional.h"
#include "cutlass/numeric_types.h"
#include "cutlass/half.h"
namespace cutlass {
////////////////////////////////////////////////////////////////////////////////////////////////////
@@ -717,6 +716,27 @@ struct multiplies<Array<T, N>> {
}
};
template <typename T, int N>
struct scale<Array<T, N>> {
T const scaling_factor_;
CUTLASS_HOST_DEVICE
scale(T scaling_factor) : scaling_factor_(scaling_factor) {
}
CUTLASS_HOST_DEVICE
Array<T, N> operator()(Array<T, N> const & rhs) const {
Array<T, N> result;
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < N; ++i) {
result[i] = rhs[i] * scaling_factor_;
}
return result;
}
};
template <typename T, int N>
struct divides<Array<T, N>> {
@@ -764,13 +784,13 @@ struct divides<Array<T, N>> {
};
template <typename T, int N>
struct maximum<Array<T, N>> {
struct maximum<Array<T, N>, false> {
CUTLASS_HOST_DEVICE
Array<T, N> operator()(Array<T, N> const &lhs, Array<T, N> const &rhs) const {
Array<T, N> result;
maximum<T> scalar_op;
maximum<T, false> scalar_op;
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < N; ++i) {
@@ -784,7 +804,7 @@ struct maximum<Array<T, N>> {
Array<T, N> operator()(Array<T, N> const &lhs, T const &scalar) const {
Array<T, N> result;
maximum<T> scalar_op;
maximum<T, false> scalar_op;
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < N; ++i) {
@@ -795,10 +815,10 @@ struct maximum<Array<T, N>> {
}
CUTLASS_HOST_DEVICE
Array<T, N> operator()( T const &scalar, Array<T, N> const &rhs) const {
Array<T, N> operator()(T const &scalar, Array<T, N> const &rhs) const {
Array<T, N> result;
maximum<T> scalar_op;
maximum<T, false> scalar_op;
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < N; ++i) {
@@ -810,18 +830,13 @@ struct maximum<Array<T, N>> {
};
template <typename T, int N>
struct minimum<Array<T, N>> {
CUTLASS_HOST_DEVICE
static T scalar_op(T const &lhs, T const &rhs) {
return (rhs < lhs ? rhs : lhs);
}
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;
minimum<T> scalar_op;
maximum<T, true> scalar_op;
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < N; ++i) {
@@ -835,7 +850,7 @@ struct minimum<Array<T, N>> {
Array<T, N> operator()(Array<T, N> const &lhs, T const &scalar) const {
Array<T, N> result;
minimum<T> scalar_op;
maximum<T, true> scalar_op;
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < N; ++i) {
@@ -846,10 +861,112 @@ struct minimum<Array<T, N>> {
}
CUTLASS_HOST_DEVICE
Array<T, N> operator()( T const &scalar, Array<T, N> const &rhs) const {
Array<T, N> operator()(T const &scalar, Array<T, N> const &rhs) const {
Array<T, N> result;
minimum<T> scalar_op;
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> {
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, false> 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, false> 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, 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;
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < N; ++i) {
@@ -1013,7 +1130,7 @@ struct plus<Array<half_t, N>> {
result_ptr[i] = __hadd2(lhs_ptr[i], rhs_ptr[i]);
}
if (N % 2) {
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 = __hadd(a_residual_ptr[N - 1], b_residual_ptr[N - 1]);
@@ -1046,7 +1163,7 @@ struct plus<Array<half_t, N>> {
result_ptr[i] = __hadd2(lhs_pair, rhs_ptr[i]);
}
if (N % 2) {
if constexpr (N % 2) {
__half const *b_residual_ptr = reinterpret_cast<__half const *>(&rhs);
__half d_residual = __hadd(reinterpret_cast<__half const &>(lhs), b_residual_ptr[N - 1]);
@@ -1078,7 +1195,7 @@ struct plus<Array<half_t, N>> {
result_ptr[i] = __hadd2(lhs_ptr[i], rhs_pair);
}
if (N % 2) {
if constexpr (N % 2) {
__half const *a_residual_ptr = reinterpret_cast<__half const *>(&lhs);
__half d_residual = __hadd(a_residual_ptr[N - 1], reinterpret_cast<__half const &>(rhs));
@@ -1113,7 +1230,7 @@ struct minus<Array<half_t, N>> {
result_ptr[i] = __hsub2(lhs_ptr[i], rhs_ptr[i]);
}
if (N % 2) {
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 = __hsub(a_residual_ptr[N - 1], b_residual_ptr[N - 1]);
@@ -1146,7 +1263,7 @@ struct minus<Array<half_t, N>> {
result_ptr[i] = __hsub2(lhs_pair, rhs_ptr[i]);
}
if (N % 2) {
if constexpr (N % 2) {
__half const *b_residual_ptr = reinterpret_cast<__half const *>(&rhs);
__half d_residual = __hsub(reinterpret_cast<__half const &>(lhs), b_residual_ptr[N - 1]);
@@ -1178,7 +1295,7 @@ struct minus<Array<half_t, N>> {
result_ptr[i] = __hsub2(lhs_ptr[i], rhs_pair);
}
if (N % 2) {
if constexpr (N % 2) {
__half const *a_residual_ptr = reinterpret_cast<__half const *>(&lhs);
__half d_residual = __hsub(a_residual_ptr[N - 1], reinterpret_cast<__half const &>(rhs));
@@ -1213,7 +1330,7 @@ struct multiplies<Array<half_t, N>> {
result_ptr[i] = __hmul2(lhs_ptr[i], rhs_ptr[i]);
}
if (N % 2) {
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 = __hmul(a_residual_ptr[N - 1], b_residual_ptr[N - 1]);
@@ -1246,7 +1363,7 @@ struct multiplies<Array<half_t, N>> {
result_ptr[i] = __hmul2(lhs_pair, rhs_ptr[i]);
}
if (N % 2) {
if constexpr (N % 2) {
__half const *b_residual_ptr = reinterpret_cast<__half const *>(&rhs);
__half d_residual = __hmul(
@@ -1281,7 +1398,7 @@ struct multiplies<Array<half_t, N>> {
result_ptr[i] = __hmul2(lhs_ptr[i], rhs_pair);
}
if (N % 2) {
if constexpr (N % 2) {
__half const *a_residual_ptr = reinterpret_cast<__half const *>(&lhs);
__half d_residual = __hmul(
@@ -1319,7 +1436,7 @@ struct divides<Array<half_t, N>> {
result_ptr[i] = __h2div(lhs_ptr[i], rhs_ptr[i]);
}
if (N % 2) {
if constexpr (N % 2) {
__half const *a_residual_ptr = reinterpret_cast<__half const *>(&lhs);
__half const *b_residual_ptr = reinterpret_cast<__half const *>(&rhs);
@@ -1355,7 +1472,7 @@ struct divides<Array<half_t, N>> {
result_ptr[i] = __h2div(lhs_pair, rhs_ptr[i]);
}
if (N % 2) {
if constexpr (N % 2) {
__half const *b_residual_ptr = reinterpret_cast<__half const *>(&rhs);
__half d_residual = __hdiv(
@@ -1390,7 +1507,7 @@ struct divides<Array<half_t, N>> {
result_ptr[i] = __h2div(lhs_ptr[i], rhs_pair);
}
if (N % 2) {
if constexpr (N % 2) {
__half const *a_residual_ptr = reinterpret_cast<__half const *>(&lhs);
__half d_residual = __hdiv(
@@ -1427,7 +1544,7 @@ struct negate<Array<half_t, N>> {
result_ptr[i] = __hneg2(source_ptr[i]);
}
if (N % 2) {
if constexpr (N % 2) {
half_t x = lhs[N - 1];
__half lhs_val = -reinterpret_cast<__half const &>(x);
result[N - 1] = reinterpret_cast<half_t const &>(lhs_val);
@@ -1468,7 +1585,7 @@ struct multiply_add<Array<half_t, N>, Array<half_t, N>, Array<half_t, N>> {
result_ptr[i] = __hfma2(a_ptr[i], b_ptr[i], c_ptr[i]);
}
if (N % 2) {
if constexpr (N % 2) {
__half const *a_residual_ptr = reinterpret_cast<__half const *>(&a);
__half const *b_residual_ptr = reinterpret_cast<__half const *>(&b);
@@ -1514,7 +1631,7 @@ struct multiply_add<Array<half_t, N>, Array<half_t, N>, Array<half_t, N>> {
result_ptr[i] = __hfma2(a_pair, b_ptr[i], c_ptr[i]);
}
if (N % 2) {
if constexpr (N % 2) {
__half const *b_residual_ptr = reinterpret_cast<__half const *>(&b);
__half const *c_residual_ptr = reinterpret_cast<__half const *>(&c);
@@ -1558,7 +1675,7 @@ struct multiply_add<Array<half_t, N>, Array<half_t, N>, Array<half_t, N>> {
result_ptr[i] = __hfma2(a_ptr[i], b_pair, c_ptr[i]);
}
if (N % 2) {
if constexpr (N % 2) {
__half const *a_residual_ptr = reinterpret_cast<__half const *>(&a);
__half const *c_residual_ptr = reinterpret_cast<__half const *>(&c);
@@ -1603,7 +1720,7 @@ struct multiply_add<Array<half_t, N>, Array<half_t, N>, Array<half_t, N>> {
result_ptr[i] = __hfma2(a_ptr[i], b_ptr[i], c_pair);
}
if (N % 2) {
if constexpr (N % 2) {
__half const *a_residual_ptr = reinterpret_cast<__half const *>(&a);
__half const *b_residual_ptr = reinterpret_cast<__half const *>(&b);
@@ -1653,7 +1770,7 @@ struct multiply_add_relu0<Array<half_t, N>, Array<half_t, N>, Array<half_t, N>>
result_ptr[i] = __hfma2_relu(a_ptr[i], b_ptr[i], c_ptr[i]);
}
if (N % 2) {
if constexpr (N % 2) {
__half const *a_residual_ptr = reinterpret_cast<__half const *>(&a);
__half const *b_residual_ptr = reinterpret_cast<__half const *>(&b);
@@ -1700,7 +1817,7 @@ struct multiply_add_relu0<Array<half_t, N>, Array<half_t, N>, Array<half_t, N>>
result_ptr[i] = __hfma2_relu(a_pair, b_ptr[i], c_ptr[i]);
}
if (N % 2) {
if constexpr (N % 2) {
__half const *b_residual_ptr = reinterpret_cast<__half const *>(&b);
__half const *c_residual_ptr = reinterpret_cast<__half const *>(&c);
@@ -1745,7 +1862,7 @@ struct multiply_add_relu0<Array<half_t, N>, Array<half_t, N>, Array<half_t, N>>
result_ptr[i] = __hfma2_relu(a_ptr[i], b_pair, c_ptr[i]);
}
if (N % 2) {
if constexpr (N % 2) {
__half const *a_residual_ptr = reinterpret_cast<__half const *>(&a);
__half const *c_residual_ptr = reinterpret_cast<__half const *>(&c);
@@ -1791,7 +1908,7 @@ struct multiply_add_relu0<Array<half_t, N>, Array<half_t, N>, Array<half_t, N>>
result_ptr[i] = __hfma2_relu(a_ptr[i], b_ptr[i], c_pair);
}
if (N % 2) {
if constexpr (N % 2) {
__half const *a_residual_ptr = reinterpret_cast<__half const *>(&a);
__half const *b_residual_ptr = reinterpret_cast<__half const *>(&b);
@@ -1820,7 +1937,7 @@ 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>> {
struct minimum<Array<half_t, N>, false> {
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;
@@ -1835,7 +1952,7 @@ struct minimum<Array<half_t, N>> {
result_ptr[i] = __hmin2(lhs_ptr[i], rhs_ptr[i]);
}
if (N % 2) {
if constexpr (N % 2) {
__half const *a_residual_ptr = reinterpret_cast<__half const *>(&lhs);
__half const *b_residual_ptr = reinterpret_cast<__half const *>(&rhs);
@@ -1871,7 +1988,7 @@ struct minimum<Array<half_t, N>> {
result_ptr[i] = __hmin2(lhs_pair, rhs_ptr[i]);
}
if (N % 2) {
if constexpr (N % 2) {
__half const *b_residual_ptr = reinterpret_cast<__half const *>(&rhs);
__half d_residual = __hmin(
@@ -1906,7 +2023,7 @@ struct minimum<Array<half_t, N>> {
result_ptr[i] = __hmin2(lhs_ptr[i], rhs_pair);
}
if (N % 2) {
if constexpr (N % 2) {
__half const *a_residual_ptr = reinterpret_cast<__half const *>(&lhs);
__half d_residual = __hmin(
@@ -1929,7 +2046,7 @@ struct minimum<Array<half_t, N>> {
};
template <int N>
struct maximum<Array<half_t, N>> {
struct maximum<Array<half_t, N>, false> {
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;
@@ -1944,7 +2061,7 @@ struct maximum<Array<half_t, N>> {
result_ptr[i] = __hmax2(lhs_ptr[i], rhs_ptr[i]);
}
if (N % 2) {
if constexpr (N % 2) {
__half const *a_residual_ptr = reinterpret_cast<__half const *>(&lhs);
__half const *b_residual_ptr = reinterpret_cast<__half const *>(&rhs);
@@ -1980,7 +2097,7 @@ struct maximum<Array<half_t, N>> {
result_ptr[i] = __hmax2(lhs_pair, rhs_ptr[i]);
}
if (N % 2) {
if constexpr (N % 2) {
__half const *b_residual_ptr = reinterpret_cast<__half const *>(&rhs);
__half d_residual = __hmax(
@@ -2015,7 +2132,7 @@ struct maximum<Array<half_t, N>> {
result_ptr[i] = __hmax2(lhs_ptr[i], rhs_pair);
}
if (N % 2) {
if constexpr (N % 2) {
__half const *a_residual_ptr = reinterpret_cast<__half const *>(&lhs);
__half d_residual = __hmax(
@@ -2063,7 +2180,7 @@ struct multiply_add<Array<bfloat16_t, N>, Array<bfloat16_t, N>, Array<bfloat16_t
);
}
if (N % 2) {
if constexpr (N % 2) {
uint16_t *result_ptr = reinterpret_cast<uint16_t *>(&result);
uint16_t const *a_residual_ptr = reinterpret_cast<uint16_t const *>(&a);
@@ -2114,7 +2231,7 @@ struct multiply_add<Array<bfloat16_t, N>, Array<bfloat16_t, N>, Array<bfloat16_t
);
}
if (N % 2) {
if constexpr (N % 2) {
uint16_t *result_ptr = reinterpret_cast<uint16_t *>(&result);
uint16_t const *a_residual_ptr = reinterpret_cast<uint16_t const *>(&a);
@@ -2165,7 +2282,7 @@ struct multiply_add<Array<bfloat16_t, N>, Array<bfloat16_t, N>, Array<bfloat16_t
);
}
if (N % 2) {
if constexpr (N % 2) {
uint16_t *result_ptr = reinterpret_cast<uint16_t *>(&result);
uint16_t const *a_residual_ptr = reinterpret_cast<uint16_t const *>(&a);
@@ -2216,7 +2333,7 @@ struct multiply_add<Array<bfloat16_t, N>, Array<bfloat16_t, N>, Array<bfloat16_t
);
}
if (N % 2) {
if constexpr (N % 2) {
uint16_t *result_ptr = reinterpret_cast<uint16_t *>(&result);
uint16_t const *a_residual_ptr = reinterpret_cast<uint16_t const *>(&a);