CUTLASS 2.8 (#363)

CUTLASS 2.8

examples/13_two_tensor_op_fusion/kernel/b2b_implicit_gemm_convolution.h

@@ -79,6 +79,10 @@ struct B2bImplicitGemmConvolution {
   using ElementAccumulator = typename EpilogueOutputOp0::ElementAccumulator;
   using ElementCompute = typename EpilogueOutputOp0::ElementCompute;
 
+  /// Scale and Bias
+  using ElementScaleBias = typename B2bMma::IteratorAccumulatorScaleBias::Element;
+  using LayoutScaleBias = typename B2bMma::IteratorAccumulatorScaleBias::Layout;
+
   using WarpMmaOperator0 = typename B2bMma::Policy0::Operator;
   using WarpMmaOperator1 = typename B2bMma::Policy1::Operator;
 
@@ -103,13 +107,14 @@ struct B2bImplicitGemmConvolution {
 
   using TensorRefA0 = typename B2bMma::IteratorA0::TensorRef;
   using TensorRefB0 = typename B2bMma::IteratorB0::TensorRef;
+  using TensorRefScaleBias0 = typename B2bMma::IteratorAccumulatorScaleBias::TensorRef;
   using TensorRefB1 = typename B2bMma::IteratorB1::TensorRef;
   using TensorRefC = cutlass::TensorRef<ElementC, LayoutC>;
 
   /// Check iterator A and B convolution dimension are the same and 
   // set device::B2bImplicitGemmConvolution::kConvDim
   static_assert(B2bMma::IteratorA0::kConvDim == B2bMma::IteratorB0::kConvDim, 
-    "Convolution on different different dimensions is not supported");
+    "Convolution on different dimensions is not supported");
   static int const kConvDim = B2bMma::IteratorA0::kConvDim;
 
   /// Conv dimension and problem size structure (Conv2d or Conv3d)
@@ -148,6 +153,8 @@ struct B2bImplicitGemmConvolution {
     TensorRefA0 ref_A0;
     TensorRefB0 ref_B0;
     TensorRefC ref_C0;
+    TensorRefScaleBias0 ref_Scale0;
+    TensorRefScaleBias0 ref_Bias0;
     TensorRefB1 ref_B1;
     TensorRefC ref_C1;
     TensorRefC ref_D1;
@@ -178,6 +185,8 @@ struct B2bImplicitGemmConvolution {
       TensorRefA0 const & ref_A0,
       TensorRefB0 const & ref_B0,
       TensorRefC const & ref_C0,
+      TensorRefScaleBias0 const & ref_Scale0,
+      TensorRefScaleBias0 const & ref_Bias0,
       TensorRefB1 const & ref_B1,
       TensorRefC const & ref_C1,
       TensorRefC const & ref_D1,
@@ -190,6 +199,8 @@ struct B2bImplicitGemmConvolution {
       ref_A0(ref_A0),
       ref_B0(ref_B0),
       ref_C0(ref_C0),
+      ref_Scale0(ref_Scale0),
+      ref_Bias0(ref_Bias0),
       ref_B1(ref_B1),
       ref_C1(ref_C1),
       ref_D1(ref_D1),
@@ -218,6 +229,8 @@ struct B2bImplicitGemmConvolution {
     typename B2bMma::IteratorB0::Element const *ptr_B0;
     typename Epilogue::OutputTileIterator::Params iterator_C0;
     typename Epilogue::OutputTileIterator::Element *ptr_C0;
+    typename B2bMma::IteratorAccumulatorScaleBias::Element *ptr_Scale0;
+    typename B2bMma::IteratorAccumulatorScaleBias::Element *ptr_Bias0;
     typename B2bMma::IteratorB1::Params iterator_B1;
     typename B2bMma::IteratorB1::Element const *ptr_B1;
     typename Epilogue::OutputTileIterator::Params iterator_C1;
@@ -252,6 +265,8 @@ struct B2bImplicitGemmConvolution {
       ptr_B0(args.ref_B0.data()),
       iterator_C0(ConvOutputIteratorParameter::layout(args.ref_C0)),
       ptr_C0(args.ref_C0.data()),
+      ptr_Scale0(args.ref_Scale0.data()),
+      ptr_Bias0(args.ref_Bias0.data()),
       iterator_B1(args.problem_size_1, args.ref_B1.layout()),
       ptr_B1(args.ref_B1.data()),
       iterator_C1(ConvOutputIteratorParameter::layout(args.ref_C1)),
@@ -350,6 +365,28 @@ struct B2bImplicitGemmConvolution {
     int warp_idx = __shfl_sync(0xffffffff, threadIdx.x / 32, 0);
     int lane_idx = threadIdx.x % 32;
 
+    // Construct iterators to accumulator scale/bias vector
+    typename B2bMma::IteratorAccumulatorScaleBias iterator_Scale0(
+      params.ptr_Scale0,
+      {1, params.problem_size_0.K},
+      thread_idx,
+      warp_idx,
+      MatrixCoord(
+        0, threadblock_tile_idx.n() * B2bMma::Shape0::kN
+      )
+    );
+
+    typename B2bMma::IteratorAccumulatorScaleBias iterator_Bias0(
+      params.ptr_Bias0,
+      {1, params.problem_size_0.K},
+      thread_idx,
+      warp_idx,
+      MatrixCoord(
+        0, threadblock_tile_idx.n() * B2bMma::Shape0::kN
+      )
+    );
+
+
     //
     // Main loop
     //
@@ -366,7 +403,8 @@ struct B2bImplicitGemmConvolution {
     accumulators.clear();
 
     // Compute threadblock-scoped matrix multiply-add
-    b2bMma(params.gemm_k_iterations_0, accumulators, iterator_A0, iterator_B0, iterator_B1, src_accum, output_op_0);
+    b2bMma(params.gemm_k_iterations_0, accumulators, iterator_A0, iterator_B0, 
+        iterator_Scale0, iterator_Bias0, iterator_B1, src_accum, output_op_0);
 
     //
     // Epilogue