Updates for CUTLASS 3.5.0 (#1468)

This commit is contained in:
Vijay Thakkar
2024-04-11 21:33:40 -04:00
committed by GitHub
parent a40e08e9d5
commit 7d49e6c7e2
171 changed files with 7526 additions and 1888 deletions
+1 -1
View File
@@ -172,7 +172,7 @@ struct ReLu<Array<T, N>> {
template <typename T>
struct Clamp {
struct Arguments {
T lower_bound = CUTLASS_STL_NAMESPACE::numeric_limits<T>::min();
T lower_bound = CUTLASS_STL_NAMESPACE::numeric_limits<T>::lowest();
T upper_bound = CUTLASS_STL_NAMESPACE::numeric_limits<T>::max();
};
@@ -63,8 +63,52 @@ template<class Op>
struct kIsHeavy_member_or_false<Op, typename cutlass::platform::enable_if<Op::kIsHeavy>::type> {
static constexpr bool value = Op::kIsHeavy;
};
} // namespace (anonymous)
/////////////////////////////////////////////////////////////////////////////////////////////////
namespace detail {
struct EmptyArguments {};
template<class T, class = void>
struct ElementwiseOpDispatcher {
using Arguments = EmptyArguments;
T op;
CUTLASS_HOST_DEVICE
ElementwiseOpDispatcher(Arguments) {}
template <typename ValueType>
CUTLASS_HOST_DEVICE
ValueType operator()(ValueType value) {
return op(value);
}
};
template<class T>
struct ElementwiseOpDispatcher<T, std::void_t<typename T::Arguments>> {
using Arguments = typename T::Arguments;
Arguments args;
T op;
CUTLASS_HOST_DEVICE
ElementwiseOpDispatcher(Arguments args_):args(args_) {}
template <typename ValueType>
CUTLASS_HOST_DEVICE
ValueType operator()(ValueType value) {
return op(value, args);
}
};
}
/////////////////////////////////////////////////////////////////////////////////////////////////
/// This base class is meant to define the concept required of the
/// EpilogueWithBroadcast::OutputOp
template <
@@ -95,9 +139,13 @@ public:
using ElementwiseOp = ElementwiseOp_;
using BinaryOp = BinaryOp_;
using ElementwiseOpDispatcher = detail::ElementwiseOpDispatcher<ElementwiseOp>;
using ElementwiseArguments = typename ElementwiseOpDispatcher::Arguments;
// Indicates that this epilogue applies only one binary operation
static bool const kIsSingleSource = true;
using FragmentAccumulator = Array<ElementAccumulator, kElementsPerAccess>;
using FragmentCompute = Array<ElementCompute, kElementsPerAccess>;
using FragmentC = Array<ElementC, kElementsPerAccess>;
@@ -127,6 +175,7 @@ public:
ElementCompute beta; ///< scales source tensor
ElementCompute const *alpha_ptr; ///< pointer to accumulator scalar - if not null, loads it from memory
ElementCompute const *beta_ptr; ///< pointer to source scalar - if not null, loads it from memory
ElementwiseArguments elementwise; ///< Arguments for elementwise operation
//
// Methods
@@ -142,8 +191,9 @@ public:
CUTLASS_HOST_DEVICE
Params(
ElementCompute alpha,
ElementCompute beta
): alpha(alpha), beta(beta), alpha_ptr(nullptr), beta_ptr(nullptr) {
ElementCompute beta,
ElementwiseArguments elementwise_ = ElementwiseArguments{}
): alpha(alpha), beta(beta), alpha_ptr(nullptr), beta_ptr(nullptr), elementwise(elementwise_) {
}
@@ -157,8 +207,9 @@ public:
CUTLASS_HOST_DEVICE
Params(
ElementCompute const *alpha_ptr,
ElementCompute const *beta_ptr
): alpha(0), beta(0), alpha_ptr(alpha_ptr), beta_ptr(beta_ptr) {
ElementCompute const *beta_ptr,
ElementwiseArguments elementwise_ = ElementwiseArguments{}
): alpha(0), beta(0), alpha_ptr(alpha_ptr), beta_ptr(beta_ptr), elementwise(elementwise_) {
}
@@ -178,6 +229,7 @@ private:
ElementCompute alpha_;
ElementCompute beta_;
ElementwiseArguments const &elementwise_;
bool skip_elementwise_;
public:
@@ -188,7 +240,7 @@ public:
/// Constructor from Params
CUTLASS_HOST_DEVICE
LinearCombinationBiasElementwise(Params const &params) {
LinearCombinationBiasElementwise(Params const &params): elementwise_(params.elementwise) {
alpha_ = (params.alpha_ptr ? *params.alpha_ptr : params.alpha);
beta_ = (params.beta_ptr ? *params.beta_ptr : params.beta);
@@ -290,7 +342,7 @@ public:
FragmentC const &frag_C,
FragmentCompute const &V) const {
ElementwiseOp elementwise_op;
ElementwiseOpDispatcher elementwise_op(elementwise_);
BinaryOp binary_op;
FragmentCompute tmp_Accum = NumericArrayConverter<ElementCompute, ElementAccumulator, kElementsPerAccess>()(AB);
@@ -322,7 +374,7 @@ public:
FragmentAccumulator const &AB,
FragmentCompute const &V) const {
ElementwiseOp elementwise_op;
ElementwiseOpDispatcher elementwise_op(elementwise_);
BinaryOp binary_op;
FragmentCompute tmp_Accum = NumericArrayConverter<ElementCompute, ElementAccumulator, kElementsPerAccess>()(AB);
@@ -432,17 +432,12 @@ public:
intermediate = mul_add_accumulator(alpha_, converted_accumulator, intermediate); // D = alpha * Accum + X
}
// Convert floats back to INT
FragmentAccumulator scaled_accumulator;
//
// Convert float => ElementOutput_ with clamping
//
NumericArrayConverter<ElementOutput, ElementCompute, kCount, Round> destination_converter;
NumericArrayConverter<int, ElementCompute, kCount, Round> compute_converter;
scaled_accumulator = compute_converter(intermediate);
// Convert to destination numeric type
NumericArrayConverter<ElementOutput, int, kCount, Round> destination_converter;
return destination_converter(scaled_accumulator);
return destination_converter(intermediate);
}
/// Computes linear scaling: D = alpha * accumulator
@@ -466,17 +461,12 @@ public:
intermediate = mul_add_accumulator(alpha_, converted_accumulator); // D = alpha * Accum
}
// Convert floats back to INT
FragmentAccumulator scaled_accumulator;
//
// Convert float => ElementOutput_ with clamping
//
NumericArrayConverter<ElementOutput, ElementCompute, kCount, Round> destination_converter;
NumericArrayConverter<int, ElementCompute, kCount, Round> compute_converter;
scaled_accumulator = compute_converter(intermediate);
// Convert to destination numeric type
NumericArrayConverter<ElementOutput, int, kCount, Round> destination_converter;
return destination_converter(scaled_accumulator);
return destination_converter(intermediate);
}
};
@@ -156,23 +156,24 @@ public:
NumericArrayConverter<ElementCompute, ElementOutput, kCount, Round> source_converter;
NumericArrayConverter<ElementCompute, ElementAccumulator, kCount, Round> accumulator_converter;
ComputeFragment converted_source(
ComputeFragment converted_source{
source_converter(source.real),
source_converter(source.imag));
source_converter(source.imag)};
ComputeFragment converted_accumulator(
ComputeFragment converted_accumulator{
accumulator_converter(accumulator.real),
accumulator_converter(accumulator.imag));
// Perform binary operations
ComputeFragment intermediate;
accumulator_converter(accumulator.imag)};
multiplies<Array<ElementCompute, kCount> > mul_op;
multiply_add<Array<ElementCompute, kCount> > mul_add_op;
// Perform binary operations
// complex multiply: I = beta * C
intermediate.real = mul_op(beta_.real(), converted_source.real);
intermediate.imag = mul_op(beta_.real(), converted_source.imag);
ComputeFragment intermediate {
mul_op(beta_.real(), converted_source.real),
mul_op(beta_.real(), converted_source.imag)
};
intermediate.real = mul_add_op(-beta_.imag(), converted_source.imag, intermediate.real);
intermediate.imag = mul_add_op( beta_.imag(), converted_source.real, intermediate.imag);
@@ -187,9 +188,9 @@ public:
// Convert to destination numeric type
NumericArrayConverter<ElementOutput, ElementCompute, kCount, Round> destination_converter;
return FragmentOutput(
return FragmentOutput{
destination_converter(intermediate.real),
destination_converter(intermediate.imag));
destination_converter(intermediate.imag)};
}
/// Computes linear scaling: D = alpha * accumulator + beta * source
@@ -200,19 +201,19 @@ public:
// Convert source to interal compute numeric type
NumericArrayConverter<ElementCompute, ElementAccumulator, kCount, Round> accumulator_converter;
ComputeFragment converted_accumulator(
ComputeFragment converted_accumulator{
accumulator_converter(accumulator.real),
accumulator_converter(accumulator.imag));
accumulator_converter(accumulator.imag)};
// Perform binary operations
ComputeFragment intermediate;
multiplies<Array<ElementCompute, kCount> > mul_op;
multiply_add<Array<ElementCompute, kCount> > mul_add_op;
// complex multiply-add: I = alpha * AB + I
intermediate.real = mul_op(alpha_.real(), converted_accumulator.real);
intermediate.imag = mul_op(alpha_.real(), converted_accumulator.imag);
ComputeFragment intermediate {
mul_op(alpha_.real(), converted_accumulator.real),
mul_op(alpha_.real(), converted_accumulator.imag)
};
intermediate.real = mul_add_op(-alpha_.imag(), converted_accumulator.imag, intermediate.real);
intermediate.imag = mul_add_op( alpha_.imag(), converted_accumulator.real, intermediate.imag);
@@ -220,9 +221,9 @@ public:
// Convert to destination numeric type
NumericArrayConverter<ElementOutput, ElementCompute, kCount, Round> destination_converter;
return FragmentOutput(
return FragmentOutput{
destination_converter(intermediate.real),
destination_converter(intermediate.imag));
destination_converter(intermediate.imag)};
}
};