[inductor] fix triton bucketize mask propagation (pytorch#159961)

See https://hud.pytorch.org/pytorch/pytorch/commit/6b414f56a4a133a428af618d8ed1553849341497

Pull Request resolved: pytorch#159961
Approved by: https://github.com/eellison

Author: v0i0

test/inductor/test_torchinductor.py

@@ -13697,6 +13697,28 @@ def f(a_list):
         print(profile_output)
         self.assertFalse("Pageable" in profile_output)
 
+    @unittest.skipIf(
+        config.cpp_wrapper,
+        "cpp_wrapper samples will lead to invalid indexing",
+    )
+    def test_inductor_triton_bucketize_respects_masking(self):
+        def fn(inp, repeats, output_size):
+            # return torch.repeat_interleave(inp, repeats, dim=0, output_size=output_size)
+            idx = torch.searchsorted(
+                repeats.cumsum(0),
+                torch.arange(0, output_size, device=repeats.device),
+                right=True,
+            )
+            return torch.index_select(inp, 0, idx)
+
+        inp = torch.arange(0, 4, device=self.device)
+        repeats = torch.tensor([1, 2, 3, 4], device=self.device)
+        output_size = repeats.sum().item()
+        args = (inp, repeats, output_size)
+        self.assertEqual(fn(*args), torch.compile(fn)(*args))
+
+    # end of class CommonTemplate - add new tests here
+
 
 @dataclasses.dataclass
 class TestFailure:

torch/_inductor/codegen/triton.py

@@ -2669,6 +2669,18 @@ def guard_cooperative_store(self, name, buffer):
         buffer.writeline(DeferredLine(name, f"if rsplit_id == ({idx} % RSPLIT):"))
         return buffer.indent()
 
+    def _combine_masks(self, *variables: Optional[CSEVariable]):
+        masks = None
+        for elem in variables:
+            if elem is None:
+                continue
+            if hasattr(elem, "mask_vars"):
+                if masks is None:
+                    masks = elem.mask_vars
+                else:
+                    masks = masks | elem.mask_vars
+        return masks
+
     def bucketize(
         self,
         values: CSEVariable,
@@ -2718,6 +2730,9 @@ def bucketize(
             dtype=indexing_dtype,  # type: ignore[attr-defined]
         )
 
+        masks = self._combine_masks(values, boundary_indices, sorter_indices)
+        result.mask_vars = masks  # type: ignore[attr-defined]
+
         return result
 
     def reduction_resize(self, value) -> str: