Add provenance annotations to output code (#988)

Author: jansel

docs/api/settings.md

@@ -190,6 +190,11 @@ See :class:`helion.autotuner.LocalAutotuneCache` for details on cache keys and b
 
    Print generated Triton code to stderr. Default is ``False``. Controlled by ``HELION_PRINT_OUTPUT_CODE=1``.
 
+.. autoattribute:: Settings.output_origin_lines
+
+   Annotate generated Triton code with ``# src[<file>:<line>]`` comments indicating the originating Helion statements.
+   Default is ``True``. Controlled by ``HELION_OUTPUT_ORIGIN_LINES`` (set to ``0`` to disable).
+
 .. autoattribute:: Settings.ignore_warnings
 
    List of warning types to suppress during compilation. Default is an empty list.
@@ -253,6 +258,7 @@ Built-in values for ``HELION_AUTOTUNER`` include ``"PatternSearch"``, ``"Differe
 | ``HELION_CACHE_DIR`` | ``LocalAutotuneCache`` | Override the on-disk directory used for cached autotuning artifacts. |
 | ``HELION_SKIP_CACHE`` | ``LocalAutotuneCache`` | When set to ``1``, ignore cached autotuning entries and rerun searches. |
 | ``HELION_PRINT_OUTPUT_CODE`` | ``print_output_code`` | Print generated Triton code to stderr for inspection. |
+| ``HELION_OUTPUT_ORIGIN_LINES`` | ``output_origin_lines`` | Include ``# src[...]`` comments in generated Triton code; set to ``0`` to disable. |
 | ``HELION_IGNORE_WARNINGS`` | ``ignore_warnings`` | Comma-separated warning names defined in ``helion.exc`` to suppress. |
 | ``HELION_ALLOW_WARP_SPECIALIZE`` | ``allow_warp_specialize`` | Permit warp-specialized code generation for ``tl.range``. |
 | ``HELION_DEBUG_DTYPE_ASSERTS`` | ``debug_dtype_asserts`` | Inject dtype assertions after each lowering step. |

helion/_compiler/ast_extension.py

@@ -2,14 +2,19 @@
 
 import ast
 import enum
+import linecache
+import os
 import re
+import textwrap
 import threading
 import typing
 from typing import TYPE_CHECKING
 from typing import TypeVar
 
 from .. import exc
+from .output_lines import OutputLines
 from .source_location import SourceLocation
+from .source_location import UnknownLocation
 from .source_location import current_location
 
 if TYPE_CHECKING:
@@ -302,6 +307,148 @@ def visit_Tuple(self, node: ast.Tuple) -> None:
         super().visit_Tuple(node)
 
 
-def unparse(ast_obj: ast.AST) -> str:
-    unparser = _TupleParensRemovedUnparser()
-    return unparser.visit(ast_obj)
+class _LocationAnnotatingOutputLines(OutputLines):
+    def __init__(self, parent: ast._Unparser) -> None:  # pyright: ignore[reportAttributeAccessIssue]
+        super().__init__(parent)
+        self._cache: dict[tuple[str, int, int], tuple[str, ...]] = {}
+        self._last_location_key: tuple[str, int, int] | None = None
+
+    def reset_last_location(self) -> None:
+        super().reset_last_location()
+        self._last_location_key = None
+
+    def insert_location_comment(self, location: object) -> None:
+        if not isinstance(location, (SourceLocation, UnknownLocation)):
+            location = UnknownLocation()
+        key = self._location_key(location)
+        if key is None or key == self._last_location_key:
+            return
+
+        comments = self._comments_for_key(key, location)
+        if comments:
+            self.insert_comments(comments)
+            self._last_location_key = key
+
+    def _location_key(
+        self, location: SourceLocation | UnknownLocation
+    ) -> tuple[str, int, int] | None:
+        if not location:
+            return ("<unknown>", 0, 0)
+        filename = location.filename
+        if not filename:
+            return None
+        start = location.lineno or 0
+        end = location.end_lineno or start
+        return (filename, start, end)
+
+    def _comments_for_key(
+        self,
+        key: tuple[str, int, int],
+        location: SourceLocation | UnknownLocation,
+    ) -> tuple[str, ...]:
+        cached = self._cache.get(key)
+        if cached is not None:
+            return cached
+
+        filename, start, end = key
+        if not location:
+            comments = ("# src[unknown]: [source unavailable]",)
+        elif start <= 0:
+            comments = (
+                f"# src[{os.path.basename(filename)}:{start}]: [source unavailable]",
+            )
+        else:
+            lines = linecache.getlines(filename)
+            if not lines:
+                linecache.checkcache(filename)
+                lines = linecache.getlines(filename)
+
+            if not lines:
+                comments = (
+                    f"# src[{os.path.basename(filename)}:{start}]: [source unavailable]",
+                )
+            else:
+                snippet_full = lines[start - 1 : end]
+                if not snippet_full:
+                    comments = (
+                        f"# src[{os.path.basename(filename)}:{start}]: [source unavailable]",
+                    )
+                else:
+                    max_lines = 3
+                    truncated = len(snippet_full) > max_lines
+                    snippet = snippet_full[:max_lines]
+                    dedented = textwrap.dedent("".join(snippet))
+                    body_list: list[str] = []
+                    base_name = os.path.basename(filename)
+                    for offset, dedented_line in enumerate(dedented.splitlines()):
+                        stripped = dedented_line.rstrip()
+                        if not stripped.strip():
+                            continue
+                        lineno = start + offset
+                        body_list.append(f"# src[{base_name}:{lineno}]: {stripped}")
+                    if truncated:
+                        range_part = f"{start}-{end}" if end != start else f"{start}"
+                        body_list.append(f"# src[{base_name}:{range_part}]: ...")
+                    comments = (
+                        tuple(body_list)
+                        if body_list
+                        else (f"# src[{base_name}:{start}]: [source unavailable]",)
+                    )
+
+        self._cache[key] = comments
+        return comments
+
+
+class _HelionUnparser(_TupleParensRemovedUnparser):
+    _indent: int
+
+    def __init__(
+        self, *args: object, output_origin_lines: bool = True, **kwargs: object
+    ) -> None:
+        super().__init__(*args, **kwargs)
+        if output_origin_lines:
+            self.output = _LocationAnnotatingOutputLines(self)
+        else:
+            self.output = OutputLines(self)
+        self._source = self.output
+        self._output_origin_lines = output_origin_lines
+
+    def visit(self, node: ast.AST) -> str:  # type: ignore[override]
+        self.output.lines.clear()
+        self.output.last_newline = 0
+        self.output.reset_last_location()
+        self.traverse(node)
+        return "".join(self.output)
+
+    def maybe_newline(self) -> None:  # type: ignore[override]
+        output = getattr(self, "output", None)
+        if output is not None and getattr(output, "_skip_next_newline", False):
+            output._skip_next_newline = False
+            return
+        super().maybe_newline()
+
+    def traverse(self, node: ast.AST | list[ast.AST]) -> None:  # pyright: ignore[reportSignatureIssue]
+        if (
+            self._output_origin_lines
+            and isinstance(node, ExtendedAST)
+            and isinstance(node, ast.stmt)
+        ):
+            if not isinstance(
+                node,
+                (
+                    ast.FunctionDef,
+                    ast.AsyncFunctionDef,
+                    ast.ClassDef,
+                    ast.Import,
+                    ast.ImportFrom,
+                ),
+            ):
+                self.output.insert_location_comment(node._location)
+        super().traverse(node)
+
+
+def unparse(ast_obj: ast.AST, *, output_origin_lines: bool = True) -> str:
+    unparser = _HelionUnparser(output_origin_lines=output_origin_lines)
+    result = unparser.visit(ast_obj)
+    del unparser.output  # break reference cycle
+    return result

helion/_compiler/output_lines.py

@@ -0,0 +1,73 @@
+from __future__ import annotations
+
+from typing import Iterable
+from typing import Iterator
+from typing import Sequence
+
+
+class OutputLines:
+    """
+    Helper to build source text while keeping track of the most recent newline so
+    callers can inject annotations before the currently buffered statement.
+    """
+
+    def __init__(self, parent: object) -> None:
+        super().__init__()
+        self.lines: list[str] = []
+        self.last_newline = 0
+        self.parent = parent
+        self._skip_next_newline = False
+
+    def extend(self, chunks: Iterable[str]) -> None:
+        """Append text while tracking the index after the last newline."""
+        concatenated = "".join(chunks)
+        if not concatenated:
+            return
+
+        new_lines = concatenated.splitlines(keepends=True)
+        self.lines.extend(new_lines)
+
+        if new_lines[-1].endswith("\n"):
+            self.last_newline = len(self.lines)
+        elif len(new_lines) > 1:
+            # Second to last line must end in newline if the last one did not.
+            assert new_lines[-2].endswith("\n")
+            self.last_newline = len(self.lines) - 1
+
+    def append(self, text: str) -> None:
+        self.extend([text])
+
+    def insert_comments(self, comments: Sequence[str]) -> None:
+        """Insert comment lines right before the current statement."""
+        if not comments:
+            return
+        if self.lines and not self.lines[-1].endswith("\n"):
+            self.lines[-1] = f"{self.lines[-1]}\n"
+            self.last_newline = len(self.lines)
+        indent = "    " * getattr(self.parent, "_indent", 0)
+        insert_at = min(max(self.last_newline, 0), len(self.lines))
+        for comment in comments:
+            assert "\n" not in comment
+            self.lines.insert(insert_at, f"{indent}{comment}\n")
+            insert_at += 1
+        self.last_newline = insert_at
+        self._skip_next_newline = True
+
+    def insert_annotation(self, annotation: str) -> None:
+        self.insert_comments((f"# {annotation}",))
+
+    def reset_last_location(self) -> None:
+        self._skip_next_newline = False
+
+    def insert_location_comment(self, location: object) -> None:
+        # Base OutputLines does not track source locations; override when needed.
+        return None
+
+    def __bool__(self) -> bool:
+        return bool(self.lines)
+
+    def __len__(self) -> int:
+        return len(self.lines)
+
+    def __iter__(self) -> Iterator[str]:
+        return iter(self.lines)

helion/_compiler/type_printer.py

@@ -4,48 +4,10 @@
 
 from .ast_extension import ExtendedAST
 from .ast_extension import _TupleParensRemovedUnparser
+from .output_lines import OutputLines
 
 if TYPE_CHECKING:
     import ast
-    from collections.abc import Iterator
-
-
-class OutputLines:
-    def __init__(self, parent: ASTPrinter) -> None:
-        super().__init__()
-        self.lines: list[str] = []
-        self.last_newline = 0
-        self.parent = parent
-
-    def extend(self, lines: list[str]) -> None:
-        """Keep track of the index right after the last newline so insert_annotation can insert in the correct spot."""
-        lines = "".join(lines).splitlines(keepends=True)
-        if not lines:
-            return
-        self.lines.extend(lines)
-        if lines[-1].endswith("\n"):
-            self.last_newline = len(self.lines)
-        elif len(lines) > 1:
-            assert lines[-2].endswith("\n")
-            self.last_newline = len(self.lines) - 1
-
-    def __bool__(self) -> bool:
-        return bool(self.lines)
-
-    def __len__(self) -> int:
-        return len(self.lines)
-
-    def __iter__(self) -> Iterator[str]:
-        return iter(self.lines)
-
-    def insert_annotation(self, annotation: str) -> None:
-        assert "\n" not in annotation
-        indent = "    " * self.parent._indent
-        self.lines.insert(self.last_newline, f"{indent}# {annotation}\n")
-        self.last_newline += 1
-
-    def append(self, text: str) -> None:
-        self.extend([text])
 
 
 class ASTPrinter(_TupleParensRemovedUnparser):

helion/_testing.py

@@ -772,8 +772,28 @@ def normalize_codegen_variants(code: str) -> str:
         # tl.sqrt( -> tl.sqrt_rn(
         return re.sub(r"\btl\.sqrt\s*\(", "tl.sqrt_rn(", code)
 
+    @staticmethod
+    def normalize_source_comment_structure(code: str) -> str:
+        pattern = re.compile(
+            r"^(?P<indent>\s*)# src\[(?P<prefix>[^:\]]+:)(?P<start>\d+|N)(?:-(?P<end>\d+|N))?]: (?P<text>.*?)(?P<newline>\r?\n|$)",
+            flags=re.MULTILINE,
+        )
+
+        def replacer(match: re.Match[str]) -> str:
+            text = match.group("text").rstrip()
+            if not text.strip():
+                return ""
+            indent = match.group("indent")
+            prefix = match.group("prefix")
+            suffix = "N-N" if match.group("end") is not None else "N"
+            newline = match.group("newline")
+            return f"{indent}# src[{prefix}{suffix}]: {text}{newline}"
+
+        return pattern.sub(replacer, code)
+
     @classmethod
     def normalize_code(cls, code: str) -> str:
+        code = cls.normalize_source_comment_structure(code)
         code = cls.normalize_tensor_descriptors(code)
         code = cls.normalize_device_name(code)
         code = cls.normalize_codegen_variants(code)

helion/runtime/kernel.py

@@ -394,7 +394,11 @@ def format_kernel_decorator(self, config: Config, settings: Settings) -> str:
         return f"@helion.kernel(config={config.__repr__()}, static_shapes={settings.static_shapes})"
 
     def to_triton_code(
-        self, config: ConfigLike | None = None, emit_repro_caller: bool = False
+        self,
+        config: ConfigLike | None = None,
+        *,
+        emit_repro_caller: bool = False,
+        output_origin_lines: bool | None = None,
     ) -> str:
         """
         Generate Triton code for the kernel based on the given configuration.
@@ -413,7 +417,11 @@ def to_triton_code(
                 config = Config(**config)  # pyright: ignore[reportArgumentType]
             self.env.config_spec.normalize(config)
             root = generate_ast(self.host_function, config, emit_repro_caller)
-            return get_needed_imports(root) + unparse(root)
+            if output_origin_lines is None:
+                output_origin_lines = self.settings.output_origin_lines
+            return get_needed_imports(root) + unparse(
+                root, output_origin_lines=output_origin_lines
+            )
 
     def compile_config(
         self, config: ConfigLike | None = None, *, allow_print: bool = True

helion/runtime/settings.py

@@ -314,6 +314,11 @@ class _Settings:
             _env_get_bool, "HELION_PRINT_OUTPUT_CODE", False
         )
     )
+    output_origin_lines: bool = dataclasses.field(
+        default_factory=functools.partial(
+            _env_get_bool, "HELION_OUTPUT_ORIGIN_LINES", True
+        )
+    )
     force_autotune: bool = dataclasses.field(
         default_factory=functools.partial(_env_get_bool, "HELION_FORCE_AUTOTUNE", False)
     )
@@ -379,6 +384,10 @@ class Settings(_Settings):
             "Set HELION_AUTOTUNE_IGNORE_ERRORS=1 to enable globally."
         ),
         "print_output_code": "If True, print the output code of the kernel to stderr.",
+        "output_origin_lines": (
+            "If True, annotate generated Triton code with source-origin comments. "
+            "Set HELION_OUTPUT_ORIGIN_LINES=0 to disable."
+        ),
         "force_autotune": "If True, force autotuning even if a config is provided.",
         "autotune_config_overrides": (
             "Dictionary of config key/value pairs forced during autotuning. "