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

CUTLASS 3.5.1 (#1623)

Browse Source 
* CUTLASS 3.5.1

* updates, optimizations, fixes
This commit is contained in:
Vijay Thakkar
2024-07-29 08:46:24 -04:00
committed by GitHub
parent 56b46e2d13
commit be60a0b272
312 changed files with 19793 additions and 6775 deletions
Download Patch File Download Diff File Expand all files Collapse all files
include/cutlass/functional.h
+215 -6
View File
@@ -37,6 +37,11 @@
#include "cutlass/cutlass.h"
#include "cutlass/numeric_types.h"
#include "cutlass/platform/platform.h"
#if defined(__CUDACC_RTC__)
#include "cutlass/floating_point_nvrtc.h"
#endif
#include <cuda_runtime.h>
@@ -262,7 +267,36 @@ struct reciprocal_approximate <float> {
CUTLASS_HOST_DEVICE
float operator()(float lhs) const {
float ret;
#if defined(__CUDA_ARCH__)
asm volatile ("rcp.approx.f32 %0, %1;\n" : "=f"(ret) : "f"(lhs));
#else
ret = 1.0f / lhs;
#endif
return ret;
}
};
/// reciprocal_approximate with ftz
template<typename T>
struct reciprocal_approximate_ftz : reciprocal_approximate<T>
{};
template <>
struct reciprocal_approximate_ftz <float> {
CUTLASS_HOST_DEVICE
float operator()(float lhs) const {
float ret;
#if defined(__CUDA_ARCH__)
asm volatile ("rcp.approx.ftz.f32 %0, %1;\n" : "=f"(ret) : "f"(lhs));
#else
if (std::fpclassify(lhs) == FP_SUBNORMAL) {
lhs = 0.0f;
}
ret = 1.0f / lhs;
if (std::fpclassify(ret) == FP_SUBNORMAL) {
ret = 0.0f;
}
#endif
return ret;
}
};
@@ -336,7 +370,7 @@ struct maximum<T, true> {
CUTLASS_HOST_DEVICE
T operator()(T const &lhs, T const &rhs) const {
#if defined(__CUDA_ARCH__)
return lhs > rhs or isnan(lhs) ? lhs : rhs;
return lhs > rhs or ::isnan(lhs) ? lhs : rhs;
#else
return lhs > rhs or std::isnan(lhs) ? lhs : rhs;
#endif
@@ -359,7 +393,7 @@ struct maximum<float, true> {
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)
asm volatile("max.NaN.f32 %0, %1, %2;\n" : "=f"(res) : "f"(lhs), "f"(rhs));
#elif defined(__CUDA_ARCH__)
res = lhs > rhs or isnan(lhs) ? lhs : rhs;
res = lhs > rhs or ::isnan(lhs) ? lhs : rhs;
#else
res = lhs > rhs or std::isnan(lhs) ? lhs : rhs;
#endif
@@ -394,7 +428,7 @@ struct minimum<T, true> {
CUTLASS_HOST_DEVICE
T operator()(T const &lhs, T const &rhs) const {
#if defined(__CUDA_ARCH__)
return lhs < rhs or isnan(lhs) ? lhs : rhs;
return lhs < rhs or ::isnan(lhs) ? lhs : rhs;
#else
return lhs < rhs or std::isnan(lhs) ? lhs : rhs;
#endif
@@ -409,6 +443,10 @@ struct minimum<float, false> {
}
};
template <typename T>
struct minimum_with_nan_propagation : minimum<T, true>
{};
template <typename T, bool PropagateNaN = false>
struct maximum_absolute_value {
CUTLASS_HOST_DEVICE
@@ -469,6 +507,83 @@ struct multiply_add_relu0 {
}
};
/// Guarded-multiply-add
template <typename A, typename B = A, typename C = A>
struct guarded_multiply_add {
CUTLASS_HOST_DEVICE
C operator()(A const &a, B const &b, C const &c) const {
if (isnan(a) || isnan(b)) {
return C(0);
}
return C(a) * C(b) + c;
}
};
/// Guarded-multiply-add
template <>
struct guarded_multiply_add<half_t, half_t, half_t> {
CUTLASS_HOST_DEVICE
half_t operator()(half_t const &a, half_t const &b, half_t const &c) const {
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
half_t result;
asm ("fma.rn.oob.f16 %0, %1, %2, %3;\n"
: "=h"(*reinterpret_cast<uint16_t*>(&result))
: "h"(*reinterpret_cast<uint16_t const*>(&a)), "h"(*reinterpret_cast<uint16_t const*>(&b)), "h"(*reinterpret_cast<uint16_t const*>(&c)));
return result;
#else
if (isnan(a) || isnan(b)) {
return half_t(0);
}
return a * b + c;
#endif
}
};
/// Guarded-multiply-add-relu0
template <typename A, typename B = A, typename C = A>
struct guarded_multiply_add_relu0 {
CUTLASS_HOST_DEVICE
C operator()(A const &a, B const &b, C const &c) const {
if (
#if defined(__CUDA_ARCH__)
::isnan(a) || ::isnan(b)
#else
std::isnan(a) || std::isnan(b)
#endif
) {
return C(0);
}
maximum<C> mx;
return mx(C(a) * C(b) + c, C(0));
}
};
template <>
struct guarded_multiply_add_relu0<half_t, half_t, half_t> {
CUTLASS_HOST_DEVICE
half_t operator()(half_t const &a, half_t const &b, half_t const &c) const {
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
half_t result;
asm ("fma.rn.oob.relu.f16 %0, %1, %2, %3;\n"
: "=h"(*reinterpret_cast<uint16_t*>(&result))
: "h"(*reinterpret_cast<uint16_t const*>(&a)), "h"(*reinterpret_cast<uint16_t const*>(&b)), "h"(*reinterpret_cast<uint16_t const*>(&c)));
return result;
#else
if (
#if defined(__CUDA_ARCH__)
::isnan(a) || ::isnan(b)
#else
std::isnan(a) || std::isnan(b)
#endif
) {
return half_t(0);
}
maximum<half_t> mx;
return mx(a * b + c, half_t(0));
#endif
}
};
/// Fused multiply-add
template <typename T>
struct and_add {
@@ -488,11 +603,99 @@ struct xor_add {
}
};
namespace detail {
// Whether namespace-unqualified conj(t) for t of type T is
// well-formed. This says whether the compiler can find
// namespace-unqualified conj(T) via argument-dependent lookup.
// If so, then CUTLASS assumes that conj(t) returns
// the complex conjugate of t.
template <typename T, typename Enable = void>
struct has_unqualified_conj : cutlass::platform::false_type
{};
template<typename T>
struct has_unqualified_conj<
T,
decltype(conj(cutlass::platform::declval<T>()), void())
> : cutlass::platform::true_type
{};
template <typename T>
constexpr bool has_unqualified_conj_v = has_unqualified_conj<T>::value;
} // namespace detail
// forward declaration (needed for conjugate below)
template<class T>
CUTLASS_HOST_DEVICE T conj(T const& z);
namespace detail {
// Whether cutlass::conj(t) for t of type T is well-formed.
// If so, then CUTLASS assumes that cutlass::conj(t)
// returns the complex conjugate of t.
template <typename T, typename Enable = void>
struct has_cutlass_conj : cutlass::platform::false_type
{};
template<typename T>
struct has_cutlass_conj<
T,
decltype(cutlass::conj(cutlass::platform::declval<T>()), void())
> : cutlass::platform::true_type
{};
template <typename T>
constexpr bool has_cutlass_conj_v = has_cutlass_conj<T>::value;
} // namespace detail
// Return the complex conjugate of the input.
//
// If the struct hasn't already been specialized for type T, then
//
// 1. for arithmetic types, return z;
//
// 2. for types where either (namespace-unqualified) conj(z) or
// cutlass::conj(z) is well formed, declare "using cutlass::conj;"
// and return conj(z); and
//
// 3. for everything else, return z.
//
// Regarding (1), the C++ Standard Library makes std::conj always
// return std::complex, even for (noncomplex) arithmetic types.
// cutlass::conj(T t) needs to return type T. This follows the
// convention of linear algebra software like the BLAS, where
// "conjugate transpose" means the same thing as "transpose" for a
// matrix of noncomplex numbers.
//
// Case (2) covers std::complex, cuda::std::complex, and non-Standard
// (including user-defined) complex number types (for which "conj(z)"
// is findable via argument-dependent lookup). cutlass::conj has a
// totally generic overload, but a more type-specific overload in any
// namespace will take precedence.
//
// Case (3) covers non-Standard non-complex number types.
//
// Users should not generally need to specialize this struct for their
// own custom complex or noncomplex types. The idiomatic way to
// identify a type T as "complex" is to make namespace-unqualified
// calls to conj(T) findable via argument-dependent lookup.
template <typename T>
struct conjugate {
CUTLASS_HOST_DEVICE
T operator()(T const &a) const {
return a;
T operator()(T const& z) const {
if constexpr (cutlass::platform::is_arithmetic_v<T>) {
return z;
}
else if constexpr (detail::has_unqualified_conj_v<T> || detail::has_cutlass_conj_v<T>) {
using cutlass::conj;
return conj(z);
}
else {
return z;
}
}
};
@@ -649,7 +852,13 @@ struct atomic_maximum<float> {
CUTLASS_DEVICE
float operator()(float *ptr, float value) const {
#if defined(__CUDA_ARCH__)
return !signbit(value) ?
// In device code, make sure that we do NOT try to use
// std::signbit, as that won't work if building with NVRTC.
// Instead, prefix "::" to call signbit from the global namespace,
// which CUDA guarantees to work in device code without including
// any headers.
//
return ! ::signbit(value) ?
__int_as_float(atomicMax((int*)ptr, __float_as_int(value))) :
__uint_as_float(atomicMin((unsigned int*)ptr, __float_as_uint(value)));
#else