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/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) {