CUTLASS 3.1 (#915)

Co-authored-by: Aniket Shivam <ashivam@nvidia.com>

include/cutlass/conv/conv2d_problem_size.h

@@ -47,13 +47,6 @@
 
 #pragma once
 
-
-#if defined(__CUDACC_RTC__)
-#include <cuda/std/cmath>
-#else
-#include <cmath>
-#endif
-
 #include "cutlass/cutlass.h"
 #include "cutlass/tensor_coord.h"
 #include "cutlass/fast_math.h"

include/cutlass/conv/convolution.h

@@ -29,18 +29,18 @@
  *
  **************************************************************************************************/
 /*! \file
-    \brief 
+    \brief
 
-This file contains definitions and utility functions for describing convolution problem sizes in terms of 
-activation (NHWC), filter (KRSC), output (NPQK), pading (pad_h, pad_w), stride (stride_h, stride_w),
-dilation (dilation_h, dilation_w).  Furthermore, it defines helper functions to map cutlass' implicit gemm 
-tensor extents, sizes, data types to that of convolutions extents, sizes, and data types. 
+This file contains definitions and utility functions for describing convolution problem sizes in terms of
+activation (NHWC), filter (KRSC), output (NPQK), padding (pad_h, pad_w), stride (stride_h, stride_w), and
+dilation (dilation_h, dilation_w).  Furthermore, it defines helper functions to map CUTLASS's implicit gemm
+tensor extents, sizes, and data types to that of the convolution's extents, sizes, and data types.
 
                         * Mapping convolutions to Gemm computation *
 
-Cutlass employs ImplicitGemm algorithm to implement convolutions. ImplicitGemm algorithm runs gemm operation 
-on convolution tensors Activation, Filter, and Output . The underlying gemm operation follows the standard 
-gemm definition:
+Cutlass implements convolutions with the Implicit Gemm algorithm.  This algorithm performs a gemm
+(general matrix-matrix multiply) on the convolution tensors Activation, Filter, and Output.
+The underlying gemm operation follows the standard gemm definition:
 
                                      C = A * B + C
 
@@ -48,22 +48,23 @@ gemm definition:
                             C is source and output matrix
 
 
-For the three convolutional operators (Fprop, Dgrad, Wgrad), ImplicitGemm matrices A, B, and C are mapped on 
-to convolution tensors Activation, Filter and Output as per the below table:
+For the three convolutional operators (Fprop, Dgrad, Wgrad), ImplicitGemm matrices A, B, and C are mapped
+to convolution tensors Activation, Filter and Output as described in the table below.
 
         ___________________________________________________________________________
-         ConvolutionalOperator |        A        |      B         |       C                           
+         ConvolutionalOperator |        A        |      B         |       C
         ___________________________________________________________________________
         |                      |                 |                |               |
-        |       Fprop          |    Activation   |    Filter      |     Output    |  
-        |       Dgrad          |     Output      |    Filter      |   Activation  |  
-        |       Wgrad          |     Output      |  Activation    |     Filter    | 
+        |       Fprop          |    Activation   |    Filter      |     Output    |
+        |       Dgrad          |     Output      |    Filter      |   Activation  |
+        |       Wgrad          |     Output      |  Activation    |     Filter    |
         ___________________________________________________________________________
 
-In convolution codebase, DO NOT mix using (A, B, C) with (Acvitation, Filter, Output).
+In convolution codebase, DO NOT mix using (A, B, C) with (Activation, Filter, Output).
 
-For example, a convolution class/function with A, B, Output is confusing and error-prone. Instead use below 
-mapping functions and adhere to using either A, B, C or Acvitation, Filter, Output. 
+For example, it's confusing and error prone to document a convolution class or function
+as operating on "A, B, Output."  Instead, use the mapping functions below,
+and adhere to using either A, B, C or Activation, Filter, Output.
 
 Map elements' data types (ImplicitGemm -> Conv): GemmToConvElementMap
 Map elements' data types (Conv -> ImplicitGemm): ConvToGemmElementMap
@@ -83,20 +84,20 @@ namespace conv {
 ////////////////////////////////////////////////////////////////////////////////////////////////////
 
 /// Convolutional operator
-enum class Operator { 
-  kFprop, 
-  kDgrad, 
-  kWgrad 
+enum class Operator {
+  kFprop,
+  kDgrad,
+  kWgrad
 };
 
-/// Distinguishes convolution  from cross correlation
-enum class Mode { 
-  kCrossCorrelation, 
-  kConvolution 
+/// Distinguishes convolution from cross correlation
+enum class Mode {
+  kCrossCorrelation,
+  kConvolution
 };
 
 /// Selects among several implementation variants trading off performance with simplicity
-enum class IteratorAlgorithm { 
+enum class IteratorAlgorithm {
   kAnalytic,      ///< functionally correct in all cases but lower performance
   kOptimized,     ///< optimized for R <= 32, S <= 32 and unity-stride dgrad
   kFixedChannels, ///< Analytic algorithm optimized for fixed channel count (C == AccessSize)
@@ -113,9 +114,9 @@ enum class StrideSupport {
 };
 
 /// Identifies split-K mode
-enum class SplitKMode { 
-  kNone, 
-  kSerial, 
+enum class SplitKMode {
+  kNone,
+  kSerial,
   kParallel
 };
 

include/cutlass/conv/kernel/default_conv2d_fprop_with_broadcast.h

@@ -99,7 +99,7 @@ struct DefaultConv2dFpropWithBroadcast {
     AlignmentB
   >::Kernel;
 
-  // Replace epilogue
+  // Define epilogue
   using Epilogue = typename cutlass::conv::kernel::detail::DefaultConvEpilogueWithBroadcastTensorOp<
     ArchTag,
     typename ImplicitGemmBase::Epilogue::Shape,

include/cutlass/conv/kernel/default_conv2d_fprop_with_reduction.h

@@ -100,7 +100,7 @@ struct DefaultConv2dFpropWithReduction {
     AlignmentB
   >::Kernel;
 
-  // Replace epilogue
+  // Define epilogue
   using Epilogue = typename cutlass::conv::kernel::detail::DefaultConvEpilogueWithReductionTensorOp<
     ArchTag,
     typename ImplicitGemmBase::Epilogue::Shape,

include/cutlass/conv/kernel/default_conv2d_group_fprop.h

@@ -222,6 +222,138 @@ struct DefaultConv2dGroupFprop <
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
+/// Defines a kernel for Conv2dGroupFprop specialization for Analytic IteratorAlgorithm and
+/// 2 stage pipeline that supports all GroupMode.
+
+template <
+  typename ElementA,
+  typename LayoutA,
+  typename ElementB,
+  typename LayoutB,
+  typename ElementC,
+  typename LayoutC,
+  typename ElementAccumulator,
+  typename ArchTag,
+  typename ThreadblockShape,
+  typename WarpShape,
+  typename InstructionShape,
+  typename EpilogueOutputOp,
+  typename ThreadblockSwizzle,
+  typename MathOperatorTag,
+  conv::GroupMode GroupMode,
+  conv::StrideSupport StrideSupport,
+  int AlignmentA,
+  int AlignmentB
+>
+struct DefaultConv2dGroupFprop <
+  ElementA,
+  LayoutA,
+  ElementB,
+  LayoutB,
+  ElementC,
+  LayoutC,
+  ElementAccumulator,
+  arch::OpClassTensorOp,
+  ArchTag,
+  ThreadblockShape,
+  WarpShape,
+  InstructionShape,
+  EpilogueOutputOp,
+  ThreadblockSwizzle,
+  2,
+  MathOperatorTag,
+  GroupMode,
+  IteratorAlgorithm::kAnalytic,
+  StrideSupport,
+  AlignmentA,
+  AlignmentB
+> {
+
+  static_assert(std::is_same<LayoutA, cutlass::layout::TensorNHWC>::value,
+    "Current group conv only support NHWC layout");
+  static_assert(std::is_same<LayoutB, cutlass::layout::TensorNHWC>::value,
+    "Current group conv only support NHWC layout");
+  static_assert(std::is_same<LayoutC, cutlass::layout::TensorNHWC>::value,
+    "Current group conv only support NHWC layout");
+
+  // Define the core components from GEMM
+  using MmaCore = typename cutlass::gemm::threadblock::DefaultMmaCore<
+      ThreadblockShape, WarpShape, InstructionShape, ElementA, layout::RowMajor,
+      ElementB, layout::ColumnMajor, ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp,
+      2, MathOperatorTag>;
+
+  // Define iterators over tiles from the A operand
+  using ThreadMapA = typename MmaCore::IteratorThreadMapA;
+  using AccessTypeA = cutlass::AlignedArray<ElementA, AlignmentA>;
+  using IteratorA =
+    cutlass::conv::threadblock::TileIterator<
+      cutlass::conv::threadblock::Conv2dFpropActivationTileAccessIteratorAnalytic<
+        cutlass::MatrixShape<ThreadblockShape::kM, ThreadblockShape::kK>,
+        ElementA, LayoutA,
+        ThreadMapA,
+        AccessTypeA,
+        GroupMode
+      >
+    >;
+
+  using SmemIteratorA = typename MmaCore::SmemIteratorA;
+
+  // Define iterators over tiles from the B operand
+  using ThreadMapB = typename MmaCore::IteratorThreadMapB;
+  using AccessTypeB = cutlass::AlignedArray<ElementB, AlignmentB>;
+  using IteratorB =
+    cutlass::conv::threadblock::TileIterator<
+      cutlass::conv::threadblock::Conv2dFpropFilterTileAccessIteratorAnalytic<
+        cutlass::MatrixShape<ThreadblockShape::kK, ThreadblockShape::kN>,
+        ElementB, LayoutB,
+        ThreadMapB,
+        AccessTypeB,
+        GroupMode
+      >
+    >;
+
+  using SmemIteratorB = typename MmaCore::SmemIteratorB;
+
+  // Warp-level GEMM components
+  using WarpMmaTensorOp = typename MmaCore::MmaTensorOp;
+  using MmaPolicy = typename MmaCore::MmaPolicy;
+
+  // Define the Mma
+  using Mma = threadblock::ImplicitGemmPipelined<
+    ThreadblockShape,
+    IteratorA,
+    SmemIteratorA,
+    IteratorB,
+    SmemIteratorB,
+    ElementC,
+    LayoutC,
+    MmaPolicy
+  >;
+
+  static const int kPartitionsK = ThreadblockShape::kK / WarpShape::kK;
+
+  // Define the epilogue
+  using Epilogue = typename detail::DefaultConvEpilogue<
+    ArchTag,
+    ThreadblockShape,
+    WarpMmaTensorOp,
+    kPartitionsK,
+    EpilogueOutputOp
+  >::Epilogue;
+
+  // Define the kernel
+  using Kernel = cutlass::conv::kernel::ImplicitGemmConvolution<
+    Mma,
+    Epilogue,
+    ThreadblockSwizzle,
+    conv::Operator::kFprop,
+    Conv2dProblemSize,
+    GroupMode
+  >;
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
 /// Defines a kernel for Conv2dGroupFprop specialization for Optimized IteratorAlgorithm and multistage
 /// pipeline that supports GroupMode::kSingleGroup.
 template <

include/cutlass/conv/threadblock/conv2d_fprop_activation_tile_access_iterator_analytic.h

@@ -303,7 +303,7 @@ public:
   static Status can_implement(Conv2dProblemSize const &problem_size) {
 
     // check alignment constraint on iterator's contiguous dimension
-    if (problem_size.C % AccessType::kElements) {
+    if ((problem_size.C / problem_size.groups) % AccessType::kElements) {
       return Status::kErrorInvalidProblem;
     }
 

include/cutlass/conv/threadblock/conv2d_fprop_activation_tile_access_iterator_optimized.h

@@ -388,7 +388,7 @@ public:
   static Status can_implement(Conv2dProblemSize const &problem_size) {
 
     // check alignment constraint on iterator's contiguous dimension
-    if (problem_size.C % AccessType::kElements) {
+    if ((problem_size.C / problem_size.groups) % AccessType::kElements) {
       return Status::kErrorInvalidProblem;
     }
 

include/cutlass/conv/threadblock/conv2d_fprop_filter_tile_access_iterator_analytic.h

@@ -290,7 +290,7 @@ public:
   static Status can_implement(Conv2dProblemSize const &problem_size) {
 
     // check alignment constraint on iterator's contiguous dimension
-    if (problem_size.K % AccessType::kElements) {
+    if ((problem_size.C / problem_size.groups) % AccessType::kElements) {
       return Status::kErrorInvalidProblem;
     }
 

include/cutlass/conv/threadblock/conv2d_fprop_filter_tile_access_iterator_optimized.h

@@ -288,7 +288,7 @@ public:
   static Status can_implement(Conv2dProblemSize const &problem_size) {
 
     // check alignment constraint on iterator's contiguous dimension
-    if (problem_size.C % AccessType::kElements) {
+    if ((problem_size.C / problem_size.groups) % AccessType::kElements) {
       return Status::kErrorInvalidProblem;
     }
 

include/cutlass/conv/threadblock/conv2d_wgrad_activation_tile_access_iterator_analytic.h

@@ -268,7 +268,7 @@ public:
   static Status can_implement(Conv2dProblemSize const &problem_size) {
 
     // check alignment constraint on iterator's contiguous dimension
-    if (problem_size.K % AccessType::kElements) {
+    if (problem_size.C % AccessType::kElements) {
       return Status::kErrorInvalidProblem;
     }
 

include/cutlass/conv/threadblock/conv2d_wgrad_activation_tile_access_iterator_optimized.h

@@ -304,7 +304,7 @@ public:
   static Status can_implement(Conv2dProblemSize const &problem_size) {
 
     // check alignment constraint on iterator's contiguous dimension
-    if (problem_size.K % AccessType::kElements) {
+    if (problem_size.C % AccessType::kElements) {
       return Status::kErrorInvalidProblem;
     }
 

include/cutlass/conv/threadblock/conv2d_wgrad_output_gradient_tile_access_iterator_analytic.h

@@ -243,7 +243,7 @@ public:
   static Status can_implement(Conv2dProblemSize const &problem_size) {
 
     // check alignment constraint on iterator's contiguous dimension
-    if (problem_size.C % AccessType::kElements) {
+    if (problem_size.K % AccessType::kElements) {
       return Status::kErrorInvalidProblem;
     }
 

include/cutlass/conv/threadblock/conv2d_wgrad_output_gradient_tile_access_iterator_optimized.h

@@ -293,7 +293,7 @@ public:
   static Status can_implement(Conv2dProblemSize const &problem_size) {
 
     // check alignment constraint on iterator's contiguous dimension
-    if (problem_size.C % AccessType::kElements) {
+    if (problem_size.K % AccessType::kElements) {
       return Status::kErrorInvalidProblem;
     }
 

include/cutlass/conv/threadblock/conv3d_wgrad_activation_tile_access_iterator_analytic.h

@@ -270,7 +270,7 @@ public:
   static Status can_implement(Conv3dProblemSize const &problem_size) {
 
     // check alignment constraint on iterator's contiguous dimension
-    if (problem_size.K % (128/sizeof_bits<Element>::value)) {
+    if (problem_size.C % (128/sizeof_bits<Element>::value)) {
       return Status::kErrorInvalidProblem;
     }
 

include/cutlass/conv/threadblock/conv3d_wgrad_activation_tile_access_iterator_optimized.h

@@ -300,7 +300,7 @@ public:
   static Status can_implement(Conv3dProblemSize const &problem_size) {
 
     // check alignment constraint on iterator's contiguous dimension
-    if (problem_size.K % (128/sizeof_bits<Element>::value)) {
+    if (problem_size.C % (128/sizeof_bits<Element>::value)) {
       return Status::kErrorInvalidProblem;
     }
 

include/cutlass/conv/threadblock/conv3d_wgrad_output_gradient_tile_access_iterator_analytic.h

@@ -248,7 +248,7 @@ public:
   static Status can_implement(Conv3dProblemSize const &problem_size) {
 
     // check alignment constraint on iterator's contiguous dimension
-    if (problem_size.C % (128/sizeof_bits<Element>::value)) {
+    if (problem_size.K % (128/sizeof_bits<Element>::value)) {
       return Status::kErrorInvalidProblem;
     }
 

include/cutlass/conv/threadblock/conv3d_wgrad_output_gradient_tile_access_iterator_optimized.h

@@ -291,7 +291,7 @@ public:
   static Status can_implement(Conv3dProblemSize const &problem_size) {
 
     // check alignment constraint on iterator's contiguous dimension
-    if (problem_size.C % (128/sizeof_bits<Element>::value)) {
+    if (problem_size.K % (128/sizeof_bits<Element>::value)) {
       return Status::kErrorInvalidProblem;
     }
 

include/cutlass/conv/threadblock/implicit_gemm_multistage.h

@@ -134,6 +134,12 @@ public:
     /// Number of cp.async instructions to load on group of operand B
     static int const kAccessesPerGroupB =
         (AsyncCopyIterationsPerStageB + Base::kWarpGemmIterations - 1) / Base::kWarpGemmIterations;
+
+    // Optional staged-accumulation (e.g., tf32x3 kernels) for improved numerical
+    // accuracy, where each mainloop iteration first accumulates into a temporary
+    // set of freshly-cleared accumulators, which are subsequently added to the
+    // final accumulator set.
+    static bool const kStagedAccumulation = arch::UseStagedAccumulation<typename Operator::MathOperator>::value;
   };
 
  private:
@@ -387,10 +393,7 @@ public:
 
     FragmentC tmp_accum;
 
-    if (platform::is_same<typename Operator::MathOperator,
-                          arch::OpMultiplyAddFastF32>::value
-      || platform::is_same<typename Operator::MathOperator,
-                           arch::OpMultiplyAddComplexFastF32>::value) {
+    if (Detail::kStagedAccumulation) {
       tmp_accum.clear();
     }
 
@@ -444,10 +447,7 @@ public:
         copy_tiles_and_advance(iterator_A, iterator_B, group_start_iteration_A,
                                group_start_iteration_B);
 
-        if (platform::is_same<typename Operator::MathOperator,
-                              arch::OpMultiplyAddFastF32>::value
-          || platform::is_same<typename Operator::MathOperator,
-                               arch::OpMultiplyAddComplexFastF32>::value) {
+        if (Detail::kStagedAccumulation) {
           warp_mma(
             tmp_accum, 
             warp_transformed_frag_A[warp_mma_k % 2],
@@ -518,10 +518,7 @@ public:
 
     }
 
-    if (platform::is_same<typename Operator::MathOperator,
-                          arch::OpMultiplyAddFastF32>::value
-      || platform::is_same<typename Operator::MathOperator,
-                           arch::OpMultiplyAddComplexFastF32>::value) {
+    if (Detail::kStagedAccumulation) {
       accum = plus_accum(accum, tmp_accum); 
     }
   

include/cutlass/conv/threadblock/threadblock_swizzle.h

@@ -107,7 +107,7 @@ struct StridedDgradHorizontalThreadblockSwizzle :
     // compute number of tiles in m dimension
     int tile_m = get_strided_dgrad_tile_m(problem_size, tile_size.m());
 
-    // compute number of tiles in n dimension
+    // compute number of tiles in n dimension 
     int tile_n = (implicit_gemm_problem_size.n() + tile_size.n() - 1) / tile_size.n();
 
     return gemm::GemmCoord(
@@ -148,7 +148,7 @@ struct StridedDgradIdentityThreadblockSwizzle :
     // compute number of tiles in m dimension
     int tile_m = get_strided_dgrad_tile_m(problem_size, tile_size.m());
 
-    // compute number of tiles in n dimension
+    // compute number of tiles in n dimension 
     int tile_n = (implicit_gemm_problem_size.n() + tile_size.n() - 1) / tile_size.n();
 
     return gemm::GemmCoord(