Updates and Bug fixes to CUTLASS 3.3 (#1232)

include/cutlass/epilogue/collective/builders/sm90_builder.inl

@@ -154,7 +154,7 @@ sm90_get_smem_store_op_for_accumulator() {
   }
   else {
     // auto-vectorizing store
-    return DefaultCopy{};
+    return AutoVectorizingCopyWithAssumedAlignment{};
   }
 }
 
@@ -175,7 +175,7 @@ sm90_get_smem_load_op_for_source() {
   }
   else {
     // auto-vectorizing load
-    return DefaultCopy{};
+    return AutoVectorizingCopyWithAssumedAlignment{};
   }
 }
 

include/cutlass/epilogue/fusion/sm90_visitor_compute_tma_warpspecialized.hpp

@@ -438,7 +438,7 @@ struct Sm90ReLUAuxStore {
         using VecType = uint_bit_t<V>;
         Tensor tC_rAux_vec = recast<VecType>(tC_rAux);
         Tensor tC_gAux_vec = recast<VecType>(tC_gAux);
-        Tensor tC_cAux_vec = tC_cAux.compose(make_layout(Int<size(tC_rAux_vec)>{}, Int<V>{}));
+        Tensor tC_cAux_vec = tC_cAux.compose(make_layout(Int<size(tC_rAux_vec)>{}, Int<V>{})); // only works if vector is logically sequential
         auto predicate_fn = [&] (auto&&... coords) { return elem_less(tC_cAux_vec(coords...), residue_mn); };
         copy_if(FunctionPredTensor(predicate_fn), tC_rAux_vec, tC_gAux_vec);
       }
@@ -662,7 +662,7 @@ struct Sm90AuxLoad<
         }
 
         if (elem_less(repeat_like(residue_mn, _0{}), residue_mn)) { // (partially) in-bounds CTA tile
-          copy(tC_gAux, tC_rAux);
+          copy_aligned(tC_gAux, tC_rAux);
         }
       }
     }
@@ -677,7 +677,7 @@ struct Sm90AuxLoad<
         }
 
         if (elem_less(repeat_like(residue_mn, _0{}), residue_mn)) {
-          copy(tC_gAux(_,_,_,epi_m,epi_n), tC_rAux);
+          copy_aligned(tC_gAux(_,_,_,epi_m,epi_n), tC_rAux);
         }
       }
     }

include/cutlass/epilogue/fusion/sm90_visitor_load_tma_warpspecialized.hpp

@@ -641,8 +641,9 @@ struct Sm90RowBroadcast {
 
       if (epi_m == 0) { // Assumes M-major subtile loop
         // Filter so we don't issue redundant copies over stride-0 modes
+        // (only works if 0-strides are in same location, which is by construction)
         int bcast_pipe_index = (load_iteration / EpiTiles) % Stages;
-        copy(filter(tCsRow(_,_,_,epi_m,epi_n,bcast_pipe_index)), filter(tCrRow));
+        copy_aligned(filter(tCsRow(_,_,_,epi_m,epi_n,bcast_pipe_index)), filter(tCrRow));
       }
     }
 
@@ -774,7 +775,8 @@ struct Sm90ColBroadcast {
       }
 
       // Filter so we don't issue redundant copies over stride-0 modes
-      copy(filter(tCgCol), filter(tCrCol));
+      // (only works if 0-strides are in same location, which is by construction)
+      copy_aligned(filter(tCgCol), filter(tCrCol));
     }
 
     template <typename ElementAccumulator, int FragmentSize>

include/cutlass/epilogue/fusion/sm90_visitor_store_tma_warpspecialized.hpp

@@ -915,8 +915,9 @@ public:
         using ElementGmem = conditional_t<FinalReduction, ElementCompute volatile, ElementCompute>;
         Tensor tCgBuf = sm90_partition_for_epilogue<ReferenceSrc>(gBuf_nl(_,_,n,l), epi_tile, tiled_copy, thread_idx);
         if (is_reduced_lane) {
-          // Filter so we don't issue redunant copies over stride-0 modes
-          copy(filter(tCrCol), recast<ElementGmem>(filter(tCgBuf)));
+          // Filter so we don't issue redundant copies over stride-0 modes
+          // (only works if 0-strides are in same location, which is by construction)
+          copy_aligned(filter(tCrCol), recast<ElementGmem>(filter(tCgBuf)));
         }
         sync_fn();
       }
@@ -934,7 +935,8 @@ public:
         Tensor tCsBuf = sm90_partition_for_epilogue<ReferenceSrc>(sBuf(_,_,get<1>(warp_mn)), epi_tile, tiled_copy, thread_idx);
         if (is_reduced_lane) {
           // Filter so we don't issue redunant copies over stride-0 modes
-          copy(filter(tCrCol), filter(tCsBuf));
+          // (only works if 0-strides are in same location, which is by construction)
+          copy_aligned(filter(tCrCol), filter(tCsBuf));
         }
         sync_fn();