Update

guangyey · guangyey · commit 62ab0c3c6bb0 · 2026-03-12T10:22:23.000Z
[ghstack-poisoned]
diff --git a/aten/src/ATen/native/cuda/SortingRadixSelect.cuh b/aten/src/ATen/native/cuda/SortingRadixSelect.cuh
@@ -521,6 +521,7 @@ __device__ __forceinline__ void countRadixAggregateCounts(
   for (uint32_t i = 0; i < RadixSize; ++i) {
     counts[i] = smem[buffer_offset + i];
   }
+  __syncthreads(); // Wait for all threads to finish reading the final counts.
 }
 
 // This function counts the distribution of all input values in a
diff --git a/test/inductor/pallas_expected_failures/CpuTests.test_max_pool2d_with_indices_backward6_cpu b/test/inductor/pallas_expected_failures/CpuTests.test_max_pool2d_with_indices_backward6_cpu
diff --git a/test/inductor/test_torchinductor.py b/test/inductor/test_torchinductor.py
@@ -10354,7 +10354,7 @@ def fn(a, b, c):
 
     # From https://github.com/pytorch/pytorch/issues/93384
     def test_max_pool2d_with_indices_backward6(self):
-        # dilation is not 1. Should still generate kernels.
+        # dilation is not 1. Should fallback
         def fn(a, b, c):
             return aten.max_pool2d_with_indices_backward(
                 a, b, [3, 2], [2, 1], [1, 1], [1, 2], False, c
@@ -10378,7 +10378,7 @@ def fn(a, b, c):
                 indices,
             ],
         )
-        assertGeneratedKernelCountEqual(self, 1)
+        assertGeneratedKernelCountEqual(self, 0)
 
     def test_issue102546(self):
         def fn(x):
diff --git a/test/inductor/test_torchinductor_codegen_dynamic_shapes.py b/test/inductor/test_torchinductor_codegen_dynamic_shapes.py
@@ -219,6 +219,9 @@ def run(*ex, **kwargs):
     "test_max_pool2d_with_indices_backward5_dynamic_shapes": TestFailure(
         ("cpu", "cuda")
     ),
+    "test_max_pool2d_with_indices_backward6_dynamic_shapes": TestFailure(
+        ("cpu", "cuda", "xpu")
+    ),
     "test_misaligned_address_issue1_dynamic_shapes": TestFailure(("cpu",)),
     "test_mm_views_dynamic_shapes": TestFailure(("cpu", "cuda", "xpu")),
     "test_new_empty_dynamic_shapes": TestFailure(("cpu", "cuda", "xpu")),
diff --git a/torch/_inductor/lowering.py b/torch/_inductor/lowering.py
@@ -5179,6 +5179,11 @@ def max_pool2d_with_indices_backward(
     is_channels_last = (x_stride is not None and x_stride[1] == 1) or (
         gO_stride is not None and gO_stride[1] == 1
     )
+    if any(d != 1 for d in dilation):
+        # dilation NYI
+        return fallback_max_pool2d_with_indices_backward(
+            grad_output, x, kernel_size, stride, padding, dilation, ceil_mode, indices
+        )
 
     *_batch, _height, width = x.get_size()
     *_, pooled_height, pooled_width = grad_output.get_size()
@@ -5187,17 +5192,13 @@ def max_pool2d_with_indices_backward(
     grad_loader = grad_output.make_loader()
     new_size = list(x.get_size())
 
-    # Effective kernel size accounts for dilation
-    effective_kh = (kernel_size[0] - 1) * dilation[0] + 1
-    effective_kw = (kernel_size[1] - 1) * dilation[1] + 1
-
     h_window_size = max(
-        max(FloorDiv(h, stride[0]) - max(0, FloorDiv(h - effective_kh, stride[0])), 1)
-        for h in range(effective_kh * 2)
+        max(FloorDiv(h, stride[0]) - max(0, FloorDiv(h - kernel_size[0], stride[0])), 1)
+        for h in range(kernel_size[0] * 2)
     )
     w_window_size = max(
-        max(FloorDiv(w, stride[1]) - max(0, FloorDiv(w - effective_kw, stride[1])), 1)
-        for w in range(effective_kw * 2)
+        max(FloorDiv(w, stride[1]) - max(0, FloorDiv(w - kernel_size[1], stride[1])), 1)
+        for w in range(kernel_size[1] * 2)
     )
 
     window_size = h_window_size * w_window_size
@@ -5216,10 +5217,10 @@ def fn(idx):
         h = h + padding[0]
         w = w + padding[1]
         phstart = ops.index_expr(
-            FloorDiv(h - effective_kh + stride[0], stride[0]), torch.int32
+            FloorDiv(h - kernel_size[0] + stride[0], stride[0]), torch.int32
         )
         pwstart = ops.index_expr(
-            FloorDiv(w - effective_kw + stride[1], stride[1]), torch.int32
+            FloorDiv(w - kernel_size[1] + stride[1], stride[1]), torch.int32
         )
         phend = ops.index_expr(FloorDiv(h, stride[0]) + 1, torch.int32)
         pwend = ops.index_expr(FloorDiv(w, stride[1]) + 1, torch.int32)

Original file line number	Diff line number	Diff line change
`@@ -521,6 +521,7 @@ __device__ __forceinline__ void countRadixAggregateCounts(`
`521`	`521`	`for (uint32_t i = 0; i < RadixSize; ++i) {`
`522`	`522`	`counts[i] = smem[buffer_offset + i];`
`523`	`523`	`}`
	`524`	`+ __syncthreads(); // Wait for all threads to finish reading the final counts.`
`524`	`525`	`}`
`525`	`526`
`526`	`527`	`// This function counts the distribution of all input values in a`