Fix eager reduction warp shuffle order to start from offset=16

ghstack-source-id: 4e8b6c8
Pull Request resolved: pytorch/pytorch#164790

Author: PaulZhang12

aten/src/ATen/native/cuda/Reduce.cuh

@@ -655,8 +655,14 @@ 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);

aten/src/ATen/native/cuda/reduction_template.cuh

@@ -466,7 +466,11 @@ 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);

test/test_decomp.py

@@ -220,6 +220,8 @@ def op_assert_ref(test_case, op, test_dtype, i, orig, decomp, ref, args, kwargs)
         (torch.bfloat16, torch.ops.aten.reflection_pad2d_backward.default): 5e-3,
         (torch.float16, torch.ops.aten.reflection_pad3d_backward.default): 5e-3,
         (torch.bfloat16, torch.ops.aten.reflection_pad3d_backward.default): 5e-2,
+        (torch.float16, torch.ops.aten._batch_norm_with_update.default): 2e-7,
+        (torch.bfloat16, torch.ops.aten._batch_norm_with_update.default): 2e-7,
         # see https://github.com/pytorch/pytorch/pull/96264
         (torch.float16, torch.ops.aten.mv.default): 1e-5,
         (torch.bfloat16, torch.ops.aten.mv.default): 1e-5,
@@ -295,6 +297,7 @@ def op_assert_equal(test_case, op, test_dtype, orig, decomp, args, kwargs):
         rtol, atol = tol_table[(decomp.dtype, op)]
     else:
         rtol, atol = _getDefaultRtolAndAtol(orig.dtype, decomp.dtype)
+
     test_case.assertEqual(
         orig,
         decomp,