CUTLASS 2.6 (#298)

CUTLASS 2.6

include/cutlass/epilogue/threadblock/predicated_tile_iterator.h

@@ -56,7 +56,7 @@ namespace threadblock {
 
 ////////////////////////////////////////////////////////////////////////////////
 
-/// Tile iterator used to load and store output tile from shared memory in epilogue.
+/// Tile iterator used to load and store output tile from global memory in epilogue.
 ///
 /// Satisfies: ReadableTileIterator | PredicatedTileIterator | ForwardTileIterator
 ///
@@ -105,6 +105,7 @@ public:
 
   /// Uses a non-template class
   struct Params : PredicatedTileIteratorParams {
+    using Base = PredicatedTileIteratorParams;
 
     CUTLASS_HOST_DEVICE
     Params() { }
@@ -115,9 +116,11 @@ public:
         layout.stride(0) * int(sizeof(AccessType)) / kElementsPerAccess,
         make_OutputTileThreadMapDesc<ThreadMap>()
       ) 
-    {
-        
-    }
+    { }
+
+    CUTLASS_HOST_DEVICE
+    Params(Base const &base) : 
+      Base(base) { }
   };
 
   /// Mask object
@@ -177,6 +180,14 @@ private:
   /// Internal state counter
   int state_[3];
  
+  //
+  // Static asserts about internal strides
+  //
+
+  static_assert(sizeof(extent_row_) == 4, "Expected 32b extents");
+  static_assert(sizeof(thread_start_row_) == 4, "Expected 32b extents");
+  static_assert(sizeof(PredicatedTileIteratorParams::stride) == 8, "Expected 64b strides");
+
 private:
 
   //
@@ -236,7 +247,7 @@ public:
 
   /// Loads a fragment from memory
   CUTLASS_DEVICE
-  void load_with_byte_offset(Fragment &frag, int64_t byte_offset) {
+  void load_with_byte_offset(Fragment &frag, int64_t byte_offset) const {
 
     uint8_t *byte_pointer = byte_pointer_;
     AccessType *frag_ptr = reinterpret_cast<AccessType *>(&frag);
@@ -292,18 +303,16 @@ public:
       }
     }
   }
-
-
   /// Loads a fragment from memory
   CUTLASS_DEVICE
-  void load(Fragment &frag) {
+  void load(Fragment &frag) const {
 
     load_with_byte_offset(frag, 0);
   }
 
   /// Stores a fragment to memory
   CUTLASS_DEVICE
-  void store_with_byte_offset(Fragment const &frag, int64_t byte_offset) {
+  void store_with_byte_offset(Fragment const &frag, int64_t byte_offset) const {
     uint8_t *byte_pointer = byte_pointer_;
     AccessType const *frag_ptr = reinterpret_cast<AccessType const *>(&frag);
 
@@ -357,7 +366,7 @@ public:
 
   /// Stores a fragment to memory
   CUTLASS_DEVICE
-  void store(Fragment const &frag) {
+  void store(Fragment const &frag) const {
 
     store_with_byte_offset(frag, 0);
   }
@@ -421,7 +430,7 @@ public:
 
 ////////////////////////////////////////////////////////////////////////////////
 
-/// Tile iterator used to load output tile from shared memory in epilogue.
+/// Tile iterator used to load output tile from global memory in epilogue.
 ///
 /// Satisfies: ReadableTileIterator | InterleavedPredicatedTileIterator | ForwardTileIterator
 ///
@@ -454,51 +463,23 @@ public:
   /// Memory access size
   using AccessType = AlignedArray<Element, ThreadMap::kElementsPerAccess>;
 
-  //
-  // Parameters struct
-  //
-
-  struct Params {
-
-    //
-    // Data members
-    //
-
-    LongIndex stride;               ///< stride in bytes between columns
-
-    LongIndex advance_row;          ///< amount to add to move to the next 'row' position
-    LongIndex advance_column;       ///< amount to add to move to the next 'column' position
-
-    //
-    // Methods
-    //
+  /// Uses a non-template class
+  struct Params : InterleavedPredicatedTileIteratorParams {
+    using Base = InterleavedPredicatedTileIteratorParams;
 
     CUTLASS_HOST_DEVICE
-    Status initialize(Index stride_) {
-
-      stride = LongIndex(stride_);
-
-      advance_row =
-          ThreadMap::Delta::kContiguous * sizeof_bits<Element>::value / 8;
-
-      advance_column = LongIndex(stride_) - ThreadMap::Iterations::kContiguous *
-                                                kElementsPerAccess *
-                                                sizeof_bits<Element>::value *
-                                                ThreadMap::kWarpSize / 8;
-
-      return Status::kSuccess;
-    }
+    Params() { }
 
     CUTLASS_HOST_DEVICE
-    Params() {
-      initialize(0);
-    }
+    Params(Layout const &layout): 
+      Base(
+        layout.stride(0) * int(sizeof(AccessType)) / kElementsPerAccess,
+        make_InterleavedPredicatedTileIteratorDesc<Element, ThreadMap>()
+      ) { }
 
     CUTLASS_HOST_DEVICE
-    Params(Layout const &layout) {
-
-      initialize(layout.stride(0) * int(sizeof(AccessType)) / kElementsPerAccess);
-    }
+    Params(Base const &base) : 
+      Base(base) { }
   };
 
   /// Mask object
@@ -705,7 +686,7 @@ public:
 
 ///////////////////////////////////////////////////////////////////////////////
 
-/// Tile iterator used to load output tile from shared memory in epilogue.
+/// Tile iterator used to load output tile from global memory in epilogue.
 ///
 /// Satisfies: ReadableTileIterator | InterleavedMaskedTileIterator | ForwardTileIterator
 ///