Add support for print(prefix_str, *tensors) (#140)

Author: yf225

helion/_compiler/device_ir.py

@@ -30,6 +30,7 @@
 from ..autotuner.config_spec import ReductionLoopSpec
 from ..language import _tracing_ops
 from ..language._decorators import args_to_proxies
+from ..language._decorators import get_device_func_replacement
 from .ast_extension import ExtendedAST
 from .ast_extension import LoopType
 from .ast_extension import NodeVisitor
@@ -788,8 +789,17 @@ def visit_Call(self, node: ast.Call) -> object:
                 kwargs.update(self._to_proxy(kwarg.value))
             else:
                 kwargs[kwarg.arg] = self._to_proxy(kwarg.value)
+
+        if isinstance(
+            (func_type_info := node.func._type_info),  # pyre-ignore[16]
+            CallableType,
+        ) and (replacement := get_device_func_replacement(func_type_info.value)):
+            func = replacement
+        else:
+            func = self.visit(node.func)
+
         # pyre-ignore[6]
-        return CheckForIndexCalls.retry_call(self.visit(node.func), args, kwargs)
+        return CheckForIndexCalls.retry_call(func, args, kwargs)
 
     def visit_Attribute(self, node: ast.Attribute) -> object:
         return getattr(self.visit(node.value), node.attr)

helion/_compiler/type_propagation.py

@@ -22,6 +22,7 @@
 
 from .. import exc
 from ..autotuner.config_spec import BlockSizeSpec
+from ..language._decorators import get_device_func_replacement
 from ..language._decorators import is_api_func
 from .ast_extension import ExtendedAST
 from .ast_extension import LoopType
@@ -1828,6 +1829,14 @@ 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)
+
+        if (
+            isinstance(func, CallableType)
+            and self.origin().is_device()
+            and (replacement := get_device_func_replacement(func.value))
+        ):
+            func = CallableType(func.origin, replacement)
+
         unhandled = []
         args = []
         kwargs = {}

helion/language/__init__.py

@@ -4,6 +4,7 @@
 from .constexpr import specialize as specialize
 from .creation_ops import full as full
 from .creation_ops import zeros as zeros
+from .device_print import device_print as device_print
 from .loops import Tile as Tile
 from .loops import grid as grid
 from .loops import register_block_size as register_block_size

helion/language/_decorators.py

@@ -289,3 +289,22 @@ def _to_proxy(arg: TypeInfo) -> object:
         return arg.proxy()
     except NotImplementedError:
         raise exc.TracedArgNotSupported(arg) from None
+
+
+# Tracks 1-1 mapping between Python functions and their Helion API counterparts within device function.
+_DEVICE_FUNC_REPLACEMENTS: dict[object, APIFunc] = {}
+
+
+def device_func_replacement(python_func: object) -> _Decorator:
+    def _impl(fn: _C) -> _C:
+        assert is_api_func(fn), (
+            f"{device_func_replacement.__qualname__} can only be used on API functions"
+        )
+        _DEVICE_FUNC_REPLACEMENTS[python_func] = fn
+        return fn  # pyre-ignore[7]
+
+    return _impl
+
+
+def get_device_func_replacement(func: object) -> APIFunc | None:
+    return _DEVICE_FUNC_REPLACEMENTS.get(func)

helion/language/device_print.py

@@ -0,0 +1,83 @@
+from __future__ import annotations
+
+import ast
+import builtins
+from typing import TYPE_CHECKING
+
+from torch.fx import has_side_effect
+
+from .. import exc
+from .._compiler.ast_extension import create
+from .._compiler.ast_extension import expr_from_string
+from . import _decorators
+
+if TYPE_CHECKING:
+    from .._compiler.inductor_lowering import CodegenState
+
+
+@has_side_effect
+@_decorators.device_func_replacement(builtins.print)
+@_decorators.api(is_device_only=False)
+def device_print(prefix: str, *values: object) -> None:
+    """
+    Print values from device code.
+
+    :param prefix: A string prefix for the print statement
+    :param values: Tensor values to print
+    """
+    raise exc.NotInsideKernel
+
+
+@_decorators.register_fake(device_print)
+def _(*values: object, sep: str = " ", end: str = "\n") -> None:
+    return None
+
+
+@_decorators.type_propagation(device_print)
+def _(*args: object, origin: object, **kwargs: object) -> object:
+    from .._compiler.type_propagation import LiteralType
+    from .._compiler.type_propagation import NoType
+    from .._compiler.type_propagation import TensorType
+
+    # Check that we have at least one argument (prefix)
+    if len(args) == 0:
+        raise ValueError("print() requires at least one argument (prefix)")
+
+    # First argument must be the prefix string
+    if not (isinstance(args[0], LiteralType) and isinstance(args[0].value, str)):
+        raise TypeError(
+            f"First argument to print() must be a string prefix, got {args[0]}"
+        )
+
+    # For compile-time values like tensor shapes, we should error out
+    for i, arg in enumerate(args[1:]):
+        if not isinstance(arg, TensorType):
+            raise TypeError(
+                f"print() only supports runtime tensor values. "
+                f"Argument {i + 1} is {arg}, not a tensor. "
+                f"Compile-time values like tensor shapes are not supported yet."
+            )
+
+    return NoType(origin=origin)
+
+
+# pyre-fixme[56]
+@_decorators.codegen(device_print)
+def _(state: CodegenState) -> None:
+    prefix = state.proxy_arg(0)
+    call_args = [create(ast.Constant, value=prefix)]
+
+    # Handle varargs
+    if len(state.proxy_args) > 1:
+        assert len(state.ast_args) > 1
+        ast_varargs = state.ast_args[1]
+        call_args.extend(ast_varargs[0])  # pyre-fixme[16]
+
+    call_expr = create(
+        ast.Call,
+        func=expr_from_string("tl.device_print"),
+        args=call_args,
+        keywords=[],
+    )
+    stmt = create(ast.Expr, value=call_expr)
+    state.add_statement(stmt)