Print statement for codegen backends

src/finchlite/codegen/__init__.py

@@ -19,6 +19,7 @@
     CKernel,
     CLibrary,
     CLowerer,
+    PrintableCFType,
 )
 from .numba_codegen import (
     NumbaBinaryOperator,
@@ -31,6 +32,7 @@
     NumbaNAryOperator,
     NumbaOperator,
     NumbaUnaryOperator,
+    PrintableNumbaFType,
 )
 
 __all__ = [
@@ -64,6 +66,8 @@
     "NumbaUnaryOperator",
     "NumpyBuffer",
     "NumpyBufferFType",
+    "PrintableCFType",
+    "PrintableNumbaFType",
     "SafeBuffer",
     "SafeBufferFType",
 ]

src/finchlite/codegen/c_codegen/__init__.py

@@ -13,11 +13,13 @@
     COperator,
     CStackFType,
     CUnaryOperator,
+    PrintableCFType,
     c_eq,
     c_function_call,
     c_function_name,
     c_hash,
     c_literal,
+    c_print,
     c_type,
     construct_from_c,
     create_shared_lib,
@@ -44,11 +46,13 @@
     "COperator",
     "CStackFType",
     "CUnaryOperator",
+    "PrintableCFType",
     "c_eq",
     "c_function_call",
     "c_function_name",
     "c_hash",
     "c_literal",
+    "c_print",
     "c_type",
     "construct_from_c",
     "create_shared_lib",

src/finchlite/codegen/c_codegen/c.py

@@ -1,3 +1,4 @@
+import builtins
 import ctypes
 import logging
 import shutil
@@ -481,6 +482,116 @@ def struct_c_type(fmt: StructFType):
 }
 
 
+ctype_print_fmt: dict[Any, str] = {
+    ctypes.c_bool: "%d",
+    ctypes.c_char: "%c",
+    ctypes.c_wchar: "%lc",
+    ctypes.c_byte: "%d",
+    ctypes.c_ubyte: "%d",
+    ctypes.c_int8: "%d",
+    ctypes.c_int16: "%d",
+    ctypes.c_int32: "%d",
+    ctypes.c_int64: "%ld",
+    ctypes.c_uint8: "%u",
+    ctypes.c_uint16: "%u",
+    ctypes.c_uint32: "%u",
+    ctypes.c_uint64: "%lu",
+    ctypes.c_char_p: "%s",
+    ctypes.c_wchar_p: "%ls",
+    ctypes.c_float: "%f",
+    ctypes.c_double: "%f",
+}
+
+
+class PrintableCFType(ABC):
+    @abstractmethod
+    def c_print(self, ctx, obj) -> str:
+        """
+        Return a C expression for this value in printf argument position.
+        """
+        ...
+
+
+def c_print(fmt: FType, ctx, obj) -> str:
+    match fmt:
+        case PrintableCFType():
+            return fmt.c_print(ctx, obj)
+        case (
+            algebra.ftypes.FDTypeNumpy()
+            | algebra.int_
+            | algebra.float_
+            | algebra.bool_
+            | algebra.str_
+        ):
+            return c_print_scalar(fmt, ctx, obj)
+        case _:
+            raise NotImplementedError(f"No C print mapping for {fmt}")
+
+
+def c_print_scalar(fmt: FType, ctx, obj) -> str:
+    ctype = c_type(fmt)
+    if ctype not in ctype_print_fmt:
+        raise NotImplementedError(f"No C print mapping for {fmt}")
+    if isinstance(fmt, algebra.ftypes.FDTypeBoolean):
+        return f"(int){obj}"
+    return obj
+
+
+def c_print_string_fallback(fmt: FType, ctx) -> str:
+    try:
+        fallback = ctx.ctype_name(c_type(fmt))
+    except NotImplementedError:
+        fallback = str(fmt)
+    return f'"{c_string_literal(fallback)}"'
+
+
+def c_printf_arg(fmt: FType, ctx, obj, conversion: str) -> str:
+    if conversion == "s":
+        try:
+            return c_print(fmt, ctx, obj)
+        except NotImplementedError:
+            return c_print_string_fallback(fmt, ctx)
+    return c_print(fmt, ctx, obj)
+
+
+def c_string_literal(value: str) -> str:
+    return value.replace("\\", "\\\\").replace('"', '\\"').replace("\n", "\\n")
+
+
+def printf_conversions(fmt: str) -> list[str]:
+    conversions = []
+    idx = 0
+    while idx < len(fmt):
+        if fmt[idx] != "%":
+            idx += 1
+            continue
+        idx += 1
+        if idx < len(fmt) and fmt[idx] == "%":
+            idx += 1
+            continue
+        while idx < len(fmt) and fmt[idx] in "-+ #0":
+            idx += 1
+        if idx < len(fmt) and fmt[idx] == "*":
+            raise NotImplementedError("dynamic printf widths are not supported")
+        while idx < len(fmt) and fmt[idx].isdigit():
+            idx += 1
+        if idx < len(fmt) and fmt[idx] == ".":
+            idx += 1
+            if idx < len(fmt) and fmt[idx] == "*":
+                raise NotImplementedError("dynamic printf precisions are not supported")
+            while idx < len(fmt) and fmt[idx].isdigit():
+                idx += 1
+        if fmt.startswith(("hh", "ll"), idx):
+            idx += 2
+        elif idx < len(fmt) and fmt[idx] in "hljztL":
+            idx += 1
+        if idx >= len(fmt):
+            raise ValueError("incomplete printf format specifier")
+        conversions.append(fmt[idx])
+        idx += 1
+    return conversions
+
+
 class CGenerator(UnvalidatedForm, CLowerer):
     def lower(self, prgm: asm.AssemblyNode) -> CCode:
         ctx = CContext()
@@ -869,6 +980,31 @@ def __call__(self, prgm: asm.AssemblyNode):
                         )
                     self(func)
                 return None
+            case asm.Print(args):
+                self.add_header("#include <stdio.h>")
+                match args:
+                    case (asm.Literal(str() as fmt), *vals):
+                        pass
+                    case _:
+                        raise TypeError("Print expects a literal format string")
+                conversions = printf_conversions(fmt)
+                if builtins.len(conversions) != builtins.len(vals):
+                    raise ValueError(
+                        f"printf format expects {builtins.len(conversions)} arguments, "
+                        f"got {builtins.len(vals)}"
+                    )
+                print_args = [
+                    c_printf_arg(val.result_type, self, self(val), conversion)
+                    for val, conversion in zip(vals, conversions, strict=False)
+                ]
+                fmt_str = c_string_literal(fmt)
+                if print_args:
+                    self.exec(
+                        f'{feed}printf("{fmt_str}", {", ".join(print_args)});'
+                    )
+                else:
+                    self.exec(f'{feed}printf("{fmt_str}");')
+                return None
             case _:
                 raise NotImplementedError(
                     f"Unrecognized assembly node type: {type(prgm)}"

src/finchlite/codegen/numba_codegen/__init__.py

@@ -12,6 +12,7 @@
     NumbaOperator,
     NumbaStackFType,
     NumbaUnaryOperator,
+    PrintableNumbaFType,
     construct_from_numba,
     deserialize_from_numba,
     numba_binary_function_call,
@@ -21,6 +22,7 @@
     numba_getattr,
     numba_jitclass_type,
     numba_nary_function_call,
+    numba_print,
     numba_setattr,
     numba_type,
     numba_unary_function_call,
@@ -45,6 +47,7 @@
     "NumbaOperator",
     "NumbaStackFType",
     "NumbaUnaryOperator",
+    "PrintableNumbaFType",
     "construct_from_numba",
     "deserialize_from_numba",
     "numba_binary_function_call",
@@ -54,6 +57,7 @@
     "numba_getattr",
     "numba_jitclass_type",
     "numba_nary_function_call",
+    "numba_print",
     "numba_setattr",
     "numba_type",
     "numba_unary_function_call",

src/finchlite/codegen/numba_codegen/numba.py

@@ -207,6 +207,23 @@ def construct_from_numba(self, numba_buffer):
         ...
 
 
+class PrintableNumbaFType(ABC):
+    @abstractmethod
+    def numba_print(self, ctx, obj) -> str:
+        """
+        Return a Numba expression for this value in printf argument position.
+        """
+        ...
+
+
+def numba_print(fmt: FType, ctx, obj) -> str:
+    match fmt:
+        case PrintableNumbaFType():
+            return fmt.numba_print(ctx, obj)
+        case _:
+            return ctx(obj)
+
+
 def to_numpy_type(t: FType) -> np.dtype:
     """Return a NumPy dtype for a Finch scalar/data type."""
     if isinstance(t, algebra.ftypes.FDTypeNumpy):
@@ -588,7 +605,15 @@ def lower(self, prgm: asm.AssemblyNode) -> NumbaCode:
 
 
 class NumbaContext(Context):
-    def __init__(self, tab="    ", indent=0, types=None, slots=None):
+    def __init__(
+        self,
+        tab="    ",
+        indent=0,
+        types=None,
+        slots=None,
+        imports=None,
+        importset=None,
+    ):
         if types is None:
             types = ScopedDict()
         if slots is None:
@@ -601,13 +626,18 @@ def __init__(self, tab="    ", indent=0, types=None, slots=None):
         self.types = types
         self.slots = slots
 
-        self.imports = [
-            "import _operator, builtins",
-            "from numba import njit",
-            "import numpy",
-            "from numpy import int64, float64",
-            "\n",
-        ]
+        if imports is None:
+            imports = [
+                "import _operator, builtins",
+                "from numba import njit",
+                "import numpy",
+                "from numpy import int64, float64",
+                "\n",
+            ]
+        if importset is None:
+            importset = set(imports)
+        self.imports = imports
+        self._importset = importset
 
     @property
     def feed(self) -> str:
@@ -622,10 +652,17 @@ def emit_global(self):
     def emit(self):
         return "\n".join([*self.preamble, *self.epilogue])
 
+    def add_import(self, import_line: str):
+        if import_line not in self._importset:
+            self.imports.insert(-1, import_line)
+            self._importset.add(import_line)
+
     def block(self) -> "NumbaContext":
         blk = super().block()
         blk.indent = self.indent
         blk.tab = self.tab
+        blk.imports = self.imports
+        blk._importset = self._importset
         blk.types = self.types
         blk.slots = self.slots
         return blk
@@ -874,6 +911,26 @@ def __call__(self, prgm: asm.AssemblyNode):
                         )
                     self(func)
                 return None
+            case asm.Print(args):
+                match args:
+                    case (asm.Literal(str() as fmt), *vals):
+                        pass
+                    case _:
+                        raise TypeError("Print expects a literal format string")
+                print_args = [
+                    numba_print(val.result_type, self, val) for val in vals
+                ]
+                tuple_expr = ", ".join(print_args)
+                if len(print_args) == 1:
+                    tuple_expr = f"{tuple_expr},"
+                fmt_var = self.freshen("fmt")
+                self.add_import("from numba import objmode")
+                self.exec(
+                    f"{feed}with objmode():\n"
+                    f"{feed}{self.tab}{fmt_var} = {fmt!r}\n"
+                    f"{feed}{self.tab}print({fmt_var} % ({tuple_expr}), end='')"
+                )
+                return None
             case node:
                 raise NotImplementedError(f"Unrecognized node: {node}")
 

src/finchlite/finch_assembly/interpreter.py

@@ -371,10 +371,15 @@ def my_func(*args_e):
                             )
                 return AssemblyInterpreterLibrary(self, kernels)
             case asm.Print(args):
-                args_value_str = ""
-                for arg in args:
-                    args_value_str = args_value_str + f"{self(arg)} "
-                print(args_value_str, file=self.stdout)
+                match args:
+                    case (asm.Literal(str() as fmt), *vals):
+                        print(
+                            fmt % tuple(self(val) for val in vals),
+                            end="",
+                            file=self.stdout,
+                        )
+                    case _:
+                        raise TypeError("Print expects a literal format string")
                 return None
             case asm.Stack(val):
                 raise NotImplementedError(