CUTLASS 3.5.0 (#1411)

This commit is contained in:
Vijay Thakkar
2024-03-19 17:51:04 -04:00
committed by GitHub
parent ffa34e7075
commit 629f4653c3
468 changed files with 48729 additions and 7252 deletions
@@ -121,7 +121,7 @@ struct CommandLine {
* Returns the commandline parameter for a given index (not including flags)
*/
template <typename value_t>
void get_cmd_line_argument(int index, value_t& val) const {
void get_cmd_line_argument(size_t index, value_t& val) const {
using namespace std;
if (index < args.size()) {
istringstream str_stream(args[index]);
@@ -63,7 +63,7 @@ using ComplexDouble = cuda::std::complex<double>;
// User could potentially define Half instead of cute::
#ifndef BLAM_HALF_TYPE
#define BLAM_HALF_TYPE 1
#include <cute/numeric/half.hpp>
#include <cute/numeric/numeric_types.hpp>
namespace blam {
using Half = cute::half_t;
}
@@ -69,7 +69,7 @@ CUTLASS_DEVICE void dump_fragment(Fragment const& frag, int N = 0, int M = 0,
return;
}
int total_elements = frag.size();
int total_elements = int(frag.size());
if (M < 0 || M > total_elements) {
if (thread_id == 0 && block_id == 0)
@@ -42,8 +42,8 @@
namespace cutlass {
__global__ void rmsnorm_twoPassAlgo_e8(float4 *output, const float4 *input,
const float4 *weight,
const int m, const int n, float epsilon) {
const float4 *weight,
const int m, const int n, float epsilon) {
const int m_idx = blockIdx.x;
const int tid = threadIdx.x;
const int bdimx = blockDim.x;
@@ -115,9 +115,9 @@ __global__ void rmsnorm_twoPassAlgo_e8(float4 *output, const float4 *input,
template<typename T>
__global__ void rmsnorm_twoPassAlgo_e1(T* output,
const T* input,
const T* weight,
const int m, const int n,
const T* input,
const T* weight,
const int m, const int n,
float epsilon)
{
const int m_idx = blockIdx.x;
@@ -156,7 +156,7 @@ void rmsnorm(cutlass::MatrixCoord tensor_size,
TensorRef<T, layout::RowMajor> ref_output,
TensorRef<T, layout::RowMajor> ref_input,
TensorRef<T, layout::RowMajor> ref_weight,
cudaStream_t stream, float epsilon = 1e-5){
cudaStream_t stream, float epsilon = 1e-5f){
const int m = tensor_size.row();
const int n = tensor_size.column();
T* output = ref_output.data();
@@ -112,9 +112,13 @@ public:
/// Example
/// int2: kBitsStoredVec = 8; kElementsPerStoredVec = 4; kNumStoragePerStoredVec = 1 uint8_t;
/// int4: kBitsStoredVec = 8; kElementsPerStoredVec = 2; kNumStoragePerStoredVec = 1 uint8_t;
static int const kBitsStoredVec = (sizeof_bits<Element>::value < 8) ? cutlass::lcm(static_cast<int>(sizeof_bits<Element>::value), 8) : sizeof_bits<Element>::value;
static int const kElementsPerStoredVec = kBitsStoredVec / sizeof_bits<Element>::value;
static int const kNumStoragePerStoredVec = kBitsStoredVec / (sizeof(Element) * 8);
static constexpr int kBitsStoredVec = (sizeof_bits<Element>::value < 8) ? cutlass::lcm(sizeof_bits<Element>::value, 8) : sizeof_bits<Element>::value;
static constexpr int kElementsPerStoredVec = kBitsStoredVec / sizeof_bits<Element>::value;
static constexpr int kNumStoragePerStoredVec = kBitsStoredVec / (sizeof(Element) * 8);
static_assert(kBitsStoredVec != 0, "kBitsStoredVec can not be zero");
static_assert(kElementsPerStoredVec != 0, "kElementsPerStoredVec can not be zero");
static_assert(kNumStoragePerStoredVec != 0, "kNumStoragePerStoredVec can not be zero");
private:
@@ -108,4 +108,354 @@ make_cute_packed_stride(cute::Stride<cute::Int<1>, IntT, int64_t> s, cute::Shape
/////////////////////////////////////////////////////////////////////////////////////////////////
// Strides for convolutions
// Output cutlass::layout::TensorNDHWC -> rank-3 stride (InT,_1,_0)
// Note: For fprop/dgrad kernel, strides are assumed to be layout right in NZPQK/NDHWC order
// and therefore can be coalesced to just q/w. For wgrad kernel, strides are assumed to be layout
// right in KTRSC order and can be coalesced to just k.
// We enforce this condition here with asserts.
template <class IntT, size_t RankT_>
cute::Stride<IntT, cute::Int<1>, cute::Int<0>>
make_cute_packed_stride(
cute::Stride<IntT, cute::Int<1>, cute::Int<0>> s,
cute::array<int32_t, RankT_> shape_output,
cute::array<IntT, RankT_> stride_output,
cutlass::conv::Operator conv_op) {
static_assert(std::is_integral_v<IntT>,
"Stride must have an integral type so it can be set dynamically. Static strides not supported.");
static_assert(RankT_ >= 3u);
constexpr static int RankT = static_cast<int>(RankT_);
assert(stride_output[RankT-1] == 1);
cute::for_each(cute::make_seq<RankT-2>{}, [&](auto i) {
assert(stride_output[i] == shape_output[i+1] * stride_output[i+1]);
});
auto s_copy = s;
cute::get<0>(s_copy) = (conv_op == cutlass::conv::Operator::kWgrad) ?
stride_output[0] :
stride_output[RankT-2];
return s_copy;
}
//
// Activation tensor ((w, h, d, n), _1) for fprop kernel
//
// Activation cutlass::layout::TensorNWC -> rank-2 stride ((W,N),_1)
template <class IntT>
cute::Stride<cute::Stride<IntT, IntT>, cute::Int<1>>
make_cute_packed_stride(
cute::Stride<cute::Stride<IntT, IntT>, cute::Int<1>> s,
cute::array<IntT, 3> stride_nwc,
conv::Operator ConvOp) {
static_assert(std::is_integral_v<IntT>,
"Stride must have an integral type so it can be set dynamically. Static strides not supported.");
assert(stride_nwc[2] == 1);
auto s_copy = s;
cute::get<0,0>(s_copy) = stride_nwc[1];
cute::get<0,1>(s_copy) = stride_nwc[0];
return s_copy;
}
// Activation cutlass::layout::TensorNHWC -> rank-2 stride ((W,H,N),_1)
template <class IntT>
cute::Stride<cute::Stride<IntT, IntT, IntT>, cute::Int<1>>
make_cute_packed_stride(
cute::Stride<cute::Stride<IntT, IntT, IntT>, cute::Int<1>> s,
cute::array<IntT, 4> stride_nhwc,
conv::Operator ConvOp) {
static_assert(std::is_integral_v<IntT>,
"Stride must have an integral type so it can be set dynamically. Static strides not supported.");
assert(stride_nhwc[3] == 1);
auto s_copy = s;
cute::for_each(cute::make_seq<3>{}, [&](auto i) {
cute::get<0,i>(s_copy) = stride_nhwc[2-i];
});
return s_copy;
}
// Activation cutlass::layout::TensorNDHWC -> rank-2 stride ((W,H,D,N),_1)
template <class IntT>
cute::Stride<cute::Stride<IntT, IntT, IntT, IntT>, cute::Int<1>>
make_cute_packed_stride(
cute::Stride<cute::Stride<IntT, IntT, IntT, IntT>, cute::Int<1>> s,
cute::array<IntT, 5> stride_ndhwc,
conv::Operator ConvOp) {
static_assert(std::is_integral_v<IntT>,
"Stride must have an integral type so it can be set dynamically. Static strides not supported.");
assert(stride_ndhwc[4] == 1);
auto s_copy = s;
cute::for_each(cute::make_seq<4>{}, [&](auto i) {
cute::get<0,i>(s_copy) = stride_ndhwc[3-i];
});
return s_copy;
}
//
// Filter tensor (k, (_1, s, r, t)) for fprop kernel
//
// Filter cutlass::layout::TensorNWC -> rank-2 stride (k, (_1, s))
template <class IntT>
cute::Stride<IntT, cute::Stride<cute::Int<1>, IntT>>
make_cute_packed_stride(
cute::Stride<IntT, cute::Stride<cute::Int<1>, IntT>> s,
cute::array<IntT, 3> stride_ksc,
conv::Operator ConvOp) {
static_assert(std::is_integral_v<IntT>,
"Stride must have an integral type so it can be set dynamically. Static strides not supported.");
assert(stride_ksc[2] == 1);
auto s_copy = s;
cute::get<0,0>(s_copy) = stride_ksc[0];
cute::get<1,1>(s_copy) = stride_ksc[1];
return s_copy;
}
// Filter cutlass::layout::TensorNHWC -> rank-2 stride (k, (_1, s, r))
template <class IntT>
cute::Stride<IntT, cute::Stride<cute::Int<1>, IntT, IntT>>
make_cute_packed_stride(
cute::Stride<IntT, cute::Stride<cute::Int<1>, IntT, IntT>> s,
cute::array<IntT, 4> stride_krsc,
conv::Operator ConvOp) {
static_assert(std::is_integral_v<IntT>,
"Stride must have an integral type so it can be set dynamically. Static strides not supported.");
assert(stride_krsc[3] == 1);
auto s_copy = s;
cute::get<0,0>(s_copy) = stride_krsc[0];
cute::for_each(cute::make_seq<2>{}, [&](auto i) {
cute::get<1,2-i>(s_copy) = stride_krsc[i+1];
});
return s_copy;
}
// Filter cutlass::layout::TensorNDHWC -> rank-2 stride (k, (_1, s, r, t))
template <class IntT>
cute::Stride<IntT, cute::Stride<cute::Int<1>, IntT, IntT, IntT>>
make_cute_packed_stride(
cute::Stride<IntT, cute::Stride<cute::Int<1>, IntT, IntT, IntT>> s,
cute::array<IntT, 5> stride_ktrsc,
conv::Operator ConvOp) {
static_assert(std::is_integral_v<IntT>,
"Stride must have an integral type so it can be set dynamically. Static strides not supported.");
assert(stride_ktrsc[4] == 1);
auto s_copy = s;
cute::get<0,0>(s_copy) = stride_ktrsc[0];
cute::for_each(cute::make_seq<3>{}, [&](auto i) {
cute::get<1,3-i>(s_copy) = stride_ktrsc[i+1];
});
return s_copy;
}
//
// Activation tensor (_1, (w, h, d, n)) for wgrad kernel
//
// It is also Filter tensor ((_1), (k, s, r, t)) for dgrad kernel
//
// Activation cutlass::layout::TensorNWC -> rank-2 stride (_1, (W,N)) in wgrad
// Filter cutlass::layout::TensorNWC -> rank-2 stride ((_1), (k, s)) in dgrad
template <class IntT>
cute::Stride<cute::Int<1>, cute::Stride<IntT, IntT>>
make_cute_packed_stride(
cute::Stride<cute::Int<1>, cute::Stride<IntT, IntT>> s,
cute::array<IntT, 3> stride_nwc,
conv::Operator ConvOp) {
static_assert(std::is_integral_v<IntT>,
"Stride must have an integral type so it can be set dynamically. Static strides not supported.");
assert(stride_nwc[2] == 1);
auto s_copy = s;
if (ConvOp == cutlass::conv::Operator::kWgrad) {
cute::get<1,0>(s_copy) = stride_nwc[1];
cute::get<1,1>(s_copy) = stride_nwc[0];
}
else if (ConvOp == cutlass::conv::Operator::kDgrad) {
// stride_nwc in dgrad is ksc.
cute::get<1,0>(s_copy) = stride_nwc[0];
cute::get<1,1>(s_copy) = stride_nwc[1];
}
return s_copy;
}
// Activation cutlass::layout::TensorNHWC -> rank-2 stride (_1, (W,H,N)) in wgrad
// Filter cutlass::layout::TensorNHWC -> rank-2 stride ((_1), (k, s, r)) in dgrad
template <class IntT>
cute::Stride<cute::Int<1>, cute::Stride<IntT, IntT, IntT>>
make_cute_packed_stride(
cute::Stride<cute::Int<1>, cute::Stride<IntT, IntT, IntT>> s,
cute::array<IntT, 4> stride_nhwc,
conv::Operator ConvOp) {
static_assert(std::is_integral_v<IntT>,
"Stride must have an integral type so it can be set dynamically. Static strides not supported.");
assert(stride_nhwc[3] == 1);
auto s_copy = s;
if (ConvOp == cutlass::conv::Operator::kWgrad) {
cute::for_each(cute::make_seq<3>{}, [&](auto i) {
cute::get<1,i>(s_copy) = stride_nhwc[2-i];
});
}
else if (ConvOp == cutlass::conv::Operator::kDgrad) {
// stride_nhwc in dgrad is krsc.
cute::get<1,0>(s_copy) = stride_nhwc[0];
cute::for_each(cute::make_seq<2>{}, [&](auto i) {
cute::get<1,2-i>(s_copy) = stride_nhwc[i+1];
});
}
return s_copy;
}
// Activation cutlass::layout::TensorNDHWC -> rank-2 stride (_1, (W,H,D,N)) in wgrad
// Filter cutlass::layout::TensorNDHWC -> rank-2 stride ((_1), (k, s, r, t)) in dgrad
template <class IntT>
cute::Stride<cute::Int<1>, cute::Stride<IntT, IntT, IntT, IntT>>
make_cute_packed_stride(
cute::Stride<cute::Int<1>, cute::Stride<IntT, IntT, IntT, IntT>> s,
cute::array<IntT, 5> stride_ndhwc,
conv::Operator ConvOp) {
static_assert(std::is_integral_v<IntT>,
"Stride must have an integral type so it can be set dynamically. Static strides not supported.");
assert(stride_ndhwc[4] == 1);
auto s_copy = s;
if (ConvOp == cutlass::conv::Operator::kWgrad) {
cute::for_each(cute::make_seq<4>{}, [&](auto i) {
cute::get<1,i>(s_copy) = stride_ndhwc[3-i];
});
}
else if (ConvOp == cutlass::conv::Operator::kDgrad) {
// stride_ndhwc in dgrad is ktrsc.
cute::get<1,0>(s_copy) = stride_ndhwc[0];
cute::for_each(cute::make_seq<3>{}, [&](auto i) {
cute::get<1,3-i>(s_copy) = stride_ndhwc[i+1];
});
}
return s_copy;
}
//
// NZPQ tensor (_1, nzpq) for wgrad kernel
//
// cutlass::layout::TensorNWC -> rank-2 stride (_1, nzpq)
template <class IntT>
cute::Stride<cute::Int<1>, IntT>
make_cute_packed_stride(
cute::Stride<cute::Int<1>, IntT> s,
cute::array<IntT, 3> stride_nqk,
conv::Operator ConvOp) {
static_assert(std::is_integral_v<IntT>,
"Stride must have an integral type so it can be set dynamically. Static strides not supported.");
assert(stride_nqk[2] == 1);
auto s_copy = s;
cute::get<1>(s_copy) = stride_nqk[1];
return s_copy;
}
// cutlass::layout::TensorNHWC -> rank-2 stride (_1, nzpq)
template <class IntT>
cute::Stride<cute::Int<1>, IntT>
make_cute_packed_stride(
cute::Stride<cute::Int<1>, IntT> s,
cute::array<IntT, 4> stride_npqk,
conv::Operator ConvOp) {
static_assert(std::is_integral_v<IntT>,
"Stride must have an integral type so it can be set dynamically. Static strides not supported.");
assert(stride_npqk[3] == 1);
auto s_copy = s;
cute::get<1>(s_copy) = stride_npqk[2];
return s_copy;
}
// cutlass::layout::TensorNDHWC -> rank-2 stride (_1, nzpq)
template <class IntT>
cute::Stride<cute::Int<1>, IntT>
make_cute_packed_stride(
cute::Stride<cute::Int<1>, IntT> s,
cute::array<IntT, 5> stride_nzpqk,
conv::Operator ConvOp) {
static_assert(std::is_integral_v<IntT>,
"Stride must have an integral type so it can be set dynamically. Static strides not supported.");
assert(stride_nzpqk[4] == 1);
auto s_copy = s;
cute::get<1>(s_copy) = stride_nzpqk[3];
return s_copy;
}
//
// Wgrad output tensor (k, (_1, s, r, t), _0)
//
// Filter cutlass::layout::TensorKCS -> rank-3 stride (k, (_1, s), _0)
template <class IntT>
cute::Stride<IntT, cute::Stride<cute::Int<1>, IntT>, cute::Int<0>>
make_cute_packed_stride(
cute::Stride<IntT, cute::Stride<cute::Int<1>, IntT>, cute::Int<0>> s,
[[maybe_unused]] cute::array<int32_t, 3> shape_output,
cute::array<IntT, 3> stride_ksc,
conv::Operator ConvOp) {
static_assert(std::is_integral_v<IntT>,
"Stride must have an integral type so it can be set dynamically. Static strides not supported.");
assert(stride_ksc[2] == 1);
auto s_copy = s;
cute::get<0,0>(s_copy) = stride_ksc[0];
cute::get<1,1>(s_copy) = stride_ksc[1];
return s_copy;
}
// Filter cutlass::layout::TensorKCSR -> rank-3 stride (k, (_1, s, r), _0)
template <class IntT>
cute::Stride<IntT, cute::Stride<cute::Int<1>, IntT, IntT>, cute::Int<0>>
make_cute_packed_stride(
cute::Stride<IntT, cute::Stride<cute::Int<1>, IntT, IntT>, cute::Int<0>> s,
[[maybe_unused]] cute::array<int32_t, 4> shape_output,
cute::array<IntT, 4> stride_krsc,
conv::Operator ConvOp) {
static_assert(std::is_integral_v<IntT>,
"Stride must have an integral type so it can be set dynamically. Static strides not supported.");
assert(stride_krsc[3] == 1);
auto s_copy = s;
cute::get<0,0>(s_copy) = stride_krsc[0];
cute::for_each(cute::make_seq<2>{}, [&](auto i) {
cute::get<1,2-i>(s_copy) = stride_krsc[i+1];
});
return s_copy;
}
// Filter cutlass::layout::TensorKCSRT -> rank-3 stride (k, (_1, s, r, t), _0)
template <class IntT>
cute::Stride<IntT, cute::Stride<cute::Int<1>, IntT, IntT, IntT>, cute::Int<0>>
make_cute_packed_stride(
cute::Stride<IntT, cute::Stride<cute::Int<1>, IntT, IntT, IntT>, cute::Int<0>> s,
[[maybe_unused]] cute::array<int32_t, 5> shape_output,
cute::array<IntT, 5> stride_ktrsc,
conv::Operator ConvOp) {
static_assert(std::is_integral_v<IntT>,
"Stride must have an integral type so it can be set dynamically. Static strides not supported.");
assert(stride_ktrsc[4] == 1);
auto s_copy = s;
cute::get<0,0>(s_copy) = stride_ktrsc[0];
cute::for_each(cute::make_seq<3>{}, [&](auto i) {
cute::get<1,3-i>(s_copy) = stride_ktrsc[i+1];
});
return s_copy;
}
/////////////////////////////////////////////////////////////////////////////////////////////////
} // namespace cutlass
+40 -13
View File
@@ -40,8 +40,9 @@
#include <cute/util/type_traits.hpp>
#include <cute/tensor.hpp>
#include <cute/numeric/half.hpp>
#include <cute/numeric/numeric_types.hpp>
#include <cute/numeric/complex.hpp>
#include <cutlass/layout/layout.h>
// The computed infinity norm does not include
@@ -233,7 +234,8 @@ print_relative_error(
T1 const& data,
T2 const& reference,
bool print_verbose = false,
bool print_error = true) {
bool print_error = true,
double error_margin = 0.00001) {
using std::abs; using std::sqrt;
// Use either double or complex<double> for error computation
@@ -252,8 +254,8 @@ print_relative_error(
double tot_norm_sq = 0;
double tot_ind_rel_err = 0;
double max_ind_rel_err = 0;
for (std::size_t i = 0; i < n; ++i)
{
double max_diff = 0;
for (std::size_t i = 0; i < n; ++i) {
error_type val = data[i];
error_type ref = reference[i];
@@ -267,6 +269,9 @@ print_relative_error(
// Maximum relative error
max_ind_rel_err = std::max(max_ind_rel_err, rel_error);
// Maximum delta in value error
max_diff = std::max(max_diff, diff);
// Total relative error
tot_error_sq += diff * diff;
tot_norm_sq += aref * aref;
@@ -276,18 +281,40 @@ print_relative_error(
}
}
printf("Vector reference norm: [%.5e]\n", sqrt(tot_norm_sq));
double ave_rel_err = tot_ind_rel_err / double(n);
if (print_error) {
printf("Average relative error: %.3e\n", ave_rel_err);
}
if (print_error) {
printf("Maximum relative error: %.3e\n", max_ind_rel_err);
}
if (print_error) {
printf("Maximum difference : %.3e\n", max_diff);
}
double tot_rel_err = sqrt(tot_error_sq/(tot_norm_sq+eps));
if (print_error)
printf("Vector relative error: [%.5e]\n", tot_rel_err);
if (print_error) {
printf("Vector relative error: %.3e\n", tot_rel_err);
}
double ave_rel_err = tot_ind_rel_err / double(n);
if (print_error)
printf("Average relative error: [%.5e]\n", ave_rel_err);
printf("Vector reference norm: %.3e\n", sqrt(tot_norm_sq));
if (print_error)
printf("Maximum relative error: [%.5e]\n", max_ind_rel_err);
return (tot_rel_err <= error_margin) ? EXIT_SUCCESS : EXIT_FAILURE;
}
return (tot_rel_err == 0.0) ? EXIT_SUCCESS : EXIT_FAILURE;
// Overload for cute::Tensor<>
template <class Engine, class Layout>
int
print_relative_error(
cute::Tensor<Engine, Layout> data,
cute::Tensor<Engine, Layout> reference,
bool print_verbose = false,
bool print_error = true,
double error_margin = 0.00001) {
assert(size(data) == size(reference));
return print_relative_error(static_cast<std::size_t>(size(data)),
data, reference,
print_verbose, print_error, error_margin);
}
@@ -68,7 +68,7 @@ template <int Rank>
struct LinearToCoordinateHelper<Rank, 0> {
CUTLASS_HOST_DEVICE
void operator()(Coord<Rank> &coord, int64_t idx, Coord<Rank> const &extent) const {
void operator()(Coord<Rank> &coord, int64_t idx, Coord<Rank> const &) const {
coord[Rank - 1] = int(idx);
}
};
@@ -134,9 +134,8 @@ struct RandomGaussianFunc {
stddev(static_cast<FloatType>(stddev_)),
int_scale(int_scale_) {
float_scale_up = FloatType(IntType(1) << int_scale);
float_scale_up += FloatType(0.5) * float_scale_up;
float_scale_down = FloatType(1) / FloatType(IntType(1) << int_scale);
float_scale_up = FloatType(IntType(2) << int_scale); // scale up to clamp low order bits
float_scale_down = FloatType(1) / FloatType(IntType(2) << int_scale);
}
};
@@ -172,8 +171,8 @@ struct RandomGaussianFunc {
Element result;
if (params.int_scale >= 0) {
rnd = FloatType(IntType(rnd * params.float_scale_up));
result = Element(rnd * params.float_scale_down);
rnd = FloatType(IntType(std::llround(rnd * params.float_scale_up)));
result = Element(IntType(rnd * params.float_scale_down));
}
else {
result = Element(rnd);
@@ -448,9 +447,8 @@ struct RandomUniformFunc {
max(static_cast<FloatType>(max_)),
int_scale(int_scale_) {
float_scale_up = FloatType(IntType(1) << int_scale);
float_scale_up += FloatType(0.5) * float_scale_up;
float_scale_down = FloatType(1) / FloatType(IntType(1) << int_scale);
float_scale_up = FloatType(IntType(2) << int_scale); // scale up to clamp low order bits
float_scale_down = FloatType(1) / FloatType(IntType(2) << int_scale);
}
};
@@ -489,8 +487,8 @@ struct RandomUniformFunc {
Element result;
if (params.int_scale >= 0) {
rnd = FloatType(IntType(rnd * params.float_scale_up));
result = Element(rnd * params.float_scale_down);
rnd = FloatType(IntType(std::llround(rnd * params.float_scale_up)));
result = Element(IntType(rnd * params.float_scale_down));
}
else {
result = Element(rnd);
@@ -774,9 +772,13 @@ struct RandomSparseMetaFunc {
MetaSizeInBits(MetaSizeInBits_) {
if (MetaSizeInBits_ == 2) {
range = 6;
} else if (MetaSizeInBits_ == 4) {
}
else if (MetaSizeInBits_ == 4) {
range = 2;
}
else {
throw std::invalid_argument("Invalid MetaSizeInBits");
}
}
};
@@ -1161,34 +1163,10 @@ struct TensorClearPartialFunc {
/// Parameters structure
struct Params {
//
// Data members
//
TensorView view;
Element element;
FillMode fill_mode;
int alignment;
/// Default ctor
CUTLASS_HOST_DEVICE
Params(): fill_mode(FillMode::kNone) { }
//
// Methods
//
/// Construction of Gaussian RNG functor.
Params(
TensorView view_,
Element element_,
FillMode fill_mode_,
int alignment_
):
view(view_), element(element_), fill_mode(fill_mode_), alignment(alignment_) {
}
TensorView view{};
Element element{};
FillMode fill_mode{FillMode::kNone};
int alignment{0};
};
//
@@ -1307,7 +1285,7 @@ void TensorClearPartial(
TensorForEach<Func, Layout::kRank, Params>(
view.extent(),
Params(view, element, fill_mode, alignment),
Params{view, element, fill_mode, alignment},
/*grid_size*/0, /*block_size*/0,
stream
);
@@ -120,7 +120,7 @@ __global__ void TensorTransformReducePartial(
ComputeType *workspace) { /// Device-side workspace for accumulating partial results. The reduced element is stored in workspace[0]
int64_t idx = threadIdx.x + blockIdx.x * blockDim.x;
int64_t size = view_A.size();
auto size = static_cast<int64_t>(view_A.size());
__shared__ ComputeType scratchpad[kBlockSize];
@@ -0,0 +1,649 @@
/***************************************************************************************************
* Copyright (c) 2023 - 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 Reference implementation for CONV in host-side code.
*/
#pragma once
/////////////////////////////////////////////////////////////////////////////////////////////////
#include "cutlass/complex.h"
#include "cutlass/numeric_conversion.h"
#include "cutlass/epilogue/thread/activation.h"
#include "cute/tensor.hpp"
/////////////////////////////////////////////////////////////////////////////////////////////////
namespace cutlass::reference::host {
/////////////////////////////////////////////////////////////////////////////////////////////////
namespace detail {
template<class EngineAct, class LayoutAct>
bool
is_activation_in_bounds(
cute::Tensor<EngineAct, LayoutAct> const& activation,
int32_t n_, int32_t d_, int32_t h_, int32_t w_, int32_t c_) {
return ((n_ >= 0 && n_ < size<4>(activation)) &&
(d_ >= 0 && d_ < size<3>(activation)) &&
(h_ >= 0 && h_ < size<2>(activation)) &&
(w_ >= 0 && w_ < size<1>(activation)) &&
(c_ >= 0 && c_ < size<0>(activation)));
}
template<class EngineAct, class LayoutAct>
bool
is_activation_in_bounds(
cute::Tensor<EngineAct, LayoutAct> const& activation,
int32_t n_, int32_t h_, int32_t w_, int32_t c_) {
return ((n_ >= 0 && n_ < size<3>(activation)) &&
(h_ >= 0 && h_ < size<2>(activation)) &&
(w_ >= 0 && w_ < size<1>(activation)) &&
(c_ >= 0 && c_ < size<0>(activation)));
}
template<class EngineAct, class LayoutAct>
bool
is_activation_in_bounds(
cute::Tensor<EngineAct, LayoutAct> const& activation,
int32_t n_, int32_t w_, int32_t c_) {
return ((n_ >= 0 && n_ < size<2>(activation)) &&
(w_ >= 0 && w_ < size<1>(activation)) &&
(c_ >= 0 && c_ < size<0>(activation)));
}
} // namespace detail
template<
class ElementAcc_,
class ElementScalar_,
class ElementCompute_,
class ElementC_,
class ElementOut_,
class TensorAlpha_,
class TensorBeta_,
class TensorBias_,
class ActivationFunctor_ = cutlass::epilogue::thread::Identity<ElementCompute_>>
struct ConvEpilogueFusionParams {
using ElementAcc = ElementAcc_;
using ElementScalar = ElementScalar_;
using ElementCompute = ElementCompute_;
using ElementC = ElementC_;
using ElementOut = ElementOut_;
using TensorAlpha = TensorAlpha_;
using TensorBeta = TensorBeta_;
using TensorBias = TensorBias_;
using ActivationFunctor = ActivationFunctor_;
ElementScalar alpha = ElementScalar(1);
ElementScalar beta = ElementScalar(0);
TensorAlpha tensor_alpha{};
TensorBeta tensor_beta{};
TensorBias tensor_bias{};
};
template<
cutlass::conv::Operator ConvOp,
int NumSpatialDims,
class TensorA,
class TensorB,
class TensorC,
class TensorD,
class ShapePadding,
class StrideTraversal,
class ShapeDilation,
class EpilogueFusionParams>
struct ConvReferenceImpl {
using ElementAcc = typename EpilogueFusionParams::ElementAcc;
using ElementC = typename EpilogueFusionParams::ElementC;
using ElementOut = typename EpilogueFusionParams::ElementOut;
using ElementScalar = typename EpilogueFusionParams::ElementScalar;
using ElementCompute = typename EpilogueFusionParams::ElementCompute;
using ElementBias = typename EpilogueFusionParams::TensorBias::value_type;
using ActivationFunctor = typename EpilogueFusionParams::ActivationFunctor;
// Input related converter
NumericConverter<ElementCompute, ElementAcc> acc_converter;
NumericConverter<ElementCompute, ElementC> residual_converter;
NumericConverter<ElementCompute, ElementBias> bias_converter;
// Scale related converter
NumericConverter<ElementCompute, ElementScalar> scale_converter;
// Output related converter
NumericConverter<ElementOut, ElementCompute> output_converter;
EpilogueFusionParams& epi_fusion_params_;
TensorA const& tensor_a_;
TensorB const& tensor_b_;
TensorC const& tensor_c_;
TensorD& tensor_d_;
ShapePadding const& padding_;
StrideTraversal const& tstride_;
ShapeDilation const& dilation_;
// Epilogue activation operation
ActivationFunctor epi_activation;
ConvReferenceImpl(
TensorA const& tensor_a,
TensorB const& tensor_b,
TensorC const& tensor_c,
TensorD& tensor_d,
ShapePadding const& padding,
StrideTraversal const& tstride,
ShapeDilation const& dilation,
EpilogueFusionParams& epi_fusion_params)
: tensor_a_(tensor_a),
tensor_b_(tensor_b),
tensor_c_(tensor_c),
tensor_d_(tensor_d),
padding_(padding),
tstride_(tstride),
dilation_(dilation),
epi_fusion_params_(epi_fusion_params) {
static_assert(rank(ShapePadding{}) == rank(ShapeDilation{}));
static_assert(rank(ShapePadding{}) == rank(StrideTraversal{}));
}
void compute_reference() {
if constexpr (ConvOp == cutlass::conv::Operator::kFprop) {
fprop_reference(cute::Int<NumSpatialDims>{});
}
else if constexpr (ConvOp == cutlass::conv::Operator::kDgrad) {
dgrad_reference(cute::Int<NumSpatialDims>{});
}
else {
wgrad_reference(cute::Int<NumSpatialDims>{});
}
}
private:
// Specialization for 1D fprop kernel
void fprop_reference(cute::Int<1> spatial_dims) {
int32_t N = size<2>(tensor_d_);
int32_t Q = size<1>(tensor_d_);
int32_t K = size<0>(tensor_d_);
int32_t S = size<1>(tensor_b_);
int32_t C = size<0>(tensor_b_);
#if defined(_OPENMP)
#pragma omp parallel for collapse(2)
#endif
for (int32_t n = 0; n < N; ++n) {
for (int32_t q = 0; q < Q; ++q) {
for (int32_t k = 0; k < K; ++k) {
auto accumulator = ElementAcc(0);
for (int32_t s = 0; s < S; ++s) {
for (int32_t c = 0; c < C; ++c) {
int32_t w = q * cute::get<0>(tstride_) - cute::get<0>(padding_) + s * cute::get<0>(dilation_);
if (detail::is_activation_in_bounds(tensor_a_, n, w, c)) {
accumulator += ElementAcc(tensor_a_(c, w, n) * tensor_b_(c, s, k));
}
}
}
ElementScalar alpha = raw_pointer_cast(epi_fusion_params_.tensor_alpha.data()) ?
epi_fusion_params_.tensor_alpha[k] : epi_fusion_params_.alpha;
ElementScalar beta = raw_pointer_cast(epi_fusion_params_.tensor_beta.data()) ?
epi_fusion_params_.tensor_beta[k] : epi_fusion_params_.beta;
ElementCompute output = scale_converter(alpha) * acc_converter(accumulator) +
scale_converter(beta) * residual_converter(tensor_c_(k, q, n));
if (raw_pointer_cast(epi_fusion_params_.tensor_bias.data())) {
output += bias_converter(epi_fusion_params_.tensor_bias[k]);
}
output = epi_activation(output);
tensor_d_(k, q, n) = output_converter(output);
}
}
}
}
// Specialization for 2D fprop kernel
void fprop_reference(cute::Int<2> spatial_dims) {
int32_t N = size<3>(tensor_d_);
int32_t P = size<2>(tensor_d_);
int32_t Q = size<1>(tensor_d_);
int32_t K = size<0>(tensor_d_);
int32_t R = size<2>(tensor_b_);
int32_t S = size<1>(tensor_b_);
int32_t C = size<0>(tensor_b_);
#if defined(_OPENMP)
#pragma omp parallel for collapse(3)
#endif
for (int32_t n = 0; n < N; ++n) {
for (int32_t p = 0; p < P; ++p) {
for (int32_t q = 0; q < Q; ++q) {
for (int32_t k = 0; k < K; ++k) {
auto accumulator = ElementAcc(0);
for (int32_t r = 0; r < R; ++r) {
for (int32_t s = 0; s < S; ++s) {
for (int32_t c = 0; c < C; ++c) {
int32_t w = q * cute::get<0>(tstride_) - cute::get<0>(padding_) + s * cute::get<0>(dilation_);
int32_t h = p * cute::get<1>(tstride_) - cute::get<1>(padding_) + r * cute::get<1>(dilation_);
if (detail::is_activation_in_bounds(tensor_a_, n, h, w, c)) {
accumulator += ElementAcc(tensor_a_(c, w, h, n) * tensor_b_(c, s, r, k));
}
}
}
}
ElementScalar alpha = raw_pointer_cast(epi_fusion_params_.tensor_alpha.data()) ?
epi_fusion_params_.tensor_alpha[k] : epi_fusion_params_.alpha;
ElementScalar beta = raw_pointer_cast(epi_fusion_params_.tensor_beta.data()) ?
epi_fusion_params_.tensor_beta[k] : epi_fusion_params_.beta;
ElementCompute output = scale_converter(alpha) * acc_converter(accumulator) +
scale_converter(beta) * residual_converter(tensor_c_(k, q, p, n));
if (raw_pointer_cast(epi_fusion_params_.tensor_bias.data())) {
output += bias_converter(epi_fusion_params_.tensor_bias[k]);
}
output = epi_activation(output);
tensor_d_(k, q, p, n) = output_converter(output);
}
}
}
}
}
// Specialization for 3D fprop kernel
void fprop_reference(cute::Int<3> spatial_dims) {
int32_t N = size<4>(tensor_d_);
int32_t Z = size<3>(tensor_d_);
int32_t P = size<2>(tensor_d_);
int32_t Q = size<1>(tensor_d_);
int32_t K = size<0>(tensor_d_);
int32_t T = size<3>(tensor_b_);
int32_t R = size<2>(tensor_b_);
int32_t S = size<1>(tensor_b_);
int32_t C = size<0>(tensor_b_);
#if defined(_OPENMP)
#pragma omp parallel for collapse(3)
#endif
for (int32_t n = 0; n < N; ++n) {
for (int32_t z = 0; z < Z; ++z) {
for (int32_t p = 0; p < P; ++p) {
for (int32_t q = 0; q < Q; ++q) {
for (int32_t k = 0; k < K; ++k) {
auto accumulator = ElementAcc(0);
for (int32_t t = 0; t < T; ++t) {
for (int32_t r = 0; r < R; ++r) {
for (int32_t s = 0; s < S; ++s) {
for (int32_t c = 0; c < C; ++c) {
int32_t w = q * cute::get<0>(tstride_) - cute::get<0>(padding_) + s * cute::get<0>(dilation_);
int32_t h = p * cute::get<1>(tstride_) - cute::get<1>(padding_) + r * cute::get<1>(dilation_);
int32_t d = z * cute::get<2>(tstride_) - cute::get<2>(padding_) + t * cute::get<2>(dilation_);
if (detail::is_activation_in_bounds(tensor_a_, n, d, h, w, c)) {
accumulator += ElementAcc(tensor_a_(c, w, h, d, n) * tensor_b_(c, s, r, t, k));
}
}
}
}
}
ElementScalar alpha = raw_pointer_cast(epi_fusion_params_.tensor_alpha.data()) ?
epi_fusion_params_.tensor_alpha[k] : epi_fusion_params_.alpha;
ElementScalar beta = raw_pointer_cast(epi_fusion_params_.tensor_beta.data()) ?
epi_fusion_params_.tensor_beta[k] : epi_fusion_params_.beta;
ElementCompute output = scale_converter(alpha) * acc_converter(accumulator) +
scale_converter(beta) * residual_converter(tensor_c_(k, q, p, z, n));
if (raw_pointer_cast(epi_fusion_params_.tensor_bias.data())) {
output += bias_converter(epi_fusion_params_.tensor_bias[k]);
}
output = epi_activation(output);
tensor_d_(k, q, p, z, n) = output_converter(output);
}
}
}
}
}
}
// Specialization for 1D dgrad kernel
void dgrad_reference(cute::Int<1> spatial_dims) {
int32_t N = size<2>(tensor_d_);
int32_t W = size<1>(tensor_d_);
int32_t C = size<0>(tensor_d_);
int32_t K = size<2>(tensor_b_);
int32_t S = size<1>(tensor_b_);
#if defined(_OPENMP)
#pragma omp parallel for collapse(2)
#endif
for (int32_t n = 0; n < N; ++n) {
for (int32_t w = 0; w < W; ++w) {
for (int32_t c = 0; c < C; ++c) {
auto accumulator = ElementAcc(0);
for (int32_t k = 0; k < K; ++k) {
for (int32_t s = 0; s < S; ++s) {
int32_t q = w + cute::get<0>(padding_) - s * cute::get<0>(dilation_);
if (q % cute::get<0>(tstride_) == 0) {
q /= cute::get<0>(tstride_);
} else {
continue;
}
if (detail::is_activation_in_bounds(tensor_a_, n, q, k)) {
accumulator += ElementAcc(tensor_a_(k, q, n) * tensor_b_(c, s, k));
}
}
}
ElementScalar alpha = raw_pointer_cast(epi_fusion_params_.tensor_alpha.data())
? epi_fusion_params_.tensor_alpha[c] : epi_fusion_params_.alpha;
ElementScalar beta = raw_pointer_cast(epi_fusion_params_.tensor_beta.data())
? epi_fusion_params_.tensor_beta[c] : epi_fusion_params_.beta;
ElementCompute output = scale_converter(alpha) * acc_converter(accumulator) +
scale_converter(beta) * residual_converter(tensor_c_(c, w, n));
if (raw_pointer_cast(epi_fusion_params_.tensor_bias.data())) {
output += bias_converter(epi_fusion_params_.tensor_bias[c]);
}
output = epi_activation(output);
tensor_d_(c, w, n) = output_converter(output);
}
}
}
}
// Specialization for 2D dgrad kernel
void dgrad_reference(cute::Int<2> spatial_dims) {
int32_t N = size<3>(tensor_d_);
int32_t H = size<2>(tensor_d_);
int32_t W = size<1>(tensor_d_);
int32_t C = size<0>(tensor_d_);
int32_t K = size<3>(tensor_b_);
int32_t R = size<2>(tensor_b_);
int32_t S = size<1>(tensor_b_);
#if defined(_OPENMP)
#pragma omp parallel for collapse(3)
#endif
for (int32_t n = 0; n < N; ++n) {
for (int32_t h = 0; h < H; ++h) {
for (int32_t w = 0; w < W; ++w) {
for (int32_t c = 0; c < C; ++c) {
auto accumulator = ElementAcc(0);
for (int32_t k = 0; k < K; ++k) {
for (int32_t r = 0; r < R; ++r) {
for (int32_t s = 0; s < S; ++s) {
int32_t q = w + cute::get<0>(padding_) - s * cute::get<0>(dilation_);
int32_t p = h + cute::get<1>(padding_) - r * cute::get<1>(dilation_);
if (q % cute::get<0>(tstride_) == 0) {
q /= cute::get<0>(tstride_);
} else {
continue;
}
if (p % cute::get<1>(tstride_) == 0) {
p /= cute::get<1>(tstride_);
} else {
continue;
}
if (detail::is_activation_in_bounds(tensor_a_, n, p, q, k)) {
accumulator += ElementAcc(tensor_a_(k, q, p, n) * tensor_b_(c, s, r, k));
}
}
}
}
ElementScalar alpha = raw_pointer_cast(epi_fusion_params_.tensor_alpha.data())
? epi_fusion_params_.tensor_alpha[c] : epi_fusion_params_.alpha;
ElementScalar beta = raw_pointer_cast(epi_fusion_params_.tensor_beta.data())
? epi_fusion_params_.tensor_beta[c] : epi_fusion_params_.beta;
ElementCompute output = scale_converter(alpha) * acc_converter(accumulator) +
scale_converter(beta) * residual_converter(tensor_c_(c, w, h, n));
if (raw_pointer_cast(epi_fusion_params_.tensor_bias.data())) {
output += bias_converter(epi_fusion_params_.tensor_bias[c]);
}
tensor_d_(c, w, h, n) = output_converter(output);
}
}
}
}
}
// Specialization for 3D dgrad kernel
void dgrad_reference(cute::Int<3> spatial_dims) {
int32_t N = size<4>(tensor_d_);
int32_t D = size<3>(tensor_d_);
int32_t H = size<2>(tensor_d_);
int32_t W = size<1>(tensor_d_);
int32_t C = size<0>(tensor_d_);
int32_t K = size<4>(tensor_b_);
int32_t T = size<3>(tensor_b_);
int32_t R = size<2>(tensor_b_);
int32_t S = size<1>(tensor_b_);
#if defined(_OPENMP)
#pragma omp parallel for collapse(3)
#endif
for (int32_t n = 0; n < N; ++n) {
for (int32_t d = 0; d < D; ++d) {
for (int32_t h = 0; h < H; ++h) {
for (int32_t w = 0; w < W; ++w) {
for (int32_t c = 0; c < C; ++c) {
auto accumulator = ElementAcc(0);
for (int32_t k = 0; k < K; ++k) {
for (int32_t t = 0; t < T; ++t) {
for (int32_t r = 0; r < R; ++r) {
for (int32_t s = 0; s < S; ++s) {
int32_t q = w + cute::get<0>(padding_) - s * cute::get<0>(dilation_);
int32_t p = h + cute::get<1>(padding_) - r * cute::get<1>(dilation_);
int32_t z = d + cute::get<2>(padding_) - t * cute::get<2>(dilation_);
if (q % cute::get<0>(tstride_) == 0) {
q /= cute::get<0>(tstride_);
} else {
continue;
}
if (p % cute::get<1>(tstride_) == 0) {
p /= cute::get<1>(tstride_);
} else {
continue;
}
if (z % cute::get<2>(tstride_) == 0) {
z /= cute::get<2>(tstride_);
} else {
continue;
}
if (detail::is_activation_in_bounds(tensor_a_, n, z, p, q, k)) {
accumulator += ElementAcc(tensor_a_(k, q, p, z, n) * tensor_b_(c, s, r, t, k));
}
}
}
}
}
ElementScalar alpha = raw_pointer_cast(epi_fusion_params_.tensor_alpha.data())
? epi_fusion_params_.tensor_alpha[c] : epi_fusion_params_.alpha;
ElementScalar beta = raw_pointer_cast(epi_fusion_params_.tensor_beta.data())
? epi_fusion_params_.tensor_beta[c] : epi_fusion_params_.beta;
ElementCompute output = scale_converter(alpha) * acc_converter(accumulator) +
scale_converter(beta) * residual_converter(tensor_c_(c, w, h, d, n));
if (raw_pointer_cast(epi_fusion_params_.tensor_bias.data())) {
output += bias_converter(epi_fusion_params_.tensor_bias[c]);
}
output = epi_activation(output);
tensor_d_(c, w, h, d, n) = output_converter(output);
}
}
}
}
}
}
// Specialization for 1D wgrad kernel
void wgrad_reference(cute::Int<1> spatial_dims) {
int32_t N = size<2>(tensor_a_);
int32_t Q = size<1>(tensor_a_);
int32_t K = size<0>(tensor_a_);
int32_t S = size<1>(tensor_d_);
int32_t C = size<0>(tensor_d_);
#if defined(_OPENMP)
#pragma omp parallel for collapse(2)
#endif
for (int32_t k = 0; k < K; ++k) {
ElementScalar alpha = raw_pointer_cast(epi_fusion_params_.tensor_alpha.data()) ?
epi_fusion_params_.tensor_alpha[k] : epi_fusion_params_.alpha;
ElementScalar beta = raw_pointer_cast(epi_fusion_params_.tensor_beta.data()) ?
epi_fusion_params_.tensor_beta[k] : epi_fusion_params_.beta;
for (int32_t s = 0; s < S; ++s) {
for (int32_t c = 0; c < C; ++c) {
auto accumulator = ElementAcc(0);
for (int32_t n = 0; n < N; ++n) {
for (int32_t q = 0; q < Q; ++q) {
int32_t w = q * cute::get<0>(tstride_) - cute::get<0>(padding_) + s * cute::get<0>(dilation_);
if (detail::is_activation_in_bounds(tensor_b_, n, w, c)) {
accumulator += ElementAcc(tensor_b_(c, w, n) * tensor_a_(k, q, n));
}
}
}
ElementCompute output = scale_converter(alpha) * acc_converter(accumulator) +
scale_converter(beta) * residual_converter(tensor_c_(c, s, k));
if (raw_pointer_cast(epi_fusion_params_.tensor_bias.data())) {
output += bias_converter(epi_fusion_params_.tensor_bias[k]);
}
output = epi_activation(output);
tensor_d_(c, s, k) = output_converter(output);
}
}
}
}
// Specialization for 2D wgrad kernel
void wgrad_reference(cute::Int<2> spatial_dims) {
int32_t N = size<3>(tensor_a_);
int32_t P = size<2>(tensor_a_);
int32_t Q = size<1>(tensor_a_);
int32_t K = size<0>(tensor_a_);
int32_t R = size<2>(tensor_d_);
int32_t S = size<1>(tensor_d_);
int32_t C = size<0>(tensor_d_);
#if defined(_OPENMP)
#pragma omp parallel for collapse(3)
#endif
for (int32_t k = 0; k < K; ++k) {
ElementScalar alpha = raw_pointer_cast(epi_fusion_params_.tensor_alpha.data()) ?
epi_fusion_params_.tensor_alpha[k] : epi_fusion_params_.alpha;
ElementScalar beta = raw_pointer_cast(epi_fusion_params_.tensor_beta.data()) ?
epi_fusion_params_.tensor_beta[k] : epi_fusion_params_.beta;
for (int32_t r = 0; r < R; ++r) {
for (int32_t s = 0; s < S; ++s) {
for (int32_t c = 0; c < C; ++c) {
auto accumulator = ElementAcc(0);
for (int32_t n = 0; n < N; ++n) {
for (int32_t p = 0; p < P; ++p) {
for (int32_t q = 0; q < Q; ++q) {
int32_t w = q * cute::get<0>(tstride_) - cute::get<0>(padding_) + s * cute::get<0>(dilation_);
int32_t h = p * cute::get<1>(tstride_) - cute::get<1>(padding_) + r * cute::get<1>(dilation_);
if (detail::is_activation_in_bounds(tensor_b_, n, h, w, c)) {
accumulator += ElementAcc(tensor_b_(c, w, h, n) * tensor_a_(k, q, p, n));
}
}
}
}
ElementCompute output = scale_converter(alpha) * acc_converter(accumulator) +
scale_converter(beta) * residual_converter(tensor_c_(c, s, r, k));
if (raw_pointer_cast(epi_fusion_params_.tensor_bias.data())) {
output += bias_converter(epi_fusion_params_.tensor_bias[k]);
}
output = epi_activation(output);
tensor_d_(c, s, r, k) = output_converter(output);
}
}
}
}
}
// Specialization for 3D wgrad kernel
void wgrad_reference(cute::Int<3> spatial_dims) {
int32_t N = size<4>(tensor_a_);
int32_t Z = size<3>(tensor_a_);
int32_t P = size<2>(tensor_a_);
int32_t Q = size<1>(tensor_a_);
int32_t K = size<0>(tensor_a_);
int32_t T = size<3>(tensor_d_);
int32_t R = size<2>(tensor_d_);
int32_t S = size<1>(tensor_d_);
int32_t C = size<0>(tensor_d_);
#if defined(_OPENMP)
#pragma omp parallel for collapse(3)
#endif
for (int32_t k = 0; k < K; ++k) {
ElementScalar alpha = raw_pointer_cast(epi_fusion_params_.tensor_alpha.data()) ?
epi_fusion_params_.tensor_alpha[k] : epi_fusion_params_.alpha;
ElementScalar beta = raw_pointer_cast(epi_fusion_params_.tensor_beta.data()) ?
epi_fusion_params_.tensor_beta[k] : epi_fusion_params_.beta;
for (int32_t t = 0; t < T; ++t) {
for (int32_t r = 0; r < R; ++r) {
for (int32_t s = 0; s < S; ++s) {
for (int32_t c = 0; c < C; ++c) {
auto accumulator = ElementAcc(0);
for (int32_t n = 0; n < N; ++n) {
for (int32_t z = 0; z < Z; ++z) {
for (int32_t p = 0; p < P; ++p) {
for (int32_t q = 0; q < Q; ++q) {
int32_t w = q * cute::get<0>(tstride_) - cute::get<0>(padding_) + s * cute::get<0>(dilation_);
int32_t h = p * cute::get<1>(tstride_) - cute::get<1>(padding_) + r * cute::get<1>(dilation_);
int32_t d = z * cute::get<2>(tstride_) - cute::get<2>(padding_) + t * cute::get<2>(dilation_);
if (detail::is_activation_in_bounds(tensor_b_, n, d, h, w, c)) {
accumulator += ElementAcc(tensor_b_(c, w, h, d, n) * tensor_a_(k, q, p, z, n));
}
}
}
}
}
ElementCompute output = scale_converter(alpha) * acc_converter(accumulator) +
scale_converter(beta) * residual_converter(tensor_c_(c, s, r, t, k));
if (raw_pointer_cast(epi_fusion_params_.tensor_bias.data())) {
output += bias_converter(epi_fusion_params_.tensor_bias[k]);
}
output = epi_activation(output);
tensor_d_(c, s, r, t, k) = output_converter(output);
}
}
}
}
}
}
};
/////////////////////////////////////////////////////////////////////////////////////////////////
} // cutlass::reference::host
/////////////////////////////////////////////////////////////////////////////////////////////////
@@ -236,6 +236,7 @@ void gett_mainloop(
acc[m_b][n_b] = fma_op(a_frag[m_b], b_frag[n_b], acc[m_b][n_b]);
}
}
}
}
@@ -39,6 +39,7 @@
#include <cstdlib>
#include <cmath>
#include <random>
#include <stdexcept>
// Cutlass includes
#include "cutlass/cutlass.h"
@@ -196,7 +197,7 @@ struct RandomGaussianFunc {
// Sample from the Gaussian distribution for a nonzero element
if (bernoulli_result) {
if (int_scale >= 0) {
rnd = double(int64_t(rnd * double(1 << int_scale))) / double(1 << int_scale);
rnd = double(std::llround(rnd * double(1 << int_scale))) / double(1 << int_scale);
result = static_cast<Element>(rnd);
}
else {
@@ -567,7 +568,7 @@ struct RandomUniformFunc {
// testing
Element result;
if (int_scale >= 0) {
rnd = double(int64_t(rnd * double(1 << int_scale))) / double(1 << int_scale);
rnd = double(std::llround(rnd * double(1 << int_scale))) / double(1 << int_scale);
result = static_cast<Element>(Real(rnd));
}
else {
@@ -1381,9 +1382,13 @@ struct RandomSparseMetaFunc {
std::srand((unsigned)seed);
if (MetaSizeInBits_ == 2) {
range = 6;
} else if (MetaSizeInBits_ == 4) {
}
else if (MetaSizeInBits_ == 4) {
range = 2;
}
else {
throw std::invalid_argument("Invalid MetaSizeInBits");
}
}
/// Compute random value and update RNG state
@@ -61,7 +61,7 @@ ComputeType TensorTransformReduce(
TransformOp transform
) {
for (int64_t idx = 0; idx < view.size(); ++idx) {
for (int64_t idx = 0; idx < int64_t(view.size()); ++idx) {
typename Layout::TensorCoord coord;
cutlass::reference::detail::LinearToCoordinate<Layout::kRank>()(coord, idx, view.extent());
@@ -94,7 +94,7 @@ ComputeType TensorTransformReduce(
throw std::runtime_error("Tensor extents must match.");
}
for (int64_t idx = 0; idx < view_A.size(); ++idx) {
for (int64_t idx = 0; idx < int64_t(view_A.size()); ++idx) {
typename Layout::TensorCoord coord;
cutlass::reference::detail::LinearToCoordinate<Layout::kRank>()(coord, idx, view_A.extent());