[inductor] add lowering for repeat_interleave.Tensor with output size specified (#147160)

v0i0 · v0i0 · commit 9625356c6866 · 2025-07-16T13:49:06.000-07:00
ghstack-source-id: fdeeea1 Pull Request resolved: #158462
diff --git a/test/inductor/test_torchinductor.py b/test/inductor/test_torchinductor.py
@@ -13626,6 +13626,18 @@ def forward(self, x):
                 FileCheck().check("cpp_fused_add_0").run(code)
             self.assertEqual(refe_out, test_out)
 
+    def test_repeat_interleave_pass(self):
+        # https://github.com/pytorch/pytorch/issues/147160
+        def f(input, repeats):
+            return torch.repeat_interleave(input, repeats, dim=0, output_size=3) + 1
+
+        input = torch.tensor([[1, 2], [3, 4]], device="cuda")
+        repeat = torch.tensor([1, 2], device="cuda")
+        f_compiled = torch.compile(f)
+        test, (code,) = run_and_get_code(f_compiled, input, repeat)
+        self.assertEqual(test, f(input, repeat))
+        self.assertFalse("repeat_interleave.Tensor" in code)
+
 
 @dataclasses.dataclass
 class TestFailure:
diff --git a/torch/_inductor/decomposition.py b/torch/_inductor/decomposition.py
@@ -1150,3 +1150,15 @@ def rrelu_with_noise_functional(
     else:
         negative_slope = (lower + upper) / 2
         return aten.leaky_relu(self, negative_slope), torch.Tensor()
+
+
+@register_decomposition(aten.repeat_interleave.Tensor)
+def repeast_interleave_Tensor(
+    repeat: torch.Tensor,
+    output_size: Optional[int] = None,
+) -> torch.Tensor:
+    if output_size is None or type(output_size) is not int:
+        return NotImplemented
+    cumsum = repeat.cumsum(0)
+    pos = torch.arange(output_size, device=repeat.device)
+    return torch.searchsorted(cumsum, pos, right=True)
