documentation: (marimo) move pipeline_html helper to utils

docs/marimo/linalg_snitch.py

@@ -257,11 +257,11 @@ def _(mo):
 
 @app.cell
 def _(MLContext, get_all_dialects):
-    ctx = MLContext()
+    linalg_ctx = MLContext()
 
     for dialect_name, dialect_factory in get_all_dialects().items():
-        ctx.register_dialect(dialect_name, dialect_factory)
-    return ctx, dialect_factory, dialect_name
+        linalg_ctx.register_dialect(dialect_name, dialect_factory)
+    return dialect_factory, dialect_name, linalg_ctx
 
 
 @app.cell
@@ -278,9 +278,10 @@ def _(
     convert_linalg_to_loops,
     convert_memref_to_riscv,
     convert_scf_to_riscv_scf,
+    linalg_ctx,
     linalg_module,
-    pipeline_accordion,
     reconcile_unrealized_casts,
+    xmo,
 ):
     lower_to_riscv = PipelinePass(
         [
@@ -293,12 +294,12 @@ def _(
         ]
     )
 
-    riscv_module, riscv_accordion = pipeline_accordion(
-        tuple(("", p) for p in lower_to_riscv.passes), linalg_module
+    riscv_ctx, riscv_module, riscv_html = xmo.pipeline_html(
+        linalg_ctx, linalg_module, tuple(("", p) for p in lower_to_riscv.passes)
     )
 
-    riscv_accordion
-    return lower_to_riscv, riscv_accordion, riscv_module
+    riscv_html
+    return lower_to_riscv, riscv_ctx, riscv_html, riscv_module
 
 
 @app.cell
@@ -318,8 +319,9 @@ def _(
     CanonicalizePass,
     PipelinePass,
     RISCVRegisterAllocation,
-    pipeline_accordion,
+    riscv_ctx,
     riscv_module,
+    xmo,
 ):
     allocate_registers = PipelinePass(
         [
@@ -328,21 +330,22 @@ def _(
         ]
     )
 
-    regalloc_module, regalloc_accordion = pipeline_accordion(
-        tuple(("", p) for p in allocate_registers.passes), riscv_module
+    regalloc_ctx, regalloc_module, regalloc_html = xmo.pipeline_html(
+        riscv_ctx, riscv_module, tuple(("", p) for p in allocate_registers.passes),
     )
 
-    regalloc_accordion
-    return allocate_registers, regalloc_accordion, regalloc_module
+    regalloc_html
+    return allocate_registers, regalloc_ctx, regalloc_html, regalloc_module
 
 
 @app.cell
 def _(
     CanonicalizePass,
     ConvertRiscvScfToRiscvCfPass,
     PipelinePass,
-    pipeline_accordion,
+    regalloc_ctx,
     regalloc_module,
+    xmo,
 ):
     lower_to_asm = PipelinePass(
         [
@@ -351,12 +354,12 @@ def _(
         ]
     )
 
-    riscv_asm_module, assembly_accordion = pipeline_accordion(
-        (("", lower_to_asm),), regalloc_module
+    riscv_asm_ctx, riscv_asm_module, assembly_html = xmo.pipeline_html(
+        regalloc_ctx, regalloc_module, (("", lower_to_asm),),
     )
 
-    assembly_accordion
-    return assembly_accordion, lower_to_asm, riscv_asm_module
+    assembly_html
+    return assembly_html, lower_to_asm, riscv_asm_ctx, riscv_asm_module
 
 
 @app.cell
@@ -396,8 +399,9 @@ def _(
     arith_add_fastmath,
     convert_linalg_to_memref_stream,
     convert_riscv_scf_for_to_frep,
+    linalg_ctx,
     linalg_module,
-    pipeline_accordion,
+    xmo,
 ):
     from xdsl.transforms.test_lower_linalg_to_snitch import LOWER_MEMREF_STREAM_TO_SNITCH_STREAM_PASSES, OPTIMISE_MEMREF_STREAM_PASSES
 
@@ -411,16 +415,17 @@ def _(
         ]
     )
 
-    snitch_stream_module, snitch_stream_accordion = pipeline_accordion(
-        tuple(("", p) for p in convert_linalg_to_snitch.passes), linalg_module
+    snitch_stream_ctx, snitch_stream_module, snitch_stream_html = xmo.pipeline_html(
+        linalg_ctx, linalg_module, tuple(("", p) for p in convert_linalg_to_snitch.passes),
     )
 
-    snitch_stream_accordion
+    snitch_stream_html
     return (
         LOWER_MEMREF_STREAM_TO_SNITCH_STREAM_PASSES,
         OPTIMISE_MEMREF_STREAM_PASSES,
         convert_linalg_to_snitch,
-        snitch_stream_accordion,
+        snitch_stream_ctx,
+        snitch_stream_html,
         snitch_stream_module,
     )
 
@@ -432,17 +437,18 @@ def _(mo):
 
 
 @app.cell
-def _(pipeline_accordion, snitch_stream_module):
+def _(snitch_stream_ctx, snitch_stream_module, xmo):
     from xdsl.transforms.test_lower_linalg_to_snitch import LOWER_SNITCH_STREAM_TO_ASM_PASSES
 
-    snitch_asm_module, snitch_asm_accordion = pipeline_accordion(
-        tuple(("", p) for p in LOWER_SNITCH_STREAM_TO_ASM_PASSES), snitch_stream_module
+    snitch_asm_ctx, snitch_asm_module, snitch_asm_html = xmo.pipeline_html(
+        snitch_stream_ctx, snitch_stream_module, tuple(("", p) for p in LOWER_SNITCH_STREAM_TO_ASM_PASSES)
     )
 
-    snitch_asm_accordion
+    snitch_asm_html
     return (
         LOWER_SNITCH_STREAM_TO_ASM_PASSES,
-        snitch_asm_accordion,
+        snitch_asm_ctx,
+        snitch_asm_html,
         snitch_asm_module,
     )
 
@@ -489,7 +495,7 @@ def _(mo):
 
 
 @app.cell
-def _(TypedPtr, a_shape, b_shape, c_shape, ctx, mo, riscv_module):
+def _(TypedPtr, a_shape, b_shape, c_shape, mo, riscv_ctx, riscv_module):
     from math import prod
 
     from xdsl.interpreter import Interpreter, OpCounter
@@ -507,7 +513,7 @@ def _(TypedPtr, a_shape, b_shape, c_shape, ctx, mo, riscv_module):
     riscv_op_counter = OpCounter()
     riscv_interpreter = Interpreter(riscv_module, listeners=(riscv_op_counter,))
 
-    register_implementations(riscv_interpreter, ctx, include_wgpu=False, include_onnx=False)
+    register_implementations(riscv_interpreter, riscv_ctx, include_wgpu=False, include_onnx=False)
 
     riscv_interpreter.call_op("matmul", (a_shaped.data_ptr.raw, b_shaped.data_ptr.raw, riscv_c_shaped.data_ptr.raw))
 
@@ -547,10 +553,10 @@ def _(
     b_shaped,
     c_len,
     c_shape,
-    ctx,
     mo,
     register_implementations,
     riscv_c_shaped,
+    snitch_stream_ctx,
     snitch_stream_module,
 ):
     snitch_op_counter = OpCounter()
@@ -560,7 +566,7 @@ def _(
 
     snitch_c_shaped = ShapedArray(TypedPtr.new_float64([0.0] * c_len), c_shape)
 
-    register_implementations(snitch_interpreter, ctx, include_wgpu=False, include_onnx=False)
+    register_implementations(snitch_interpreter, snitch_stream_ctx, include_wgpu=False, include_onnx=False)
 
     snitch_interpreter.call_op(
         "matmul",
@@ -671,46 +677,5 @@ def format_row(key: str, *values: str):
     )
 
 
-@app.cell
-def _():
-    from collections import Counter
-    return (Counter,)
-
-
-@app.cell
-def _(Counter, ModuleOp, ModulePass, PipelinePass, ctx, mo, xmo):
-    def spec_str(p: ModulePass) -> str:
-        if isinstance(p, PipelinePass):
-            return ",".join(str(c.pipeline_pass_spec()) for c in p.passes)
-        else:
-            return str(p.pipeline_pass_spec())
-
-    def pipeline_accordion(
-        passes: tuple[tuple[mo.Html, ModulePass], ...], module: ModuleOp
-    ) -> tuple[ModuleOp, mo.Html]:
-        res = module.clone()
-        d = []
-        total_key_count = Counter(spec_str(p) for _, p in passes)
-        d_key_count = Counter()
-        for text, p in passes:
-            p.apply(ctx, res)
-            spec = spec_str(p)
-            d_key_count[spec] += 1
-            if total_key_count[spec] != 1:
-                header = f"{spec} ({d_key_count[spec]})"
-            else:
-                header = spec
-            html_res = xmo.module_html(res)
-            d.append(mo.vstack(
-                (
-                    header,
-                    text,
-                    html_res,
-                )
-            ))
-        return (res, mo.carousel(d))
-    return pipeline_accordion, spec_str
-
-
 if __name__ == "__main__":
     app.run()

docs/marimo/onnx/onnx_demo.py

@@ -170,11 +170,14 @@ def _(
     ConvertOnnxToLinalgPass,
     EmptyTensorToAllocTensorPass,
     MLIROptPass,
+    ctx,
     init_module,
     mo,
-    pipeline_accordion,
+    xmo,
 ):
-    bufferized_module, linalg_accordion = pipeline_accordion(
+    bufferized_ctx, bufferized_module, linalg_html = xmo.pipeline_html(
+        ctx,
+        init_module,
         (
             (
                 mo.md(
@@ -211,12 +214,11 @@ def _(
                     ]
                 )
             )
-        ),
-        init_module
+        )
     )
 
-    linalg_accordion
-    return bufferized_module, linalg_accordion
+    linalg_html
+    return bufferized_module, linalg_html
 
 
 @app.cell
@@ -265,42 +267,5 @@ def _():
     return (xmo,)
 
 
-@app.cell(hide_code=True)
-def _():
-    from collections import Counter
-    return (Counter,)
-
-
-@app.cell(hide_code=True)
-def _(Counter, ModuleOp, ModulePass, PipelinePass, ctx, mo, xmo):
-    def spec_str(p: ModulePass) -> str:
-        if isinstance(p, PipelinePass):
-            return ",".join(str(c.pipeline_pass_spec()) for c in p.passes)
-        else:
-            return str(p.pipeline_pass_spec())
-
-    def pipeline_accordion(passes: tuple[tuple[mo.Html, ModulePass], ...], module: ModuleOp) -> tuple[ModuleOp, mo.Html]:
-        res = module.clone()
-        d = {}
-        total_key_count = Counter(spec_str(p) for _, p in passes)
-        d_key_count = Counter()
-        for text, p in passes:
-            p.apply(ctx, res)
-            spec = spec_str(p)
-            d_key_count[spec] += 1
-            if total_key_count[spec] != 1:
-                header = f"{spec} ({d_key_count[spec]})"
-            else:
-                header = spec
-            html_res = xmo.module_html(res)
-            d[header] = mo.vstack((
-                text,
-                # mo.plain_text(f"Pass: {p.pipeline_pass_spec()}"),
-                html_res
-            ))
-        return (res, mo.accordion(d))
-    return pipeline_accordion, spec_str
-
-
 if __name__ == "__main__":
     app.run()

xdsl/utils/marimo.py

@@ -1,6 +1,11 @@
+from collections import Counter
+from collections.abc import Sequence
+
 import marimo as mo
 
+from xdsl.context import MLContext
 from xdsl.dialects.builtin import ModuleOp
+from xdsl.passes import ModulePass, PipelinePass
 
 
 def asm_html(asm: str) -> mo.Html:
@@ -15,3 +20,53 @@ def module_html(module: ModuleOp) -> mo.Html:
     Returns a Marimo-optimised representation of the module passed in.
     """
     return mo.ui.code_editor(str(module), language="javascript", disabled=True)
+
+
+def _spec_str(p: ModulePass) -> str:
+    """
+    A string representation of the pass passed in, to display to the user.
+    """
+    if isinstance(p, PipelinePass):
+        return ",".join(str(c.pipeline_pass_spec()) for c in p.passes)
+    else:
+        return str(p.pipeline_pass_spec())
+
+
+def pipeline_html(
+    ctx: MLContext, module: ModuleOp, passes: Sequence[tuple[mo.Html, ModulePass]]
+) -> tuple[MLContext, ModuleOp, mo.Html]:
+    """
+    Returns a tuple of the resulting context and module after applying the passes, and the
+    Marimo-optimised representation of the modules throughout compilation.
+
+    The input is the input MLContext, a sequence of tuples of
+    (pass description, module pass).
+
+    Marimo's reactive mechanism relies on a graph of values defined in one cell and used
+    in another, and cannot detect mutation by reference, hence the new values instead of
+    the usual mutation.
+    """
+    res = module.clone()
+    ctx = ctx.clone()
+    d: list[mo.Html] = []
+    total_key_count = Counter(_spec_str(p) for _, p in passes)
+    d_key_count = Counter[str]()
+    for text, p in passes:
+        p.apply(ctx, res)
+        spec = _spec_str(p)
+        d_key_count[spec] += 1
+        if total_key_count[spec] != 1:
+            header = f"{spec} ({d_key_count[spec]})"
+        else:
+            header = spec
+        html_res = module_html(res)
+        d.append(
+            mo.vstack(
+                (
+                    header,
+                    text,
+                    html_res,
+                )
+            )
+        )
+    return (ctx, res, mo.carousel(d))