Add block_sync_lds() between blockwise reduction again

e6dce0db · Qianfeng Zhang · 3ac68bcc · e6dce0db
Commit e6dce0db authored Nov 09, 2022 by Qianfeng Zhang
Show whitespace changes
Inline Side-by-side

Showing with 2 additions and 0 deletions

include/ck/tensor_operation/gpu/grid/gridwise_batchnorm_backward_blockwise_welford.hpp ...pu/grid/gridwise_batchnorm_backward_blockwise_welford.hpp +2 -0

No files found.
--- a/include/ck/tensor_operation/gpu/grid/gridwise_batchnorm_backward_blockwise_welford.hpp
+++ b/include/ck/tensor_operation/gpu/grid/gridwise_batchnorm_backward_blockwise_welford.hpp
@@ -478,6 +478,8 @@ struct GridwiseBatchNormBackwardWithBlockwiseWelford
        };
        static_for<0, MThreadSliceSize, 1>{}([&](auto I) {
+            if constexpr(I > 0)
+                block_sync_lds();
            BlockwiseReduce::Reduce(reduce_work_buf, dscale_thread_buf(I));
            block_sync_lds();
            BlockwiseReduce::Reduce(reduce_work_buf, dbias_thread_buf(I));