CUTLASS 2.10 updates (#622)

Co-authored-by: Aniket Shivam <ashivam@nvidia.com>
This commit is contained in:
ANIKET SHIVAM
2022-09-12 18:26:30 -07:00
committed by GitHub
parent beae168f90
commit e773429f7e
96 changed files with 8365 additions and 1667 deletions

View File

@@ -69,7 +69,7 @@ struct global_load;
/////////////////////////////////////////////////////////////////////////////////////////////////
// The redundant mov PTX instruction is used to enforce the compiler to
// initialize data to zero before ld.global
// keep the initializing code before ld.global
template <typename AccessType>
struct global_load<AccessType,
32

View File

@@ -59,6 +59,38 @@ struct Identity {
}
};
/////////////////////////////////////////////////////////////////////////////////////////////////
template <typename T>
struct LinearCombinationGenericParams {
T alpha; ///< scales accumulators
T beta; ///< scales source tensor
T const *alpha_ptr; ///< pointer to accumulator scalar - if not null, loads it from memory
T const *beta_ptr; ///< pointer to source scalar - if not null, loads it from memory
//
// Methods
//
CUTLASS_HOST_DEVICE
LinearCombinationGenericParams():
alpha(T(1)),
beta(T(0)),
alpha_ptr(nullptr),
beta_ptr(nullptr) { }
CUTLASS_HOST_DEVICE
LinearCombinationGenericParams(
T alpha,
T beta = T(0)
): alpha(alpha), beta(beta), alpha_ptr(nullptr), beta_ptr(nullptr) { }
CUTLASS_HOST_DEVICE
LinearCombinationGenericParams(
T const *alpha_ptr,
T const *beta_ptr = nullptr
): alpha(0), beta(0), alpha_ptr(alpha_ptr), beta_ptr(beta_ptr) { }
};
/////////////////////////////////////////////////////////////////////////////////////////////////
/// ReLu operator - propagates NaNs
@@ -79,6 +111,14 @@ struct ReLu {
return mx(value, T(0));
}
/// Host-constructable parameters structure
using Params = LinearCombinationGenericParams<T>;
CUTLASS_HOST_DEVICE
T operator()(T value, Params const &params_) const {
return this->operator()(value);
}
};
template <typename T, int N>
@@ -96,20 +136,74 @@ struct ReLu<Array<T, N>> {
maximum<Array<T, N> > mx;
return mx(frag, T(0));
}
/// Host-constructable parameters structure
using Params = LinearCombinationGenericParams<T>;
CUTLASS_HOST_DEVICE
Array<T, N> operator()(Array<T, N> const &frag, Params const &params_) const {
return this->operator()(frag);
}
};
// Leaky Relu operator
template <typename T>
struct LeakyReLU {
struct Params: LinearCombinationGenericParams<T> {
T leaky_alpha; ///< leaky_alpha
// Methods
using LinearCombinationGenericParams<T>::LinearCombinationGenericParams;
CUTLASS_HOST_DEVICE
Params():
LinearCombinationGenericParams<T>(),
leaky_alpha(T(1)) {}
CUTLASS_HOST_DEVICE
Params(
T alpha,
T beta,
T leaky_alpha = T(1)
): LinearCombinationGenericParams<T>(alpha, beta), leaky_alpha(leaky_alpha) {}
};
CUTLASS_HOST_DEVICE
T operator()(T const &value, T const & alpha_recip) const {
T res = value > T(0) ? value : value * alpha_recip;
return res;
}
CUTLASS_HOST_DEVICE
T operator()(T const &value, Params const &params_) const {
this->operator()(value, params_.leaky_alpha);
}
};
template <typename T, int N>
struct LeakyReLU<Array<T, N> > {
struct Params: LinearCombinationGenericParams<T> {
T leaky_alpha; ///< leaky_alpha
using LinearCombinationGenericParams<T>::LinearCombinationGenericParams;
// Methods
CUTLASS_HOST_DEVICE
Params():
LinearCombinationGenericParams<T>(),
leaky_alpha(T(1)) {}
CUTLASS_HOST_DEVICE
Params(
T alpha,
T beta,
T leaky_alpha = T(1)
): LinearCombinationGenericParams<T>(alpha, beta), leaky_alpha(leaky_alpha) {}
};
CUTLASS_HOST_DEVICE
Array<T, N> operator()(Array<T, N> const &rhs, T const & alpha_recip) const {
Array<T, N> y;
@@ -122,6 +216,11 @@ struct LeakyReLU<Array<T, N> > {
return y;
}
CUTLASS_HOST_DEVICE
Array<T, N> operator()(Array<T, N> const &rhs, Params const &params_) const {
return this->operator()(rhs, params_.leaky_alpha);
}
};
// Tanh operator
@@ -131,6 +230,13 @@ struct Tanh {
T operator()(T const &scalar) const {
return fast_tanh(scalar);
}
using Params = LinearCombinationGenericParams<T>;
CUTLASS_HOST_DEVICE
T operator()(T const &scalar, Params const &params_) const {
return this->operator()(scalar);
}
};
template <typename T, int N>
@@ -147,6 +253,13 @@ struct Tanh<Array<T, N> > {
return y;
}
using Params = LinearCombinationGenericParams<T>;
CUTLASS_HOST_DEVICE
Array<T, N> operator()(Array<T, N> const &rhs, Params const &params_) const {
return this->operator()(rhs);
}
};
template <int N>
@@ -159,6 +272,13 @@ struct Tanh<Array<half_t, N>> {
return tanh(z);
}
using Params = LinearCombinationGenericParams<T>;
CUTLASS_HOST_DEVICE
Array<T, N> operator()(Array<T, N> const &rhs, Params const &params_) const {
return this->operator()(rhs);
}
};
// Sigmoid operator
@@ -168,6 +288,13 @@ struct Sigmoid {
T operator()(T const &scalar) const {
return T(1) / (T(1) + fast_exp(-scalar));
}
using Params = LinearCombinationGenericParams<T>;
CUTLASS_HOST_DEVICE
T operator()(T const &scalar, Params const &params_) const {
return this->operator()(scalar);
}
};
template <typename T, int N>
@@ -184,6 +311,13 @@ struct Sigmoid<Array<T, N> > {
return y;
}
using Params = LinearCombinationGenericParams<T>;
CUTLASS_HOST_DEVICE
Array<T, N> operator()(Array<T, N> const &rhs, Params const &params_) const {
return this->operator()(rhs);
}
};
template <int N>
@@ -208,6 +342,12 @@ struct Sigmoid<Array<half_t, N>> {
fast_exp(neg(z))));
#endif
}
using Params = LinearCombinationGenericParams<T>;
Array<T, N> operator()(Array<T, N> const &z, Params const &params_) const {
return this->operator()(z);
}
};
// SiLu (swish) operator introduced by Elfwing et al. in the following paper
@@ -222,6 +362,13 @@ struct SiLu {
Sigmoid<T> sigmoid;
return scalar * sigmoid(scalar);
}
using Params = LinearCombinationGenericParams<T>;
CUTLASS_HOST_DEVICE
T operator()(T const &scalar, Params const &params_) const {
return this->operator()(scalar);
}
};
template <typename T, int N>
@@ -232,6 +379,13 @@ struct SiLu<Array<T, N>> {
multiplies<Array<T, N>> mul;
return mul(rhs, sigmoid_op(rhs));
}
using Params = LinearCombinationGenericParams<T>;
CUTLASS_HOST_DEVICE
Array<T, N> operator()(Array<T, N> const &rhs, Params const &params_) const {
return this->operator()(rhs);
}
};
// Hardswish operator introduced by Howard et al. in the following paper
@@ -248,6 +402,13 @@ struct HardSwish {
T relu6 = mn(mx(x + T(3), T(0)), T(6));
return x * relu6 / T(6);
}
using Params = LinearCombinationGenericParams<T>;
CUTLASS_HOST_DEVICE
T operator()(T const &x, Params const &params_) const {
return this->operator()(x);
}
};
template <>
@@ -261,6 +422,13 @@ struct HardSwish<float> {
T relu6 = mn(mx(x + T(3), T(0)), T(6));
return x * relu6 * 0.16666667f;
}
using Params = LinearCombinationGenericParams<T>;
CUTLASS_HOST_DEVICE
T operator()(T const &x, Params const &params_) const {
return this->operator()(x);
}
};
template <typename T, int N>
@@ -277,6 +445,13 @@ struct HardSwish<Array<T, N> > {
return y;
}
using Params = LinearCombinationGenericParams<T>;
CUTLASS_HOST_DEVICE
Array<T, N> operator()(Array<T, N> const &x, Params const &params_) const {
return this->operator()(x);
}
};
template <int N>
@@ -292,6 +467,13 @@ struct HardSwish<Array<half_t, N> > {
return mul(mul(mn(mx(add(rhs, T(3)), T(0)), T(6)), rhs), T(0.16666667f));
}
using Params = LinearCombinationGenericParams<T>;
CUTLASS_HOST_DEVICE
Array<T, N> operator()(Array<T, N> const &x, Params const &params_) const {
return this->operator()(x);
}
};
//
@@ -311,6 +493,13 @@ struct GELU {
return T(cutlass::constants::half<T>() * scalar *
(cutlass::constants::one<T>() + (T)erff((float)(scalar / cutlass::constants::root_two<T>()))));
}
using Params = LinearCombinationGenericParams<T>;
CUTLASS_HOST_DEVICE
T operator()(T const &scalar, Params const &params_) const {
return this->operator()(scalar);
}
};
template <>
@@ -320,6 +509,13 @@ struct GELU<float> {
return cutlass::constants::half<float>() * scalar *
(cutlass::constants::one<float>() + erff( scalar / cutlass::constants::root_two<float>() ));
}
using Params = LinearCombinationGenericParams<float>;
CUTLASS_HOST_DEVICE
float operator()(float const &scalar, Params const &params_) const {
return this->operator()(scalar);
}
};
template <>
@@ -329,6 +525,13 @@ struct GELU<double> {
return cutlass::constants::half<double>() * scalar *
(cutlass::constants::one<double>() + erf( scalar / cutlass::constants::root_two<double>() ));
}
using Params = LinearCombinationGenericParams<double>;
CUTLASS_HOST_DEVICE
double operator()(double const &scalar, Params const &params_) const {
return this->operator()(scalar);
}
};
template <typename T, int N>
@@ -345,6 +548,13 @@ struct GELU<Array<T, N> > {
return y;
}
using Params = LinearCombinationGenericParams<T>;
CUTLASS_HOST_DEVICE
Array<T, N> operator()(Array<T, N> const &rhs, Params const &params_) const {
return this->operator()(rhs);
}
};
// GELU operator implemented using the Taylor series approximation
@@ -360,6 +570,9 @@ struct GELU_taylor {
return T(cutlass::constants::half<T>() * z *
(cutlass::constants::one<T>() + fast_tanh(k0 * z * (cutlass::constants::one<T>() + k1 * z * z))));
}
using Params = LinearCombinationGenericParams<T>;
};
template <int N>
@@ -386,6 +599,8 @@ struct GELU_taylor<Array<half_t, N> > {
return y;
}
using Params = LinearCombinationGenericParams<half_t>;
};
template <typename T, int N>
@@ -403,6 +618,8 @@ struct GELU_taylor<Array<T, N> > {
return y;
}
using Params = LinearCombinationGenericParams<T>;
};
/// Computes backwards pass for GELU operator assuming d_t is the layer gradient and

View File

@@ -40,6 +40,7 @@
#include "cutlass/functional.h"
#include "cutlass/numeric_conversion.h"
#include "cutlass/epilogue/thread/scale_type.h"
#include "cutlass/epilogue/thread/linear_combination_params.h"
/////////////////////////////////////////////////////////////////////////////////////////////////
@@ -76,22 +77,23 @@ public:
using FragmentAccumulator = Array<ElementAccumulator, kCount>;
using ComputeFragment = Array<ElementCompute, kCount>;
using ParamsBase = LinearCombinationParams;
static FloatRoundStyle const kRound = Round;
/// Host-constructable parameters structure
struct Params {
struct Params : ParamsBase{
ElementCompute alpha; ///< scales accumulators
ElementCompute beta; ///< scales source tensor
ElementCompute const *alpha_ptr; ///< pointer to accumulator scalar - if not null, loads it from memory
ElementCompute const *beta_ptr; ///< pointer to source scalar - if not null, loads it from memory
//
// Methods
//
CUTLASS_HOST_DEVICE
Params():
ParamsBase(
ElementCompute(1),
ElementCompute(0)
),
alpha(ElementCompute(1)),
beta(ElementCompute(0)),
alpha_ptr(nullptr),
@@ -101,30 +103,43 @@ public:
Params(
ElementCompute alpha,
ElementCompute beta
): alpha(alpha), beta(beta), alpha_ptr(nullptr), beta_ptr(nullptr) {
}
):
ParamsBase(alpha, beta),
alpha(alpha), beta(beta), alpha_ptr(nullptr), beta_ptr(nullptr) { }
CUTLASS_HOST_DEVICE
Params(
ElementCompute alpha
): alpha(alpha), beta(0), alpha_ptr(nullptr), beta_ptr(nullptr) {
}
):
ParamsBase(alpha, ElementCompute(0)),
alpha(alpha), beta(0), alpha_ptr(nullptr), beta_ptr(nullptr) { }
CUTLASS_HOST_DEVICE
Params(
ElementCompute const *alpha_ptr,
ElementCompute const *beta_ptr
): alpha(0), beta(0), alpha_ptr(alpha_ptr), beta_ptr(beta_ptr) {
}
):
ParamsBase(*alpha_ptr, *beta_ptr),
alpha(0), beta(0), alpha_ptr(alpha_ptr), beta_ptr(beta_ptr) { }
CUTLASS_HOST_DEVICE
Params(
ElementCompute const *alpha_ptr
): alpha(0), beta(0), alpha_ptr(alpha_ptr), beta_ptr(nullptr) {
):
ParamsBase(*alpha_ptr, ElementCompute(0)),
alpha(0), beta(0), alpha_ptr(alpha_ptr), beta_ptr(nullptr) { }
CUTLASS_HOST_DEVICE
Params(
ParamsBase const& base
): ParamsBase(base), alpha_ptr(nullptr), beta_ptr(nullptr) {
#if defined(__CUDA_ARCH__)
alpha = reinterpret_cast<ElementCompute const&>(base.alpha_data);
beta = reinterpret_cast<ElementCompute const&>(base.beta_data);
#else
memcpy( alpha, base.alpha_data, sizeof(ElementCompute) );
memcpy( beta, base.alpha_data, sizeof(ElementCompute) );
#endif
}
};
@@ -142,7 +157,6 @@ public:
/// Constructs the function object, possibly loading from pointers in host memory
CUTLASS_HOST_DEVICE
LinearCombination(Params const &params) {
alpha_ = (params.alpha_ptr ? *params.alpha_ptr : params.alpha);
beta_ = (params.beta_ptr ? *params.beta_ptr : params.beta);
}

View File

@@ -83,40 +83,7 @@ public:
static FloatRoundStyle const kRound = Round;
/// Host-constructable parameters structure
struct Params {
ElementCompute alpha; ///< scales accumulators
ElementCompute beta; ///< scales source tensor
ElementCompute const *alpha_ptr; ///< pointer to accumulator scalar - if not null, loads it from memory
ElementCompute const *beta_ptr; ///< pointer to source scalar - if not null, loads it from memory
//
// Methods
//
CUTLASS_HOST_DEVICE
Params():
alpha(ElementCompute(1)),
beta(ElementCompute(0)),
alpha_ptr(nullptr),
beta_ptr(nullptr) { }
CUTLASS_HOST_DEVICE
Params(
ElementCompute alpha,
ElementCompute beta = ElementCompute(0)
): alpha(alpha), beta(beta), alpha_ptr(nullptr), beta_ptr(nullptr) {
}
CUTLASS_HOST_DEVICE
Params(
ElementCompute const *alpha_ptr,
ElementCompute const *beta_ptr = nullptr
): alpha(0), beta(0), alpha_ptr(alpha_ptr), beta_ptr(beta_ptr) {
}
};
using Params = typename ActivationFunctor<FragmentCompute>::Params;
private:
@@ -124,8 +91,7 @@ private:
// Data members
//
ElementCompute alpha_;
ElementCompute beta_;
Params params_;
bool skip_elementwise_;
public:
@@ -133,9 +99,9 @@ public:
/// Constructs the function object, possibly loading from pointers in host memory
CUTLASS_HOST_DEVICE
LinearCombinationGeneric(Params const &params) {
alpha_ = (params.alpha_ptr ? *params.alpha_ptr : params.alpha);
beta_ = (params.beta_ptr ? *params.beta_ptr : params.beta);
params_ = params;
params_.alpha = (params.alpha_ptr ? *params.alpha_ptr : params.alpha);
params_.beta = (params.beta_ptr ? *params.beta_ptr : params.beta);
skip_elementwise_ = false;
}
@@ -148,14 +114,14 @@ public:
if (Scale == ScaleType::Nothing) return false;
return beta_ != ElementCompute(0);
return params_.beta != ElementCompute(0);
}
/// Functionally required for serial reduction in the epilogue
CUTLASS_HOST_DEVICE
void set_k_partition(int k_partition, int k_partition_count) {
if (k_partition) {
beta_ = ElementCompute(1);
params_.beta = ElementCompute(1);
}
if (k_partition != k_partition_count - 1) {
@@ -186,15 +152,15 @@ public:
if (Scale == ScaleType::NoBetaScaling) {
intermediate = converted_source;
intermediate = mul_add_accumulator(alpha_, converted_accumulator, intermediate); // D = alpha * Accum + X
intermediate = mul_add_accumulator(params_.alpha, converted_accumulator, intermediate); // D = alpha * Accum + X
} else if (Scale == ScaleType::Nothing) {
intermediate = converted_accumulator;
} else {
intermediate = mul_add_source(beta_, converted_source); // X = beta * C + uniform
intermediate = mul_add_accumulator(alpha_, converted_accumulator, intermediate); // D = alpha * Accum + X
intermediate = mul_add_source(params_.beta, converted_source); // X = beta * C + uniform
intermediate = mul_add_accumulator(params_.alpha, converted_accumulator, intermediate); // D = alpha * Accum + X
}
intermediate = skip_elementwise_ ? intermediate : activation(intermediate);
intermediate = skip_elementwise_ ? intermediate : activation(intermediate, params_);
// Convert to destination numeric type
NumericArrayConverter<ElementOutput, ElementCompute, kCount, Round> destination_converter;
@@ -222,10 +188,10 @@ public:
if (Scale == ScaleType::Nothing) {
intermediate = converted_accumulator;
} else {
intermediate = mul_add_accumulator(alpha_, converted_accumulator); // D = alpha * Accum
intermediate = mul_add_accumulator(params_.alpha, converted_accumulator); // D = alpha * Accum
}
intermediate = skip_elementwise_ ? intermediate : activation(intermediate);
intermediate = skip_elementwise_ ? intermediate : activation(intermediate, params_);
// Convert to destination numeric type
NumericArrayConverter<ElementOutput, ElementCompute, kCount, Round> destination_converter;

View File

@@ -0,0 +1,75 @@
/***************************************************************************************************
* Copyright (c) 2017 - 2022 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
*/
#pragma once
/////////////////////////////////////////////////////////////////////////////////////////////////
namespace cutlass {
namespace epilogue {
namespace thread {
/////////////////////////////////////////////////////////////////////////////////////////////////
struct LinearCombinationParams {
uint64_t alpha_data[2];
uint64_t beta_data[2];
CUTLASS_HOST_DEVICE
LinearCombinationParams()
: alpha_data {0lu, 0lu}, beta_data {0lu, 0lu}
{ }
template <typename ElementCompute>
CUTLASS_HOST_DEVICE
LinearCombinationParams(ElementCompute alpha, ElementCompute beta)
: alpha_data {0lu, 0lu}, beta_data {0lu, 0lu}
{
#if defined(__CUDA_ARCH__)
reinterpret_cast<ElementCompute&>(alpha_data) = alpha;
reinterpret_cast<ElementCompute&>(beta_data) = beta;
#else
memcpy( alpha_data, &alpha, sizeof(ElementCompute) );
memcpy( beta_data, &beta, sizeof(ElementCompute) );
#endif
}
};
/////////////////////////////////////////////////////////////////////////////////////////////////
} // namespace thread
} // namespace epilogue
} // namespace cutlass
/////////////////////////////////////////////////////////////////////////////////////////////////

View File

@@ -0,0 +1,156 @@
/***************************************************************************************************
* Copyright (c) 2017 - 2022 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
*/
#pragma once
#include "cutlass/cutlass.h"
#include "cutlass/layout/matrix.h"
#include "cutlass/fast_math.h"
/////////////////////////////////////////////////////////////////////////////////////////////////
namespace cutlass {
namespace epilogue {
namespace threadblock {
/////////////////////////////////////////////////////////////////////////////////////////////////
template <
int Rank
>
struct PredicatedTileIteratorAffineLayoutRankNParams {
using Layout = layout::AffineRankN<Rank>;
using TensorCoord = typename Layout::TensorCoord;
static bool const kBigEndian = false;
//
// Data members
//
Layout layout;
/// Stride in units of bytes along M modes
Coord<Layout::kRank/2, typename Layout::LongIndex> stride_m;
/// Stride in units of bytes along N modes
Coord<Layout::kRank/2, typename Layout::LongIndex> stride_n;
/// Fast divmod objects divided by tensor extents
FastDivmod divmod_m[(Layout::kRank == 2) ? 1 : (Layout::kRank/2 - 1)];
/// Fast divmod objects divided by tensor extents
FastDivmod divmod_n[(Layout::kRank == 2) ? 1 : (Layout::kRank/2 - 1)];
int64_t rank2_inc_col;
int64_t rank2_inc_row;
//
// Methods
//
CUTLASS_HOST_DEVICE
PredicatedTileIteratorAffineLayoutRankNParams() { }
CUTLASS_HOST_DEVICE
PredicatedTileIteratorAffineLayoutRankNParams(TensorCoord const &extent,
Layout const &layout_,
int64_t element_sizeof_bits)
: layout(layout_)
{
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < Layout::kRank / 2; ++i) {
stride_m[i] = OffsetBytes(layout_.stride()[i], element_sizeof_bits);
stride_n[i] = OffsetBytes(layout_.stride()[i + Layout::kRank / 2], element_sizeof_bits);
}
if (kBigEndian) {
// "Big Endian" scheme
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < Layout::kRank / 2 - 1; ++i) {
divmod_m[i] = FastDivmod(extent[i + 1]);
divmod_n[i] = FastDivmod(extent[i + Layout::kRank / 2 + 1]);
}
}
else {
// "Little Endian" scheme
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < Layout::kRank / 2 - 1; ++i) {
divmod_m[i] = FastDivmod(extent[i]);
divmod_n[i] = FastDivmod(extent[i + Layout::kRank / 2]);
}
}
#if 0
//
// Debug print statements to verify extents and strides are passed correctly.
//
printf("PredicatedTileIteratorAffine::Params() entered\n");
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < Layout::kRank; ++i) {
printf(" extent[%d]: %d\n", i, extent[i]);
}
for (int i = 0; i < Layout::kRank; ++i) {
printf(" stride[%d]: %ld\n", i, layout_.stride()[i]);
}
printf("PredicatedTileIteratorAffine::Params() returning\n");
#endif
}
CUTLASS_HOST_DEVICE
PredicatedTileIteratorAffineLayoutRankNParams(Layout const &layout_,
int32_t threadmap_delta_kColumn,
int32_t threadmap_delta_kRow,
int64_t element_sizeof_bits)
: layout(layout_)
{
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < Layout::kRank / 2; ++i) {
stride_m[i] = OffsetBytes(layout_.stride()[i], element_sizeof_bits);
stride_n[i] = OffsetBytes(layout_.stride()[i + Layout::kRank / 2], element_sizeof_bits);
}
rank2_inc_col = threadmap_delta_kColumn * stride_n[0];
rank2_inc_row = threadmap_delta_kRow * stride_m[0];
}
};
/////////////////////////////////////////////////////////////////////////////////////////////////
} // namespace threadblock
} // namespace epilogue
} // namespace cutlass
////////////////////////////////////////////////////////////////////////////////

View File

@@ -561,6 +561,17 @@ CUTLASS_HOST_DEVICE int64_t OffsetBytes(int64_t index) {
}
}
CUTLASS_HOST_DEVICE int64_t OffsetBytes(int64_t index, int64_t element_sizeof_bits) {
if (element_sizeof_bits >= 8) {
return index * (element_sizeof_bits / 8);
}
else {
int64_t const kElementsPerByte = ((8 / element_sizeof_bits) + ((element_sizeof_bits >= 8) ? 1 : 0));
return index / kElementsPerByte;
}
}
/////////////////////////////////////////////////////////////////////////////////////////////////
// Min/Max
/////////////////////////////////////////////////////////////////////////////////////////////////

View File

@@ -88,12 +88,12 @@ public:
using ElementA = typename MapArguments::ElementA;
using LayoutA = typename MapArguments::LayoutA;
static ComplexTransform const kTransformA = MapArguments::kTransformA;
static int const kAlignmentA = GemmKernel::kAlignmentA;
static int const kAlignmentA = MapArguments::kAlignmentA;
using ElementB = typename MapArguments::ElementB;
using LayoutB = typename MapArguments::LayoutB;
static ComplexTransform const kTransformB = MapArguments::kTransformB;
static int const kAlignmentB = GemmKernel::kAlignmentB;
static int const kAlignmentB = MapArguments::kAlignmentB;
using ElementC = typename GemmKernel::ElementC;
using LayoutC = typename MapArguments::LayoutC;

View File

@@ -40,7 +40,7 @@
#include "cutlass/cutlass.h"
#include "cutlass/numeric_types.h"
#include "cutlass/transform/thread/unaryOp.h"
#include "cutlass/transform/thread/unary_op.h"
#include "cutlass/array.h"
#include "cutlass/half.h"
@@ -1273,6 +1273,40 @@ struct NumericArrayConverter<uint8_t, int, N, Round> {
/////////////////////////////////////////////////////////////////////////////////////////////////
/// Partial specialization for Array<int8_t> <= Array<float>
/// Conversion is performed with saturation regardless of setting of
/// the `Round` template parameter.
template <
int N,
FloatRoundStyle Round
>
struct NumericArrayConverter<int8_t, float, N, Round> {
using result_type = Array<int8_t, N>;
using source_type = Array<float, N>;
static FloatRoundStyle const round_style = Round;
CUTLASS_HOST_DEVICE
static result_type convert(source_type const & source) {
// Convert float to int
Array<int32_t, N> temporary;
NumericArrayConverter<int, float, N, Round> compute_converter;
temporary = compute_converter(source);
// Convert to int to int8_t
NumericArrayConverter<int8_t, int32_t, N, Round> destination_converter;
return destination_converter(temporary);
}
CUTLASS_HOST_DEVICE
result_type operator()(source_type const &s) {
return convert(s);
}
};
/////////////////////////////////////////////////////////////////////////////////////////////////
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 750) && \
((__CUDACC_VER_MAJOR__ > 10) || \
((__CUDACC_VER_MAJOR__ >= 10) && (__CUDACC_VER_MINOR__ >= 2)))