Updates to fused epilogue (#383)
* Enhancements and fixes to fused GEMM and Convolution epilogue. * Need to explicitly list cudart as unit test library dependency.
This commit is contained in:
@@ -58,6 +58,7 @@ struct Identity {
|
||||
/// ReLu operator - propagates NaNs
|
||||
template <typename T>
|
||||
struct ReLu {
|
||||
static const bool kIsHeavy=false;
|
||||
CUTLASS_HOST_DEVICE
|
||||
T operator()(T const & threshold, T value) const {
|
||||
if (value < threshold) {
|
||||
@@ -76,6 +77,7 @@ struct ReLu {
|
||||
|
||||
template <typename T, int N>
|
||||
struct ReLu<Array<T, N>> {
|
||||
static const bool kIsHeavy=false;
|
||||
CUTLASS_HOST_DEVICE
|
||||
Array<T, N> operator()(T const & threshold, Array<T, N> const &frag) const {
|
||||
Array<T, N> result;
|
||||
|
||||
@@ -201,8 +201,8 @@ public:
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int i = 0; i < kElementsPerAccess; ++i) {
|
||||
ElementCompute z = binary_op(alpha_ * tmp_Accum[i] + beta_ * tmp_C[i], V[i]);
|
||||
result_Z[i] = z;
|
||||
result_T[i] = skip_elementwise_ ? z : elementwise_op(z);
|
||||
result_T[i] = z;
|
||||
result_Z[i] = skip_elementwise_ ? z : elementwise_op(z);
|
||||
}
|
||||
|
||||
NumericArrayConverter<ElementZ, ElementCompute, kElementsPerAccess> convert_z;
|
||||
@@ -230,8 +230,8 @@ public:
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int i = 0; i < kElementsPerAccess; ++i) {
|
||||
ElementCompute z = binary_op(alpha_ * tmp_Accum[i], V[i]);
|
||||
result_Z[i] = z;
|
||||
result_T[i] = skip_elementwise_ ? z : elementwise_op(z);
|
||||
result_T[i] = z;
|
||||
result_Z[i] = skip_elementwise_ ? z : elementwise_op(z);
|
||||
}
|
||||
|
||||
NumericArrayConverter<ElementZ, ElementCompute, kElementsPerAccess> convert_z;
|
||||
|
||||
Reference in New Issue
Block a user