Do not create a new variable for tile assignments since tiles are immutable (#334)

Author: oulgen

helion/_compiler/device_ir.py

@@ -748,7 +748,9 @@ def visit_Assign(self, node: ast.Assign) -> None:
             value = self.visit(node.value)
             # For simple variable assignments like `a = b`, we need to create a new
             # variable to avoid phi node issues when the source variable gets mutated
-            if isinstance(node.value, ast.Name) and isinstance(value, torch.Tensor):
+            if isinstance(node.value, ast.Name) and (
+                isinstance(value, torch.Tensor) and not isinstance(value, Tile)
+            ):
                 value = _new_var(value)
             self._assign(target, value)
             return None

test/test_misc.expected

@@ -43,6 +43,29 @@ def kernel(a_list, b_dict, b_tuple, c_named_tuple, d_dataclass, *, _launcher=_de
     _launcher(_kernel_kernel, (triton.cdiv(a0.size(0), _BLOCK_SIZE_0),), a0, o0, o1, a0.size(0), a0.stride(0), o0.stride(0), o1.stride(0), _BLOCK_SIZE_0, num_warps=4, num_stages=3)
     return [o0, o1]
 
+--- assertExpectedJournal(TestMisc.test_propagate_tile)
+from __future__ import annotations
+
+import torch
+import triton
+import triton.language as tl
+from helion.runtime import default_launcher as _default_launcher
+
+@triton.jit
+def _copy_kernel_kernel(a, out, a_size_0, a_stride_0, out_stride_0, _BLOCK_SIZE_0: tl.constexpr):
+    pid_0 = tl.program_id(0)
+    offset_0 = pid_0 * _BLOCK_SIZE_0
+    indices_0 = (offset_0 + tl.arange(0, _BLOCK_SIZE_0)).to(tl.int32)
+    mask_0 = indices_0 < a_size_0
+    load = tl.load(a + indices_0 * a_stride_0, mask_0, other=0)
+    tl.store(out + indices_0 * out_stride_0, load, mask_0)
+
+def copy_kernel(a: torch.Tensor, *, _launcher=_default_launcher):
+    out = torch.empty_like(a)
+    _BLOCK_SIZE_0 = 4
+    _launcher(_copy_kernel_kernel, (triton.cdiv(a.size(0), _BLOCK_SIZE_0),), a, out, a.size(0), a.stride(0), out.stride(0), _BLOCK_SIZE_0, num_warps=4, num_stages=3)
+    return out
+
 --- assertExpectedJournal(TestMisc.test_scalar_tensor_item_method)
 from __future__ import annotations
 

test/test_misc.py

@@ -390,6 +390,22 @@ def tuple_unpack_kernel(inp_tuple) -> torch.Tensor:
 
         self.assertExpectedJournal(code)
 
+    def test_propagate_tile(self):
+        @helion.kernel
+        def copy_kernel(a: torch.Tensor) -> torch.Tensor:
+            out = torch.empty_like(a)
+
+            for tile in hl.tile(a.size(0), block_size=4):
+                t1 = tile
+                t2 = tile
+                out[t2] = a[t1]
+            return out
+
+        args = (torch.randn(16, device=DEVICE, dtype=torch.bfloat16),)
+        code, result = code_and_output(copy_kernel, args)
+        torch.testing.assert_close(result, args[0])
+        self.assertExpectedJournal(code)
+
 
 if __name__ == "__main__":
     unittest.main()