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

v4.5.1 update. (#3237)

Browse Source
This commit is contained in:
Junkai-Wu
2026-05-18 22:35:08 -04:00
committed by GitHub
parent e406c186f5
commit 982cb9e718
42 changed files with 6487 additions and 336 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tools/util/include/cutlass/util/reference/device/tensor_fill.h
+128 -8
View File
@@ -554,6 +554,112 @@ struct RandomUniformFunc {
}
};
/// Computes an exponent-uniform random distribution for UE8M0 scale factors.
template <>
struct RandomUniformFunc<float_ue8m0_t> {
using Element = float_ue8m0_t;
using FloatType = float;
/// Parameters structure
struct Params {
//
// Data members
//
uint64_t seed;
int exp_min;
int exp_range;
int int_scale; ///< Retained for Params compatibility; exponent is integral.
double pnan;
int exclude_zero; ///< Retained for Params compatibility; unused for UE8M0.
/// Default ctor
CUTLASS_HOST_DEVICE
Params() { }
//
// Methods
//
CUTLASS_HOST_DEVICE
static int closest_log2_exp(FloatType value) {
using CUTLASS_CMATH_NAMESPACE :: log2;
using CUTLASS_CMATH_NAMESPACE :: nearbyint;
// UE8M0 scale factors are strictly positive. Keep invalid lower bounds
// finite so callers using the generic [0, max] default do not produce NaN.
FloatType min_scale = FloatType(Element::bitcast(0x01));
FloatType positive_value = value > FloatType(0) ? value : min_scale;
return int(nearbyint(log2(positive_value)));
}
/// Construction of uniform RNG functor.
Params(
uint64_t seed_ = 0,
FloatType max_ = FloatType(1),
FloatType min_ = FloatType(0),
int int_scale_ = -1,
double pnan_ = 0,
int exclude_zero_ = -1
):
seed(seed_),
exp_min(closest_log2_exp(min_)),
exp_range(closest_log2_exp(max_) - closest_log2_exp(min_)),
int_scale(int_scale_),
pnan(pnan_),
exclude_zero(exclude_zero_) {
}
};
//
// Data members
//
/// Parameters object
Params params;
/// RNG state object
curandState_t rng_state;
//
// Methods
//
/// Device-side initialization of RNG
CUTLASS_DEVICE
RandomUniformFunc(Params const &params): params(params) {
uint64_t gtid = threadIdx.x + blockIdx.x * blockDim.x;
curand_init(params.seed, gtid, 0, &rng_state);
}
/// Compute random value and update RNG state
CUTLASS_DEVICE
Element operator()() {
// Draw random float in [0.0, 1.0] to determine if element should be NaN.
if constexpr (std::numeric_limits<Element>::has_quiet_NaN) {
if (params.pnan > 0 && (curand_uniform(&rng_state) < (params.pnan))) {
return Element(NAN);
}
}
using CUTLASS_CMATH_NAMESPACE :: pow;
FloatType rnd = random_uniform_float<FloatType>(&rng_state);
int exponent_count = params.exp_range + 1;
int exponent_offset = int(rnd * FloatType(exponent_count));
exponent_offset = exponent_offset < exponent_count ? exponent_offset : exponent_count - 1;
FloatType exp = FloatType(params.exp_min + exponent_offset);
FloatType sf = FloatType(pow(FloatType(2), exp));
return Element(sf);
}
};
/// Computes a random Gaussian distribution
template <typename Real>
struct RandomUniformFunc<complex<Real>> {
@@ -763,6 +869,16 @@ struct TensorFillRandomUniformFunc {
}
};
template <typename Element>
struct UniformDistributionValueType {
using Type = typename RealType<Element>::Type;
};
template <>
struct UniformDistributionValueType<float_ue8m0_t> {
using Type = float;
};
} // namespace detail
///////////////////////////////////////////////////////////////////////////////////////////////////
@@ -774,8 +890,10 @@ template <
void TensorFillRandomUniform(
TensorView<Element, Layout> view, ///< destination tensor
uint64_t seed, ///< seed for RNG
typename RealType<Element>::Type max = Element(1), ///< upper bound of distribution
typename RealType<Element>::Type min = Element(0), ///< lower bound for distribution
typename detail::UniformDistributionValueType<Element>::Type max =
typename detail::UniformDistributionValueType<Element>::Type(1), ///< upper bound of distribution
typename detail::UniformDistributionValueType<Element>::Type min =
typename detail::UniformDistributionValueType<Element>::Type(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.
@@ -805,8 +923,8 @@ void BlockFillRandomUniform(
Element *ptr,
size_t capacity,
uint64_t seed, ///< seed for RNG
typename RealType<Element>::Type max, ///< upper bound of distribution
typename RealType<Element>::Type min, ///< lower bound for distribution
typename detail::UniformDistributionValueType<Element>::Type max, ///< upper bound of distribution
typename detail::UniformDistributionValueType<Element>::Type min, ///< 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.
@@ -1768,6 +1886,7 @@ void TensorFillRandom(
) {
using Real = typename RealType<Element>::Type;
using UniformReal = typename detail::UniformDistributionValueType<Element>::Type;
if (dist.kind == Distribution::Gaussian) {
TensorFillRandomGaussian<Element, Layout>(
@@ -1782,8 +1901,8 @@ void TensorFillRandom(
TensorFillRandomUniform<Element, Layout>(
view,
seed,
static_cast<Real>(dist.uniform.max),
static_cast<Real>(dist.uniform.min),
static_cast<UniformReal>(dist.uniform.max),
static_cast<UniformReal>(dist.uniform.min),
dist.int_scale,
dist.uniform.pnan,
exclude_zero,
@@ -1830,6 +1949,7 @@ void BlockFillRandom(
cudaStream_t stream = nullptr) {
using Real = typename RealType<Element>::Type;
using UniformReal = typename detail::UniformDistributionValueType<Element>::Type;
if (dist.kind == Distribution::Gaussian) {
BlockFillRandomGaussian<Element>(
@@ -1846,8 +1966,8 @@ void BlockFillRandom(
ptr,
capacity,
seed,
static_cast<Real>(dist.uniform.max),
static_cast<Real>(dist.uniform.min),
static_cast<UniformReal>(dist.uniform.max),
static_cast<UniformReal>(dist.uniform.min),
dist.int_scale,
dist.uniform.pnan,
stream);
tools/util/include/cutlass/util/reference/host/tensor_fill.h
+68
View File
@@ -658,6 +658,74 @@ public:
}
};
/// Computes an exponent-uniform random distribution for UE8M0 scale factors.
template <>
struct RandomUniformFunc<float_ue8m0_t> {
using Element = float_ue8m0_t;
uint64_t seed;
int exp_min;
int exp_range;
int int_scale; ///< Retained for Params compatibility; exponent is integral.
double pnan;
private:
using engine_type = std::mt19937;
public:
engine_type bernoulli_rnd;
std::bernoulli_distribution bernoulli_dist;
bool exclude_zero; ///< Retained for Params compatibility; unused for UE8M0.
static int closest_log2_exp(double value) {
// UE8M0 scale factors are strictly positive. Keep invalid lower bounds
// finite so callers using the generic [0, max] default do not produce NaN.
double min_scale = double(Element::bitcast(0x01));
double positive_value = value > 0.0 ? value : min_scale;
return int(std::nearbyint(std::log2(positive_value)));
}
RandomUniformFunc(
uint64_t seed_ = 0,
double max = 1,
double min_ = 0,
int int_scale_ = -1,
double pnan_ = 0,
bool exclude_zero_ = false
):
seed(seed_),
exp_min(closest_log2_exp(min_)),
exp_range(closest_log2_exp(max) - closest_log2_exp(min_)),
int_scale(int_scale_),
pnan(pnan_),
bernoulli_rnd{static_cast<engine_type::result_type>(seed_)},
bernoulli_dist(pnan_),
exclude_zero(exclude_zero_)
{
std::srand((unsigned)seed);
}
/// Compute random value and update RNG state
Element operator()() {
// Sample from NaN distribution.
if constexpr (std::numeric_limits<Element>::has_quiet_NaN) {
if (pnan > 0 && bernoulli_dist(bernoulli_rnd)) {
return Element(NAN);
}
}
double rnd = double(std::rand()) / double(RAND_MAX);
int exponent_count = exp_range + 1;
int exponent_offset = int(rnd * double(exponent_count));
exponent_offset = exponent_offset < exponent_count ? exponent_offset : exponent_count - 1;
double sf = std::pow(2.0, double(exp_min + exponent_offset));
return Element(sf);
}
};
/// Partial specialization for initializing a complex value.
template <typename Element>
struct RandomUniformFunc<complex<Element> > {