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

CUTLASS 3.2 (#1024)

Browse Source 
* CUTLASS 3.2
This commit is contained in:
ANIKET SHIVAM
2023-08-07 20:50:32 -04:00
committed by GitHub
parent a0d787b746
commit 4575443d44
392 changed files with 47559 additions and 7940 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tools/util/include/cutlass/util/distribution.h
+9 -2
View File
@@ -57,6 +57,10 @@ struct Distribution {
struct {
double mean;
double stddev;
double pnz;
double pnzA;
double pnzB;
double pnzC;
} gaussian;
/// Elements are linear combination of row and column index
@@ -88,10 +92,11 @@ struct Distribution {
}
/// Configures distribution as Gaussian distribution
Distribution &set_gaussian(double _mean, double _stddev, int _int_scale = 0) {
Distribution &set_gaussian(double _mean, double _stddev, int _int_scale = 0, double _pnz = 100.0) {
kind = Gaussian;
gaussian.mean = _mean;
gaussian.stddev = _stddev;
gaussian.pnz = _pnz;
int_scale = _int_scale;
return *this;
}
@@ -123,7 +128,9 @@ inline std::ostream &operator<<(std::ostream &out, cutlass::Distribution const &
out << "uniform, min: " << dist.uniform.min << ", max: " << dist.uniform.max;
break;
case cutlass::Distribution::Gaussian:
out << "gaussian, mean: " << dist.gaussian.mean << ", stddev: " << dist.gaussian.stddev;
out << "gaussian, mean: " << dist.gaussian.mean << ", stddev: " << dist.gaussian.stddev
<< ", pnzA: " << dist.gaussian.pnzA << ", pnzB: "
<< dist.gaussian.pnzB << ", pnzC: " << dist.gaussian.pnzC;
break;
case cutlass::Distribution::Identity:
out << "identity";
tools/util/include/cutlass/util/host_tensor.h
+23 -8
View File
@@ -47,6 +47,7 @@
#include "cutlass/cutlass.h"
#include "cutlass/tensor_ref.h"
#include "cutlass/tensor_view.h"
#include "cutlass/fast_math.h"
#include "device_memory.h"
@@ -103,8 +104,17 @@ public:
/// Constant reference to element in tensor
using ConstReference = typename ConstTensorRef::Reference;
/// Used to handle packing of subbyte elements
static int const kElementsPerStoredItem = (sizeof_bits<Element>::value < 8 ? (8 / sizeof_bits<Element>::value) : 1);
/// Note: Below is used to handle packing of subbyte elements
/// kBitsStoredVec : The bits of store vec that could be divisiable by the element
/// kElementsPerStoredVec : The number of elements could be stored in per store vec
/// kNumStoragePerStoredVec : How much storage(i.e. sizeof(element storage)) the store vec needs to consume.
/// Usually the element storage of subbyte is uint8_t.
/// 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(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);
private:
@@ -170,8 +180,7 @@ public:
device_.reset();
host_.clear();
count /= kElementsPerStoredItem;
count = count / kElementsPerStoredVec * kNumStoragePerStoredVec;
host_.resize(count);
// Allocate memory
@@ -217,7 +226,7 @@ public:
LongIndex new_size = size_t(layout_.capacity(extent_));
if (static_cast<decltype(host_.size())>(new_size) > host_.size()) {
reserve(new_size);
reserve(new_size, device_backed_);
}
}
@@ -232,7 +241,7 @@ public:
/// Returns the number of elements stored in the host tensor
size_t size() const {
return host_.size() * kElementsPerStoredItem;
return host_.size() / kNumStoragePerStoredVec * kElementsPerStoredVec;
}
/// Returns the logical capacity based on extent and layout. May differ from size().
@@ -254,6 +263,9 @@ public:
/// Gets pointer to host data
Element const * host_data() const { return host_.data(); }
/// Gets pointer to host data with a pointer offset
Element const * host_data_ptr_offset(LongIndex ptr_element_offset) const { return &ReferenceFactory<Element>::get(host_.data(), ptr_element_offset); }
/// Gets a constant reference to an element in host memory
ConstReference host_data(LongIndex idx) const {
return ReferenceFactory<Element const>::get(host_data(), idx);
@@ -262,11 +274,14 @@ public:
/// Gets pointer to device data
Element * device_data() { return device_.get(); }
/// Gets pointer to device data
Element const * device_data() const { return device_.get(); }
/// Gets pointer to device data with a pointer offset
Element * device_data_ptr_offset(LongIndex ptr_element_offset) { return &ReferenceFactory<Element>::get(device_data(), ptr_element_offset); }
/// Gets pointer to device data
Element const * device_data() const { return device_.get(); }
/// Gets pointer to device data with a pointer offset
Element const * device_data_ptr_offset(LongIndex ptr_element_offset) const { return &ReferenceFactory<Element>::get(device_data(), ptr_element_offset); }
/// Accesses the tensor reference pointing to data
TensorRef host_ref(LongIndex ptr_element_offset=0) { return TensorRef(host_data_ptr_offset(ptr_element_offset), layout_); }
tools/util/include/cutlass/util/packed_stride.hpp
+31 -25
View File
@@ -29,7 +29,7 @@
*
**************************************************************************************************/
/*! \file
\brief Utilities for packing a rank-X shape into a rank-(X-1) stride in CuTe.
\brief Utilities for packing constructing canonical CuTe stride types for 3.x mainloop params.
*/
#pragma once
@@ -38,25 +38,29 @@
/////////////////////////////////////////////////////////////////////////////////////////////////
namespace cutlass {
/////////////////////////////////////////////////////////////////////////////////////////////////
// Strides without batch mode
template <class StrideIntT>
cute::Stride<StrideIntT, cute::Int<1>>
make_cute_packed_stride(cute::Stride<StrideIntT, cute::Int<1>> s, cute::Shape<int,int,int> shape_MKL) {
static_assert(std::is_integral_v<StrideIntT>,
template <class IntT>
cute::Stride<IntT, cute::Int<1>>
make_cute_packed_stride(cute::Stride<IntT, cute::Int<1>> s, cute::Shape<int,int,int> shape_MKL) {
static_assert(std::is_integral_v<IntT>,
"Stride must have an integral type so it can be set dynamically. Static strides not supported.");
auto s_copy = s;
cute::get<0>(s_copy) = static_cast<StrideIntT>(cute::get<1>(shape_MKL));
cute::get<0>(s_copy) = static_cast<IntT>(cute::get<1>(shape_MKL));
return s_copy;
}
template <class StrideIntT>
cute::Stride<cute::Int<1>, StrideIntT>
make_cute_packed_stride(cute::Stride<cute::Int<1>, StrideIntT> s, cute::Shape<int,int,int> shape_MKL) {
static_assert(std::is_integral_v<StrideIntT>,
template <class IntT>
cute::Stride<cute::Int<1>, IntT>
make_cute_packed_stride(cute::Stride<cute::Int<1>, IntT> s, cute::Shape<int,int,int> shape_MKL) {
static_assert(std::is_integral_v<IntT>,
"Stride must have an integral type so it can be set dynamically. Static strides not supported.");
auto s_copy = s;
cute::get<1>(s_copy) = static_cast<StrideIntT>(cute::get<0>(shape_MKL));
cute::get<1>(s_copy) = static_cast<IntT>(cute::get<0>(shape_MKL));
return s_copy;
}
@@ -64,38 +68,40 @@ make_cute_packed_stride(cute::Stride<cute::Int<1>, StrideIntT> s, cute::Shape<in
// Strides with batch mode
template <class StrideIntT>
cute::Stride<StrideIntT, cute::Int<1>, int64_t>
make_cute_packed_stride(cute::Stride<StrideIntT, cute::Int<1>, int64_t> s, cute::Shape<int,int,int> shape_MKL) {
static_assert(std::is_integral_v<StrideIntT>,
template <class IntT>
cute::Stride<IntT, cute::Int<1>, int64_t>
make_cute_packed_stride(cute::Stride<IntT, cute::Int<1>, int64_t> s, cute::Shape<int,int,int> shape_MKL) {
static_assert(std::is_integral_v<IntT>,
"Stride must have an integral type so it can be set dynamically. Static strides not supported.");
auto s_copy = s;
cute::get<0>(s_copy) = static_cast<StrideIntT>(cute::get<1>(shape_MKL));
cute::get<0>(s_copy) = static_cast<IntT>(cute::get<1>(shape_MKL));
int batch_count = cute::get<2>(shape_MKL);
if (batch_count > 1) {
cute::get<2>(s_copy) = static_cast<StrideIntT>(cute::get<0>(shape_MKL) * cute::get<1>(shape_MKL));
cute::get<2>(s_copy) = static_cast<IntT>(cute::get<0>(shape_MKL) * cute::get<1>(shape_MKL));
}
else {
cute::get<2>(s_copy) = static_cast<StrideIntT>(0);
cute::get<2>(s_copy) = static_cast<IntT>(0);
}
return s_copy;
}
template <class StrideIntT>
cute::Stride<cute::Int<1>, StrideIntT, int64_t>
make_cute_packed_stride(cute::Stride<cute::Int<1>, StrideIntT, int64_t> s, cute::Shape<int,int,int> shape_MKL) {
static_assert(std::is_integral_v<StrideIntT>,
template <class IntT>
cute::Stride<cute::Int<1>, IntT, int64_t>
make_cute_packed_stride(cute::Stride<cute::Int<1>, IntT, int64_t> s, cute::Shape<int,int,int> shape_MKL) {
static_assert(std::is_integral_v<IntT>,
"Stride must have an integral type so it can be set dynamically. Static strides not supported.");
auto s_copy = s;
cute::get<1>(s_copy) = static_cast<StrideIntT>(cute::get<0>(shape_MKL));
cute::get<1>(s_copy) = static_cast<IntT>(cute::get<0>(shape_MKL));
int batch_count = cute::get<2>(shape_MKL);
if (batch_count > 1) {
cute::get<2>(s_copy) = static_cast<StrideIntT>(cute::get<0>(shape_MKL) * cute::get<1>(shape_MKL));
cute::get<2>(s_copy) = static_cast<IntT>(cute::get<0>(shape_MKL) * cute::get<1>(shape_MKL));
}
else {
cute::get<2>(s_copy) = static_cast<StrideIntT>(0);
cute::get<2>(s_copy) = static_cast<IntT>(0);
}
return s_copy;
}
/////////////////////////////////////////////////////////////////////////////////////////////////
} // namespace cutlass
tools/util/include/cutlass/util/reference/device/tensor_fill.h
+157 -66
View File
@@ -354,19 +354,22 @@ template <
void TensorFillRandomGaussian(
TensorView<Element, Layout> view, ///< destination tensor
uint64_t seed, ///< seed for RNG
Element mean = Element(0), ///< Gaussian distribution's mean
Element stddev = Element(1), ///< Gaussian distribution's standard deviation
int bits = -1) { ///< If non-negative, specifies number of fractional bits that
typename RealType<Element>::Type mean = Element(0), ///< Gaussian distribution's mean
typename RealType<Element>::Type stddev = Element(1), ///< Gaussian distribution's standard deviation
int bits = -1, ///< If non-negative, specifies number of fractional bits that
/// are not truncated to zero. Permits reducing precision of
/// data.
cudaStream_t stream = nullptr) {
using RandomFunc = detail::RandomGaussianFunc<Element>;
using Func = detail::TensorFillRandomGaussianFunc<Element, Layout>;
using Params = typename Func::Params;
TensorForEach<Func, Layout::kRank, Params>(
view.extent(),
Params(view, typename RandomFunc::Params(seed, mean, stddev, bits))
Params(view, typename RandomFunc::Params(seed, mean, stddev, bits)),
/*grid_size*/0, /*block_size*/0,
stream
);
}
@@ -380,15 +383,16 @@ void BlockFillRandomGaussian(
uint64_t seed, ///< seed for RNG
typename RealType<Element>::Type mean, ///< Gaussian distribution's mean
typename RealType<Element>::Type stddev, ///< Gaussian distribution's standard deviation
int bits = -1) { ///< If non-negative, specifies number of fractional bits that
int bits = -1, ///< If non-negative, specifies number of fractional bits that
/// are not truncated to zero. Permits reducing precision of
/// data.
cudaStream_t stream = nullptr) {
using RandomFunc = detail::RandomGaussianFunc<Element>;
typename RandomFunc::Params params(seed, mean, stddev, bits);
BlockForEach<Element, RandomFunc>(ptr, capacity, params);
BlockForEach<Element, RandomFunc>(ptr, capacity, params, /*grid_size*/0, /*block_size*/0, stream);
}
///////////////////////////////////////////////////////////////////////////////////////////////////
@@ -686,12 +690,13 @@ template <
void TensorFillRandomUniform(
TensorView<Element, Layout> view, ///< destination tensor
uint64_t seed, ///< seed for RNG
Element max = Element(1), ///< upper bound of distribution
Element min = Element(0), ///< lower bound for distribution
int bits = -1) { ///< If non-negative, specifies number of fractional bits that
typename RealType<Element>::Type max = Element(1), ///< upper bound of distribution
typename RealType<Element>::Type min = Element(0), ///< lower bound for distribution
int bits = -1, ///< If non-negative, specifies number of fractional bits that
/// are not truncated to zero. Permits reducing precision of
/// data.
/// data.
cudaStream_t stream = nullptr) {
using RandomFunc = detail::RandomUniformFunc<Element>;
using Func = detail::TensorFillRandomUniformFunc<Element, Layout>;
using Params = typename Func::Params;
@@ -700,7 +705,9 @@ void TensorFillRandomUniform(
TensorForEach<Func, Layout::kRank, Params>(
view.extent(),
Params(view, random)
Params(view, random),
/*grid_size*/0, /*block_size*/0,
stream
);
}
@@ -714,15 +721,16 @@ void BlockFillRandomUniform(
uint64_t seed, ///< seed for RNG
typename RealType<Element>::Type max, ///< upper bound of distribution
typename RealType<Element>::Type min, ///< lower bound for distribution
int bits = -1) { ///< If non-negative, specifies number of fractional bits that
int bits = -1, ///< If non-negative, specifies number of fractional bits that
/// are not truncated to zero. Permits reducing precision of
/// data.
/// data.
cudaStream_t stream = nullptr) {
using RandomFunc = detail::RandomUniformFunc<Element>;
typename RandomFunc::Params params(seed, max, min, bits);
BlockForEach<Element, RandomFunc>(ptr, capacity, params);
BlockForEach<Element, RandomFunc>(ptr, capacity, params, /*grid_size*/0, /*block_size*/0, stream);
}
///////////////////////////////////////////////////////////////////////////////////////////////////
@@ -899,10 +907,9 @@ template <
void TensorFillRandomSparseMeta(
TensorView<Element, Layout> view, ///< destination tensor
uint64_t seed, ///< seed for RNG
int MetaSizeInBits = 2) { ///< If non-negative, specifies number of fractional bits that
/// are not truncated to zero. Permits reducing precision of
/// data.
int MetaSizeInBits = 2, ///< meta data size
cudaStream_t stream = nullptr) {
using RandomFunc = detail::RandomSparseMetaFunc<Element>;
using Func = detail::TensorFillRandomUniformFunc<Element, Layout>;
using Params = typename Func::Params;
@@ -911,7 +918,9 @@ void TensorFillRandomSparseMeta(
TensorForEach<Func, Layout::kRank, Params>(
view.extent(),
Params(view, random)
Params(view, random),
/*grid_size*/0, /*block_size*/0,
stream
);
}
@@ -923,13 +932,14 @@ void BlockFillRandomSparseMeta(
Element *ptr,
size_t capacity,
uint64_t seed, ///< seed for RNG
int MetaSizeInBits = 2) { ///< meta data size
int MetaSizeInBits = 2, ///< meta data size
cudaStream_t stream = nullptr) {
using RandomFunc = detail::RandomSparseMetaFunc<Element>;
typename RandomFunc::Params params(seed, MetaSizeInBits);
BlockForEach<Element, RandomFunc>(ptr, capacity, params);
BlockForEach<Element, RandomFunc>(ptr, capacity, params, /*grid_size*/0, /*block_size*/0, stream);
}
///////////////////////////////////////////////////////////////////////////////////////////////////
@@ -1245,14 +1255,17 @@ template <
void TensorFillDiagonal(
TensorView<Element, Layout> view, ///< destination tensor
Element diag = Element(1), ///< value to write in the diagonal
Element other = Element(0)) { ///< value to write off the diagonal
Element other = Element(0), ///< value to write off the diagonal
cudaStream_t stream = nullptr) {
typedef detail::TensorFillDiagonalFunc<Element, Layout> Func;
typedef typename Func::Params Params;
TensorForEach<Func, Layout::kRank, Params>(
view.extent(),
Params(view, diag, other)
Params(view, diag, other),
/*grid_size*/0, /*block_size*/0,
stream
);
}
@@ -1264,18 +1277,20 @@ template <
void TensorFillPartial(
TensorView<Element, Layout> view, ///< destination tensor
Element element,
FillMode fill_mode) {
FillMode fill_mode,
cudaStream_t stream = nullptr) {
typedef detail::TensorFillPartialFunc<Element, Layout> Func;
typedef typename Func::Params Params;
TensorForEach<Func, Layout::kRank, Params>(
view.extent(),
Params(view, element, fill_mode)
Params(view, element, fill_mode),
stream
);
}
/// Clears a tensor partially depending on fill mode and alignment. Elements on the wrong-side
/// Clears a tensor partially depending on fill mode and alignment. Elements on the wrong-side
/// of fillmode (upto the alignment) are overwritten with the user supplied element (typically zeros)
template <
typename Element, ///< Element type
@@ -1284,14 +1299,17 @@ void TensorClearPartial(
TensorView<Element, Layout> view, ///< destination tensor
Element element,
FillMode fill_mode,
int alignment) {
int alignment,
cudaStream_t stream = nullptr) {
typedef detail::TensorClearPartialFunc<Element, Layout> Func;
typedef typename Func::Params Params;
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
);
}
@@ -1302,10 +1320,11 @@ template <
typename Element, ///< Element type
typename Layout> ///< Layout function
void TensorFill(
TensorView<Element, Layout> view, ///< destination tensor
Element val = Element(0)) { ///< value to uniformly fill it with
TensorView<Element, Layout> view, ///< destination tensor
Element val = Element(0), ///< value to uniformly fill it with
cudaStream_t stream = nullptr) {
TensorFillDiagonal(view, val, val);
TensorFillDiagonal(view, val, val, stream);
}
///////////////////////////////////////////////////////////////////////////////////////////////////
@@ -1315,9 +1334,10 @@ template <
typename Element, ///< Element type
typename Layout> ///< Layout function
void TensorFillIdentity(
TensorView<Element, Layout> view) { ///< destination tensor
TensorView<Element, Layout> view, ///< destination tensor
cudaStream_t stream = nullptr) {
TensorFillDiagonal(view, Element(1), Element(0));
TensorFillDiagonal(view, Element(1), Element(0), stream);
}
///////////////////////////////////////////////////////////////////////////////////////////////////
@@ -1415,14 +1435,17 @@ template <
typename Layout> ///< Layout function
void TensorUpdateDiagonal(
TensorView<Element, Layout> view, ///< destination tensor
Element diag = Element(1)) {
Element diag = Element(1),
cudaStream_t stream = nullptr) {
typedef detail::TensorUpdateDiagonalFunc<Element, Layout> Func;
typedef typename Func::Params Params;
TensorForEach<Func, Layout::kRank, Params>(
view.extent(),
Params(view, diag)
Params(view, diag),
/*grid_size*/0, /*block_size*/0,
stream
);
}
@@ -1521,14 +1544,17 @@ template <
typename Layout> ///< Layout function
void TensorUpdateOffDiagonal(
TensorView<Element, Layout> view, ///< destination tensor
Element other = Element(1)) {
Element other = Element(1),
cudaStream_t stream = nullptr) {
typedef detail::TensorUpdateOffDiagonalFunc<Element, Layout> Func;
typedef typename Func::Params Params;
TensorForEach<Func, Layout::kRank, Params>(
view.extent(),
Params(view, other)
Params(view, other),
/*grid_size*/0, /*block_size*/0,
stream
);
}
@@ -1602,11 +1628,30 @@ struct TensorFillLinearFunc {
/// Compute random value and update RNG state
CUTLASS_DEVICE
void operator()(TensorCoord const &coord) {
Element sum = params.s;
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < Layout::kRank; ++i) {
sum += params.v[i] * Element(coord[i]);
if constexpr (is_complex<Element>::value) {
if constexpr (sizeof_bits<Element>::value <= 32) {
sum = Element(static_cast<complex<float>>(sum) +
static_cast<complex<float>>(params.v[i]) * static_cast<complex<float>>(coord[i]));
}
}
else if constexpr (sizeof_bits<Element>::value <= 32) {
if constexpr (std::numeric_limits<Element>::is_integer) {
sum = Element(static_cast<int32_t>(sum) +
static_cast<int32_t>(params.v[i]) * static_cast<int32_t>(coord[i]));
}
else {
sum = Element(static_cast<float>(sum) +
static_cast<float>(params.v[i]) * static_cast<float>(coord[i]));
}
}
else {
sum += params.v[i] * coord[i];
}
}
params.view.at(coord) = sum;
@@ -1624,20 +1669,57 @@ template <
void TensorFillLinear(
TensorView<Element, Layout> view, ///< destination tensor
Array<Element, Layout::kRank> const & v,
Element s = Element(0)) {
Element s = Element(0),
cudaStream_t stream = nullptr) {
using Func = detail::TensorFillLinearFunc<Element, Layout>;
using Params = typename Func::Params;
TensorForEach<Func, Layout::kRank, Params>(
view.extent(),
Params(view, v, s)
Params(view, v, s),
/*grid_size*/0, /*block_size*/0,
stream
);
}
///////////////////////////////////////////////////////////////////////////////////////////////////
///////////////////////////////////////////////////////////////////////////////////////////////////
/// Fills a tensor with random values from a distribution.
template <
typename Element, ///< Element type
typename Layout> ///< Layout function
void TensorFillRandom(
TensorView<Element, Layout> view, ///< destination tensor
uint64_t seed,
Distribution dist,
cudaStream_t stream = nullptr) {
using Real = typename RealType<Element>::Type;
if (dist.kind == Distribution::Gaussian) {
TensorFillRandomGaussian<Element, Layout>(
view,
seed,
static_cast<Real>(dist.gaussian.mean),
static_cast<Real>(dist.gaussian.stddev),
dist.int_scale,
stream);
} else if (dist.kind == Distribution::Uniform) {
TensorFillRandomUniform<Element, Layout>(
view,
seed,
static_cast<Real>(dist.uniform.max),
static_cast<Real>(dist.uniform.min),
dist.int_scale,
stream);
}
}
///////////////////////////////////////////////////////////////////////////////////////////////////
///////////////////////////////////////////////////////////////////////////////////////////////////
/// Fills a block of data with sequential elements
template <
typename Element
@@ -1670,27 +1752,30 @@ void BlockFillRandom(
Element *ptr,
size_t capacity,
uint64_t seed,
Distribution dist) {
Distribution dist,
cudaStream_t stream = nullptr) {
using Real = typename RealType<Element>::Type;
if (dist.kind == Distribution::Gaussian) {
BlockFillRandomGaussian<Element>(
ptr,
capacity,
seed,
static_cast<Real>(dist.gaussian.mean),
static_cast<Real>(dist.gaussian.stddev),
dist.int_scale);
ptr,
capacity,
seed,
static_cast<Real>(dist.gaussian.mean),
static_cast<Real>(dist.gaussian.stddev),
dist.int_scale,
stream);
}
else if (dist.kind == Distribution::Uniform) {
BlockFillRandomUniform<Element>(
ptr,
capacity,
seed,
ptr,
capacity,
seed,
static_cast<Real>(dist.uniform.max),
static_cast<Real>(dist.uniform.min),
dist.int_scale);
static_cast<Real>(dist.uniform.min),
dist.int_scale,
stream);
}
}
@@ -1786,14 +1871,17 @@ template <
typename Layout> ///< Layout function
void TensorCopyDiagonalIn(
TensorView<Element, Layout> view, ///< destination tensor
Element const *ptr) { ///< dense buffer of elements
Element const *ptr, ///< dense buffer of elements
cudaStream_t stream = nullptr) {
using Func = detail::TensorCopyDiagonalInFunc<Element, Layout>;
using Params = typename Func::Params;
TensorForEach<Func, Layout::kRank, Params>(
view.extent(),
Params(view, ptr)
Params(view, ptr),
/*grid_size*/0, /*block_size*/0,
stream
);
}
@@ -1890,14 +1978,17 @@ template <
typename Layout> ///< Layout function
void TensorCopyDiagonalOut(
Element *ptr, ///< dense buffer of elements
TensorView<Element, Layout> view) { ///< source tensor
TensorView<Element, Layout> view, ///< source tensor
cudaStream_t stream = nullptr) {
using Func = detail::TensorCopyDiagonalOutFunc<Element, Layout>;
using Params = typename Func::Params;
TensorForEach<Func, Layout::kRank, Params>(
view.extent(),
Params(view, ptr)
Params(view, ptr),
/*grid_size*/0, /*block_size*/0,
stream
);
}
tools/util/include/cutlass/util/reference/device/tensor_foreach.h
+16 -8
View File
@@ -45,7 +45,10 @@ template <typename Func, int Rank, typename Params>
struct TensorForEach {
/// Constructor performs the operation.
TensorForEach(Coord<Rank> size, Params params = Params(), int grid_size = 0, int block_size = 0) {
TensorForEach(
Coord<Rank> size, Params params = Params(),
int grid_size = 0, int block_size = 0,
cudaStream_t stream = nullptr) {
if (!grid_size || !block_size) {
@@ -67,7 +70,7 @@ struct TensorForEach {
dim3 grid(grid_size, 1, 1);
dim3 block(block_size, 1, 1);
kernel::TensorForEach<Func, Rank, Params><<< grid, block >>>(size, params);
kernel::TensorForEach<Func, Rank, Params><<< grid, block, 0, stream >>>(size, params);
}
};
@@ -78,7 +81,10 @@ template <typename Func, int Rank, typename Params>
struct TensorDiagonalForEach {
/// Constructor performs the operation
TensorDiagonalForEach(Coord<Rank> size, Params params = Params(), int start = 0, int end = -1, int block_size = 128) {
TensorDiagonalForEach(
Coord<Rank> size, Params params = Params(),
int start = 0, int end = -1,
int block_size = 128, cudaStream_t stream = nullptr) {
if (end < 0) {
end = size.min();
@@ -87,7 +93,8 @@ struct TensorDiagonalForEach {
dim3 block(block_size, 1, 1);
dim3 grid((end - start + block_size - 1) / block_size, 1, 1);
kernel::TensorDiagonalForEach<Func, Rank, Params><<< grid, block >>>(size, params, start, end);
kernel::TensorDiagonalForEach<Func, Rank, Params><<< grid, block, 0, stream >>>(
size, params, start, end);
}
};
@@ -99,11 +106,12 @@ struct BlockForEach {
/// Constructor performs the operation.
BlockForEach(
Element *ptr,
Element *ptr,
size_t capacity,
typename Func::Params params = typename Func::Params(),
int grid_size = 0,
int block_size = 0) {
int grid_size = 0,
int block_size = 0,
cudaStream_t stream = nullptr) {
if (!grid_size || !block_size) {
@@ -125,7 +133,7 @@ struct BlockForEach {
dim3 grid(grid_size, 1, 1);
dim3 block(block_size, 1, 1);
kernel::BlockForEach<Element, Func><<< grid, block >>>(ptr, capacity, params);
kernel::BlockForEach<Element, Func><<< grid, block, 0, stream >>>(ptr, capacity, params);
}
};
tools/util/include/cutlass/util/reference/host/gemm.h
+78
View File
@@ -372,6 +372,84 @@ struct Gemm<ElementA, LayoutA, ElementB, LayoutB, ElementC, LayoutC, ScalarType,
}
};
/// Partial specialization for AND-popc
template <typename ElementA, typename LayoutA, typename ElementB,
typename LayoutB, typename ElementC, typename LayoutC,
typename ScalarType, typename ComputeType>
struct Gemm<ElementA, LayoutA, ElementB, LayoutB, ElementC, LayoutC, ScalarType,
ComputeType, arch::OpAndPopc> {
void operator()(gemm::GemmCoord problem_size, ScalarType alpha,
TensorRef<ElementA, LayoutA> tensor_a,
TensorRef<ElementB, LayoutB> tensor_b, ScalarType beta,
TensorRef<ElementC, LayoutC> tensor_c,
ComputeType initial_accum = ComputeType(0)) {
static_assert(
LayoutA::kRank == 2 && LayoutB::kRank == 2 && LayoutC::kRank == 2,
"Tensors must be of rank 2");
compute_gemm<ElementA, LayoutA, ElementB, LayoutB, ElementC, LayoutC,
ScalarType, ComputeType, and_add<ComputeType>>(
problem_size, alpha, tensor_a, tensor_b, beta, tensor_c, initial_accum);
}
void operator()(gemm::GemmCoord problem_size, ScalarType alpha,
TensorRef<ElementA, LayoutA> tensor_a,
TensorRef<ElementB, LayoutB> tensor_b, ScalarType beta,
TensorRef<ElementC, LayoutC> tensor_c,
TensorRef<ElementC, LayoutC> tensor_d,
ComputeType initial_accum = ComputeType(0)) {
static_assert(
LayoutA::kRank == 2 && LayoutB::kRank == 2 && LayoutC::kRank == 2,
"Tensors must be of rank 2");
compute_gemm<ElementA, LayoutA, ElementB, LayoutB, ElementC, LayoutC,
ScalarType, ComputeType, and_add<ComputeType>>(
problem_size, alpha, tensor_a, tensor_b, beta, tensor_c, tensor_d, initial_accum);
}
};
////////////////////////////////////////////////////////////////////////////////////////////////////
/// Partial specialization for multiply-add
template <typename ElementA, typename LayoutA, typename ElementB,
typename LayoutB, typename ElementC, typename LayoutC,
typename ScalarType, typename ComputeType>
struct Gemm<ElementA, LayoutA, ElementB, LayoutB, ElementC, LayoutC, ScalarType,
ComputeType, arch::OpMultiplyAddFastF32> {
void operator()(gemm::GemmCoord problem_size, ScalarType alpha,
TensorRef<ElementA, LayoutA> tensor_a,
TensorRef<ElementB, LayoutB> tensor_b, ScalarType beta,
TensorRef<ElementC, LayoutC> tensor_c,
ComputeType initial_accum = ComputeType(0)) {
static_assert(
LayoutA::kRank == 2 && LayoutB::kRank == 2 && LayoutC::kRank == 2,
"Tensors must be of rank 2");
compute_gemm<ElementA, LayoutA, ElementB, LayoutB, ElementC, LayoutC,
ScalarType, ComputeType, multiply_add<ComputeType>>(
problem_size, alpha, tensor_a, tensor_b, beta, tensor_c, initial_accum);
}
void operator()(gemm::GemmCoord problem_size, ScalarType alpha,
TensorRef<ElementA, LayoutA> tensor_a,
TensorRef<ElementB, LayoutB> tensor_b, ScalarType beta,
TensorRef<ElementC, LayoutC> tensor_c,
TensorRef<ElementC, LayoutC> tensor_d,
ComputeType initial_accum = ComputeType(0)) {
static_assert(
LayoutA::kRank == 2 && LayoutB::kRank == 2 && LayoutC::kRank == 2,
"Tensors must be of rank 2");
compute_gemm<ElementA, LayoutA, ElementB, LayoutB, ElementC, LayoutC,
ScalarType, ComputeType, multiply_add<ComputeType>>(
problem_size, alpha, tensor_a, tensor_b, beta, tensor_c, tensor_d, initial_accum);
}
};
////////////////////////////////////////////////////////////////////////////////////////////////////
////////////////////////////////////////////////////////////////////////////////////////////////////
//
// Batched GEMM
tools/util/include/cutlass/util/reference/host/gett.hpp
+129 -31
View File
@@ -73,23 +73,29 @@ struct GettMainloopParams {
template<
class ElementScalar_,
class ElementScalingFactor_,
class ElementAccumulator_,
class ElementCompute_,
class TensorC_, // (M, N, L)
class TensorD_, // (M, N, L)
class TensorBias_, // (M, 1)
class TensorT_, // (M, N, L)
class VectorBias_ = TensorD_, // (M, 1)
class TensorAux_ = TensorD_, // (M, N, L)
class VectorAlpha_ = TensorD_, // (M, 1)
class VectorBeta_ = VectorAlpha_, // (M, 1)
class ActivationFunctor_ = cutlass::epilogue::thread::Identity<ElementCompute_>,
class BiasBinaryOp_ = cutlass::plus<ElementCompute_>
>
struct GettEpilogueParams {
using ElementScalar = ElementScalar_;
using ElementScalingFactor = ElementScalingFactor_;
using ElementAccumulator = ElementAccumulator_;
using ElementCompute = ElementCompute_;
using TensorC = TensorC_;
using TensorD = TensorD_;
using TensorBias = TensorBias_;
using TensorT = TensorT_;
using TensorAux = TensorAux_;
using VectorBias = VectorBias_;
using VectorAlpha = VectorAlpha_;
using VectorBeta = VectorBeta_;
using ActivationFunctor = ActivationFunctor_;
using BiasBinaryOp = BiasBinaryOp_;
@@ -97,17 +103,27 @@ struct GettEpilogueParams {
using LayoutC = typename TensorC::layout_type;
using EngineD = typename TensorD::engine_type;
using LayoutD = typename TensorD::layout_type;
using EngineBias = typename TensorBias::engine_type;
using LayoutBias = typename TensorBias::layout_type;
using EngineT = typename TensorT::engine_type;
using LayoutT = typename TensorT::layout_type;
ElementScalar alpha = ElementScalar(1);
ElementScalar beta = ElementScalar(0);
TensorC C{};
TensorD D{};
TensorBias Bias{};
TensorT T{};
VectorBias Bias{};
TensorAux Aux{};
VectorAlpha Valpha{};
VectorBeta Vbeta{};
ElementAccumulator* abs_max_D = nullptr;
ElementAccumulator* abs_max_Aux = nullptr;
ElementScalingFactor scale_a = ElementScalingFactor(1);
ElementScalingFactor scale_b = ElementScalingFactor(1);
ElementScalingFactor scale_c = ElementScalingFactor(1);
ElementScalingFactor scale_d = ElementScalingFactor(1);
ElementScalingFactor scale_aux = ElementScalingFactor(1);
bool beta_per_channel_scaling = false;
};
/////////////////////////////////////////////////////////////////////////////////////////////////
@@ -125,7 +141,9 @@ void Gett(
static int constexpr kBlockM = 64;
static int constexpr kBlockN = 64;
#if defined(_OPENMP)
#pragma omp parallel for collapse(3)
#endif
for (int64_t l = 0; l < cute::size<2>(mainloop_params.A.layout()); ++l) {
for (int64_t m = 0; m < cute::size<0>(mainloop_params.A.layout()); m += kBlockM) {
for (int64_t n = 0; n < cute::size<0>(mainloop_params.B.layout()); n += kBlockN) {
@@ -152,8 +170,8 @@ void gett_mainloop(
static_assert(cute::rank(typename MainloopParams::LayoutA{}) == 3, "M, K, B");
static_assert(cute::rank(typename MainloopParams::LayoutB{}) == 3, "N, K, B");
using ElementA = typename MainloopParams::EngineA::value_type;
using ElementB = typename MainloopParams::EngineB::value_type;
using ElementA = typename MainloopParams::TensorA::value_type;
using ElementB = typename MainloopParams::TensorB::value_type;
using RingOp = multiply_add<ElementAccumulator, ElementAccumulator, ElementAccumulator>;
RingOp fma_op;
@@ -217,16 +235,23 @@ void gett_epilogue(
static_assert(cute::rank(typename EpilogueParams::LayoutD{}) == 3, "N, K, B");
using ElementCompute = typename EpilogueParams::ElementCompute;
using ElementC = typename EpilogueParams::EngineC::value_type;
using ElementD = typename EpilogueParams::EngineD::value_type;
using ElementBias = typename EpilogueParams::EngineBias::value_type;
using ElementT = typename EpilogueParams::EngineT::value_type;
using ElementC = typename EpilogueParams::TensorC::value_type;
using ElementD = typename EpilogueParams::TensorD::value_type;
using ElementAux = typename EpilogueParams::TensorAux::value_type;
using ElementBias = typename EpilogueParams::VectorBias::value_type;
using ElementScalar = typename EpilogueParams::ElementScalar;
using ElementScalingFactor = typename EpilogueParams::ElementScalingFactor;
using ActivationFunctor = typename EpilogueParams::ActivationFunctor;
using BiasBinaryOp = typename EpilogueParams::BiasBinaryOp;
constexpr bool IsScalingAndAmaxOutputNeeded =
std::is_same_v<ElementD, cutlass::float_e4m3_t> or
std::is_same_v<ElementD, cutlass::float_e5m2_t>;
constexpr bool IsScalingAndAmaxAuxOutputNeeded =
std::is_same_v<ElementAux, cutlass::float_e4m3_t> or
std::is_same_v<ElementAux, cutlass::float_e5m2_t>;
// Input related converter
NumericConverter<ElementCompute, ElementAccumulator> accumulator_converter;
NumericConverter<ElementCompute, ElementC> source_converter;
@@ -234,9 +259,15 @@ void gett_epilogue(
// Scale related converter
NumericConverter<ElementCompute, ElementScalar> scale_converter;
NumericConverter<ElementCompute, ElementScalingFactor> scaling_factor_converter;
// Abs max converter
[[maybe_unused]] NumericConverter<ElementAccumulator, ElementCompute> abs_max_output_converter;
// Output related converter
NumericConverter<ElementD, ElementCompute> destination_converter;
NumericConverter<ElementT, ElementCompute> temporary_converter;
NumericConverter<ElementAux, ElementCompute> aux_destination_converter;
// Epilogue operations
multiply_add<ElementCompute, ElementCompute, ElementCompute> epilogue_fma;
multiplies<ElementCompute> mul;
@@ -250,13 +281,30 @@ void gett_epilogue(
// Do conversion
ElementCompute converted_alpha = scale_converter(epilogue_params.alpha);
ElementCompute converted_beta = scale_converter(epilogue_params.beta);
ElementCompute converted_scale_a = scaling_factor_converter(epilogue_params.scale_a);
ElementCompute converted_scale_b = scaling_factor_converter(epilogue_params.scale_b);
ElementCompute converted_scale_c = scaling_factor_converter(epilogue_params.scale_c);
ElementCompute converted_scale_d = scaling_factor_converter(epilogue_params.scale_d);
ElementCompute converted_scale_aux = scaling_factor_converter(epilogue_params.scale_aux);
// Init local var
[[maybe_unused]] ElementCompute local_abs_max_output = ElementCompute(0);
[[maybe_unused]] ElementCompute local_abs_max_aux_output = ElementCompute(0);
converted_alpha = mul(converted_alpha, mul(converted_scale_a, converted_scale_b));
converted_beta = mul(converted_beta, converted_scale_c);
for (int n_b = 0; n_b < kBlockN; ++n_b) {
for (int m_b = 0; m_b < kBlockM; ++m_b) {
if (m + m_b < cute::size<0>(epilogue_params.D.layout()) && n + n_b < cute::size<1>(epilogue_params.D.layout())) {
// Convert every type to ElementCompute first, do compute, convert to output type, write it out
ElementCompute converted_acc = accumulator_converter(acc[m_b][n_b]);
// per-row alpha
if (epilogue_params.Valpha.data()) {
converted_alpha = scale_converter(epilogue_params.Valpha(m + m_b));
}
ElementCompute output = mul(converted_alpha, converted_acc);
if (epilogue_params.Bias.data()) {
ElementCompute converted_bias = bias_converter(epilogue_params.Bias(m + m_b));
output = bias_op(output, converted_bias);
@@ -264,20 +312,54 @@ void gett_epilogue(
if (epilogue_params.C.data()) {
ElementCompute converted_src = source_converter(epilogue_params.C(m + m_b, n + n_b, l));
// per-row beta
if (epilogue_params.Vbeta.data()) {
converted_beta = scale_converter(epilogue_params.Vbeta(m + m_b));
}
output = epilogue_fma(converted_beta, converted_src, output);
}
if (epilogue_params.T.data()) {
// Store intermediate output
epilogue_params.T(m + m_b, n + n_b, l) = temporary_converter(output);
if (epilogue_params.Aux.data()) {
auto aux_output = output;
if constexpr (IsScalingAndAmaxAuxOutputNeeded) {
maximum_absolute_value_reduction<ElementCompute, true> amax_op;
local_abs_max_aux_output = amax_op(local_abs_max_aux_output, aux_output);
aux_output = epilogue_fma(converted_scale_aux, aux_output, ElementCompute(0));
}
epilogue_params.Aux(m + m_b, n + n_b, l) = aux_destination_converter(aux_output);
}
output = activation(output);
if constexpr (IsScalingAndAmaxOutputNeeded) {
maximum_absolute_value_reduction<ElementCompute, true> amax_op;
local_abs_max_output = amax_op(local_abs_max_output, output);
output = epilogue_fma(converted_scale_d, output, ElementCompute(0));
}
epilogue_params.D(m + m_b, n + n_b, l) = destination_converter(output);
}
}
}
#if defined(_OPENMP)
#pragma omp critical(Abs_Max_Data_Update)
#endif
{
if constexpr (IsScalingAndAmaxOutputNeeded) {
if (epilogue_params.abs_max_D) {
*epilogue_params.abs_max_D = maximum_with_nan_propogation<ElementAccumulator>{}(
*epilogue_params.abs_max_D, abs_max_output_converter(local_abs_max_output));
}
}
if constexpr (IsScalingAndAmaxAuxOutputNeeded) {
if (epilogue_params.abs_max_Aux) {
*epilogue_params.abs_max_Aux = maximum_with_nan_propogation<ElementAccumulator>{}(
*epilogue_params.abs_max_Aux, abs_max_output_converter(local_abs_max_aux_output));
}
}
}
}
/////////////////////////////////////////////////////////////////////////////////////////////////
@@ -310,15 +392,19 @@ void Gemm3x(
Layout layout_B = make_layout_rank3(mainloop_params.B);
Layout layout_C = make_layout_rank3(epilogue_params.C);
Layout layout_D = make_layout_rank3(epilogue_params.D);
Layout layout_Aux = make_layout_rank3(epilogue_params.Aux);
Layout layout_Bias = make_layout_rank3(epilogue_params.Bias);
Layout layout_T = make_layout_rank3(epilogue_params.T);
Layout layout_Valpha = make_layout_rank3(epilogue_params.Valpha);
Layout layout_Vbeta = make_layout_rank3(epilogue_params.Vbeta);
auto TensorA = make_tensor(mainloop_params.A.data(), layout_A);
auto TensorB = make_tensor(mainloop_params.B.data(), layout_B);
auto TensorC = make_tensor(epilogue_params.C.data(), layout_C);
auto TensorD = make_tensor(epilogue_params.D.data(), layout_D);
auto TensorBias = make_tensor(epilogue_params.Bias.data(), layout_Bias);
auto TensorT = make_tensor(epilogue_params.T.data(), layout_T);
auto TensorAux = make_tensor(epilogue_params.Aux.data(), layout_Aux);
auto VectorBias = make_tensor(epilogue_params.Bias.data(), layout_Bias);
auto VectorAlpha = make_tensor(epilogue_params.Valpha.data(), layout_Valpha);
auto VectorBeta = make_tensor(epilogue_params.Vbeta.data(), layout_Vbeta);
// Reconstruct mainloop params
GettMainloopParams<typename MainloopParams::ElementAccumulator,
@@ -331,19 +417,31 @@ void Gemm3x(
// Reconstruct epilogue params
GettEpilogueParams<typename EpilogueParams::ElementScalar,
typename EpilogueParams::ElementScalingFactor,
typename EpilogueParams::ElementAccumulator,
typename EpilogueParams::ElementCompute,
decltype(TensorC),
decltype(TensorD),
decltype(TensorBias),
decltype(TensorT)
decltype(VectorBias),
decltype(TensorAux),
decltype(VectorAlpha),
decltype(VectorBeta)
>
epilogue_params_converted{epilogue_params.alpha,
epilogue_params.beta,
TensorC,
TensorD,
TensorBias,
TensorT
VectorBias,
TensorAux,
VectorAlpha,
VectorBeta,
epilogue_params.abs_amax_D,
epilogue_params.abs_amax_Aux,
epilogue_params.scale_a,
epilogue_params.scale_b,
epilogue_params.scale_c,
epilogue_params.scale_d,
epilogue_params.scale_aux
};
Gett(mainloop_params_converted, epilogue_params_converted);
tools/util/include/cutlass/util/reference/host/tensor_fill.h
+135 -54
View File
@@ -38,6 +38,7 @@
#include <utility>
#include <cstdlib>
#include <cmath>
#include <random>
// Cutlass includes
#include "cutlass/cutlass.h"
@@ -156,6 +157,7 @@ struct RandomGaussianFunc {
double stddev;
int int_scale;
double pi;
double pnz;
//
// Methods
@@ -164,9 +166,10 @@ struct RandomGaussianFunc {
uint64_t seed_ = 0,
double mean_ = 0,
double stddev_ = 1,
int int_scale_ = -1
int int_scale_ = -1,
double pnz_ = 100.0
):
seed(seed_), mean(mean_), stddev(stddev_), int_scale(int_scale_), pi(std::acos(-1)) {
seed(seed_), mean(mean_), stddev(stddev_), int_scale(int_scale_), pi(std::acos(-1)), pnz(pnz_) {
std::srand((unsigned)seed);
}
@@ -184,12 +187,24 @@ struct RandomGaussianFunc {
// Scale and convert final result
Element result;
if (int_scale >= 0) {
rnd = double(int64_t(rnd * double(1 << int_scale))) / double(1 << int_scale);
result = static_cast<Element>(rnd);
// Sample from the Bernoulli distribution, and use the result to sample from the Gaussian
std::random_device rnd_device;
std::mt19937 bernoulli_rnd(rnd_device());
std::bernoulli_distribution bernoulli_dist(pnz / 100);
bool bernoulli_result = bernoulli_dist(bernoulli_rnd);
// 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);
result = static_cast<Element>(rnd);
}
else {
result = static_cast<Element>(rnd);
}
}
else {
result = static_cast<Element>(rnd);
result = static_cast<Element>(0);
}
return result;
@@ -205,6 +220,7 @@ struct RandomGaussianFunc<complex<Element> > {
double stddev;
int int_scale;
double pi;
double pnz;
//
// Methods
@@ -213,9 +229,10 @@ struct RandomGaussianFunc<complex<Element> > {
uint64_t seed_ = 0,
double mean_ = 0,
double stddev_ = 1,
int int_scale_ = -1
int int_scale_ = -1,
double pnz_ = 100.0
):
seed(seed_), mean(mean_), stddev(stddev_), int_scale(int_scale_), pi(std::acos(-1)) {
seed(seed_), mean(mean_), stddev(stddev_), int_scale(int_scale_), pi(std::acos(-1)), pnz(pnz_) {
std::srand((unsigned)seed);
}
@@ -228,14 +245,28 @@ struct RandomGaussianFunc<complex<Element> > {
detail::BoxMullerFunc func;
func(rnd, mean, stddev, pi);
if (int_scale >= 0) {
rnd[0] = double(int(rnd[0] * double(1 << int_scale)));
rnd[1] = double(int(rnd[1] * double(1 << int_scale)));
reals[0] = from_real<Element>(rnd[0] / double(1 << int_scale));
reals[1] = from_real<Element>(rnd[1] / double(1 << int_scale));
} else {
reals[0] = from_real<Element>(rnd[0]);
reals[1] = from_real<Element>(rnd[1]);
// Sample from the Bernoulli distribution, and use the result to sample from the Gaussian
std::random_device rnd_device;
std::mt19937 bernoulli_rnd(rnd_device());
std::bernoulli_distribution bernoulli_dist(pnz / 100);
bool bernoulli_result = bernoulli_dist(bernoulli_rnd);
// Sample from the Gaussian distribution for a nonzero element
if (bernoulli_result) {
if (int_scale >= 0) {
rnd[0] = double(int(rnd[0] * double(1 << int_scale)));
rnd[1] = double(int(rnd[1] * double(1 << int_scale)));
reals[0] = from_real<Element>(rnd[0] / double(1 << int_scale));
reals[1] = from_real<Element>(rnd[1] / double(1 << int_scale));
}
else {
reals[0] = from_real<Element>(rnd[0]);
reals[1] = from_real<Element>(rnd[1]);
}
}
else {
reals[0] = from_real<Element>(0);
reals[1] = from_real<Element>(0);
}
return complex<Element>(reals[0], reals[1]);
@@ -251,6 +282,7 @@ struct RandomGaussianFunc<Quaternion<Element> > {
double stddev;
int int_scale;
double pi;
double pnz;
//
// Methods
@@ -259,9 +291,10 @@ struct RandomGaussianFunc<Quaternion<Element> > {
uint64_t seed_ = 0,
double mean_ = 0,
double stddev_ = 1,
int int_scale_ = -1
int int_scale_ = -1,
double pnz_ = 100.0
):
seed(seed_), mean(mean_), stddev(stddev_), int_scale(int_scale_), pi(std::acos(-1)) {
seed(seed_), mean(mean_), stddev(stddev_), int_scale(int_scale_), pi(std::acos(-1)), pnz(pnz_) {
std::srand((unsigned)seed);
}
@@ -276,21 +309,37 @@ struct RandomGaussianFunc<Quaternion<Element> > {
func(rnd1, mean, stddev, pi);
func(rnd2, mean, stddev, pi);
if (int_scale >= 0) {
rnd1[0] = double(int(rnd1[0] * double(1 << int_scale)));
rnd1[1] = double(int(rnd1[1] * double(1 << int_scale)));
rnd2[0] = double(int(rnd2[0] * double(1 << int_scale)));
rnd2[1] = double(int(rnd2[1] * double(1 << int_scale)));
// Sample from the Bernoulli distribution, and use the result to sample from the Gaussian
std::random_device rnd_device;
std::mt19937 bernoulli_rnd(rnd_device());
std::bernoulli_distribution bernoulli_dist(pnz / 100);
bool bernoulli_result = bernoulli_dist(bernoulli_rnd);
reals[0] = from_real<Element>(rnd1[0] / double(1 << int_scale));
reals[1] = from_real<Element>(rnd1[1] / double(1 << int_scale));
reals[2] = from_real<Element>(rnd2[0] / double(1 << int_scale));
reals[3] = from_real<Element>(rnd2[1] / double(1 << int_scale));
} else {
reals[0] = from_real<Element>(rnd1[0]);
reals[1] = from_real<Element>(rnd1[1]);
reals[2] = from_real<Element>(rnd2[0]);
reals[3] = from_real<Element>(rnd2[1]);
// Sample from the Gaussian distribution for a nonzero element
if (bernoulli_result) {
if (int_scale >= 0) {
rnd1[0] = double(int(rnd1[0] * double(1 << int_scale)));
rnd1[1] = double(int(rnd1[1] * double(1 << int_scale)));
rnd2[0] = double(int(rnd2[0] * double(1 << int_scale)));
rnd2[1] = double(int(rnd2[1] * double(1 << int_scale)));
reals[0] = from_real<Element>(rnd1[0] / double(1 << int_scale));
reals[1] = from_real<Element>(rnd1[1] / double(1 << int_scale));
reals[2] = from_real<Element>(rnd2[0] / double(1 << int_scale));
reals[3] = from_real<Element>(rnd2[1] / double(1 << int_scale));
}
else {
reals[0] = from_real<Element>(rnd1[0]);
reals[1] = from_real<Element>(rnd1[1]);
reals[2] = from_real<Element>(rnd2[0]);
reals[3] = from_real<Element>(rnd2[1]);
}
}
else {
reals[0] = from_real<Element>(0);
reals[1] = from_real<Element>(0);
reals[2] = from_real<Element>(0);
reals[3] = from_real<Element>(0);
}
return Quaternion<Element>(reals[0], reals[1], reals[2], reals[3]);
@@ -389,11 +438,11 @@ void TensorFillRandomGaussian(
uint64_t seed, ///< seed for RNG
double mean = 0, ///< Gaussian distribution's mean
double stddev = 1, ///< Gaussian distribution's standard deviation
int bits = -1) { ///< If non-negative, specifies number of fractional bits that
/// are not truncated to zero. Permits reducing precision of
int bits = -1, ///< If non-negative, specifies number of fractional bits that
double pnz = 100.0) { /// are not truncated to zero. Permits reducing precision of
/// data.
detail::RandomGaussianFunc<Element> random_func(seed, mean, stddev, bits);
detail::RandomGaussianFunc<Element> random_func(seed, mean, stddev, bits, pnz);
detail::TensorFillGaussianFunc<Element, Layout> func(
dst,
@@ -411,16 +460,16 @@ template <
typename Element, ///< Element type
typename Layout> ///< Layout function
void TensorFillRandomGaussian(
TensorViewPlanarComplex<Element, Layout> dst, ///< destination tensor
uint64_t seed, ///< seed for RNG
double mean = 0, ///< Gaussian distribution's mean
double stddev = 1, ///< Gaussian distribution's standard deviation
int bits = -1) { ///< If non-negative, specifies number of fractional bits that
/// are not truncated to zero. Permits reducing precision of
/// data.
TensorViewPlanarComplex<Element, Layout> dst, ///< destination tensor
uint64_t seed, ///< seed for RNG
double mean = 0, ///< Gaussian distribution's mean
double stddev = 1, ///< Gaussian distribution's standard deviation
int bits = -1, ///< If non-negative, specifies number of fractional bits that
double pnz = 100.0) { /// are not truncated to zero. Permits reducing precision of
/// data.
TensorFillRandomGaussian(dst.view_real(), seed, mean, stddev, bits);
TensorFillRandomGaussian(dst.view_imag(), ~seed, mean, stddev, bits);
TensorFillRandomGaussian(dst.view_real(), seed, mean, stddev, bits, pnz);
TensorFillRandomGaussian(dst.view_imag(), ~seed, mean, stddev, bits, pnz);
}
///////////////////////////////////////////////////////////////////////////////////////////////////
@@ -434,11 +483,11 @@ void TensorFillSymmetricRandomGaussian(
cutlass::FillMode fill_mode, ///< FillMode for symmetric matrices
double mean = 0, ///< Gaussian distribution's mean
double stddev = 1, ///< Gaussian distribution's standard deviation
int bits = -1) { ///< If non-negative, specifies number of fractional bits that
/// are not truncated to zero. Permits reducing precision of
int bits = -1, ///< If non-negative, specifies number of fractional bits that
double pnz = 100.0) { /// are not truncated to zero. Permits reducing precision of
/// data.
detail::RandomGaussianFunc<Element> random_func(seed, mean, stddev, bits);
detail::RandomGaussianFunc<Element> random_func(seed, mean, stddev, bits, pnz);
detail::TensorFillSymmetricGaussianFunc<Element, Layout> func(
dst,
@@ -464,12 +513,12 @@ void BlockFillRandomGaussian(
uint64_t seed, ///< seed for RNG
double mean = 0, ///< Gaussian distribution's mean
double stddev = 1, ///< Gaussian distribution's standard deviation
int bits = -1) { ///< If non-negative, specifies number of fractional bits that
/// are not truncated to zero. Permits reducing precision of
int bits = -1, ///< If non-negative, specifies number of fractional bits that
double pnz = 100.0) { /// are not truncated to zero. Permits reducing precision of
/// data.
detail::RandomGaussianFunc<Element> random_func(seed, mean, stddev, bits);
detail::RandomGaussianFunc<Element> random_func(seed, mean, stddev, bits, pnz);
for (size_t i = 0; i < capacity; ++i) {
ReferenceFactory<Element>::get(ptr, i) = random_func();
@@ -801,10 +850,10 @@ void TensorFillRandomUniform(
uint64_t seed, ///< seed for RNG
double max = 1, ///< upper bound of distribution
double min = 0, ///< lower bound for distribution
int bits = -1) { ///< If non-negative, specifies number of fractional bits that
int bits = -1) { ///< If non-negative, specifies number of fractional bits that
/// are not truncated to zero. Permits reducing precision of
/// data.
/// data.
TensorFillRandomUniform(dst.view_real(), seed, max, min, bits);
TensorFillRandomUniform(dst.view_imag(), ~seed, max, min, bits);
}
@@ -1189,6 +1238,37 @@ void TensorFillSequential(
///////////////////////////////////////////////////////////////////////////////////////////////////
///////////////////////////////////////////////////////////////////////////////////////////////////
/// Fills a tensor with random values from a distribution.
template <
typename Element, ///< Element type
typename Layout> ///< Layout function
void TensorFillRandom(
TensorView<Element, Layout> view, ///< destination tensor
uint64_t seed,
Distribution dist) {
using Real = typename RealType<Element>::Type;
if (dist.kind == Distribution::Gaussian) {
TensorFillRandomGaussian(
view,
seed,
static_cast<Real>(dist.gaussian.mean),
static_cast<Real>(dist.gaussian.stddev),
dist.int_scale);
} else if (dist.kind == Distribution::Uniform) {
TensorFillRandomUniform(
view,
seed,
static_cast<Real>(dist.uniform.max),
static_cast<Real>(dist.uniform.min),
dist.int_scale);
}
}
///////////////////////////////////////////////////////////////////////////////////////////////////
///////////////////////////////////////////////////////////////////////////////////////////////////
/// Fills a block of data with sequential elements
template <
typename Element
@@ -1250,7 +1330,8 @@ void BlockFillRandom(
seed,
dist.gaussian.mean,
dist.gaussian.stddev,
dist.int_scale);
dist.int_scale,
dist.gaussian.pnz);
}
else if (dist.kind == Distribution::Uniform) {
BlockFillRandomUniform<Element>(