squall/cutlass
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

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

Browse Source 
Co-authored-by: Ali Hassani <ahassanijr@gmail.com>
This commit is contained in:
Ali Hassani 2023-12-05 15:35:41 -05:00 committed by GitHub
parent e9e30c2304
commit a75b4ac483
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 4 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
include/cutlass/epilogue/threadblock/epilogue_with_broadcast.h
View File

@ -946,13 +946,13 @@ private:
// //
if (OutputOp::kStoreZ) { if (OutputOp::kStoreZ) {
destination_iterator += reduce_fragment_idx;
destination_iterator.store(frag_Z); destination_iterator.store(frag_Z);
++destination_iterator;
} }
if (OutputOp::kStoreT) { if (OutputOp::kStoreT) {
tensor_iterator += reduce_fragment_idx;
tensor_iterator.store(frag_T); tensor_iterator.store(frag_T);
++tensor_iterator;
} }
} }
}; };
@ -1698,13 +1698,13 @@ private:
// //
if (OutputOp::kStoreZ) { if (OutputOp::kStoreZ) {
destination_iterator += reduce_fragment_idx;
destination_iterator.store(frag_Z); destination_iterator.store(frag_Z);
++destination_iterator;
} }
if (OutputOp::kStoreT) { if (OutputOp::kStoreT) {
tensor_iterator += reduce_fragment_idx;
tensor_iterator.store(frag_T); tensor_iterator.store(frag_T);
++tensor_iterator;
} }
} }
}; };