diff --git a/example/README.md b/example/README.md index 93269e5dc1..2b721049d9 100644 --- a/example/README.md +++ b/example/README.md @@ -5,6 +5,7 @@ This folder contains examples of how to use the Kirin library. Each example is a ## List of Examples - `simple.py`: A simple example that demonstrates how to create a simple Kirin dialect group and its kernel. +- `ir_inspector.py`: Generates an interactive HTML view of a compiled kernel; the same view can be rendered inline in a notebook. - `food`: A more sophisticated example but without any domain specifics. It demonstrates how to create a new Kirin dialect and combine it with existing dialects with custom analysis and rewrites. - `pauli`: An example that implements a dialect with rewrites that simplifies products of Pauli matrices. diff --git a/example/ir_inspector.py b/example/ir_inspector.py new file mode 100644 index 0000000000..8e8c2bbfb8 --- /dev/null +++ b/example/ir_inspector.py @@ -0,0 +1,22 @@ +"""Generate an interactive HTML view of a compiled Kirin kernel.""" + +from pathlib import Path + +from kirin.prelude import basic + + +@basic +def add_one(x: int) -> int: + return x + 1 + + +if __name__ == "__main__": + output = add_one.visualize(Path("add_one.ir.html")) + print(f"Wrote interactive IR inspector to {output.resolve()}") + +# In a Jupyter notebook, render the same inspector inline instead: +# +# from IPython.display import HTML, display +# from kirin import ir_to_html +# +# display(HTML(ir_to_html(add_one))) diff --git a/src/kirin/__init__.py b/src/kirin/__init__.py index bc01bee243..48f1d2d63e 100644 --- a/src/kirin/__init__.py +++ b/src/kirin/__init__.py @@ -1,5 +1,14 @@ # re-exports the public API of the kirin package from . import ir as ir, types as types, stdlib as stdlib, lowering as lowering from .exception import enable_stracetrace, disable_stracetrace +from .visualize import to_html as ir_to_html, write_html as write_ir_html -__all__ = ["ir", "types", "lowering", "enable_stracetrace", "disable_stracetrace"] +__all__ = [ + "ir", + "types", + "lowering", + "enable_stracetrace", + "disable_stracetrace", + "ir_to_html", + "write_ir_html", +] diff --git a/src/kirin/ir/method.py b/src/kirin/ir/method.py index 4ad9ea36ad..3516e7dece 100644 --- a/src/kirin/ir/method.py +++ b/src/kirin/ir/method.py @@ -13,6 +13,11 @@ from kirin.print.printable import Printable if typing.TYPE_CHECKING: + from collections.abc import Mapping + from pathlib import Path + from typing import Any + + from kirin.ir.ssa import SSAValue from kirin.serialization.base.serializer import Serializer from kirin.serialization.base.deserializer import Deserializer from kirin.serialization.core.serializationunit import SerializationUnit @@ -161,6 +166,19 @@ def __repr__(self) -> str: def print_impl(self, printer: Printer) -> None: return printer.print(self.code) + def visualize( + self, + file: str | Path, + *, + analysis: Mapping[SSAValue, Any] | None = None, + title: str | None = None, + ) -> Path: + """Write an interactive HTML inspector for this method's IR.""" + + from kirin.visualize import write_html + + return write_html(self, file, analysis=analysis, title=title) + def similar(self, dialects: typing.Optional["DialectGroup"] = None): return Method( dialects=dialects or self.dialects, diff --git a/src/kirin/visualize.py b/src/kirin/visualize.py new file mode 100644 index 0000000000..3876d3b7de --- /dev/null +++ b/src/kirin/visualize.py @@ -0,0 +1,518 @@ +"""Interactive HTML inspection for Kirin IR. + +The renderer is intentionally dependency-free so an exported file can be +opened locally without a running server. It is a debugging view: it does not +attempt to edit or execute the IR. +""" + +from __future__ import annotations + +import html +import json +from typing import TYPE_CHECKING, Any +from pathlib import Path +from collections.abc import Mapping + +from kirin.ir.ssa import SSAValue, ResultValue, BlockArgument +from kirin.idtable import IdTable +from kirin.ir.nodes.stmt import Statement +from kirin.ir.nodes.block import Block +from kirin.ir.nodes.region import Region + +if TYPE_CHECKING: + from kirin.source import SourceInfo + from kirin.ir.method import Method + + +def to_html( + node: Method | Statement | Region | Block, + *, + analysis: Mapping[SSAValue, Any] | None = None, + title: str | None = None, +) -> str: + """Render Kirin IR as a self-contained interactive HTML document. + + Hovering an SSA value shows its producer, type, uses, hints, and an + optional analysis fact. Hovering a statement shows its IR class, dialect, + operands, results, attributes, traits, and source span. + """ + + return _Inspector(node, analysis=analysis, title=title).render() + + +def write_html( + node: Method | Statement | Region | Block, + file: str | Path, + *, + analysis: Mapping[SSAValue, Any] | None = None, + title: str | None = None, +) -> Path: + """Write an interactive HTML inspection page and return its path.""" + + path = Path(file) + path.write_text(to_html(node, analysis=analysis, title=title), encoding="utf-8") + return path + + +class _Inspector: + def __init__( + self, + node: Method | Statement | Region | Block, + *, + analysis: Mapping[SSAValue, Any] | None, + title: str | None, + ) -> None: + self.node = node + self.analysis = analysis + self.title = title or self._default_title(node) + self.ssa_labels: IdTable[SSAValue] = IdTable() + self.ssa_ids: dict[SSAValue, str] = {} + self.statement_ids: dict[Statement, str] = {} + self.block_ids: dict[Block, str] = {} + self.ssa_details: dict[str, dict[str, Any]] = {} + self.statement_details: dict[str, dict[str, Any]] = {} + + @staticmethod + def _default_title(node: Method | Statement | Region | Block) -> str: + sym_name = getattr(node, "sym_name", None) + if sym_name is not None: + return f"Kirin IR: {sym_name or ''}" + return "Kirin IR Inspector" + + def render(self) -> str: + body = self._render_node(self.node) + data = json.dumps( + {"ssa": self.ssa_details, "statements": self.statement_details}, + ensure_ascii=True, + ) + return _PAGE.format( + title=html.escape(self.title), + body=body, + data=data.replace(" str: + if isinstance(node, Statement): + return self._render_statement(node, 0) + if isinstance(node, Region): + return self._render_region(node, 0) + if isinstance(node, Block): + return self._render_block(node, 0) + return self._render_statement(node.code, 0) + + def _ssa_id(self, value: SSAValue) -> str: + if value not in self.ssa_ids: + self.ssa_ids[value] = f"ssa-{len(self.ssa_ids)}" + return self.ssa_ids[value] + + def _ssa_label(self, value: SSAValue) -> str: + return self.ssa_labels[value] + + def _statement_id(self, statement: Statement) -> str: + if statement not in self.statement_ids: + self.statement_ids[statement] = f"stmt-{len(self.statement_ids)}" + return self.statement_ids[statement] + + def _block_id(self, block: Block) -> str: + if block not in self.block_ids: + self.block_ids[block] = f"block-{len(self.block_ids)}" + return self.block_ids[block] + + def _render_ssa(self, value: SSAValue) -> str: + key = self._ssa_id(value) + label = self._ssa_label(value) + self._add_ssa_details(value) + return ( + f'' + f"{html.escape(label)}" + ) + + def _render_statement(self, statement: Statement, depth: int) -> str: + statement_id = self._statement_id(statement) + self._add_statement_details(statement) + results = ", ".join(self._render_ssa(value) for value in statement.results) + prefix = f'{results} = ' if results else "" + operation = html.escape(self._operation_name(statement)) + args = self._render_arguments(statement) + attributes = self._render_attributes(statement) + result_types = ", ".join( + html.escape(_safe_repr(value.type)) for value in statement.results + ) + type_suffix = ( + f' : {result_types}' if result_types else "" + ) + source = self._source_label(statement.source) + source_suffix = ( + f' {html.escape(source)}' if source else "" + ) + indent = depth * 20 + line = ( + f'
{prefix}' + f'{operation}({args}){attributes}' + f"{type_suffix}{source_suffix}
" + ) + if not statement.regions: + return line + + regions = "".join( + self._render_region(region, depth + 1) for region in statement.regions + ) + return f'{line}
{regions}
' + + def _render_region(self, region: Region, depth: int) -> str: + if not region.blocks: + return '
{{}}
' + return "".join(self._render_block(block, depth) for block in region.blocks) + + def _render_block(self, block: Block, depth: int) -> str: + block_id = self._block_id(block) + args = ", ".join( + f'{self._render_ssa(arg)} : ' + f"{html.escape(_safe_repr(arg.type))}" + for arg in block.args + ) + header = "" + has_multiple_blocks = block.parent is not None and len(block.parent.blocks) > 1 + if args or has_multiple_blocks: + header = ( + f'
' + f'^{block_id}({args}):
' + ) + statements = "".join( + self._render_statement(statement, depth + 1) for statement in block.stmts + ) + return f"{header}{statements}" + + def _render_arguments(self, statement: Statement) -> str: + names = self._argument_names(statement) + values: list[str] = [] + for index, value in enumerate(statement.args): + values.append( + f'{html.escape(names.get(index, f"arg{index}"))}=' + f"{self._render_ssa(value)}" + ) + return ", ".join(values) + + def _render_attributes(self, statement: Statement) -> str: + if not statement.attributes: + return "" + rendered = ", ".join( + f"{html.escape(name)}={html.escape(_safe_repr(value))}" + for name, value in statement.attributes.items() + ) + return f' {{{rendered}}}' + + @staticmethod + def _argument_names(statement: Statement) -> dict[int, str]: + names: dict[int, str] = {} + for name, slice_ in statement._name_args_slice.items(): + if isinstance(slice_, int): + names[slice_] = name + else: + start, stop, step = slice_.indices(len(statement.args)) + for index in range(start, stop, step): + names[index] = f"{name}[{index - start}]" + return names + + @staticmethod + def _operation_name(statement: Statement) -> str: + dialect = statement.dialect.name if statement.dialect else "" + return f"{dialect + '.' if dialect else ''}{statement.name}" + + def _add_ssa_details(self, value: SSAValue) -> None: + key = self._ssa_id(value) + if key in self.ssa_details: + return + + details: dict[str, Any] = { + "Kind": type(value).__name__, + "Type": _safe_repr(value.type), + "Hints": {name: _safe_repr(hint) for name, hint in value.hints.items()}, + } + if isinstance(value, ResultValue): + details["Producer"] = { + "statement": self._statement_id(value.owner), + "operation": self._operation_name(value.owner), + "result index": value.index, + } + elif isinstance(value, BlockArgument): + details["Producer"] = { + "block": self._block_id(value.owner), + "argument index": value.index, + } + else: + details["Producer"] = _safe_repr(value.owner) + + uses = sorted( + value.uses, key=lambda use: (self._statement_id(use.stmt), use.index) + ) + details["Uses"] = [ + { + "statement": self._statement_id(use.stmt), + "operation": self._operation_name(use.stmt), + "operand index": use.index, + } + for use in uses + ] + if self.analysis is not None and value in self.analysis: + details["Analysis"] = _safe_repr(self.analysis[value]) + self.ssa_details[key] = details + + def _add_statement_details(self, statement: Statement) -> None: + key = self._statement_id(statement) + if key in self.statement_details: + return + + argument_names = self._argument_names(statement) + details: dict[str, Any] = { + "Class": f"{type(statement).__module__}.{type(statement).__qualname__}", + "Dialect": statement.dialect.name if statement.dialect else None, + "Operation": self._operation_name(statement), + "Operands": [ + { + "name": argument_names.get(index, f"arg{index}"), + "ssa": self._ssa_label(value), + "type": _safe_repr(value.type), + } + for index, value in enumerate(statement.args) + ], + "Results": [ + { + "ssa": self._ssa_label(value), + "type": _safe_repr(value.type), + "uses": len(value.uses), + } + for value in statement.results + ], + "Attributes": { + name: _safe_repr(value) for name, value in statement.attributes.items() + }, + "Traits": [type(trait).__name__ for trait in statement.traits], + "Regions": len(statement.regions), + "Successors": [self._block_id(block) for block in statement.successors], + } + if source := self._source_details(statement.source): + details["Source"] = source + self.statement_details[key] = details + + @staticmethod + def _source_label(source: SourceInfo | None) -> str | None: + if source is None: + return None + line = source.lineno + source.lineno_begin + location = f"{source.file or ''}:{line}:{source.col_offset + source.col_indent}" + return location + + def _source_details(self, source: SourceInfo | None) -> dict[str, Any] | None: + label = self._source_label(source) + if source is None or label is None: + return None + details: dict[str, Any] = {"location": label} + if snippet := _read_source_snippet(source): + details["snippet"] = snippet + return details + + +def _safe_repr(value: object) -> str: + try: + return repr(value) + except Exception: + return f"<{type(value).__name__}>" + + +def _read_source_snippet(source: SourceInfo) -> str | None: + if not source.file: + return None + try: + lines = Path(source.file).read_text(encoding="utf-8").splitlines() + except (OSError, UnicodeError): + return None + + first = source.lineno + source.lineno_begin + last = (source.end_lineno or source.lineno) + source.lineno_begin + if first < 1 or last < first or first > len(lines): + return None + return "\n".join(lines[first - 1 : min(last, len(lines))]) + + +_PAGE = """ + + + + +{title} + + + +
{title}
+
+
{body}
+ +
+ + + + +""" diff --git a/test/visualize/test_html.py b/test/visualize/test_html.py new file mode 100644 index 0000000000..428fd76dbe --- /dev/null +++ b/test/visualize/test_html.py @@ -0,0 +1,32 @@ +from kirin.prelude import basic +from kirin.visualize import to_html, write_html + + +@basic +def add_one(x: int): + return x + 1 + + +def test_html_inspector_includes_ssa_and_statement_metadata(tmp_path): + result = add_one.callable_region.blocks[0].stmts.at(0).result + page = to_html(add_one, analysis={result: "constant one"}) + + assert 'data-ssa="ssa-' in page + assert 'data-stmt="stmt-' in page + assert '"Producer"' in page + assert '"Uses"' in page + assert '"Class"' in page + assert '"Attributes"' in page + assert "constant one" in page + + path = add_one.visualize(tmp_path / "add_one.html") + assert path.read_text(encoding="utf-8").startswith("") + + +def test_write_html_accepts_a_statement(tmp_path): + statement = add_one.callable_region.blocks[0].stmts.at(0) + path = write_html(statement, tmp_path / "statement.html") + + page = path.read_text(encoding="utf-8") + assert "Kirin IR Inspector" in page + assert statement.name in page