Remove @helion.jit usage and advise use of @helion.kernel (#1116)

Author: yf225

examples/all_gather_matmul.py

@@ -79,7 +79,7 @@ def copy_engine_all_gather_w_progress(
 
 
 # %%
-@helion.jit(
+@helion.kernel(
     config=helion.Config(
         block_sizes=[128, 256, 64],
         num_warps=8,

examples/all_reduce.py

@@ -81,7 +81,7 @@ def dev_array_to_tensor_short(
 
 
 # %%
-@helion.jit(
+@helion.kernel(
     config=helion.Config(
         block_sizes=[8192],
         num_warps=32,

test/test_associative_scan.py

@@ -92,9 +92,9 @@ def cumsum_helper(x: torch.Tensor) -> torch.Tensor:
     return hl.associative_scan(add_combine_fn, x, dim=0)
 
 
-@helion.jit
+@helion.kernel
 def jit_add_combine_fn(x, y):
-    """Addition combine function with @helion.jit decorator (should be ignored)."""
+    """Addition combine function with @helion.kernel decorator (should be ignored)."""
     return x + y
 
 
@@ -496,10 +496,10 @@ def test_codegen_kernel(x: torch.Tensor) -> torch.Tensor:
         self.assertNotIn("placeholder", code)
 
     @skipIfRefEager(
-        "torch._higher_order_ops.associative_scan with nested @helion.jit is not supported by ref eager mode yet"
+        "torch._higher_order_ops.associative_scan with nested @helion.kernel is not supported by ref eager mode yet"
     )
     def test_associative_scan_jit_decorator_ignored(self):
-        """Test that @helion.jit decorator on combine functions is ignored."""
+        """Test that @helion.kernel decorator on combine functions is ignored."""
 
         @helion.kernel(autotune_effort="none")
         def test_jit_kernel(x: torch.Tensor) -> torch.Tensor:
@@ -521,8 +521,8 @@ def test_jit_kernel(x: torch.Tensor) -> torch.Tensor:
         self.assertIn("def jit_add_combine_fn_", code)
         self.assertIn("tl.associative_scan", code)
         self.assertIn("param_0 + param_1", code)
-        # Verify @helion.jit decorator doesn't appear in generated code
-        self.assertNotIn("@helion.jit", code)
+        # Verify @helion.kernel decorator doesn't appear in generated code
+        self.assertNotIn("@helion.kernel", code)
 
     @skipIfRefEager(
         "torch._higher_order_ops.associative_scan with tuple arg is not supported by ref eager mode yet"

test/test_reduce.py

@@ -73,9 +73,9 @@ def argmax_combine_unpacked_fn(left_value, left_index, right_value, right_index)
     return max_value, max_index
 
 
-@helion.jit
+@helion.kernel
 def jit_add_combine_fn(x, y):
-    """Addition combine function with @helion.jit decorator (should be ignored)."""
+    """Addition combine function with @helion.kernel decorator (should be ignored)."""
     return x + y
 
 
@@ -221,7 +221,7 @@ def test_reduce_product_kernel(x: torch.Tensor) -> torch.Tensor:
         torch.testing.assert_close(result, expected)
 
     def test_reduce_jit_combine_fn(self):
-        """Test reduce with @helion.jit decorated combine function."""
+        """Test reduce with @helion.kernel decorated combine function."""
 
         @helion.kernel(autotune_effort="none")
         def test_reduce_jit_kernel(x: torch.Tensor) -> torch.Tensor: