Add epilogue functor for residual block fusion (#391)

* Add epilogue functor for residual block fusion

* Do not run split-k tests when ActivationOp is not Identity

* explain TestSplitK param

* return early

test/unit/conv/device/conv2d_fprop_with_broadcast_sm75.cu

@@ -28,15 +28,16 @@
 
 #include "../../common/cutlass_unit_test.h"
 #include "cutlass/cutlass.h"
-
+#include "cutlass/array.h"
 #include "cutlass/epilogue/thread/linear_combination_bias_elementwise.h"
-#include "cutlass/epilogue/thread/linear_combination_bias_relu.h"
+#include "cutlass/epilogue/thread/linear_combination_residual_block.h"
+#include "cutlass/epilogue/thread/activation.h"
 
 #include "cutlass/conv/kernel/default_conv2d_fprop_with_broadcast.h"
 #include "cutlass/conv/device/implicit_gemm_convolution.h"
 
 #include "conv2d_with_broadcast_testbed.h"
-        
+
 #if defined(CUTLASS_ARCH_MMA_SM75_SUPPORTED)
 
 TEST(SM75_Device_Conv2d_Fprop_With_Broadcast_Analytic_ImplicitGemm_f16nhwc_f16nhwc_f32nhwc_tensor_op_f32,
@@ -83,6 +84,87 @@ TEST(SM75_Device_Conv2d_Fprop_With_Broadcast_Analytic_ImplicitGemm_f16nhwc_f16nh
   EXPECT_TRUE(test::conv::device::TestAllConv2dWithBroadcast<Conv2dFprop>());
 }
 
+// Test residual block fusion: UnaryOp(BinaryOp(ActivationOp(Conv2d(X) + bias), residual))
+// LinearCombinationResidualBlock does not support the split-k mode unless ActivationOp is Identity.
+// This is because the activation needs to be applied to the fully accumulated output of the Conv2d op,
+// which only the last thread block would have an access to, before applying BinaryOp.
+// The epilogue functor in the last thread block would have to be given three inputs, namely
+// partial outputs, bias, and residual, but this is not supported in the current interface.
+// Set TestSplitK = false to skip split-k tests with non-trivial ActivationOp.
+template <
+ typename ElementAccumulator,
+ template<typename T> class ActivationOp,
+ template<typename T> class BinaryOp,
+ template<typename T> class UnaryOp,
+ bool TestSplitK = true
+>
+void TestResidaulBlock() {
+  using ElementA = cutlass::half_t;
+  using ElementB = cutlass::half_t;
+  using ElementC = cutlass::half_t;
+  using ElementD = ElementC;
+  using ElementCompute = ElementAccumulator;
+
+  using EpilogueOutputOp = cutlass::epilogue::thread::LinearCombinationResidualBlock<
+    ElementD,
+    ElementAccumulator,
+    ElementCompute,
+    ElementC,
+    8,
+    ActivationOp,
+    BinaryOp,
+    UnaryOp
+  >;
+
+  using Conv2dFpropKernel = typename cutlass::conv::kernel::DefaultConv2dFpropWithBroadcast<
+    ElementA, cutlass::layout::TensorNHWC,
+    ElementB, cutlass::layout::TensorNHWC,
+    ElementC, cutlass::layout::TensorNHWC,
+    ElementAccumulator,
+    cutlass::arch::OpClassTensorOp,
+    cutlass::arch::Sm75,
+    cutlass::gemm::GemmShape<128, 128, 32>,
+    cutlass::gemm::GemmShape<64, 64, 32>,
+    cutlass::gemm::GemmShape<16, 8, 8>,
+    EpilogueOutputOp,
+    cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>,
+    2,
+    cutlass::arch::OpMultiplyAdd,
+    cutlass::conv::IteratorAlgorithm::kAnalytic
+  >::Kernel;
+
+  using Conv2dFprop = cutlass::conv::device::ImplicitGemmConvolution<Conv2dFpropKernel>;
+
+  struct ReferenceOp {
+    using OutputOp = typename Conv2dFprop::EpilogueOutputOp;
+    using ElementZ = typename OutputOp::ElementZ;
+
+    ActivationOp<ElementCompute> activation;
+    BinaryOp<ElementCompute> binary_op;
+    UnaryOp<ElementCompute> unary_op;
+
+    void operator()(ElementZ &Z, ElementZ&, ElementCompute conv2d, ElementCompute residual) {
+      Z = ElementZ(unary_op(binary_op(activation(conv2d), residual)));
+    }
+  };
+
+  bool passed = test::conv::device::TestAllConv2dWithBroadcast<Conv2dFprop, ReferenceOp, true, TestSplitK>();
+  EXPECT_TRUE(passed);
+}
+
+TEST(SM75_Device_Conv2d_Fprop_With_Residual_Block_Plus_Analytic_ImplicitGemm_f16nhwc_f16nhwc_f32nhwc_tensor_op_f32,
+     128x128_32x2_64x64x32) {
+  // Resnet
+  TestResidaulBlock<cutlass::half_t, cutlass::epilogue::thread::Identity, cutlass::plus, cutlass::epilogue::thread::ReLu>();
+}
+
+TEST(SM75_Device_Conv2d_Fprop_With_Residual_Block_Multiply_Analytic_ImplicitGemm_f16nhwc_f16nhwc_f32nhwc_tensor_op_f32,
+     128x128_32x2_64x64x32) {
+  // EfficientNet V2
+  // Do not run split-K tests since the activation op is not Identity.
+  TestResidaulBlock<float, cutlass::epilogue::thread::Sigmoid, cutlass::multiplies, cutlass::epilogue::thread::Identity, false>();
+}
+
 ////////////////////////////////////////////////////////////////////////////////
 
 #endif  // CUTLASS_ARCH_MMA_SM75_SUPPORTED

test/unit/conv/device/conv2d_with_broadcast_testbed.h

@@ -95,7 +95,8 @@ struct Conv2dWithBroadcastReferenceOp {
 
 template <
   typename Conv2d,
-  typename ReferenceOp = Conv2dWithBroadcastReferenceOp<Conv2d>
+  typename ReferenceOp,
+  bool AddBroadcastFirst = false
 >
 class TestbedConv2dWithBroadcast {
 public:
@@ -113,7 +114,8 @@ public:
   using ElementT = typename EpilogueOutputOp::ElementT;
 
   static cutlass::conv::Operator const kConvolutionalOperator = Conv2d::kConvolutionalOperator;
-
+  static const bool kAddBroadcastFirst = AddBroadcastFirst;
+  static const bool kStoreT = EpilogueOutputOp::kStoreT;
 public:
 
   /// Initialization
@@ -270,7 +272,7 @@ public:
     cutlass::conv::Conv2dProblemSize const &problem_size,
     cutlass::conv::SplitKMode const &split_k_mode = cutlass::conv::SplitKMode::kSerial,
     ElementCompute alpha = ElementCompute(1),
-    ElementCompute beta = ElementCompute(0)) {
+    ElementCompute beta = ElementCompute(1)) {
 
     // Waive test if insufficient CUDA device
     if (!sufficient()) {
@@ -300,7 +302,7 @@ public:
       {alpha, beta},
       split_k_mode,
       tensor_Broadcast.device_data(),
-      tensor_T_computed.device_data(),
+      kStoreT ? tensor_T_computed.device_data() : nullptr,
       0,         // This must be zero
       implicit_gemm_tensor_c_extent(kConvolutionalOperator, problem_size).c()
     );
@@ -338,7 +340,8 @@ public:
     //
     // Reference check
     //
-
+    // When kAddBroadcastFirst is true, add bias on the host
+    ElementCompute beta_ref = kAddBroadcastFirst ? ElementCompute(0) : beta;
 #if CUTLASS_CONV_TEST_UNIT_REFERENCE_DEVICE_ENABLED
 
     cutlass::reference::device::Conv2d<
@@ -358,7 +361,7 @@ public:
       tensor_C_reference.device_ref(),
       tensor_Y_reference.device_ref(),
       alpha, 
-      beta);
+      beta_ref);
 
     // sync host (copy device data to host) for dumping error output in case of mismatches
     tensor_Y_reference.sync_host();
@@ -382,7 +385,7 @@ public:
       tensor_C_reference.host_ref(),
       tensor_Y_reference.host_ref(),
       alpha, 
-      beta);
+      beta_ref);
 
 #endif
     ReferenceOp reference_op;
@@ -395,9 +398,16 @@ public:
   
             ElementZ z;
             ElementT t;
-    
-            reference_op(z, t, tensor_Y_reference.at({n, p, q, k}), tensor_Broadcast.at({0, 0, 0, k}));
-    
+	    ElementCompute accum = tensor_Y_reference.at({n, p, q, k});
+	    ElementCompute bias = ElementCompute(tensor_Broadcast.at({0, 0, 0, k}));
+
+            if (kAddBroadcastFirst) {
+              reference_op(z, t, accum + bias,
+                           beta * ElementCompute(tensor_C_reference.at({n, p, q, k})));
+            } else {
+              reference_op(z, t, accum, bias);
+            }
+
             tensor_Z_reference.at({n, p, q, k}) = z;
             tensor_T_reference.at({n, p, q, k}) = t;
           }
@@ -405,11 +415,11 @@ public:
       }
     }
 
-    passed = cutlass::reference::host::TensorEquals(
-      tensor_T_computed.host_view(), 
-      tensor_T_reference.host_view());
-
-    EXPECT_TRUE(passed);
+    if (kStoreT) {
+      passed = cutlass::reference::host::TensorEquals(
+          tensor_T_computed.host_view(), tensor_T_reference.host_view());
+      EXPECT_TRUE(passed);
+    }
 
     passed = cutlass::reference::host::TensorEquals(
       tensor_Z_computed.host_view(), 
@@ -479,10 +489,13 @@ public:
 // Additionaly, each conv2d test can provide conv problem sizes (conv_test_sizes) and blacklist of sizes 
 // (conv_blacklist_sizes)
 /////////////////////////////////////////////////////////////////////////////////////////////////////////////
-template <typename ImplicitGemm>
+template <typename ImplicitGemm,
+          typename ReferenceOp = Conv2dWithBroadcastReferenceOp<ImplicitGemm>,
+          bool AddBroadcastFirst = false,
+	  bool TestSplitK = true>
 bool TestAllConv2dWithBroadcast(
-  const Conv2dProblemVector & conv_test_sizes = Conv2dProblemVector(),
-  const Conv2dProblemVector & conv_blacklist_sizes = Conv2dProblemVector()) {
+  const Conv2dProblemVector &conv_test_sizes = Conv2dProblemVector(),
+  const Conv2dProblemVector &conv_blacklist_sizes = Conv2dProblemVector()) {
 
   bool passed = true;
 
@@ -490,7 +503,7 @@ bool TestAllConv2dWithBroadcast(
   // Testbed object
   //
 
-  TestbedConv2dWithBroadcast<ImplicitGemm> testbed;
+  TestbedConv2dWithBroadcast<ImplicitGemm, ReferenceOp, AddBroadcastFirst> testbed;
 
   //
   // Get conv problem sizes to run conv operator 
@@ -597,6 +610,9 @@ bool TestAllConv2dWithBroadcast(
     return passed;
   }
 
+  if (!TestSplitK)
+    return passed;
+
   // Sweep split-k-slice using serial and prallel reduction with non-unity alpha and non-zero beta for 
   // a single conv2d problem size. Convolution unit tests take a long time to run so only sweep parameters 
   // which are abolutely neccessary to catch functional bugs. The below code does provide option to sweep