diff --git a/examples/all_gather_matmul.py b/examples/all_gather_matmul.py index 001d67a3c..1f4a0445c 100644 --- a/examples/all_gather_matmul.py +++ b/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, diff --git a/examples/all_reduce.py b/examples/all_reduce.py index 17616c275..7396e9b9c 100644 --- a/examples/all_reduce.py +++ b/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, diff --git a/test/test_associative_scan.py b/test/test_associative_scan.py index f5c070d79..030def297 100644 --- a/test/test_associative_scan.py +++ b/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" diff --git a/test/test_reduce.py b/test/test_reduce.py index 5895d442c..224d1cf6f 100644 --- a/test/test_reduce.py +++ b/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: