CUTLASS 3.5.0 (#1411)
This commit is contained in:
@@ -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,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());
|
||||
|
||||
Reference in New Issue
Block a user