Add Snake activation functor for EVT (#3184)
Introduces cutlass::epilogue::thread::Snake, a two-operand activation functor implementing Snake_a(x) = x + (1/a) * sin^2(a*x) from Ziyin et al. 2020 (arXiv:2006.08195). The per-channel learnable frequency `a` flows through an EVT child (e.g. Sm90RowBroadcast), composing into Sm90EVT<Sm90Compute<Snake, ...>, x_node, alpha_node> for fused GEMM+Snake epilogues used in neural vocoders. Adds unit tests in test/unit/epilogue/thread/activation.cu covering f32 and bf16 paths, validated against float64 reference goldens. Closes #3141
This commit is contained in:
@@ -905,6 +905,41 @@ struct ElementwiseFilter<Array<T, N> > {
|
||||
}
|
||||
};
|
||||
|
||||
// Snake activation: Snake_a(x) = x + (1/a) * sin^2(a*x)
|
||||
// Introduced in Ziyin, Hartwig, Ueda, "Neural Networks Fail to Learn
|
||||
// Periodic Functions and How to Fix It," NeurIPS 2020 (arXiv:2006.08195).
|
||||
// The per-channel learnable frequency `a` is passed as the second operand
|
||||
// (intended to flow through an EVT child such as Sm90RowBroadcast).
|
||||
// Caller must ensure a != 0 (the formula is singular at a = 0).
|
||||
template <typename T>
|
||||
struct Snake {
|
||||
static const bool kIsHeavy = true;
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
T operator()(T const& x, T const& alpha) const {
|
||||
float xf = float(x);
|
||||
float af = float(alpha);
|
||||
float s = fast_sin(af * xf);
|
||||
return T(xf + s * s / af);
|
||||
}
|
||||
};
|
||||
|
||||
template <typename T, int N>
|
||||
struct Snake<Array<T, N>> {
|
||||
static const bool kIsHeavy = true;
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
Array<T, N> operator()(Array<T, N> const& x, Array<T, N> const& alpha) const {
|
||||
Array<T, N> result;
|
||||
Snake<T> scalar_op;
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int i = 0; i < N; ++i){
|
||||
result[i] = scalar_op(x[i], alpha[i]);
|
||||
}
|
||||
return result;
|
||||
}
|
||||
};
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace thread
|
||||
|
||||
Reference in New Issue
Block a user