diff --git a/docs/tutorials/descript++_notebook.py b/docs/tutorials/descript++_notebook.py
new file mode 100644
index 000000000..c4e18647a
--- /dev/null
+++ b/docs/tutorials/descript++_notebook.py
@@ -0,0 +1,811 @@
+import marimo
+
+__generated_with = "0.20.0"
+app = marimo.App()
+
+
+@app.cell
+def _():
+ import marimo as mo
+
+ return (mo,)
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ # === Utility functions for the tutorial ===
+ from io import StringIO
+ from contextlib import redirect_stderr, redirect_stdout
+ import traceback
+ from typing import Any
+ import queue
+ import time
+ import multiprocessing as mp
+ import exec_utils
+ import matplotlib.pyplot as plt
+
+ def get_backend_import(backend_name: str) -> str:
+ """Return the import statement for the selected backend."""
+ if backend_name == "MLIR":
+ return "from xtc.backends.mlir import Backend"
+ else:
+ return "from xtc.backends.tvm import Backend"
+
+ def get_print_opts_str(output_option: str) -> str:
+ """Return the compiler print options string for the selected output."""
+ opts_map = {
+ "Source IR": "print_source_ir=True",
+ "Transformed IR": "print_transformed_ir=True",
+ "Lowered IR": "print_lowered_ir=True",
+ "Assembly": "print_assembly=True",
+ }
+ return opts_map.get(output_option, "")
+
+ def get_print_opts_dict(output_option: str) -> dict:
+ """Return the compiler print options as a dictionary."""
+ opts_map = {
+ "Source IR": {"print_source_ir": True},
+ "Transformed IR": {"print_transformed_ir": True},
+ "Lowered IR": {"print_lowered_ir": True},
+ "Assembly": {"print_assembly": True},
+ }
+ return opts_map.get(output_option, {})
+
+ def get_backend_class(backend_name: str):
+ """Return the Backend class for the selected backend."""
+ if backend_name == "MLIR":
+ from xtc.backends.mlir import Backend
+ else:
+ from xtc.backends.tvm import Backend
+ return Backend
+
+ def get_output_options(backend_name: str) -> list:
+ """Return available output options based on backend (TVM doesn't support Lowered IR)."""
+ if backend_name == "TVM":
+ return ["Source IR", "Transformed IR", "Assembly"]
+ else:
+ return ["Source IR", "Transformed IR", "Lowered IR", "Assembly"]
+
+ def create_backend_radio(label: str = "Backend:"):
+ """Create a radio button for backend selection."""
+ return mo.ui.radio(options=["MLIR", "TVM"], value="TVM", label=label)
+
+ def create_output_radio(backend_name: str, label: str = "Output options:"):
+ """Create a radio button for output options based on backend."""
+ return mo.ui.radio(
+ options=get_output_options(backend_name),
+ value="Assembly",
+ label=label
+ )
+
+ def execute_editor_code(editor_value: str, display_results_fn=None, initial_namespace=None):
+ """
+ Execute editor code with stdout/stderr capture.
+ Returns (success, output, captured_data).
+ - If display_results_fn is provided, it's injected as 'display_results' in the namespace.
+ - If initial_namespace is provided, those values are injected before execution.
+ - captured_data contains any data captured via display_results.
+ """
+ captured = {"perf": 0.0}
+
+ def _display_results(perf):
+ captured["perf"] = perf
+
+ namespace = dict(initial_namespace) if initial_namespace else {}
+ if display_results_fn is not None:
+ namespace["display_results"] = _display_results
+
+ code_stderr = StringIO()
+ code_stdout = StringIO()
+
+ try:
+ with redirect_stderr(code_stderr), redirect_stdout(code_stdout):
+ exec(editor_value, namespace)
+ output = code_stderr.getvalue() + code_stdout.getvalue()
+ return True, output, captured
+ except Exception:
+ return False, traceback.format_exc(), captured
+
+ def render_editor_output(success: bool, output: str, captured: dict):
+ """Render the output of an editor execution as marimo elements."""
+ if not success:
+ return mo.md(f"**Code error:**\n```\n{output}\n```")
+
+ perf_display = mo.md(f"**Performance:** {captured['perf']:.2f}% of peak")
+ code_content = mo.md(f"```asm\n{output}\n```") if output else mo.md("*No IR output.*")
+ code_accordion = mo.accordion({"Generated Code": code_content})
+ return mo.vstack([perf_display, code_accordion])
+
+ def run_exploration(generator, get_info=None):
+ """
+ Run exploration with progress bar. Returns sorted results.
+ Generator should yield (index, total, sample_or_name, perf) tuples.
+ """
+ results = []
+ best_sample = None
+ best_perf = 0.0
+ captured_info = {}
+
+ first_result = next(generator, None)
+ if first_result is None:
+ return [], {}
+
+ idx, total, sample, perf = first_result
+ results.append({"sample": sample, "perf": perf})
+ if perf > best_perf:
+ best_perf = perf
+ best_sample = sample
+
+ with mo.status.progress_bar(total=total, title="Exploring schedules...", remove_on_exit=False) as progress:
+ progress.update(
+ title=f"Exploring schedules... (best: {best_perf:.1f}%)",
+ subtitle=f"Sample {idx + 1}/{total}: {sample} -> {perf:.1f}%"
+ )
+
+ for idx, total, sample, perf in generator:
+ results.append({"sample": sample, "perf": perf})
+
+ if perf > best_perf:
+ best_perf = perf
+ best_sample = sample
+
+ progress.update(
+ title=f"Exploring schedules... (best: {best_perf:.1f}%)",
+ subtitle=f"Sample {idx + 1}/{total}: {sample} -> {perf:.1f}%"
+ )
+
+ if get_info:
+ captured_info.update(get_info())
+
+ sorted_results = sorted(results, key=lambda x: x["perf"], reverse=True)
+ return sorted_results, captured_info, results
+
+ def start_streaming_execution(
+ *,
+ code: str,
+ out: Any,
+ throttle_s: float = 0.5,
+ ) -> Any:
+ """
+ Start a marimo Thread that launches a subprocess and streams its output.
+ Cancellation: if the spawning cell is invalidated, thread.should_exit becomes True,
+ and we terminate the subprocess.
+ """
+ def target():
+ import marimo as mo
+ thread = mo.current_thread()
+
+ ctx = mp.get_context("spawn")
+ out_q: "mp.Queue" = ctx.Queue()
+ p = ctx.Process(
+ target=exec_utils._child_exec,
+ args=(code, out_q),
+ daemon=True
+ )
+ p.start()
+
+ buf: list[str] = []
+ last = 0.0
+
+ def render(force: bool = False):
+ nonlocal last
+ now = time.time()
+ if force or (now - last) >= throttle_s:
+ out.replace(
+ mo.md(
+ f"**Output:**\n\n```text\n{''.join(buf)}\n```"
+ )
+ )
+ last = now
+
+ try:
+ out.replace(mo.md("**Output:**\n\n```text\n\n```"))
+ render(force=True)
+
+ while True:
+ # Cancel requested? (cell invalidated by Cancel click / rerun / interrupt)
+ if thread.should_exit:
+ buf.append("\n[Cancelled]\n")
+ render(force=True)
+ if p.is_alive():
+ p.terminate()
+ p.join(timeout=1)
+ break
+
+ # Process ended?
+ if not p.is_alive():
+ # Drain remaining queue chunks
+ while True:
+ try:
+ kind, payload = out_q.get_nowait()
+ except queue.Empty:
+ break
+ if kind == "chunk":
+ buf.append(payload)
+ render(force=True)
+ break
+
+ # Get output chunk (non-blocking-ish)
+ try:
+ kind, payload = out_q.get(timeout=0.1)
+ except queue.Empty:
+ continue
+
+ if kind == "chunk":
+ buf.append(payload)
+ render()
+ elif kind == "done":
+ # allow loop to observe process exit / drain remaining data
+ continue
+ finally:
+ if p.is_alive():
+ p.terminate()
+ p.join(timeout=1)
+
+ t = mo.Thread(target=target, daemon=True)
+ t.start()
+ return t
+
+ return (
+ create_backend_radio,
+ create_output_radio,
+ execute_editor_code,
+ get_backend_class,
+ get_print_opts_dict,
+ plt,
+ render_editor_output,
+ run_exploration,
+ start_streaming_execution,
+ traceback,
+ )
+
+
+@app.cell
+def _(mo):
+ mo.md(r"""
+ # Descript
+
+ XTC allows you to describe target loop structures for operators using a small DSL called `descript`. Instead of manually specifying each transformation step, you declare the desired final loop structure, and XTC automatically infers the sequence of transformations needed to achieve it.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Define a schedule declaratively
+
+ `descript` can be used to describe a completed loop structure, with annotations marking optimisations.
+
+ For example, the following loop structure:
+ ```
+ for i in ... // parallelized
+ for k in ...
+ for j in ...
+ for i1 in range(8): // unrolled
+ for j1 in range(16): // vectorized
+ ```
+
+ Can be described as:
+ ```python
+ '''
+ i: parallelize
+ k:
+ j:
+ i#8: unroll
+ j#16: vectorize
+ '''
+ ```
+
+ The loop order follows the given structure (outer to inner), `j#16` creates a tile of size 16 on `j`, and the `vectorize` attribute marks that inner loop for vectorization. Similarly, `parallelize` and `unroll` mark their respective loops for parallelization and (full) unrolling.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ **Practice.** The code below lets you define a schedule specification and see the generated assembly. Try replicating the good-enough schedule you discovered in the previous section!
+ """)
+ return
+
+
+@app.cell
+def _(mo):
+
+ descript_spec = mo.ui.code_editor(
+ value='''# Schedule specification
+ schedule_spec = """
+ i: parallelize
+ k:
+ j:
+ i#8: unroll
+ j#16: vectorize
+ """''',
+ language="python",
+ label="")
+ descript_editor = mo.ui.code_editor(
+ value='''import xtc.graphs.xtc.op as O
+ from xtc.graphs.xtc.graph import XTCGraph
+ from xtc.schedules.descript import descript_scheduler
+ from xtc.runtimes.host import HostRuntime
+
+ rt = HostRuntime.get()
+
+ # Problem setup
+ I, J, K, dtype = 16, 32, 512, "float32"
+
+ def matmul_graph(I: int, J: int, K: int, dtype: str) -> XTCGraph:
+ """Create a graph computing C = A @ B."""
+ a = O.tensor((I, K), dtype, name="A")
+ b = O.tensor((K, J), dtype, name="B")
+ with O.graph(name="matmul") as gb:
+ O.matmul(a, b, name="C")
+ return gb.graph
+
+ graph = matmul_graph(I, J, K, dtype)
+ backend = Backend(graph)
+
+ # Compile
+ scheduler = backend.get_scheduler()
+ descript_scheduler(
+ scheduler=scheduler,
+ node_name="C",
+ abstract_dims=["i", "j", "k"],
+ abstract_matrix=["A", "B", "C"],
+ spec=schedule_spec
+ )
+ schedule = scheduler.schedule()
+
+ compiler = backend.get_compiler(dump_file="matmul", shared_lib=True, **print_opts)
+ module = compiler.compile(schedule)
+
+ # Evaluate and display results
+ peak_flops = rt.evaluate_flops(dtype)
+ evaluator = module.get_evaluator()
+ results, _, _ = evaluator.evaluate()
+ perf = (I * J * K) / min(results) / peak_flops * 100
+
+ display_results(perf)
+ ''',
+ language="python",
+ label=""
+ )
+ mo.vstack([descript_spec, descript_editor])
+ return descript_editor, descript_spec
+
+
+@app.cell
+def _(create_backend_radio):
+ descript_backend_radio = create_backend_radio()
+ return (descript_backend_radio,)
+
+
+@app.cell
+def _(create_output_radio, descript_backend_radio, mo):
+ descript_output_radio = create_output_radio(descript_backend_radio.value)
+ mo.hstack([descript_backend_radio, descript_output_radio], justify="start", gap=4)
+ return (descript_output_radio,)
+
+
+@app.cell
+def _(
+ descript_backend_radio,
+ descript_editor,
+ descript_output_radio,
+ descript_spec,
+ execute_editor_code,
+ get_backend_class,
+ get_print_opts_dict,
+ mo,
+ render_editor_output,
+):
+ _namespace = {
+ "Backend": get_backend_class(descript_backend_radio.value),
+ "print_opts": get_print_opts_dict(descript_output_radio.value),
+ }
+ exec(descript_spec.value, _namespace)
+ _success, _output, _captured = execute_editor_code(descript_editor.value, display_results_fn=True, initial_namespace=_namespace)
+ mo.stop(not _success, mo.md(f"**Code error:**\n```\n{_output}\n```"))
+ render_editor_output(_success, _output, _captured)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Experimenting with Multiple Schedules
+
+ Performance engineering is often about exploring different optimization strategies. Different schedules can have dramatically different performance depending on:
+ - **Problem size**: Small matrices may not benefit from parallelization overhead
+ - **Hardware**: Cache sizes, vector width, and core count affect optimal tiling
+ - **Data layout**: Memory access patterns influence cache efficiency
+
+ In this section, we'll write a simple loop to try several schedule configurations and compare their performance.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ `descript` can be used to describe loop structures without specifying every tile size. This allows you to easily write a set of schedules to explore.
+
+ For example:
+ ```python
+ i:
+ j:
+ k:
+ j#16:
+ k#32:
+ i#i1: unroll=i_r
+ j#j1: vectorize=j_v
+ ```
+ leaves the sizes of the inner kernel as variables to explore.
+
+ On top of that, variables can also be used to explore configurations for optimizations, for instance: `i#i1: unroll=i_r` allows exploring different unroll sizes, and `j#j1: vectorize=j_v` allows toggling the vectorization.
+
+ Finally, it is also possible to use additional constraints to limit the search space with expert intuitions. For example, limiting the size of the kernel to fit in the CPU registers: `j1+j1*i1 < 128`
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ **Practice.** The code below defines several schedule configurations using our declarative scheduler. You can try different tile sizes, loop orderings, or optimization combinations. The pre-built search space (`schedule_spec`) may be poorly designed...
+ """)
+ return
+
+
+@app.cell
+def _(mo):
+ explore_schedules_intro = mo.ui.code_editor(
+ value=
+ '''import xtc.graphs.xtc.op as O
+ from xtc.search.strategies import Strategy_Descript as Strategy
+ from xtc.graphs.xtc.graph import XTCGraph
+ from xtc.backends.tvm import Backend as TVM_Backend
+ from xtc.backends.mlir import Backend as MLIR_Backend
+ from xtc.runtimes.host import HostRuntime
+ from io import StringIO
+ from contextlib import redirect_stderr
+
+ rt = HostRuntime.get()
+
+ # Problem setup
+ I, J, K, dtype = 16, 32, 512, "float32"
+
+ def matmul_graph(I: int, J: int, K: int, dtype: str) -> XTCGraph:
+ """Create a graph computing C = A @ B."""
+ a = O.tensor((I, K), dtype, name="A")
+ b = O.tensor((K, J), dtype, name="B")
+ with O.graph(name="matmul") as gb:
+ O.matmul(a, b, name="C")
+ return gb.graph
+
+ graph = matmul_graph(I=I, J=J, K=K, dtype=dtype)
+ peak_flops = rt.evaluate_flops(dtype)
+
+ # Evaluation helpers
+ def apply_schedule(strategy: Strategy, graph, backend_cls, sample):
+ """Apply a declarative schedule specification and compile."""
+ backend = backend_cls(graph)
+ scheduler = backend.get_scheduler()
+ strategy.generate(scheduler, sample)
+ schedule = scheduler.schedule()
+ comp = backend.get_compiler(dump_file="test_mlir", shared_lib=True)
+ code = StringIO()
+ with redirect_stderr(code):
+ module = comp.compile(schedule)
+ return module, code.getvalue()
+
+ def evaluate(module, peak_flops, nfmadds):
+ """Evaluate module performance as percentage of peak."""
+ evaluator = module.get_evaluator()
+ results, _, _ = evaluator.evaluate()
+ result = min(results)
+ time_flops = nfmadds / result
+ perf = time_flops / peak_flops * 100
+ return perf
+
+ def acquire(strategy: Strategy, sample: dict, backend):
+ """Evaluate a single configuration and return its performance."""
+ module, _ = apply_schedule(strategy, graph, backend, sample)
+ return evaluate(module, peak_flops, I*J*K)
+
+ # Run exploration (displays a progress bar and returns sorted results)
+ def get_info():
+ return {"dims": f"{I}x{J}x{K}", "dtype": dtype}''',
+ language="python",
+ label=""
+ )
+
+ explore_schedules_code = mo.ui.code_editor(
+ value=
+ '''def explore():
+ """Generator that yields (index, total, name, perf) for each evaluation."""
+ schedule_spec = """
+ i:
+ j:
+ k:
+ i#i1: unroll=i_u
+ j#j1: vectorize=j_v
+ constraints:
+ - j1+j1*i1<128
+ - i1*j1 > 16
+ """
+ strategy = Strategy(graph, schedule_spec, partial_unrolls=False, partial_tiles=False)
+ backend = MLIR_Backend
+ configurations = list(strategy.sample(1000))
+ total = len(configurations)
+
+ for idx, sample in enumerate(configurations):
+ perf = acquire(strategy=strategy, sample=sample, backend=backend)
+ yield (idx, total, f"{sample}", perf)
+
+ results = run_exploration(explore(), get_info)
+ ''',
+ language="python",
+ label=""
+ )
+ run_explore_button = mo.ui.run_button(label="Run exploration")
+ mo.vstack([explore_schedules_intro, explore_schedules_code, run_explore_button])
+ return explore_schedules_code, explore_schedules_intro, run_explore_button
+
+
+@app.cell
+def _(
+ explore_schedules_code,
+ explore_schedules_intro,
+ mo,
+ plt,
+ run_exploration,
+ run_explore_button,
+ traceback,
+):
+ mo.stop(not run_explore_button.value, mo.md("*Click 'Run exploration' to execute the code.*"))
+
+ # Execute the user's code with run_exploration available
+ _namespace = {"run_exploration": run_exploration}
+ try:
+ exec(explore_schedules_intro.value, _namespace)
+ exec(explore_schedules_code.value, _namespace)
+ except Exception:
+ mo.stop(True, mo.md(f"**Code error:**\n```\n{traceback.format_exc()}\n```"))
+
+ _results, _captured_info, _unsorted_results = _namespace.get("results", ([], {}))
+
+ # Build results summary
+ if not _results:
+ mo.stop(True, mo.md("**Error:** No results returned. Make sure to call `results = run_exploration(explore(), get_info)`"))
+
+ _dims = _captured_info.get("dims", "?")
+ _dtype = _captured_info.get("dtype", "?")
+ _baseline_perf = _results[-1]["perf"] if _results else 1.0
+ _best = _results[0]
+
+ _summary_lines = [
+ f"### Schedule Exploration Results",
+ f"",
+ f"- **Problem:** {_dims} matmul, {_dtype}",
+ f"",
+ f"**Total configurations evaluated:** {len(_results)}",
+ f"",
+ f"#### Top 10 configurations:",
+ f"",
+ f"| Rank | Configuration | Performance | vs Baseline |",
+ f"|------|---------------|-------------|-------------|",
+ ]
+
+ for _rank, _r in enumerate(_results[:10], 1):
+ _speedup = _r["perf"] / _baseline_perf if _baseline_perf > 0 else 0
+ _summary_lines.append(f"| {_rank} | {_r['sample']} | {_r['perf']:.2f}% | {_speedup:.2f}x |")
+
+ _summary_lines.extend([
+ f"",
+ f"#### Best configuration: **{_best['sample']}** ({_best['perf']:.2f}% of peak)",
+ ])
+
+ max_perf = _unsorted_results[0]["perf"]
+ cumul = [max_perf]
+ for d in _unsorted_results[1:]:
+ max_perf = max(max_perf, d["perf"])
+ if d["perf"] > max_perf:
+ raise Exception("HEIN")
+ cumul.append(max_perf)
+
+
+ plt.plot(cumul)
+ plt.xlabel('Number of samples')
+ plt.ylabel('% of peak performance')
+ plt.title('Cumulative best performance')
+ plt.grid(True)
+
+ ax = mo.ui.matplotlib(plt.gca())
+
+ mo.vstack([mo.md("\n".join(_summary_lines)), ax])
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Using descript
+
+ XTC presents two methods to use a `descript` schedule.
+
+ The `descript_scheduler` function allows directly applying a specification to a schedule. Using it also requires specifying the name of the root name, and the axis names used. (See the first example.)
+
+ `Strategy_Descript` allows using a specification with parameters to define a strategy that explores the value space it describes. It works like other XTC strategies, and has two parameters: `partial_unrolls` and `partial_tiles` that, when set to True, allow exploring tile and unroll sizes that do not divide their outer tile.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Transformations
+
+ As shown above, `descript` can be used to explore configurations for optimizations other than loop interchange and tiling. This section will show the supported transformations and the corresponding syntax and exploration parameters.
+
+ These transformations are written on the same line as the tile they apply to. For instance, to unroll a tile of size 16 on the axis `i` by a factor of 4, one would write: ` i#16: unroll=4`.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Tiling
+
+ Tile sizes are given with the syntax `axis#size: arguments`. The size can be fixed with an integer, or marked as a parameter to explore with a string.
+
+ By default, tile sizes divide their outer tile sizes. This can be changed by setting `partial_tile=True` in the strategy's arguments (to do so on every tile), or by adding `partial` to the tile's arguments (to only do so on that tile).
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Splitting
+
+ Splitting is given by two possible syntaxes: `axis[start:end]:` and `axis[:size:]:`.
+ `start`, `end` and `size` can be integers or parameters. `start` and `end` can be ommited on the first/last split of an axis.
+
+ The loops inside of the split need to be tabulated, to indicate where the split ends.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ Using splitting and tiling to write a `descript` specification,
+ ```
+ for j in ...:
+ for i in ...:
+ for k in ...:
+ for i1 in range(8):
+ for j1 in range(8):
+ for i2 in range(4):
+ for i3 in range(8, 17):
+ for j2 in range(8):
+ for i4 in range(3):
+ ```
+
+ Can be described as:
+ ```python
+ j:
+ i:
+ k:
+ i[:8]:
+ j#8:
+ i#4:
+ i[8:16]:
+ j#8:
+ i#3:
+ ```
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Unroll, parallelize, and vectorize
+
+ `unroll` by itself mean full loop unrolling. An unroll factor can also be used as seen above. That unroll factor can also be replaced by a variable, which indicates that the unroll factor is a parameter to explore.
+
+ By default, unroll factors divide their respective tile size. This can be changed by setting `partial_unrolls=True` in the strategy's arguments.
+
+ | Unroll syntax | Description |
+ |-------------------------------------|-------------------------------------------------------------------|
+ | `unroll` | Full unroll on the corresponding tile |
+ | `unroll=8` | Unroll of factor 8 on the corresponding tile |
+ | `unroll=i_unroll` | Unroll factor to be explored (named i_unroll in the constraints) |
+
+ `parallelize` and `vectorize` can only be applied on the outermost (for `parallelize`) or innermost (for `vectorize`) tiles.
+
+ They cannot be up to a factor, but a variable can be used to explore the presence or abscence of the transformation.
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ### Read/Write Buffers
+
+ `descript` also allows specifying read and write buffers, and gives two syntaxes to do so.
+
+ The first syntax follows the XTC API and has two optimizations: `buffer` for write buffers, and `pack` for read buffers.
+
+ `buffer` takes one optional argument: the buffer memory type for the allocation, which can be a string, or `None` to use XTC's default.
+
+ `pack` takes three arguments, in order: the index of the input to pack, the buffer memory type, and wether or not to pad the buffer.
+
+ The second syntax doesn't distiguish between buffer and pack. It is writen like an axis, but using the name of the tensor to buffer instead of an axis: `A: pack`. In the case of packing, `pad` can be added to enable padding. Using this syntax with `descript_scheduler` also requires specifying the same of the matrices (see the example).
+
+ | Buffer syntax | Description |
+ |-------------------------------------|-------------------------------------------------------------------|
+ | `buffer` | Write buffer |
+ | `pack=(0, None, False)` | Read buffer for the first input, with default memory and no padding |
+ | `B: pack pad` | Read buffer for the second input, with padding |
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ## Sandbox
+
+ This place is your playground. It is where you will achieve the greatest challenges,
+ using XTC from scratch ;)
+ """)
+ return
+
+
+@app.cell
+def _(mo):
+ sandbox_editor = mo.ui.code_editor(
+ value=
+ '''\
+ # Implement your challenge here!
+ print("Hello XTC!")''',
+ language="python",
+ label=""
+ )
+ run = mo.ui.run_button(label="Run sandbox")
+ cancel = mo.ui.button(label="Stop execution", kind="danger")
+
+ mo.vstack(
+ [
+ sandbox_editor,
+ mo.hstack([run, cancel]),
+ ]
+ )
+ return cancel, run, sandbox_editor
+
+
+@app.cell
+def _(cancel, mo, run, sandbox_editor, start_streaming_execution):
+ # This output placeholder is what the background thread updates.
+ out = mo.output
+ out
+
+ # If cancel is clicked, this cell reruns; that invalidates the previous run-cell,
+ # making the old background thread's `should_exit` flip to True, and it will terminate.
+ if cancel.value:
+ out.replace(mo.md("Cancelled (if something was running, it will stop)."))
+ else:
+ mo.stop(not run.value, mo.md("*Click 'Run sandbox' to execute the code, and 'Stop execution' to cancel long runs.*"))
+ # Start background execution. It will keep streaming until done or cancelled.
+ start_streaming_execution(code=sandbox_editor.value, out=out, throttle_s=0.05)
+ return
+
+
+if __name__ == "__main__":
+ app.run()
diff --git a/docs/tutorials/hw_counters_introduction_v2.py b/docs/tutorials/hw_counters_introduction_v2.py
new file mode 100644
index 000000000..a5e18636c
--- /dev/null
+++ b/docs/tutorials/hw_counters_introduction_v2.py
@@ -0,0 +1,1085 @@
+import marimo
+
+__generated_with = "0.19.6"
+app = marimo.App(width="full")
+
+with app.setup:
+ import os
+ import sys
+ import marimo as mo
+ from sys import platform
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ # Hardware Performance Counters & Top-down Analysis
+
+ Optimizing code requires understanding exactly how the CPU executes it.
+
+ In this Marimo, we will use **XTC** to compile a Matrix Multiplication and evaluate it using hardware counters and Top-down Microarchitecture Analysis (TMA)
+ """)
+ return
+
+
+@app.cell(hide_code=True)
+def _(mo, platform):
+ _warnings = []
+
+ if platform == "darwin":
+ _warnings.append(mo.callout(mo.md("**MacOS Detected:** Hardware counters are restricted. The Marimo will run, but TMA metrics might be unavailable."), kind="danger"))
+
+ _warnings.append(mo.callout(mo.md(r"""
+ **Hardware Counters Prerequisites:**
+ To allow userspace applications (ring 3) to read CPU PMU counters on Linux, please run the following commands in your terminal:
+ ```bash
+ sudo sysctl kernel.perf_event_paranoid=-1
+ echo 0 | sudo tee /proc/sys/kernel/nmi_watchdog
+ ```
+
+ *Note : Reloading this page may be required*
+ """), kind="warn"))
+
+ prerequisites_ui = mo.vstack(_warnings)
+ prerequisites_ui
+ return prerequisites_ui,
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ mo.md(r"""
+ ---
+
+ ### For a better experience, it's recommended to be familiar with the XTC's `descript` syntax and the general usage of XTC.
+
+ ```bash
+ marimo run docs/tutorials/xtc_101.py
+ marimo run docs/tutorials/descript++_notebook.py
+ ```
+ ---
+ """)
+ return
+
+
+@app.cell
+def _(supported_arch_msg):
+ supported_arch_msg
+ return
+
+
+# µpipe Sankey diagram
+@app.cell(hide_code=True)
+def _(mo):
+ def plot_upipe(l1_res, l2_res=None, l3_res=None):
+ """Generates a hierarchical Sankey diagram for TMA."""
+ if not l1_res or len(l1_res) < 4 or l1_res[0] < 0:
+ return mo.md("⚠️ **TMA Data not available.**")
+
+ try:
+ import plotly.graph_objects as go
+ except ImportError:
+ return mo.callout(mo.md(" **Please install `plotly` (`pip install plotly`) to visualize the Sankey diagram**"), kind="warn")
+
+ try:
+ # Color palettes (L1, L2, L3)
+ c_ret = "#4ade80" # Green
+ c_bad = "#f87171" # Red
+ c_fe = "#60a5fa" # Blue
+ c_be = "#c084fc" # Purple
+
+ desc_l1 = {
+ "Retiring": "Fraction of slots utilized by useful work.",
+ "Bad Spec": "Fraction of slots wasted due to incorrect speculations.",
+ "Frontend": "Fraction of slots where the Frontend undersupplies the Backend.",
+ "Backend": "Fraction of slots where no uops are delivered due to a lack of Backend resources."
+ }
+ desc_l2 = {
+ "Light Ops": "Retiring typical, single-uop instructions.",
+ "Heavy Ops": "Retiring complex, multi-uop or microcoded instructions.",
+ "Clears": "Wasted slots due to pipeline flushes.",
+ "Branch Misp": "Wasted slots due to incorrect branch predictions.",
+ "Fetch BW": "Frontend cannot decode or deliver enough instructions per cycle.",
+ "Fetch Lat": "Frontend is completely starved and delivering nothing.",
+ "Core Bound": "Execution stalled waiting for execution units (ALU/FPU).",
+ "Mem Bound": "Execution stalled waiting for data from the memory (L1/L2/L3/RAM)."
+ }
+ desc_l3 = [
+ "Stalls due to branch resteers.",
+ "Cycles where Divider unit was active.",
+ "Stalled on external memory (DRAM).",
+ "Limited by Decoded Stream Buffer pipeline.",
+ "Stalls due to switching from DSB to MITE.",
+ "Instructions decoded into 2 or more uops.",
+ "Floating-point (FP) operations executed.",
+ "One uop representing multiple contiguous instructions.",
+ "Stalls due to instruction cache misses.",
+ "Stalls due to ITLB misses.",
+ "Stalled without missing the L1D cache.",
+ "Stalled due to L2 cache accesses.",
+ "Stalled due to L3 cache accesses or Core contention.",
+ "Stalls due to Length Changing Prefixes.",
+ "Memory load or store uops retired.",
+ "Uops fetched by the Microcode Sequencer.",
+ "Limited by MITE pipeline (legacy decode).",
+ "Stalls due to switching uop delivery to MS unit.",
+ "Branch instructions that were not fused.",
+ "Remaining light uops not covered by sibling nodes.",
+ "Stalls due to other misprediction cases.",
+ "Machine Clears not related to memory ordering.",
+ "Stalled on accesses to Persistent Memory Modules.",
+ "Limited due to execution ports saturation.",
+ "Issue-pipeline stalled due to serializing operations.",
+ "Stalled due to store memory accesses and RFO requests."
+ ]
+
+ nodes_label = []
+ nodes_color = []
+ nodes_desc = []
+ links_source = []
+ links_target = []
+ links_value = []
+ links_color = []
+
+ def hex_to_rgba(hex_color, alpha=0.4):
+ h = hex_color.lstrip('#')
+ return f"rgba({int(h[0:2], 16)}, {int(h[2:4], 16)}, {int(h[4:6], 16)}, {alpha})"
+
+ def add_node(name, pct, color, desc=""):
+ nodes_label.append(f"{name} ({pct:.1f}%)")
+ nodes_color.append(color)
+ nodes_desc.append(desc)
+ return len(nodes_label) - 1
+
+ def add_link(source, target, value, color):
+ links_source.append(source)
+ links_target.append(target)
+ links_value.append(max(value, 0.1)) # 0.1 ensures 0% metrics remain visible as thin lines
+ links_color.append(hex_to_rgba(color, 0.4))
+
+ # TopdownL1
+ n_ret = add_node("Retiring", l1_res[0], c_ret, desc_l1["Retiring"])
+ n_bad = add_node("Bad Spec", l1_res[1], c_bad, desc_l1["Bad Spec"])
+ n_fe = add_node("Frontend", l1_res[2], c_fe, desc_l1["Frontend"])
+ n_be = add_node("Backend", l1_res[3], c_be, desc_l1["Backend"])
+
+ # TopdownL2, TopdownL3
+ if l2_res and len(l2_res) >= 8 and l2_res[0] >= 0:
+ def process_l2(l2_name, l2_idx, parent_node, color, desc, l3_mappings=None):
+ l2_val = l2_res[l2_idx]
+ n_l2 = add_node(l2_name, l2_val, color, desc)
+ add_link(parent_node, n_l2, l2_val, color)
+
+ if l3_res and l3_mappings and len(l3_res) > 0 and l3_res[0] >= 0:
+ for l3_name, l3_idx in l3_mappings:
+ if l3_idx < len(l3_res):
+ l3_val = l3_res[l3_idx]
+ desc_text = desc_l3[l3_idx] if len(desc_l3) > l3_idx else ""
+ n_l3 = add_node(l3_name, l3_val, color, desc_text)
+ add_link(n_l2, n_l3, l3_val, color)
+
+ map_light = map_heavy = map_clears = map_misp = map_fbw = map_flat = map_core = map_mem = None
+
+ if l3_res and len(l3_res) >= 26:
+ map_light = [("FP Arith", 6), ("Mem Ops", 14), ("Fused", 7), ("Non-Fused", 18), ("Other Light", 19)]
+ map_heavy = [("Few Uops", 5), ("Microcode", 15)]
+ map_clears = [("Other Nukes", 21)]
+ map_misp = [("Branch Rest.", 0), ("Other Misp.", 20)]
+ map_fbw = [("MITE", 16), ("DSB", 3), ("DSB Swit.", 4), ("LCP", 13)]
+ map_flat = [("ICache Miss", 8), ("ITLB Miss", 9), ("MS Switches", 17)]
+ map_core = [("Divider", 1), ("Ports Util.", 23), ("Serializing", 24)]
+ map_mem = [("L1 Bound", 10), ("L2 Bound", 11), ("L3 Bound", 12), ("DRAM Bound", 2), ("Store Bound", 25), ("PMM Bound", 22)]
+ elif l3_res and len(l3_res) >= 5: # Support for TopdownL3_Mem
+ map_mem = [("L1 Bound", 0), ("L2 Bound", 1), ("L3 Bound", 2), ("DRAM Bound", 3), ("Store Bound", 4)]
+
+ process_l2("Light Ops", 0, n_ret, c_ret, desc_l2["Light Ops"], map_light)
+ process_l2("Heavy Ops", 1, n_ret, c_ret, desc_l2["Heavy Ops"], map_heavy)
+ process_l2("Clears", 2, n_bad, c_bad, desc_l2["Clears"], map_clears)
+ process_l2("Branch Misp", 3, n_bad, c_bad, desc_l2["Branch Misp"], map_misp)
+ process_l2("Fetch BW", 4, n_fe, c_fe, desc_l2["Fetch BW"], map_fbw)
+ process_l2("Fetch Lat", 5, n_fe, c_fe, desc_l2["Fetch Lat"], map_flat)
+ process_l2("Core Bound", 6, n_be, c_be, desc_l2["Core Bound"], map_core)
+ process_l2("Mem Bound", 7, n_be, c_be, desc_l2["Mem Bound"], map_mem)
+
+ fig = go.Figure(data=[go.Sankey(
+ node = dict(
+ pad = 15,
+ thickness = 20,
+ line = dict(color = "black", width = 0.5),
+ label = nodes_label,
+ color = nodes_color,
+ customdata = nodes_desc,
+ hovertemplate = "%{label}
%{customdata}"
+ ),
+ link = dict(
+ source = links_source,
+ target = links_target,
+ value = links_value,
+ color = links_color,
+ hovertemplate = "%{source.label} → %{target.label}
Value: %{value:.1f}%"
+ )
+ )])
+
+ fig.update_layout(
+ title_text="Top-down Microarchitecture Pipeline µpipe",
+ font_size=12,
+ height=650 if (l3_res and len(l3_res) > 5) else 450,
+ margin=dict(t=40, l=20, r=20, b=20)
+ )
+
+ return fig
+
+ except Exception as e:
+ import traceback
+ error_trace = traceback.format_exc()
+ return mo.md(f"**Error generating Sankey diagram:**\n```python\n{error_trace}\n```")
+
+ return plot_upipe,
+
+
+
+
+@app.cell(hide_code=True)
+def _(mo):
+ btn_all = mo.ui.run_button(label="Run All Experiments", kind="success")
+
+ run_all_ui = mo.vstack([
+ mo.md("---"),
+ mo.md("## Evaluate All Steps"),
+ mo.md("Click the button below to compile and evaluate all steps at once."),
+ btn_all
+ ])
+
+ run_all_ui
+ return btn_all, run_all_ui
+
+
+
+@app.cell(hide_code=True)
+def _(mo, platform, plot_upipe):
+ def run_experiment(schedule_spec):
+ import xtc.graphs.xtc.op as O
+ from xtc.backends.mlir import Backend
+ from xtc.schedules.descript import descript_scheduler
+ from xtc.runtimes.host import HostRuntime
+ import subprocess
+
+ # Fixed matrix size for all tests
+ I, J, K, dtype = 512, 512, 512, "float32"
+
+ a = O.tensor((I, K), dtype, name="A")
+ b = O.tensor((K, J), dtype, name="B")
+ with O.graph(name="matmul") as gb:
+ O.matmul(a, b, name="C")
+
+ backend = Backend(gb.graph)
+ scheduler = backend.get_scheduler()
+
+ descript_scheduler(
+ scheduler=scheduler,
+ node_name="C",
+ abstract_dims=["i", "j", "k"],
+ abstract_matrix=["A", "B", "C"],
+ spec=schedule_spec
+ )
+
+ compiler = backend.get_compiler(
+ dump_file="matmul_mlir",
+ shared_lib=True,
+ print_assembly=True
+ )
+ module = compiler.compile(scheduler.schedule())
+
+ rt = HostRuntime.get()
+ peak_flops = rt.evaluate_flops(dtype)
+
+ # Evaluate pure execution time
+ eval_time = module.get_evaluator(validate=False)
+ time_results, _, _ = eval_time.evaluate()
+ exec_time = min(time_results)
+
+ # Calculate percentage of theoretical peak (MACs / Time / Peak_MACs)
+ peak_perf_pct = (I * J * K) / exec_time / peak_flops * 100
+
+ # Visual progress bar for Peak Perf
+ bar_color = "#ef4444" if peak_perf_pct < 10 else ("#f59e0b" if peak_perf_pct < 50 else "#22c55e")
+ fill_width = max(min(peak_perf_pct, 100), 2) # Ensure at least 2% so the text is visible
+
+ perf_ui = mo.Html(f'''
+
+
Computational Efficiency (Theoretical Peak):
+
+
+ {peak_perf_pct:.1f}%
+
+
+