squall/cutlass
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity
be692b48b0
cutlass/include/cutlass/functional.h
Yujia Zhai cc3c29a81a
CUTLASS 3.6.0 (#1850) 
* 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>
2024-10-09 15:33:27 -04:00

931 lines
24 KiB
C++
Raw Blame History

/***************************************************************************************************
* Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: BSD-3-Clause
*
* Redistribution and use in source and binary forms, with or without
* modification, are permitted provided that the following conditions are met:
*
* 1. Redistributions of source code must retain the above copyright notice, this
* list of conditions and the following disclaimer.
*
* 2. Redistributions in binary form must reproduce the above copyright notice,
* this list of conditions and the following disclaimer in the documentation
* and/or other materials provided with the distribution.
*
* 3. Neither the name of the copyright holder nor the names of its
* contributors may be used to endorse or promote products derived from
* this software without specific prior written permission.
*
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
* DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
* DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
* SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
* CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
* OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*
**************************************************************************************************/
/*! \file
\brief Define basic numeric operators
This is inspired by the Standard Library's <functional> header.
*/
#pragma once
#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>
#if defined(CUTLASS_ARCH_WMMA_ENABLED)
#include <mma.h>
#endif // defined(CUTLASS_ARCH_WMMA_ENABLED)
#ifdef _MSC_VER
// Provides support for alternate operators such as 'and', 'or', ...
#include <iso646.h>
#endif // _MSC_VER
namespace cutlass {
/////////////////////////////////////////////////////////////////////////////////////////////////
template <typename T>
struct absolute_value_op {
CUTLASS_HOST_DEVICE
T operator()(T lhs) const {
return abs(lhs);
}
};
template <>
struct absolute_value_op<float> {
CUTLASS_HOST_DEVICE
float operator()(float lhs) const { return fabs(lhs); }
};
template <typename T>
struct plus {
CUTLASS_HOST_DEVICE
T operator()(T lhs, T const &rhs) const {
lhs += rhs;
return lhs;
}
};
template <typename T>
struct minus {
CUTLASS_HOST_DEVICE
T operator()(T lhs, T const &rhs) const {
lhs -= rhs;
return lhs;
}
};
template <typename T>
struct multiplies {
CUTLASS_HOST_DEVICE
T operator()(T lhs, T const &rhs) const {
lhs *= rhs;
return lhs;
}
};
template <typename T>
struct scale {
T const scaling_factor_;
CUTLASS_HOST_DEVICE
scale(float scaling_factor) : scaling_factor_(scaling_factor) {
}
T operator()(T const &rhs) const {
T result = rhs * scaling_factor_;
return result;
}
};
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 530
/// Partial specializations needed when __CUDA_NO_HALF2_OPERATORS__ is set
template<>
struct plus<__half2> {
CUTLASS_HOST_DEVICE
__half2 operator()(__half2 lhs, __half2 const &rhs) const {
return __hadd2(lhs, rhs);
}
};
template<>
struct minus<__half2> {
CUTLASS_HOST_DEVICE
__half2 operator()(__half2 lhs, __half2 const &rhs) const {
return __hsub2(lhs, rhs);
}
};
template<>
struct multiplies<__half2> {
CUTLASS_HOST_DEVICE
__half2 operator()(__half2 lhs, __half2 const &rhs) const {
return __hmul2(lhs, rhs);
}
};
/// Partial specializations needed when __CUDA_NO_HALF_OPERATORS__ is set
template<>
struct plus<__half> {
CUTLASS_HOST_DEVICE
__half operator()(__half lhs, __half const &rhs) const {
return __hadd(lhs, rhs);
}
};
template<>
struct minus<__half> {
CUTLASS_HOST_DEVICE
__half operator()(__half lhs, __half const &rhs) const {
return __hsub(lhs, rhs);
}
};
template<>
struct multiplies<__half> {
CUTLASS_HOST_DEVICE
__half operator()(__half lhs, __half const &rhs) const {
return __hmul(lhs, rhs);
}
};
#endif // defined(__CUDA_ARCH__)
/// Squares with optional conversion
template <typename T, typename Output = T>
struct square {
CUTLASS_HOST_DEVICE
Output operator()(T lhs) const {
multiplies<Output> mul_op;
Output y = Output(lhs);
return mul_op(y, y);
}
};
/// Returns the magnitude squared of an element.
template <typename T, typename Output = T>
struct magnitude_squared {
CUTLASS_HOST_DEVICE
Output operator()(T lhs) const {
multiplies<Output> mul_op;
Output y = Output(lhs);
return mul_op(y, y);
}
};
/// Computes the square of a difference with optional conversion
template <typename T, typename Output = T>
struct square_difference {
CUTLASS_HOST_DEVICE
Output operator()(T lhs, T rhs) const {
multiplies<Output> mul_op;
Output y = Output(lhs) - Output(rhs);
return mul_op(y, y);
}
};
/// Computes the square of a difference with optional conversion
template <typename T, typename Output = T>
struct magnitude_squared_difference {
CUTLASS_HOST_DEVICE
Output operator()(T lhs, T rhs) const {
multiplies<Output> mul_op;
Output y = Output(lhs) - Output(rhs);
return mul_op(y, y);
}
};
// Computes the reciprocal square root
template <typename T>
struct inverse_square_root;
template <>
struct inverse_square_root<float> {
CUTLASS_HOST_DEVICE
float operator()(float const &lhs) const {
#if defined(__CUDA_ARCH__)
return rsqrtf(lhs);
#else
return 1.f / std::sqrt(lhs);
#endif
}
};
template <>
struct inverse_square_root<half_t> {
CUTLASS_HOST_DEVICE
half_t operator()(half_t const &lhs) const {
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ > 520)
auto result = hrsqrt(reinterpret_cast<__half const &>(lhs));
return reinterpret_cast<half_t const &>(result);
#else
return half_t(1.f / std::sqrt(half_t::convert(lhs)));
#endif
}
};
/// Divides
template <typename T>
struct divides {
CUTLASS_HOST_DEVICE
T operator()(T lhs, T const &rhs) const {
lhs /= rhs;
return lhs;
}
};
/// reciprocal_approximate
template <typename T>
struct reciprocal_approximate {
CUTLASS_HOST_DEVICE
T operator()(T lhs) const {
return divides<T>{}(T(1), lhs);
}
};
template <>
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;
}
};
/// Negate
template <typename T>
struct negate {
CUTLASS_HOST_DEVICE
T operator()(T lhs) const {
return -lhs;
}
};
/// Greater equal
template <typename T>
struct greater_equal {
CUTLASS_HOST_DEVICE
bool operator()(T const &lhs, T const &rhs) const {
return (lhs >= rhs);
}
};
/// Greater
template <typename T>
struct greater {
CUTLASS_HOST_DEVICE
bool operator()(T const &lhs, T const &rhs) const {
return (lhs > rhs);
}
};
/// Less equal
template <typename T>
struct less_equal {
CUTLASS_HOST_DEVICE
bool operator()(T const &lhs, T const &rhs) const {
return (lhs <= rhs);
}
};
/// Less
template <typename T>
struct less {
CUTLASS_HOST_DEVICE
bool operator()(T const &lhs, T const &rhs) const {
return (lhs < rhs);
}
};
template <typename T, bool PropagateNaN = false>
struct maximum {
CUTLASS_HOST_DEVICE
T operator()(T const &lhs, T const &rhs) const {
if constexpr (PropagateNaN && cutlass::platform::is_floating_point<T>::value) {
using CUTLASS_CMATH_NAMESPACE :: isnan;
// Call isnan unqualified, so argument-dependent lookup (ADL)
// will find overloads such as cutlass::isnan(half_t).
// Calling ::isnan or std::isnan directly would force
// implicit conversions to float of custom number types
// in the cutlass namespace (e.g., cutlass::half_t).
return lhs > rhs || isnan(lhs) ? lhs : rhs;
}
else {
return (lhs < rhs ? rhs : lhs);
}
}
};
// This is a subclass and not an alias
// in order to work around a known Clang issue,
// where a template template parameter with one template parameter
// does not match classes that take multiple template parameters
// but have defaults for all but the first.
template<typename T>
struct maximum_with_default_nan_propagation : public maximum<T>
{};
template <>
struct maximum<float, false> {
CUTLASS_HOST_DEVICE
float operator()(float const &lhs, float const &rhs) const {
return fmaxf(lhs, rhs);
}
};
template <>
struct maximum<float, true> {
CUTLASS_HOST_DEVICE
float operator()(float lhs, float rhs) const {
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)
float res;
asm volatile("max.NaN.f32 %0, %1, %2;\n" : "=f"(res) : "f"(lhs), "f"(rhs));
return res;
#else
using CUTLASS_CMATH_NAMESPACE :: isnan;
return lhs > rhs || isnan(lhs) ? lhs : rhs;
#endif
}
};
// This is a subclass and not an alias
// in order to work around a known Clang issue,
// where a template template parameter with one template parameter
// does not match classes that take multiple template parameters
// but have defaults for all but the first.
template <typename T>
struct maximum_with_nan_propagation : maximum<T, true>
{};
// This alias exists for backwards compatibility only.
// Please use the correctly spelled class template above.
template <typename T>
using maximum_with_nan_propogation = maximum_with_nan_propagation<T>;
template <typename T, bool PropagateNaN = false>
struct minimum {
CUTLASS_HOST_DEVICE
T operator()(T const &lhs, T const &rhs) const {
if constexpr (PropagateNaN && cutlass::platform::is_floating_point<T>::value) {
using CUTLASS_CMATH_NAMESPACE :: isnan;
return lhs < rhs || isnan(lhs) ? lhs : rhs;
}
else {
return (rhs < lhs ? rhs : lhs);
}
}
};
template <>
struct minimum<float, false> {
CUTLASS_HOST_DEVICE
float operator()(float const &lhs, float const &rhs) const {
return fminf(lhs, rhs);
}
};
template <>
struct minimum<float, true> {
CUTLASS_HOST_DEVICE
float operator()(float lhs, float rhs) const {
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)
float res;
asm volatile("min.NaN.f32 %0, %1, %2;\n" : "=f"(res) : "f"(lhs), "f"(rhs));
return res;
#else
// No need for ADL; call std::isnan(float) on host and ::isnan(float) on device.
return lhs < rhs || (CUTLASS_CMATH_NAMESPACE :: isnan(lhs)) ? lhs : rhs;
#endif
}
};
template <typename T>
struct minimum_with_nan_propagation : minimum<T, true>
{};
template <typename T, bool PropagateNaN = false>
struct maximum_absolute_value {
CUTLASS_HOST_DEVICE
float operator()(T const &lhs, T const &rhs) const {
absolute_value_op<T> abs_op;
maximum<T, PropagateNaN> max_op;
return max_op(abs_op(lhs), abs_op(rhs));
}
};
// assumes the left operand is already an absolute value
template <typename T, bool PropagateNaN = false>
struct maximum_absolute_value_reduction {
CUTLASS_HOST_DEVICE
float operator()(T const &lhs, T const &rhs) const {
absolute_value_op<T> abs_op;
maximum<T, PropagateNaN> max_op;
return max_op(lhs, abs_op(rhs));
}
};
/// Fused multiply-add
template <typename A, typename B = A, typename C = A>
struct multiply_add {
CUTLASS_HOST_DEVICE
C operator()(A const &a, B const &b, C const &c) const {
return C(a) * C(b) + c;
}
};
template <typename T>
struct square_and_plus {
CUTLASS_HOST_DEVICE
T operator()(T lhs, T const &rhs) const {
multiply_add<T> multiply_add_op;
return multiply_add_op(rhs, rhs, lhs);
}
};
// Fused multiply-add that takes exactly one template parameter.
// This is useful for working around a known Clang issue,
// where a template template parameter with one template parameter
// does not match classes that take multiple template parameters
// but have defaults for all but the first.
template <typename A>
struct homogeneous_multiply_add : public multiply_add<A, A, A>
{};
/// Fused multiply-add
template <typename A, typename B = A, typename C = A>
struct multiply_add_relu0 {
CUTLASS_HOST_DEVICE
C operator()(A const &a, B const &b, C const &c) const {
maximum<C> mx;
return mx(C(a) * C(b) + c, C(0));
}
};
/// 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 {
using CUTLASS_CMATH_NAMESPACE :: isnan;
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
// Namespace-qualifying isnan as cutlass::isnan saves the compiler
// the trouble of argument-dependent lookup. Calling std::isnan or
// ::isnan here would result in unwanted implicit conversion to float.
if (cutlass::isnan(a) || cutlass::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 {
using CUTLASS_CMATH_NAMESPACE :: isnan;
if (isnan(a) || isnan(b)) {
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 (cutlass::isnan(a) || cutlass::isnan(b)) {
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 {
CUTLASS_HOST_DEVICE
T operator()(T const &a, T const &b, T const &c) const {
return ((a & b) + c);
}
};
/// Fused multiply-add
template <typename T>
struct xor_add {
CUTLASS_HOST_DEVICE
T operator()(T const &a, T const &b, T const &c) const {
return ((a ^ b) + c);
}
};
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& 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;
}
}
};
template <typename T>
struct first {
CUTLASS_HOST_DEVICE
T operator()(T const & first, T const &...) const {
return first;
}
CUTLASS_HOST_DEVICE
T operator()(T const & first) const {
return first;
}
};
/////////////////////////////////////////////////////////////////////////////////////////////////
template <typename T>
struct logical_and {
CUTLASS_HOST_DEVICE
T operator()(T const &a, T const &b) const {
return ((static_cast<bool>(a) && static_cast<bool>(b)) ? T(1) : T());
}
};
template <typename T>
struct logical_or {
CUTLASS_HOST_DEVICE
T operator()(T const &a, T const &b) const {
return ((static_cast<bool>(a) || static_cast<bool>(b)) ? T(1) : T());
}
};
template <typename T>
struct logical_not {
CUTLASS_HOST_DEVICE
T operator()(T const &a) const {
return T(!(a));
}
};
/////////////////////////////////////////////////////////////////////////////////////////////////
template <typename T>
struct bit_and {
CUTLASS_HOST_DEVICE
T operator()(T const &a, T const &b) const {
return a & b;
}
};
template <typename T>
struct bit_or {
CUTLASS_HOST_DEVICE
T operator()(T const &a, T const &b) const {
return a | b;
}
};
template <typename T>
struct bit_not {
CUTLASS_HOST_DEVICE
T operator()(T const &a) const {
return ~a;
}
};
template <typename T>
struct bit_xor {
CUTLASS_HOST_DEVICE
T operator()(T const &a, T const &b) const {
return a ^ b;
}
};
//////////////////////////////////////////////////////////////////////////////////////////////////
/// Atomic reductions
template <typename T>
struct atomic_add
{
CUTLASS_DEVICE
void operator()(T *ptr, const T &data)
{
#if defined(__CUDA_ARCH__)
atomicAdd(ptr, data);
#else
CUTLASS_UNUSED(ptr);
CUTLASS_UNUSED(data);
CUTLASS_NOT_IMPLEMENTED();
#endif
}
};
template<>
struct atomic_add<double>
{
CUTLASS_DEVICE
void operator()(double *ptr, const double &data)
{
#if !defined(__CUDA_ARCH__)
CUTLASS_UNUSED(ptr);
CUTLASS_UNUSED(data);
CUTLASS_NOT_IMPLEMENTED();
#elif (__CUDA_ARCH__ >= 600)
atomicAdd(ptr, data);
#else
// Use CAS loop
unsigned long long int* ptr_int = reinterpret_cast<unsigned long long int*>(ptr);
unsigned long long int old_int = *ptr_int;
unsigned long long int assumed_int;
do {
double update = data + __longlong_as_double(old_int);
assumed_int = old_int;
old_int = atomicCAS(ptr_int, assumed_int, __double_as_longlong(update));
} while (assumed_int != old_int);
#endif // (__CUDA_ARCH__ >= 600)
}
};
template<>
struct atomic_add<half2>
{
CUTLASS_DEVICE
void operator()(half2 *ptr, const half2 &data)
{
#if !defined(__CUDA_ARCH__) || (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 600))
CUTLASS_UNUSED(ptr);
CUTLASS_UNUSED(data);
CUTLASS_NOT_IMPLEMENTED();
#else
// Vector-2 atomic reduction requires .target sm_60 or higher
uint32_t word = reinterpret_cast<const uint32_t&>(data);
asm volatile ("red.gpu.global.add.noftz.f16x2 [%0], %1;\n" : : "l"(ptr), "r"(word));
#endif // (__CUDA_ARCH__ >= 600)
}
};
template <typename T>
using red [[deprecated("use atomic_add instead")]] = atomic_add<T>;
template <typename T>
struct atomic_maximum {
CUTLASS_DEVICE
T operator()(T *ptr, T value) const {
#if defined(__CUDA_ARCH__)
return atomicMax(ptr, value);
#else
CUTLASS_UNUSED(ptr);
CUTLASS_UNUSED(value);
CUTLASS_NOT_IMPLEMENTED();
return 0;
#endif
}
};
template <>
struct atomic_maximum<float> {
CUTLASS_DEVICE
float operator()(float *ptr, float value) const {
#if defined(__CUDA_ARCH__)
// 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
CUTLASS_UNUSED(ptr);
CUTLASS_UNUSED(value);
CUTLASS_NOT_IMPLEMENTED();
return 0;
#endif
}
};
// is_atomic
template <class Fn>
struct is_atomic : platform::false_type {};
template <class T>
struct is_atomic<atomic_add<T>> : platform::true_type {};
template <class T>
struct is_atomic<atomic_maximum<T>> : platform::true_type {};
/////////////////////////////////////////////////////////////////////////////////////////////////
//
// Partial specializations for nvcuda::wmma::fragment<Use, m, n, k, T, Layout>
//
/////////////////////////////////////////////////////////////////////////////////////////////////
#if defined(CUTLASS_ARCH_WMMA_ENABLED)
template<typename Use, int m, int n, int k, typename T, typename Layout>
struct plus<nvcuda::wmma::fragment<Use, m, n, k, T, Layout>>
{
using Fragment = nvcuda::wmma::fragment<Use, m, n, k, T, Layout>;
using ElementType = typename Fragment::element_type;
CUTLASS_HOST_DEVICE
Fragment operator()(Fragment const &lhs, Fragment const &rhs) const
{
Fragment result;
plus<ElementType> scalar_op;
ElementType *result_elts = reinterpret_cast<ElementType*>(&result);
const ElementType *lhs_elts = reinterpret_cast<const ElementType*>(&lhs);
const ElementType *rhs_elts = reinterpret_cast<const ElementType*>(&rhs);
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < Fragment::num_elements; i++) {
result_elts[i] = scalar_op(lhs_elts[i], rhs_elts[i]);
}
return result;
}
};
#endif // defined(CUTLASS_ARCH_WMMA_ENABLED)
/////////////////////////////////////////////////////////////////////////////////////////////////
} // namespace cutlass
/////////////////////////////////////////////////////////////////////////////////////////////////
Reference in New Issue View Git Blame Copy Permalink