Revert "[inductor] Cooperative reductions (pytorch#137756)"

This reverts commit fed37db.

Reverted pytorch#137756 on behalf of https://github.com/jeanschmidt due to ROCM tests are timing out :( ([comment](pytorch#137756 (comment)))

Author: pytorchmergebot

test/inductor/test_cooperative_reductions.py

@@ -501,9 +501,6 @@ def f(x, scale, amax_keep_dim):
         expected_numel = (
             1 + hidden_size * 2 + 4 * 2048 * hidden_size * 2 + 4 * 2048 * 2 + 1
         )
-        if config.triton.cooperative_reductions:
-            expected_numel = 134225922
-
         self.assertExpectedInline(count_numel(f, *inp, True), str(expected_numel))
         self.assertExpectedInline(count_numel(f, *inp, False), str(expected_numel))
 

test/inductor/test_perf.py

@@ -11747,34 +11747,30 @@ def fn(a: torch.Tensor) -> torch.Tensor:
                 return torch.sum(a)
 
             kernels = self.get_kernels(fn, [torch.randn([256, 256], device=GPU_TYPE)])
-            expected_divisible = {
-                # kernel0 reduces from 256 to (xnumel=8, rnumel=8192), which means it reduces 256 by 256 into an array of
-                # size 8 by accumulating 8192 elements at once note that rnumel is equal to 512 * 16, so rnumel which is
-                # at slot 3 should be in the divisible by 16 descriptor
-                0: (0, 1, 3),
-                # kernel1 reduces from 8 elements to a single scalar.
-                # Since multi-kernel generate 2 variants for each kernel. The second
-                # persistent-reduction has index 2.
-                1: (0, 1),
-            }
             if config.triton.multi_kernel:
-                self.assertEqual(len(kernels), 4)
-                expected_divisible[2] = expected_divisible.pop(1)
-            elif config.triton.cooperative_reductions:
-                self.assertEqual(len(kernels), 1)
-                expected_divisible = {
-                    # one kernel, with extra workspace/semaphore args
-                    0: (0, 1, 2, 3, 5),
-                }
+                self.assertTrue(
+                    len(kernels) == 4,
+                    "SUM should result in four kernels when multi-kernel is enabled",
+                )
             else:
-                self.assertEqual(len(kernels), 2)
+                self.assertTrue(len(kernels) == 2, "SUM should result in two kernels")
 
-            for kernel_id, expected in expected_divisible.items():
-                divisible_by_16 = (
-                    kernels[kernel_id].triton_meta["configs"][0].divisible_by_16
-                )
-                self.assertEqual(divisible_by_16, expected)
+            # kernel0 reduces from 256 to (xnumel=8, rnumel=8192), which means it reduces 256 by 256 into an array of
+            # size 8 by accumulating 8192 elements at once note that rnumel is equal to 512 * 16, so rnumel which is
+            # at slot 3 should be in the divisible by 16 descriptor
+            arguments_that_are_divisible_by_16_in_kernel0 = (
+                kernels[0].triton_meta["configs"][0].divisible_by_16
+            )
+            self.assertEqual(arguments_that_are_divisible_by_16_in_kernel0, (0, 1, 3))
 
+            # kernel1 reduces from 8 elements to a single scalar.
+            # Since multi-kernel generate 2 variants for each kernel. The second
+            # persistent-reduction has index 2.
+            kernel1_index = 2 if config.triton.multi_kernel else 1
+            arguments_that_are_divisible_by_16_in_kernel1 = (
+                kernels[kernel1_index].triton_meta["configs"][0].divisible_by_16
+            )
+            self.assertEqual(arguments_that_are_divisible_by_16_in_kernel1, (0, 1))
             torch._dynamo.reset()
 
         @config.patch(assume_aligned_inputs=False)

test/inductor/test_torchinductor.py

@@ -315,11 +315,6 @@ def test_reduction(
         full = torch.randn(full_size).to(device)
         view = torch.as_strided(full, view_size, full.stride())
 
-        if num_triton_kernels == 2 and config.triton.cooperative_reductions:
-            # fewer kernels with cooperative reductions
-            num_triton_kernels = 1
-            num_block_pointers -= 2
-
         # Expect at least 1 block pointer for the input.
         # Add 2 more if we generate 2 kernels.
         result, (code,) = self.run_and_compare(

test/inductor/test_torchinductor_strided_blocks.py

@@ -1457,7 +1457,7 @@ def semaphores(self, min_size: sympy.Expr):
         arg = WorkspaceArg(
             count=min_size,
             zero_mode=WorkspaceZeroMode.ZERO_PER_GRAPH,
-            dtype=torch.uint32,
+            dtype=torch.int32,
             inner_name="sem_ptr",
             outer_name=f"semaphores_{current_device.type}_{current_device.index}",
             device=current_device,

torch/_inductor/codegen/common.py

@@ -1660,7 +1660,7 @@ class HalideScheduling(SIMDScheduling):
     int32_type = "hl.Int(32)"
     # TODO(jansel): Halide doesn't actually support 64 bit indexing...
     int64_type = "hl.Int(64)"
-    kernel_type = HalideKernel  # type: ignore[arg-type,assignment]
+    kernel_type = HalideKernel  # type: ignore[arg-type]
 
     @classmethod
     def get_backend_features(cls, device: torch.device):

torch/_inductor/codegen/halide.py

@@ -330,7 +330,6 @@ def __init__(
         pid_cache=None,
         reduction_hint=ReductionHint.DEFAULT,
         override_persistent_reduction=None,
-        override_cooperative_reduction=None,
     ) -> None:
         if pid_cache is None:
             pid_cache = {}
@@ -349,11 +348,6 @@ def __init__(
         self.index_dtype: str = index_dtype
         self.last_usage: OrderedSet[str] = OrderedSet()
         self.buf_accesses: DefaultDict[str, List[Dep]] = collections.defaultdict(list)
-        self.cooperative_reduction: bool = (
-            override_cooperative_reduction
-            if override_cooperative_reduction is not None
-            else self.should_use_cooperative_reduction()
-        )
         self.persistent_reduction: bool = (
             override_persistent_reduction
             if override_persistent_reduction is not None
@@ -427,9 +421,6 @@ def store_reduction(self, name: str, index: sympy.Expr, value: CSEVariable):
         finally:
             self.inside_reduction = prior
 
-    def should_use_cooperative_reduction(self) -> bool:
-        return False  # defined in subclass
-
     def should_use_persistent_reduction(self) -> bool:
         return False  # defined in subclass
 
@@ -515,7 +506,7 @@ def set_last_usage(self, nodes):
         )
 
     def disable_reduction(self):
-        should_flush = self.range_trees[-1].is_loop or self.cooperative_reduction
+        should_flush = self.range_trees[-1].is_loop
 
         @contextlib.contextmanager
         def ctx():
@@ -1334,7 +1325,6 @@ def get_kernel_args(self, node_schedule, numel, reduction_numel):
     def codegen_node_schedule(
         self, node_schedule, buf_accesses, numel, reduction_numel
     ):
-        from torch._inductor.codegen.triton import TritonKernel
         from torch._inductor.codegen.triton_split_scan import TritonSplitScanKernel
 
         tiled_groups = self.select_tiling(node_schedule, numel, reduction_numel)
@@ -1344,8 +1334,7 @@ def codegen_node_schedule(
             index_dtype,
         ) = self.get_kernel_args(node_schedule, numel, reduction_numel)
 
-        is_scan = schedule_contains_op(node_schedule, "scan")
-        is_split_scan = is_scan and any(
+        is_split_scan = any(
             isinstance(node, BaseSchedulerNode) and node.is_split_scan()
             for node in node_schedule
         )
@@ -1360,10 +1349,6 @@ def codegen_node_schedule(
             index_dtype=index_dtype,
         )
 
-        if is_scan and kernel_type == TritonKernel:
-            # TODO(jansel): scan does not yet work with cooperative reductions
-            kernel_kwargs["override_cooperative_reduction"] = False
-
         # ops.sort only works with persistent reduction, and is not bandwidth bound anyway
         # so taking the hit of non-coalesced loads is okay
         if has_sort := schedule_contains_op(node_schedule, "sort"):