Fix Stream-K reduce bug in epilogue with broadcast (#1224)

alihassanijr · Ali Hassani · web-flow · commit a75b4ac48316 · 2023-12-05T15:35:41.000-05:00
Co-authored-by: Ali Hassani &lt;ahassanijr@gmail.com&gt;
diff --git a/include/cutlass/epilogue/threadblock/epilogue_with_broadcast.h b/include/cutlass/epilogue/threadblock/epilogue_with_broadcast.h
@@ -946,13 +946,13 @@ class EpilogueWithBroadcast<
       //
 
       if (OutputOp::kStoreZ) {
+        destination_iterator += reduce_fragment_idx;
         destination_iterator.store(frag_Z);
-        ++destination_iterator;
       }
 
       if (OutputOp::kStoreT) {
+        tensor_iterator += reduce_fragment_idx;
         tensor_iterator.store(frag_T);
-        ++tensor_iterator;
       }
     }
 };
@@ -1698,13 +1698,13 @@ class EpilogueWithBroadcast<
       //
 
       if (OutputOp::kStoreZ) {
+        destination_iterator += reduce_fragment_idx;
         destination_iterator.store(frag_Z);
-        ++destination_iterator;
       }
 
       if (OutputOp::kStoreT) {
+        tensor_iterator += reduce_fragment_idx;
         tensor_iterator.store(frag_T);
-        ++tensor_iterator;
       }
     }
 };

Original file line number	Diff line number	Diff line change
`@@ -946,13 +946,13 @@ class EpilogueWithBroadcast<`
`946`	`946`	`//`
`947`	`947`
`948`	`948`	`if (OutputOp::kStoreZ) {`
	`949`	`+ destination_iterator += reduce_fragment_idx;`
`949`	`950`	`destination_iterator.store(frag_Z);`
`950`		`- ++destination_iterator;`
`951`	`951`	`}`
`952`	`952`
`953`	`953`	`if (OutputOp::kStoreT) {`
	`954`	`+ tensor_iterator += reduce_fragment_idx;`
`954`	`955`	`tensor_iterator.store(frag_T);`
`955`		`- ++tensor_iterator;`
`956`	`956`	`}`
`957`	`957`	`}`
`958`	`958`	`};`
`@@ -1698,13 +1698,13 @@ class EpilogueWithBroadcast<`
`1698`	`1698`	`//`
`1699`	`1699`
`1700`	`1700`	`if (OutputOp::kStoreZ) {`
	`1701`	`+ destination_iterator += reduce_fragment_idx;`
`1701`	`1702`	`destination_iterator.store(frag_Z);`
`1702`		`- ++destination_iterator;`
`1703`	`1703`	`}`
`1704`	`1704`
`1705`	`1705`	`if (OutputOp::kStoreT) {`
	`1706`	`+ tensor_iterator += reduce_fragment_idx;`
`1706`	`1707`	`tensor_iterator.store(frag_T);`
`1707`		`- ++tensor_iterator;`
`1708`	`1708`	`}`
`1709`	`1709`	`}`
`1710`	`1710`	`};`