:Add ct.print to suport fstring and python-tyle print

Signed-off-by: Ziheng Deng <zihengd@nvidia.com>

Author: gunnersdeng

changelog.d/ct-print.md

@@ -0,0 +1,4 @@
+<!--- SPDX-FileCopyrightText: Copyright (c) <2026> NVIDIA CORPORATION & AFFILIATES. All rights reserved. -->
+<!--- SPDX-License-Identifier: Apache-2.0 -->
+
+- Added `print()` and `ct.print()` to support python-style f-strings and positional arguments

docs/source/operations.rst

@@ -189,6 +189,7 @@ Utility
    :nosignatures:
 
    printf
+   print
    assert_
 
 

src/cuda/tile/__init__.py

@@ -126,6 +126,7 @@
     permute,
     pow,
     printf,
+    print,
     prod,
     reduce,
     reshape,
@@ -267,6 +268,7 @@
     "permute",
     "pow",
     "printf",
+    "print",
     "prod",
     "reduce",
     "reshape",

src/cuda/tile/_ir/op_impl.py

@@ -428,19 +428,36 @@ def validate_dtype(cls, dtype: DType, specifier: str) -> bool:
         else:
             return False
 
+    # Placeholder emitted by ast2hir for type-inferred format specifiers
+    _TYPE_INFER = '\x01'
+
     @classmethod
     def parse_format(cls, format: str, arg_types: Tuple[Union[TileTy, DType], ...]) -> str:
         last_pos = pos = 0
         arg_idx = 0
         tokens = []
         while pos < len(format):
-            if format[pos] == "%":
+            ch = format[pos]
+            if ch == cls._TYPE_INFER:
+                tokens.append(format[last_pos:pos])
+                if arg_idx >= len(arg_types):
+                    raise TileTypeError("Not enough arguments for format string")
+                dtype = get_dtype(arg_types[arg_idx])
+                if is_boolean(dtype) or is_integral(dtype):
+                    tokens.append('%d')
+                elif is_float(dtype) or is_restricted_float(dtype):
+                    tokens.append('%f')
+                else:
+                    raise TileTypeError(f"Cannot infer format for dtype {dtype}")
+                arg_idx += 1
+                pos += 1
+                last_pos = pos
+            elif ch == "%":
                 tokens.append(format[last_pos:pos])
                 last_pos = pos
                 # escape "%%"
                 if (pos + 1 < len(format) and format[pos + 1] == "%"):
                     pos += 2
-                    continue
                 elif (m := cls.pattern.match(format, pos)):
                     # get a format match
                     _, _, _, _, sp = m.groups()
@@ -458,10 +475,10 @@ def parse_format(cls, format: str, arg_types: Tuple[Union[TileTy, DType], ...])
                     pos = m.end()
                     tokens.append(format[last_pos:pos])
                     last_pos = pos
-                    continue
                 else:
                     raise TileTypeError("Invalid format string")
-            pos += 1
+            else:
+                pos += 1
         tokens.append(format[last_pos:pos])
         if arg_idx < len(arg_types):
             raise TileTypeError("Too many arguments for format string")

src/cuda/tile/_passes/ast2hir.py

@@ -3,6 +3,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import ast
+import builtins as _builtins
 import inspect
 import itertools
 import operator
@@ -18,6 +19,7 @@
 from cuda.tile._ir.type import ClosureDefaultPlaceholder
 from cuda.tile._passes.ast_util import ast_get_all_local_names
 from cuda.tile._stub import static_eval, static_assert, static_iter
+import cuda.tile._stub as _ct_stub
 
 
 @lru_cache
@@ -248,8 +250,8 @@ def decorate(f):
 _expr_handlers: Dict[Type[ast.AST], Callable] = {}
 
 
-_KEYWORD_LIKE_FUNCS = (static_eval, static_assert, static_iter)
-_KEYWORD_LIKE_FUNC_NAMES = ("static_eval", "static_assert", "static_iter")
+_KEYWORD_LIKE_FUNCS = (static_eval, static_assert, static_iter, _ct_stub.print, _builtins.print)
+_KEYWORD_LIKE_FUNC_NAMES = ("static_eval", "static_assert", "static_iter", "print", "print")
 
 
 @_register(_expr_handlers, ast.Call)
@@ -277,6 +279,8 @@ def _call_expr(call: ast.Call, ctx: _Context) -> hir.Value:
         elif kwd_func == "static_iter":
             raise TileSyntaxError("static_iter() is only allowed as iterable in a `for` loop,"
                                   " i.e. `for i in ct.static_iter(...)`")
+        elif kwd_func == "print":
+            return _handle_ct_print(call, ctx)
         else:
             raise TileSyntaxError(f"{kwd_func} is not expected here")
     else:
@@ -378,6 +382,131 @@ def _is_cuda_module(value: ast.expr, ctx: _Context) -> bool:
     return ctx.frozen_globals.get(value.id) is cuda
 
 
+# ================================
+# ct.print() helper functions
+# ================================
+
+def _escape_format_str(s: str) -> str:
+    """Escape a literal string for use in a C printf format (replace % with %%)."""
+    return s.replace('%', '%%')
+
+
+def _python_spec_to_c_printf(py_spec: str, ctx: _Context) -> str:
+    """Convert a Python format spec string to a C printf format specifier."""
+    import re
+    m = re.fullmatch(
+        r'(?P<align>[<>^])?'
+        r'(?P<sign>[+ -])?'
+        r'(?P<alt>\#)?'
+        r'(?P<zero>0)?'
+        r'(?P<width>[0-9]+)?'
+        r'(?:\.(?P<precision>[0-9]+))?'
+        r'(?P<type>[diouxXeEfFgGaA])?',
+        py_spec)
+    if m is None or m.group(0) != py_spec:
+        raise ctx.syntax_error(f"ct.print(): unsupported format spec '{py_spec}'")
+    align = m.group('align')
+    sign = m.group('sign')
+    alt = m.group('alt')
+    zero = m.group('zero')
+    width = m.group('width') or ''
+    precision = ('.' + m.group('precision')) if m.group('precision') is not None else ''
+    typ = m.group('type') or ''
+
+    flags = ''
+    if align == '<':
+        flags += '-'
+    if sign in ('+', ' '):
+        flags += sign
+    if alt:
+        flags += '#'
+    if zero and align != '<':
+        flags += '0'
+
+    return f'%{flags}{width}{precision}{typ}'
+
+
+def _extract_format_spec(spec_node, ctx: _Context):
+    """Extract explicit format spec from a FormattedValue's format_spec.
+    Returns a C printf format specifier (e.g. '%.2f') or None for type-inferred."""
+    if spec_node is None:
+        return None
+    if not isinstance(spec_node, ast.JoinedStr):
+        raise ctx.syntax_error("ct.print(): internal error: unexpected format_spec node")
+    if len(spec_node.values) == 0:
+        return None
+    if len(spec_node.values) != 1 or not isinstance(spec_node.values[0], ast.Constant):
+        raise ctx.syntax_error(
+            "ct.print() f-string: dynamic format specs (e.g. {x:{width}}) are not supported")
+    py_spec = str(spec_node.values[0].value)
+    return _python_spec_to_c_printf(py_spec, ctx)
+
+
+def _process_fstring(node: ast.JoinedStr, format_parts: list, tile_var_hirs: list,
+                     ctx: _Context) -> None:
+    """Decompose a JoinedStr (f-string) into format template parts and HIR vars."""
+    for part in node.values:
+        if isinstance(part, ast.Constant):
+            format_parts.append(_escape_format_str(str(part.value)))
+        elif isinstance(part, ast.FormattedValue):
+            if part.conversion != -1:
+                raise ctx.syntax_error(
+                    "ct.print() f-string: !r, !s, !a conversions are not supported")
+            c_spec = _extract_format_spec(part.format_spec, ctx)
+            if c_spec is not None:
+                format_parts.append(c_spec)
+            else:
+                format_parts.append('\x01')
+            tile_var_hirs.append(_expr(part.value, ctx))
+        else:
+            raise ctx.syntax_error("ct.print(): unsupported f-string component")
+
+
+def _require_str_constant(node: ast.expr, ctx: _Context, kwarg_name: str) -> str:
+    """Require a keyword argument to be a string constant at AST level."""
+    if not isinstance(node, ast.Constant) or not isinstance(node.value, str):
+        raise ctx.syntax_error(
+            f"ct.print(): keyword argument '{kwarg_name}' must be a string constant")
+    return node.value
+
+
+def _handle_ct_print(call: ast.Call, ctx: _Context) -> hir.Value:
+    """Handle ct.print() calls by decomposing f-strings and building HIR."""
+    sep = ' '
+    end = '\n'
+    for kw in call.keywords:
+        if kw.arg == 'sep':
+            sep = _require_str_constant(kw.value, ctx, 'sep')
+        elif kw.arg == 'end':
+            end = _require_str_constant(kw.value, ctx, 'end')
+        else:
+            raise ctx.syntax_error(
+                f"ct.print() got unexpected keyword argument '{kw.arg}'")
+
+    format_parts = []
+    tile_var_hirs = []
+    first = True
+
+    for arg_node in call.args:
+        if not first:
+            format_parts.append(_escape_format_str(sep))
+        first = False
+
+        if isinstance(arg_node, ast.JoinedStr):
+            _process_fstring(arg_node, format_parts, tile_var_hirs, ctx)
+        elif isinstance(arg_node, ast.Constant) and isinstance(arg_node.value, str):
+            format_parts.append(_escape_format_str(arg_node.value))
+        else:
+            format_parts.append('\x01')
+            tile_var_hirs.append(_expr(arg_node, ctx))
+
+    format_parts.append(_escape_format_str(end))
+    format_template = ''.join(format_parts)
+
+    template_hir = ctx.call(hir.identity, (format_template,))
+    return ctx.call(_ct_stub.printf, (template_hir, *tile_var_hirs))
+
+
 @_register(_expr_handlers, ast.Name)
 def _name_expr(name: ast.Name, ctx: Any) -> hir.Value:
     if not isinstance(name.ctx, ast.Load):

src/cuda/tile/_stub.py

@@ -2035,6 +2035,38 @@ def my_print_kernel():
     """
 
 
+@function
+def print(*args, sep: str = ' ', end: str = '\n') -> None:
+    """Print values at runtime from the device using Python-style syntax.
+
+    Supports Python f-strings and positional arguments similar to Python's
+    built-in ``print()`` function.
+
+    Args:
+        *args: Values to print. Each argument can be:
+            - A string literal or f-string
+            - A tile value (format inferred from dtype: int→``%d``, float→``%f``)
+        sep (str): Separator inserted between arguments (default: ``' '``)
+        end (str): String appended after the last argument (default: ``'\\n'``)
+
+    Examples:
+
+        >>> tile = ct.arange(4, dtype=ct.int32)
+        >>> ct.print(f"tile={tile}")
+        >>> ct.print(f"x={tile:.5f}", end='')
+        >>> ct.print("tile:", tile, sep='=')
+
+    Notes:
+        This operation has significant overhead, and should only be used
+        for debugging purposes.
+
+        F-string expressions must evaluate to tile values. Constant compile-time
+        values are supported as string-formatted segments.
+
+        Use ``opt_level=0`` to prevent block-level output interleaving.
+    """
+
+
 @function
 def assert_(cond, /, message=None) -> None:
     """Assert that all elements of the given tile are True.