Move ct.print logic out of AST

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

Author: gunnersdeng

src/cuda/tile/_ir/hir.py

@@ -241,6 +241,7 @@ def if_else(cond, then_block, else_block, /): ...
 def loop(body, iterable, /): ...  # infinite if `iterable` is None
 def static_foreach(body, items, /): ...
 def build_tuple(*items): ...  # Makes a tuple (i.e. returns `items`)
+def build_formatted_string(format, *values): ...  # Creates a FormattedStringTy value
 def unpack(iterable, expected_len, /): ...
 def identity(x): ...   # Identity function (i.e. returns `x`)
 def store_var(name, value, /): ...  # Store into a named variable

src/cuda/tile/_ir/ir.py

@@ -243,6 +243,15 @@ def as_tuple(self) -> tuple["Var", ...]:
         return self.items
 
 
+@dataclass
+class FormattedStringValue(AggregateValue):
+    format: "Any"  # StringFormat from type.py
+    values: tuple  # tuple of Var
+
+    def as_tuple(self) -> tuple:
+        return self.values
+
+
 @dataclass
 class RangeValue(AggregateValue):
     start: Var

src/cuda/tile/_ir/op_impl.py

@@ -426,6 +426,15 @@ class PrintfValidator:
     specifiers = r"([diuoxXeEfFgGaAcspn])"
     pattern = re.compile("%" + flags + width + precision + length + specifiers)
 
+    @classmethod
+    def infer_format(cls, dtype: DType) -> str:
+        if is_boolean(dtype) or is_integral(dtype):
+            return '%d'
+        elif is_float(dtype) or is_restricted_float(dtype):
+            return '%f'
+        else:
+            raise TileTypeError(f"print(): cannot infer format for dtype {dtype}")
+
     @classmethod
     def validate_dtype(cls, dtype: DType, specifier: str) -> bool:
         if is_boolean(dtype) or is_integral(dtype):
@@ -435,8 +444,58 @@ def validate_dtype(cls, dtype: DType, specifier: str) -> bool:
         else:
             return False
 
-    # Placeholder emitted by ast2hir for type-inferred format specifiers
-    _TYPE_INFER = '\x01'
+    # Python format spec regex: [align][sign][alt][zero][width][.precision][type]
+    _py_spec_pattern = re.compile(
+        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])?'
+    )
+
+    @staticmethod
+    def escape_str(s: str) -> str:
+        """Escape a literal string for use in a C printf format (replace % with %%)."""
+        return s.replace('%', '%%')
+
+    @classmethod
+    def apply_python_spec(cls, py_spec: str, dtype: DType) -> str:
+        """Convert a Python format spec to a complete C printf specifier for the given dtype.
+
+        If py_spec omits the type character, it is inferred from dtype.
+        If py_spec includes a type character, it is validated against dtype.
+        Raises TileTypeError on type mismatch; ValueError on unrecognised spec syntax.
+        """
+        m = cls._py_spec_pattern.fullmatch(py_spec)
+        if m is None or m.group(0) != py_spec:
+            raise ValueError(f"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')
+
+        flags = ''
+        if align == '<':
+            flags += '-'
+        if sign in ('+', ' '):
+            flags += sign
+        if alt:
+            flags += '#'
+        if zero and align != '<':
+            flags += '0'
+
+        if typ is None:
+            typ = cls.infer_format(dtype)[1:]  # inferred type char, e.g. 'd' from '%d'
+        elif not cls.validate_dtype(dtype, typ):
+            raise TileTypeError(
+                f"print(): format spec '{py_spec}' is incompatible with dtype {dtype}")
+        return f'%{flags}{width}{precision}{typ}'
 
     @classmethod
     def parse_format(cls, format: str, arg_types: Tuple[Union[TileTy, DType], ...]) -> str:
@@ -445,21 +504,7 @@ def parse_format(cls, format: str, arg_types: Tuple[Union[TileTy, DType], ...])
         tokens = []
         while pos < len(format):
             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 == "%":
+            if ch == "%":
                 tokens.append(format[last_pos:pos])
                 last_pos = pos
                 # escape "%%"

src/cuda/tile/_ir/ops.py

@@ -1,6 +1,7 @@
 # SPDX-FileCopyrightText: Copyright (c) <2025> NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 #
 # SPDX-License-Identifier: Apache-2.0
+import builtins
 import enum
 import math
 import operator
@@ -23,7 +24,8 @@
     add_operation, Builder,
     enter_nested_block, nested_block, PhiState, LoopVarState,
     TupleValue, make_aggregate, RangeValue, BoundMethodValue, ArrayValue, ConstantState,
-    ListValue, TiledViewValue, ClosureValue, MemoryEffect, attribute, operand, BlockRestriction
+    ListValue, TiledViewValue, ClosureValue, MemoryEffect, attribute, operand,
+    BlockRestriction, FormattedStringValue,
 )
 from .type import PointerTy
 from . import hir
@@ -57,7 +59,8 @@
     PartitionViewTy, TupleTy, TileTy, NoneType, BoundMethodTy, ArrayTy,
     ListTy, make_tile_ty, SliceType, DTypeConstructor, RangeIterType, Type,
     NONE, ModuleTy, TypeTy, LooselyTypedScalar, DTypeSpec, StringTy, InvalidType,
-    array_size_type, ClosureTy, LiveCapturedScope, TokenTy, TiledViewTy
+    array_size_type, ClosureTy, LiveCapturedScope, TokenTy, TiledViewTy, FormattedStringTy,
+    StringFormat, FormattedPiece,
 )
 from cuda.tile._datatype import (
     DType, is_integral, is_float, is_signed, is_boolean, is_restricted_float,
@@ -1318,6 +1321,39 @@ def build_tuple(items: tuple[Var, ...]) -> Var:
     return res
 
 
+@impl(hir.build_formatted_string)
+def build_formatted_string_impl(format: StringFormat, values: tuple[Var, ...]) -> Var:
+    new_pieces = []
+    new_values = []
+    for piece in format.pieces:
+        if isinstance(piece, str):
+            new_pieces.append(piece)
+        else:
+            val_var = values[piece.value_idx]
+            val_ty = val_var.get_type()
+            if isinstance(val_ty, FormattedStringTy):
+                if piece.format_spec is not None:
+                    raise TileTypeError(
+                        "f-string: cannot apply format spec to a formatted string value",
+                        val_var.loc)
+                inner_val = val_var.get_aggregate()
+                assert isinstance(inner_val, FormattedStringValue)
+                offset = len(new_values)
+                for inner_piece in val_ty.format.pieces:
+                    if isinstance(inner_piece, str):
+                        new_pieces.append(inner_piece)
+                    else:
+                        new_pieces.append(FormattedPiece(
+                            offset + inner_piece.value_idx, inner_piece.format_spec))
+                new_values.extend(inner_val.values)
+            else:
+                new_pieces.append(FormattedPiece(len(new_values), piece.format_spec))
+                new_values.append(val_var)
+    new_fmt = StringFormat(tuple(new_pieces))
+    ty = FormattedStringTy(new_fmt, tuple(v.get_type() for v in new_values))
+    return make_aggregate(FormattedStringValue(new_fmt, tuple(new_values)), ty)
+
+
 @impl(hir.unpack)
 def unpack_impl(iterable: Var, expected_len: Var) -> Var:
     ty = iterable.get_type()
@@ -3638,6 +3674,50 @@ def printf_impl(format: Var, args: Tuple[Var, ...]) -> None:
     add_operation(TilePrintf, (), format=parsed_format, args=args)
 
 
+@impl(ct.print)
+@impl(builtins.print)
+def print_impl(args: Tuple[Var, ...], sep: Var, end: Var) -> None:
+    sep_str = PrintfValidator.escape_str(require_constant_str(sep))
+    end_str = PrintfValidator.escape_str(require_constant_str(end))
+
+    format_parts = []
+    leaf_vars = []
+    first = True
+
+    for arg_var in args:
+        if not first:
+            format_parts.append(sep_str)
+        else:
+            first = False
+
+        arg_ty = arg_var.get_type()
+        if isinstance(arg_ty, FormattedStringTy):
+            fmt_val = arg_var.get_aggregate()
+            assert isinstance(fmt_val, FormattedStringValue)
+            for piece in arg_ty.format.pieces:
+                if isinstance(piece, str):
+                    format_parts.append(PrintfValidator.escape_str(piece))
+                else:
+                    value_ty = arg_ty.value_types[piece.value_idx]
+                    dtype = get_dtype(value_ty)
+                    if piece.format_spec is not None:
+                        format_parts.append(PrintfValidator.apply_python_spec(
+                            piece.format_spec, dtype))
+                    else:
+                        format_parts.append(PrintfValidator.infer_format(dtype))
+                    leaf_vars.append(fmt_val.values[piece.value_idx])
+        elif isinstance(arg_ty, StringTy):
+            format_parts.append(PrintfValidator.escape_str(arg_ty.value))
+        else:
+            tile_ty = require_tile_type(arg_var)
+            format_parts.append(PrintfValidator.infer_format(get_dtype(tile_ty)))
+            leaf_vars.append(arg_var)
+
+    format_parts.append(end_str)
+    final_format = ''.join(format_parts)
+    add_operation(TilePrintf, (), format=final_format, args=tuple(leaf_vars))
+
+
 @dataclass(eq=False)
 class TileAssert(Operation, opcode="assert", memory_effect=MemoryEffect.STORE):
     message: str = attribute()

src/cuda/tile/_ir/type.py

@@ -210,6 +210,50 @@ def map(self, unwrap: Callable[[Type], Any]) -> Tuple[Any, ...]:
         return tuple(unwrap(t) for t in self.value_types)
 
 
+# ============== Formatted String Type ===============
+
+@dataclass(frozen=True)
+class FormattedPiece:
+    """A single typed placeholder in a formatted string."""
+    value_idx: int          # index into FormattedStringTy.value_types
+    format_spec: str | None  # None = type-inferred; otherwise Python format spec (e.g. '.2f')
+
+
+@dataclass(frozen=True)
+class StringFormat:
+    """Immutable format template for a formatted string value."""
+    pieces: tuple[str | FormattedPiece, ...]
+
+
+@dataclass(frozen=True)
+class FormattedStringTy(Type):
+    format: "StringFormat"
+    value_types: tuple
+
+    def is_aggregate(self) -> bool:
+        return True
+
+    def aggregate_item_types(self) -> tuple:
+        return self.value_types
+
+    def make_aggregate_value(self, items: tuple) -> "AggregateValue":
+        from .ir import FormattedStringValue
+        return FormattedStringValue(self.format, items)
+
+    def __str__(self):
+        parts = []
+        for piece in self.format.pieces:
+            if isinstance(piece, str):
+                parts.append(piece)
+            else:
+                ty = self.value_types[piece.value_idx]
+                if piece.format_spec is not None:
+                    parts.append(f"{{<{ty}>:{piece.format_spec}}}")
+                else:
+                    parts.append(f"{{<{ty}>}}")
+        return 'FormattedString<"' + "".join(parts) + '">'
+
+
 def size_to_bytecode(s: Optional[int]) -> int:
     return bc.DYNAMIC_SHAPE if s is None else s
 

src/cuda/tile/_ir/typing_support.py

@@ -116,6 +116,7 @@ def wrapped(handler: TypeHandler):
     float: lambda x=0, /: None,
     int: lambda x=0, /: None,
     bool: lambda x=False, /: None,
+    print: lambda *args, sep=' ', end='\n': None,
 }