Fix duplicate argument handling in inductor lowering (#222)

Author: jansel

helion/_compiler/inductor_lowering.py

@@ -105,9 +105,20 @@ def prepare_node_lowering(
         node.meta["lowering"] = SympyExprLowering(val._sympy_())
         return
 
+    # Track arguments to reuse names for duplicates
+    arg_to_name: dict[Node, str] = {}
+
     def convert_arg(arg: Node) -> TensorBox:
         example = arg.meta["val"]
-        input_names.append(name := f"{node.name}_input{len(input_names)}")
+
+        # Reuse existing name for duplicate arguments
+        if arg in arg_to_name:
+            name = arg_to_name[arg]
+        else:
+            name = f"{node.name}_input{len(input_names)}"
+            arg_to_name[arg] = name
+            input_names.append(name)
+
         if isinstance(example, (torch.SymInt, torch.SymFloat, torch.SymBool)):
             dtype = {
                 torch.SymInt: torch.int64,

test/test_misc.py

@@ -16,6 +16,25 @@
 
 
 class TestMisc(TestCase):
+    def test_binary_operation_duplicate_args(self):
+        """Test case to reproduce issue #221: binary operations with duplicate tensor references"""
+
+        @helion.kernel(use_default_config=True)
+        def kernel_with_duplicate_refs(x: torch.Tensor) -> torch.Tensor:
+            result = torch.empty_like(x)
+            for tile in hl.tile(x.shape):
+                val = x[tile]
+                result[tile] = (
+                    val * val + val
+                )  # Multiple uses of same variable - triggers the bug
+            return result
+
+        x = torch.randn([16, 16], device=DEVICE)
+        expected = x * x + x
+
+        code, result = code_and_output(kernel_with_duplicate_refs, (x,))
+        torch.testing.assert_close(result, expected)
+
     def test_torch_alloc(self):
         @helion.kernel(config={"block_sizes": [64, 64]})
         def fn(x: torch.Tensor) -> torch.Tensor: