CUTLASS 2.3 initial commit (#134)

CUTLASS 2.3 adds GEMMs targeting Sparse Tensor Cores on the NVIDIA Ampere Architecture, fast SGEMM, and small matrix classes, bug fixes, and performance enhancements.

include/cutlass/transform/pitch_linear_thread_map.h

@@ -44,13 +44,17 @@ namespace transform {
 
 ////////////////////////////////////////////////////////////////////////////////
 
-/// Strip-mines a pitch-linear tile among a given number of threads, first along the contiguous
-/// dimension then along the strided dimension.
+/// Strip-mines a pitch-linear tile among a given number of threads, first along
+/// the contiguous dimension then along the strided dimension.
 ///
-/// The tile must be divisible by the thread count such that all threads may execute the same
-/// number of iterations with the same delta to exhaustively cover the tile.
+/// The tile must be divisible by the thread count such that all threads may
+/// execute the same number of iterations with the same delta to exhaustively
+/// cover the tile.
 ///
 /// This class satisfies the "RegularThreadMapping" concept.
+///
+/// This ThreadMap is used by SIMT kernels and operand E of the sparse tensor
+/// kernels.
 template <
   typename Shape_,
   int Threads,
@@ -96,16 +100,17 @@ struct PitchLinearStripminedThreadMap {
 
   /// Number of iterations by each thread
   using Iterations = typename platform::conditional<
-    Threads >= Detail::ShapeVec::kContiguous,
-    layout::PitchLinearShape<
-      1,
-      (Threads >= Detail::ShapeVec::kContiguous ? Detail::ShapeVec::kStrided / (kThreads / Detail::ShapeVec::kContiguous) : 0)
-    >,
-    layout::PitchLinearShape<
-      Detail::ShapeVec::kContiguous / kThreads,
-      Detail::ShapeVec::kStrided
-    >
-  >::type;
+      Threads >= Detail::ShapeVec::kContiguous,
+      layout::PitchLinearShape<
+          1,
+          // Redo the comparison here to work around divide by zero compiler
+          // error.  The compiler evaluates both path of platform::conditional.
+          (Threads >= Detail::ShapeVec::kContiguous
+               ? Detail::ShapeVec::kStrided /
+                     (kThreads / Detail::ShapeVec::kContiguous)
+               : 0)>,
+      layout::PitchLinearShape<Detail::ShapeVec::kContiguous / kThreads,
+                               Detail::ShapeVec::kStrided>>::type;
 
   /// Interval between accesses along each dimension of the tensor's logical coordinate space
   /// (in units of Elements)
@@ -125,13 +130,13 @@ struct PitchLinearStripminedThreadMap {
   /// (in units of Elements)
   CUTLASS_HOST_DEVICE
   static TensorCoord initial_offset(int thread_id) {
-    
     return TensorCoord(
       (thread_id % Detail::ShapeVec::kContiguous) * kElementsPerAccess, 
       thread_id / Detail::ShapeVec::kContiguous);
   }
 };
 
+/// This ThreadMap is used by GEMV
 template <
   typename Shape,
   int Threads,
@@ -196,6 +201,8 @@ struct PitchLinearTilePolicyStripminedThreadStrided
 
 /// Policy defining a warp-raked arrangement in which a shape is partitioned into contiguous
 /// elements.
+///
+/// This ThreadMap is used by tensor core kernels.
 template <
   typename Shape_,
   int Threads,
@@ -241,6 +248,14 @@ struct PitchLinearWarpRakedThreadMap {
       Shape::kStrided
     >;
 
+    static_assert(
+      !(ShapeInAccesses::kContiguous % WarpThreadArrangement::kContiguous),
+      "ShapeInAccesses must be divisible by WarpThreadArrangement.");
+
+    static_assert(
+      !(ShapeInAccesses::kStrided % WarpThreadArrangement::kStrided),
+      "ShapeInAccesses must be divisible by WarpThreadArrangement.");
+
     // compute number of warp-level accesses total
     using WarpAccessIterations = layout::PitchLinearShape<
       ShapeInAccesses::kContiguous / WarpThreadArrangement::kContiguous,
@@ -672,16 +687,17 @@ struct PitchLinear2DThreadTileStripminedThreadMap <Shape_, Threads, cutlass::lay
 
   /// Number of iterations by each thread
   using Iterations = typename platform::conditional<
-    Threads >= Detail::ShapeVec::kContiguous,
-    layout::PitchLinearShape<
-      1,
-      (Threads >= Detail::ShapeVec::kContiguous ? Detail::ShapeVec::kStrided / (kThreads / Detail::ShapeVec::kContiguous) : 0)
-    >,
-    layout::PitchLinearShape<
-      Detail::ShapeVec::kContiguous / kThreads,
-      Detail::ShapeVec::kStrided
-    >
-  >::type;
+      Threads >= Detail::ShapeVec::kContiguous,
+      layout::PitchLinearShape<
+          1,
+          // Redo the comparison here to work around divide by zero compiler
+          // error.  The compiler evaluates both path of platform::conditional.
+          (Threads >= Detail::ShapeVec::kContiguous
+               ? Detail::ShapeVec::kStrided /
+                     (kThreads / Detail::ShapeVec::kContiguous)
+               : 0)>,
+      layout::PitchLinearShape<Detail::ShapeVec::kContiguous / kThreads,
+                               Detail::ShapeVec::kStrided>>::type;
 
   /// Interval between accesses along each dimension of the tensor's logical coordinate space
   /// (in units of Elements)

include/cutlass/transform/threadblock/predicated_tile_access_iterator.h

@@ -292,7 +292,7 @@ class PredicatedTileAccessIterator<Shape_, Element_, layout::PitchLinear,
         min(threadblock_offset.strided() + residue_size, extent_.strided())
       );
     } else {
-      
+
       Index residue_size = (extent_[kAdvanceRank] - threadblock_offset.contiguous()) % Shape::kContiguous;
       if (!residue_size) {
         residue_size = Shape::kContiguous;
@@ -852,7 +852,7 @@ class PredicatedTileAccessIterator<Shape_, Element_, layout::RowMajor,
 
 ////////////////////////////////////////////////////////////////////////////////
 
-/// Specialization of PredicatedTileAccessIterator for interleaved-32 data.  It
+/// Specialization of PredicatedTileAccessIterator for interleaved data.  It
 /// is mapped to the congruous layout.
 ///
 /// Satisfies: ForwardTileIteratorConcept |
@@ -1032,7 +1032,7 @@ class PredicatedTileAccessIterator<Shape_, Element_,
 
 ////////////////////////////////////////////////////////////////////////////////
 
-/// Specialization of PredicatedTileAccessIterator for interleaved-32 data.  It
+/// Specialization of PredicatedTileAccessIterator for interleaved data.  It
 /// is mapped to the congruous layout.
 ///
 /// Satisfies: ForwardTileIteratorConcept |

include/cutlass/transform/threadblock/regular_tile_access_iterator_tensor_op.h

@@ -451,6 +451,10 @@ class RegularTileAccessIterator<Shape_, Element_,
 
   using ThreadMap = ThreadMap_;
 
+  static_assert(!(ThreadMap::Delta::kContiguous % kCrosswise),
+                "kCrosswise is the smallest unit in the contiguous dimension "
+                "for shared memory swizzling.");
+
   /// Internal details made public to facilitate introspection
   struct Detail {
     /// This iterator is specialized for an access size that is 128 bits in
@@ -554,8 +558,10 @@ class RegularTileAccessIterator<Shape_, Element_,
 
     int access_offset =
         stride_idx * ThreadMap::Delta::kStrided * stride_ / Layout::kFactor +
-        iteration_contiguous_ * ThreadMap::Delta::kContiguous /
-            ThreadMap::kElementsPerAccess;
+        // kCrosswise elements in the contiguous dimension would span to a
+        // shared memory cache line.
+        iteration_contiguous_ * (ThreadMap::Delta::kContiguous / kCrosswise) *
+            Layout::TileShape::kContiguous;
     char *access_byte_ptr =
         reinterpret_cast<char *>(access_ptr + access_offset);
     return reinterpret_cast<AccessType *>(access_byte_ptr + byte_offset_);

include/cutlass/transform/threadblock/regular_tile_iterator_pitch_linear.h

@@ -49,7 +49,8 @@ namespace threadblock {
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
-/// Regular tile iterator specialized for pitch-linear
+/// Regular tile iterator specialized for pitch-linear.  This one is used by 2-stage SIMT kernels
+/// and sparse tensor core meta data.
 template <
   typename Shape_,
   typename Element_,
@@ -139,7 +140,8 @@ public:
       for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) {
 
         int idx = c + s * ThreadMap::Iterations::kContiguous;
-        frag_ptr[idx] = access_ptr[c * ThreadMap::Delta::kContiguous];
+        frag_ptr[idx] = access_ptr[c * ThreadMap::Delta::kContiguous /
+                                   ThreadMap::kElementsPerAccess];
       }
 
       if (s + 1 < ThreadMap::Iterations::kStrided) {
@@ -180,7 +182,8 @@ public:
       for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) {
 
         int idx = c + s * ThreadMap::Iterations::kContiguous;
-        access_ptr[c * ThreadMap::Delta::kContiguous] = frag_ptr[idx];
+        access_ptr[c * ThreadMap::Delta::kContiguous /
+                   ThreadMap::kElementsPerAccess] = frag_ptr[idx];
       }
 
       if (s + 1 < ThreadMap::Iterations::kStrided) {