Updates for CUTLASS 3.5.0 (#1468)
This commit is contained in:
@@ -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 ¶ms) {
|
||||
LinearCombinationBiasElementwise(Params const ¶ms): 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)};
|
||||
}
|
||||
};
|
||||
|
||||
|
||||
Reference in New Issue
Block a user