Better error message for calling Helion kernel from another kernel (#1008)

Author: yf225

helion/_compiler/type_propagation.py

@@ -30,6 +30,7 @@
 from ..language.stack_tensor import StackTensor
 from ..language.tile_proxy import Tile
 from ..language.tile_proxy import _CheckForIndexCalls
+from ..runtime.kernel import Kernel
 from .ast_extension import ExtendedAST
 from .ast_extension import LoopType
 from .ast_extension import create
@@ -105,6 +106,8 @@ def _get(self, name: str) -> TypeInfo:
                 return TypeInfo.from_example(value, origin)
 
             origin = self.function.global_scope_origin(name)
+            if isinstance(value, Kernel):
+                return TypeInfo.from_example(value, origin)
             if not isinstance(
                 value,
                 (types.ModuleType, types.FunctionType, types.BuiltinFunctionType),
@@ -1973,9 +1976,12 @@ def visit_Compare(self, node: ast.Compare) -> TypeInfo:
 
     def visit_Call(self, node: ast.Call) -> TypeInfo:
         # TODO(jansel): test handling if *args and **kwargs
-        # TODO(jansel): check for calling a Kernel here
         func = self.visit(node.func)
 
+        # Check for calling a Helion kernel from within another Helion kernel
+        if isinstance(func, CallableType) and isinstance(func.value, Kernel):
+            raise exc.NestedKernelCallsNotSupported
+
         if (
             isinstance(func, CallableType)
             and self.origin().is_device()

helion/exc.py

@@ -452,3 +452,12 @@ class NoDeviceLoopsInKernel(BaseError):
         "Kernel contains no device loops. Add an hl.tile(...) or hl.grid(...) loop "
         "around your device computations."
     )
+
+
+class NestedKernelCallsNotSupported(BaseError):
+    message = (
+        "Calling a Helion kernel from within another Helion kernel is not supported. "
+        "Helion kernels can only be called from outside of @helion.kernel functions. "
+        "If you need to share code between kernels, consider extracting the shared logic "
+        "into a regular Python function that can be called from within both kernels."
+    )

test/test_errors.py

@@ -16,6 +16,22 @@
 import helion.language as hl
 
 
+@helion.kernel()
+def _test_inner_kernel(x: torch.Tensor) -> torch.Tensor:
+    out = torch.empty_like(x)
+    for tile in hl.tile(x.shape):
+        out[tile] = x[tile] * 2
+    return out
+
+
+@helion.kernel()
+def _test_outer_kernel_calling_inner(x: torch.Tensor) -> torch.Tensor:
+    out = torch.empty_like(x)
+    for tile in hl.tile(x.shape):
+        out[tile] = _test_inner_kernel(x[tile])
+    return out
+
+
 class TestErrors(RefEagerTestDisabled, TestCase):
     def test_autotune_no_valid_configs(self):
         class FakeKernel:
@@ -452,6 +468,15 @@ def fn(x: torch.Tensor) -> torch.Tensor:
         with self.assertRaises(helion.exc.TileOfTile):
             code_and_output(fn, (torch.randn(8, device=DEVICE),))
 
+    def test_nested_kernel_calls(self):
+        with self.assertRaisesRegex(
+            helion.exc.NestedKernelCallsNotSupported,
+            r"Calling a Helion kernel from within another Helion kernel is not supported",
+        ):
+            code_and_output(
+                _test_outer_kernel_calling_inner, (torch.randn(8, device=DEVICE),)
+            )
+
 
 if __name__ == "__main__":
     unittest.main()