frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-20 12:54:11 +08:00
Code Issues Packages Projects Releases Wiki Activity

Revert "[ATen] Fix CUDA reduction warp shuffle order (#164790)"

Browse Source 
This reverts commit 36371b8ec7a1baed255c18451b2c716386a54c95.

Reverted https://github.com/pytorch/pytorch/pull/164790 on behalf of https://github.com/clee2000 due to was reverted due to failing internal tests after merge D84992607 ([comment](https://github.com/pytorch/pytorch/pull/164790#issuecomment-3420373755))
This commit is contained in:
PyTorch MergeBot
2025-10-20 03:06:52 +00:00
parent 47804ce467
commit 602ace5eb4
3 changed files with 1 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
aten/src/ATen/native/cuda/Reduce.cuh
View File

@ -653,14 +653,8 @@ struct ReduceOp {
}
__syncthreads();
// Intra-warp reduction, fix CUDA to have offset decreasing for better numerics
// matching Triton, etc.
// todo for AMD
#ifdef USE_ROCM
for (int offset = 1; offset < dim_x; offset <<= 1) {
#else
for (int offset = dim_x >> 1; offset > 0; offset >>= 1) {
#endif
#pragma unroll
for (int i = 0; i < output_vec_size; i++) {
arg_t other = ops.warp_shfl_down(value[i], offset);

4
aten/src/ATen/native/cuda/reduction_template.cuh
View File

@ -466,11 +466,7 @@ struct ReduceJitOp {
__syncthreads();
#ifdef USE_ROCM
for (int offset = 1; offset < dim_x; offset <<= 1) {
#else
for (int offset = dim_x >> 1; offset > 0; offset >>= 1) {
#endif
#pragma unroll
for (int i = 0; i < output_vec_size; i++) {
arg_t other = reducer::warp_shfl_down(value[i], offset);