[llvm-branch-commits] [mlir] [mlir] mlir-opt-repl v2: save, bookmark, verify, pass-pipeline syntax, tab completion (PR #203803)
Maksim Levental via llvm-branch-commits
llvm-branch-commits at lists.llvm.org
Sun Jun 14 16:21:59 PDT 2026
https://github.com/makslevental created https://github.com/llvm/llvm-project/pull/203803
New features:
- save: write current IR to a file
- bookmark: name history steps for easy rewind-by-name
- verify: run --verify-diagnostics on current IR
- pass-pipeline syntax: 'run builtin.module(canonicalize,cse)' works
- tab completion: commands and pass names auto-complete with Tab
>From 47a4da17a1269573a78996e3b2bc4d653e36366e Mon Sep 17 00:00:00 2001
From: makslevental <m_levental at apple.com>
Date: Sun, 14 Jun 2026 16:13:58 -0700
Subject: [PATCH] [mlir] mlir-opt-repl v2: save, bookmark, verify,
pass-pipeline syntax, tab completion
New features:
- save: write current IR to a file
- bookmark: name history steps for easy rewind-by-name
- verify: run --verify-diagnostics on current IR
- pass-pipeline syntax: 'run builtin.module(canonicalize,cse)' works
- tab completion: commands and pass names auto-complete with Tab
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply at anthropic.com>
---
mlir/test/mlir-opt-repl/conftest.py | 5 +
.../mlir-opt-repl/test_cli_and_edge_cases.py | 104 ++++++++++-
mlir/test/mlir-opt-repl/test_mcp.py | 134 ++++++++++++++
mlir/test/mlir-opt-repl/test_repl.py | 74 ++++++++
mlir/tools/mlir-opt-repl/README.md | 54 ++++--
.../mlir-opt-repl/src/mlir_opt_repl/engine.py | 104 ++++++++++-
.../mlir-opt-repl/src/mlir_opt_repl/mcp.py | 42 ++++-
.../mlir-opt-repl/src/mlir_opt_repl/repl.py | 173 ++++++++++++++----
8 files changed, 629 insertions(+), 61 deletions(-)
diff --git a/mlir/test/mlir-opt-repl/conftest.py b/mlir/test/mlir-opt-repl/conftest.py
index 39972267255a4..59b4c44119206 100644
--- a/mlir/test/mlir-opt-repl/conftest.py
+++ b/mlir/test/mlir-opt-repl/conftest.py
@@ -14,6 +14,7 @@
import mlir_opt_repl.engine as engine
from mlir_opt_repl.mcp import mcp_main
from mlir_opt_repl.repl import interactive_main
+from mlir_opt_repl import repl as repl_module
SAMPLE_MLIR = "func.func @test(%arg0: f32, %arg1: f32) -> f32 { %0 = arith.addf %arg0, %arg1 : f32 return %0 : f32 }"
INIT_MSG = {"jsonrpc": "2.0", "id": 0, "method": "initialize", "params": {}}
@@ -23,9 +24,13 @@
def reset_engine():
engine.current_ir = None
engine.ir_history = []
+ engine.bookmarks = {}
+ repl_module.bookmarks = {}
yield
engine.current_ir = None
engine.ir_history = []
+ engine.bookmarks = {}
+ repl_module.bookmarks = {}
@contextmanager
diff --git a/mlir/test/mlir-opt-repl/test_cli_and_edge_cases.py b/mlir/test/mlir-opt-repl/test_cli_and_edge_cases.py
index f92fc171fe9a6..c8e891bace3c1 100644
--- a/mlir/test/mlir-opt-repl/test_cli_and_edge_cases.py
+++ b/mlir/test/mlir-opt-repl/test_cli_and_edge_cases.py
@@ -4,14 +4,23 @@
import subprocess
from unittest.mock import patch
+import pytest
from click.testing import CliRunner
import mlir_opt_repl.engine as engine
-from conftest import capture_stdio
+from conftest import (
+ INIT_MSG,
+ capture_stdio,
+ parse_responses,
+ run_mcp,
+ run_repl,
+ tool_call,
+)
+from mlir_opt_repl import repl as repl_module
from mlir_opt_repl.__main__ import cli
from mlir_opt_repl.diff import render_side_by_side, render_unified_diff
from mlir_opt_repl.mcp import mcp_main
-from mlir_opt_repl.repl import interactive_main
+from mlir_opt_repl.repl import _completer, _get_pass_names, interactive_main
class TestClickCLI:
@@ -58,7 +67,6 @@ def test_mlir_opt_not_found(self):
engine.MLIR_OPT = old
def test_check_mlir_opt_exits(self):
- import pytest
old = engine.MLIR_OPT
engine.MLIR_OPT = "/nonexistent/mlir-opt"
@@ -113,6 +121,96 @@ def test_load_stdin_eof(self):
assert "Loaded" in stdout.getvalue()
+class TestCompleter:
+ def test_command_completion(self):
+
+ with patch("mlir_opt_repl.repl.readline") as mock_rl:
+ mock_rl.get_line_buffer.return_value = "lo"
+ result = _completer("lo", 0)
+ assert result == "load "
+
+ def test_pass_completion(self):
+
+ with patch("mlir_opt_repl.repl.readline") as mock_rl:
+ mock_rl.get_line_buffer.return_value = "run convert-arith"
+ result = _completer("convert-arith", 0)
+ assert result is not None
+ assert "convert-arith" in result
+
+ def test_no_match(self):
+
+ with patch("mlir_opt_repl.repl.readline") as mock_rl:
+ mock_rl.get_line_buffer.return_value = "zzz"
+ result = _completer("zzz", 0)
+ assert result is None
+
+ def test_bookmark_completion(self):
+
+ repl_module.bookmarks = {"mymark": 0}
+ with patch("mlir_opt_repl.repl.readline") as mock_rl:
+ mock_rl.get_line_buffer.return_value = "rewind my"
+ result = _completer("my", 0)
+ assert result == "mymark "
+
+ def test_state_out_of_range(self):
+
+ with patch("mlir_opt_repl.repl.readline") as mock_rl:
+ mock_rl.get_line_buffer.return_value = "lo"
+ result = _completer("lo", 99)
+ assert result is None
+
+ def test_other_command_no_completions(self):
+ with patch("mlir_opt_repl.repl.readline") as mock_rl:
+ mock_rl.get_line_buffer.return_value = "save "
+ result = _completer("", 0)
+ assert result is None
+
+
+class TestBookmarkInvalidIndex:
+ def test_rewind_to_invalid_bookmark_mcp(self):
+
+ engine.ir_history = [("initial", "module {}")]
+ engine.current_ir = "module {}"
+ engine.bookmarks = {"stale": 99}
+ result = engine.handle_tool_call("rewind", {"target": "stale"})
+ assert result["isError"] is True
+ assert "invalid index" in result["content"][0]["text"]
+
+ def test_rewind_to_invalid_bookmark_repl(self):
+
+ engine.ir_history = [("initial", "module {}")]
+ engine.current_ir = "module {}"
+ repl_module.bookmarks = {"stale": 99}
+ output = run_repl("rewind stale\nquit\n")
+ assert "invalid index" in output
+
+
+class TestVerifyInvalid:
+ def test_verify_fails_on_bad_ir(self):
+ engine.current_ir = "func.func @f() -> i32 { return }"
+ engine.ir_history = [("initial", engine.current_ir)]
+ result = engine.handle_tool_call("verify", {})
+ assert result["isError"] is True
+ assert "Verification failed" in result["content"][0]["text"]
+
+
+class TestBookmarkNoBookmarksMCP:
+ def test_list_empty(self):
+ output = run_mcp(
+ INIT_MSG,
+ tool_call(
+ 1,
+ "run_pipeline",
+ {"mlir": "func.func @f() { return }", "passes": ["canonicalize"]},
+ ),
+ tool_call(2, "bookmark"),
+ )
+ assert (
+ "(no bookmarks)"
+ in parse_responses(output)[2]["result"]["content"][0]["text"]
+ )
+
+
class TestDiffEdgeCases:
def test_non_pretty_side_by_side(self):
result = render_side_by_side(
diff --git a/mlir/test/mlir-opt-repl/test_mcp.py b/mlir/test/mlir-opt-repl/test_mcp.py
index 1a70aabe8921e..d49c986fc5cf7 100644
--- a/mlir/test/mlir-opt-repl/test_mcp.py
+++ b/mlir/test/mlir-opt-repl/test_mcp.py
@@ -27,6 +27,9 @@ def test_all_tools_present(self):
"reset",
"list_passes",
"rewind",
+ "bookmark",
+ "save",
+ "verify",
"history",
}
@@ -279,3 +282,134 @@ def test_notification_no_response(self):
INIT_MSG, {"jsonrpc": "2.0", "method": "notifications/initialized"}
)
assert len(parse_responses(output)) == 1
+
+
+class TestBookmarkMCP:
+ def test_bookmark_and_rewind(self):
+ output = run_mcp(
+ INIT_MSG,
+ tool_call(
+ 1, "run_pipeline", {"mlir": SAMPLE_MLIR, "passes": ["canonicalize"]}
+ ),
+ tool_call(2, "bookmark", {"name": "pre-lower"}),
+ tool_call(3, "chain_pipeline", {"passes": ["convert-arith-to-llvm"]}),
+ tool_call(4, "rewind", {"target": "pre-lower"}),
+ )
+ responses = parse_responses(output)
+ assert (
+ "Bookmarked [1] as 'pre-lower'"
+ in responses[2]["result"]["content"][0]["text"]
+ )
+ assert (
+ "Rewound to bookmark 'pre-lower'"
+ in responses[4]["result"]["content"][0]["text"]
+ )
+ assert "arith.addf" in responses[4]["result"]["content"][0]["text"]
+
+ def test_bookmark_list(self):
+ output = run_mcp(
+ INIT_MSG,
+ tool_call(
+ 1, "run_pipeline", {"mlir": SAMPLE_MLIR, "passes": ["canonicalize"]}
+ ),
+ tool_call(2, "bookmark", {"name": "snap"}),
+ tool_call(3, "bookmark"),
+ )
+ responses = parse_responses(output)
+ assert "snap" in responses[3]["result"]["content"][0]["text"]
+
+ def test_bookmark_no_history(self):
+ output = run_mcp(INIT_MSG, tool_call(1, "bookmark", {"name": "x"}))
+ assert parse_responses(output)[1]["result"]["isError"] is True
+
+ def test_bookmark_shown_in_history(self):
+ output = run_mcp(
+ INIT_MSG,
+ tool_call(
+ 1, "run_pipeline", {"mlir": SAMPLE_MLIR, "passes": ["canonicalize"]}
+ ),
+ tool_call(2, "bookmark", {"name": "marked"}),
+ tool_call(3, "history"),
+ )
+ assert "marked" in parse_responses(output)[3]["result"]["content"][0]["text"]
+
+
+class TestSaveMCP:
+ def test_save(self, tmp_path):
+ path = str(tmp_path / "out.mlir")
+ output = run_mcp(
+ INIT_MSG,
+ tool_call(
+ 1,
+ "run_pipeline",
+ {"mlir": "func.func @f() { return }", "passes": ["canonicalize"]},
+ ),
+ tool_call(2, "save", {"path": path}),
+ )
+ assert "Saved to" in parse_responses(output)[2]["result"]["content"][0]["text"]
+ assert "func.func @f" in open(path).read()
+
+ def test_save_no_ir(self):
+ output = run_mcp(INIT_MSG, tool_call(1, "save", {"path": "/tmp/x.mlir"}))
+ assert parse_responses(output)[1]["result"]["isError"] is True
+
+ def test_save_no_path(self):
+ output = run_mcp(
+ INIT_MSG,
+ tool_call(
+ 1,
+ "run_pipeline",
+ {"mlir": "func.func @f() { return }", "passes": ["canonicalize"]},
+ ),
+ tool_call(2, "save", {"path": ""}),
+ )
+ assert parse_responses(output)[2]["result"]["isError"] is True
+
+
+class TestVerifyMCP:
+ def test_valid(self):
+ output = run_mcp(
+ INIT_MSG,
+ tool_call(
+ 1,
+ "run_pipeline",
+ {"mlir": "func.func @f() { return }", "passes": ["canonicalize"]},
+ ),
+ tool_call(2, "verify"),
+ )
+ assert (
+ "IR is valid" in parse_responses(output)[2]["result"]["content"][0]["text"]
+ )
+
+ def test_no_ir(self):
+ output = run_mcp(INIT_MSG, tool_call(1, "verify"))
+ assert parse_responses(output)[1]["result"]["isError"] is True
+
+
+class TestPassPipelineMCP:
+ def test_pipeline_string(self):
+ output = run_mcp(
+ INIT_MSG,
+ tool_call(
+ 1,
+ "run_pipeline",
+ {"mlir": SAMPLE_MLIR, "passes": ["builtin.module(canonicalize,cse)"]},
+ ),
+ )
+ text = parse_responses(output)[1]["result"]["content"][0]["text"]
+ assert "arith.addf" in text
+
+ def test_chain_pipeline_string(self):
+ output = run_mcp(
+ INIT_MSG,
+ tool_call(
+ 1, "run_pipeline", {"mlir": SAMPLE_MLIR, "passes": ["canonicalize"]}
+ ),
+ tool_call(
+ 2,
+ "chain_pipeline",
+ {"passes": ["builtin.module(convert-arith-to-llvm)"]},
+ ),
+ )
+ text = parse_responses(output)[2]["result"]["content"][0]["text"]
+ assert "llvm.fadd" in text
diff --git a/mlir/test/mlir-opt-repl/test_repl.py b/mlir/test/mlir-opt-repl/test_repl.py
index be4f214ae0cff..221c2c4d9f869 100644
--- a/mlir/test/mlir-opt-repl/test_repl.py
+++ b/mlir/test/mlir-opt-repl/test_repl.py
@@ -1,5 +1,6 @@
"""Tests for the interactive REPL."""
+import mlir_opt_repl.engine as engine
from conftest import SAMPLE_MLIR, run_repl
LOAD_SAMPLE = f"load -\n{SAMPLE_MLIR}\n\n"
@@ -201,3 +202,76 @@ def test_empty_lines_ignored(self):
def test_exit_alias(self):
output = run_repl("exit\n")
assert "mlir-opt-repl" in output
+
+
+class TestSave:
+ def test_save_to_file(self, tmp_path):
+ f = tmp_path / "out.mlir"
+ output = run_repl(LOAD_SAMPLE + f"save {f}\nquit\n")
+ assert "Saved to" in output
+ assert f.exists()
+ assert "arith.addf" in f.read_text()
+
+ def test_save_no_arg(self):
+ output = run_repl(LOAD_SAMPLE + "save\nquit\n")
+ assert "Usage:" in output
+
+ def test_save_no_ir(self):
+ output = run_repl("save /tmp/x.mlir\nquit\n")
+ assert "No IR to save" in output
+
+
+class TestBookmark:
+ def test_bookmark_and_rewind(self):
+ output = run_repl(
+ LOAD_SAMPLE
+ + "run canonicalize\nbookmark pre-lower\nrun convert-arith-to-llvm\nrewind pre-lower\nquit\n"
+ )
+ assert "Bookmarked [1] as 'pre-lower'" in output
+ assert "Rewound to bookmark 'pre-lower'" in output
+ assert "arith.addf" in output
+
+ def test_bookmark_list(self):
+ output = run_repl(
+ LOAD_SAMPLE + "run canonicalize\nbookmark mymark\nbookmark\nquit\n"
+ )
+ assert "mymark" in output
+ assert "[1]" in output
+
+ def test_bookmark_no_arg_empty(self):
+ output = run_repl("bookmark\nquit\n")
+ assert "(no bookmarks)" in output
+
+ def test_bookmark_shown_in_history(self):
+ output = run_repl(
+ LOAD_SAMPLE + "run canonicalize\nbookmark snap\nhistory\nquit\n"
+ )
+ assert "snap" in output
+
+
+class TestVerify:
+ def test_valid_ir(self):
+ output = run_repl(LOAD_SAMPLE + "run canonicalize\nverify\nquit\n")
+ assert "IR is valid" in output
+
+ def test_no_ir(self):
+ output = run_repl("verify\nquit\n")
+ assert "No IR loaded" in output
+
+ def test_invalid_ir(self):
+ engine.current_ir = "func.func @f() -> i32 { return }"
+ engine.ir_history = [("initial", engine.current_ir)]
+ output = run_repl("verify\nquit\n")
+ assert "error" in output.lower()
+
+
+class TestPassPipeline:
+ def test_pipeline_string(self):
+ output = run_repl(LOAD_SAMPLE + "run builtin.module(canonicalize,cse)\nquit\n")
+ assert "arith.addf" in output
+
+ def test_pipeline_in_history(self):
+ output = run_repl(
+ LOAD_SAMPLE + "run builtin.module(canonicalize)\nhistory\nquit\n"
+ )
+ assert "--pass-pipeline" in output
diff --git a/mlir/tools/mlir-opt-repl/README.md b/mlir/tools/mlir-opt-repl/README.md
index 9eec15b3c3e06..c0cbd074e267b 100644
--- a/mlir/tools/mlir-opt-repl/README.md
+++ b/mlir/tools/mlir-opt-repl/README.md
@@ -24,7 +24,10 @@ export MLIR_OPT=/path/to/llvm-project/build/bin/mlir-opt
mlir-opt-repl
```
-Commands: `load`, `run`, `ir`, `history`, `diff`, `sbs`, `rewind`, `reset`, `passes`, `quit`.
+Commands: `load`, `run`, `ir`, `history`, `diff`, `sbs`, `rewind`, `bookmark`,
+`save`, `verify`, `reset`, `passes`, `help`, `quit`.
+
+Tab completion is supported for commands and pass names.
Diffs are rendered with ANSI colors (red/green for changes, dim for context).
```
@@ -33,7 +36,17 @@ mlir-opt-repl> run canonicalize
mlir-opt-repl> run convert-arith-to-llvm
mlir-opt-repl> diff
mlir-opt-repl> sbs 0 2
-mlir-opt-repl> rewind 1
+mlir-opt-repl> bookmark pre-lower
+mlir-opt-repl> run convert-func-to-llvm
+mlir-opt-repl> rewind pre-lower
+mlir-opt-repl> save output.mlir
+mlir-opt-repl> verify
+```
+
+Pass-pipeline syntax is also supported:
+
+```
+mlir-opt-repl> run builtin.module(canonicalize,cse,convert-arith-to-llvm)
```
## MCP Server (for Claude Code)
@@ -79,11 +92,14 @@ For development without installing:
| Tool | Description |
|------|-------------|
-| `run_pipeline` | Run MLIR source through passes, stores result as current state |
+| `run_pipeline` | Run MLIR source through passes (supports pass-pipeline strings), stores result |
| `chain_pipeline` | Apply additional passes to the current IR state |
| `get_current_ir` | Show the current IR without running passes |
-| `reset` | Clear current IR state and history |
-| `rewind` | Undo the last N pass applications |
+| `reset` | Clear current IR state, history, and bookmarks |
+| `rewind` | Undo the last N steps, or rewind to a named bookmark |
+| `bookmark` | Bookmark the current history step (or list bookmarks) |
+| `save` | Save current IR to a file |
+| `verify` | Verify current IR is valid |
| `history` | Show pass application timeline, with unified or side-by-side diffs |
| `list_passes` | List available mlir-opt passes with optional filter |
@@ -93,17 +109,23 @@ For development without installing:
run_pipeline(mlir="func.func @f(...) ...", passes=["canonicalize"])
→ canonicalized IR (saved as state)
-chain_pipeline(passes=["convert-arith-to-llvm"])
- → arith ops converted to LLVM dialect
+chain_pipeline(passes=["builtin.module(convert-arith-to-llvm)"])
+ → arith ops converted via pass-pipeline syntax
+
+bookmark(name="pre-func-lower")
+ → bookmarks current step
+
+chain_pipeline(passes=["convert-func-to-llvm"])
+ → fully lowered
-history(format="unified")
- → shows unified diff between each step
+rewind(target="pre-func-lower")
+ → back to bookmarked state
-rewind(steps=1)
- → back to post-canonicalize state, try a different path
+save(path="output.mlir")
+ → writes current IR to file
-history(format="side_by_side", width=120)
- → two-column comparison of each step
+verify()
+ → confirms IR is valid
```
## Project Structure
@@ -113,9 +135,9 @@ pyproject.toml
src/mlir_opt_repl/
__init__.py
__main__.py — Click CLI: mlir-opt-repl [mcp|repl]
- engine.py — Core state + logic: run_mlir_opt, list_passes, handle_tool_call
+ engine.py — Core state + logic: run_mlir_opt, handle_tool_call, bookmarks
mcp.py — MCP protocol: TOOLS schema, send/recv, dispatch
- repl.py — Interactive terminal REPL
+ repl.py — Interactive terminal REPL with tab completion
diff.py — Side-by-side and unified diff renderers
render.py — ANSI color constants
```
@@ -127,7 +149,7 @@ Tests live in `mlir/test/mlir-opt-repl/` and use pytest with 100% line coverage:
```bash
PYTHONPATH=mlir/tools/mlir-opt-repl/src \
MLIR_OPT=/path/to/build/bin/mlir-opt \
- python3 -m pytest mlir/test/mlir-opt-repl \
+ python3 -m pytest mlir/test/mlir-opt-repl -n auto \
--cov=mlir_opt_repl \
--cov-config=mlir/tools/mlir-opt-repl/pyproject.toml \
--cov-fail-under=100
diff --git a/mlir/tools/mlir-opt-repl/src/mlir_opt_repl/engine.py b/mlir/tools/mlir-opt-repl/src/mlir_opt_repl/engine.py
index 72f2b1a06f353..e4ce192ec3430 100644
--- a/mlir/tools/mlir-opt-repl/src/mlir_opt_repl/engine.py
+++ b/mlir/tools/mlir-opt-repl/src/mlir_opt_repl/engine.py
@@ -21,6 +21,17 @@ def check_mlir_opt():
current_ir = None
ir_history = []
+bookmarks = {}
+
+
+def _build_pass_args(passes, extra_args=None):
+ if len(passes) == 1 and "(" in passes[0]:
+ args = [f"--pass-pipeline={passes[0]}"]
+ else:
+ args = ["--" + p.lstrip("-") for p in passes]
+ if extra_args:
+ args += ["--" + a.lstrip("-") for a in extra_args]
+ return args
def run_mlir_opt(ir_text, args):
@@ -69,13 +80,12 @@ def list_passes():
def handle_tool_call(name, arguments):
- global current_ir, ir_history
+ global current_ir, ir_history, bookmarks
if name == "run_pipeline":
mlir = arguments["mlir"]
- passes = ["--" + p.lstrip("-") for p in arguments["passes"]]
- extra = ["--" + a.lstrip("-") for a in arguments.get("extra_args", [])]
- output, err = run_mlir_opt(mlir, passes + extra)
+ passes = _build_pass_args(arguments["passes"], arguments.get("extra_args"))
+ output, err = run_mlir_opt(mlir, passes)
if err:
return {
"content": [{"type": "text", "text": f"Error:\n{err}"}],
@@ -97,9 +107,8 @@ def handle_tool_call(name, arguments):
],
"isError": True,
}
- passes = ["--" + p.lstrip("-") for p in arguments["passes"]]
- extra = ["--" + a.lstrip("-") for a in arguments.get("extra_args", [])]
- output, err = run_mlir_opt(current_ir, passes + extra)
+ passes = _build_pass_args(arguments["passes"], arguments.get("extra_args"))
+ output, err = run_mlir_opt(current_ir, passes)
if err:
return {
"content": [
@@ -119,6 +128,7 @@ def handle_tool_call(name, arguments):
elif name == "reset":
current_ir = None
ir_history = []
+ bookmarks = {}
return {"content": [{"type": "text", "text": "IR state cleared."}]}
elif name == "rewind":
@@ -127,7 +137,31 @@ def handle_tool_call(name, arguments):
"content": [{"type": "text", "text": "Error: no history to rewind."}],
"isError": True,
}
+ target = arguments.get("target")
steps = arguments.get("steps", 1)
+ if target and target in bookmarks:
+ idx = bookmarks[target]
+ if idx >= len(ir_history):
+ return {
+ "content": [
+ {
+ "type": "text",
+ "text": f"Error: bookmark '{target}' points to invalid index.",
+ }
+ ],
+ "isError": True,
+ }
+ ir_history = ir_history[: idx + 1]
+ desc, ir = ir_history[-1]
+ current_ir = ir
+ return {
+ "content": [
+ {
+ "type": "text",
+ "text": f"Rewound to bookmark '{target}' ({desc}).\n\n{ir}",
+ }
+ ]
+ }
if steps >= len(ir_history):
desc, ir = ir_history[0]
ir_history = [ir_history[0]]
@@ -152,6 +186,58 @@ def handle_tool_call(name, arguments):
]
}
+ elif name == "bookmark":
+ if not ir_history:
+ return {
+ "content": [{"type": "text", "text": "Error: no history to bookmark."}],
+ "isError": True,
+ }
+ bm_name = arguments.get("name", "")
+ if not bm_name:
+ if not bookmarks:
+ return {"content": [{"type": "text", "text": "(no bookmarks)"}]}
+ lines = []
+ for n, idx in sorted(bookmarks.items(), key=lambda x: x[1]):
+ desc = ir_history[idx][0] if idx < len(ir_history) else "?"
+ lines.append(f" {n} -> [{idx}] {desc}")
+ return {"content": [{"type": "text", "text": "\n".join(lines)}]}
+ idx = len(ir_history) - 1
+ bookmarks[bm_name] = idx
+ return {
+ "content": [{"type": "text", "text": f"Bookmarked [{idx}] as '{bm_name}'"}]
+ }
+
+ elif name == "save":
+ if current_ir is None:
+ return {
+ "content": [{"type": "text", "text": "Error: no IR to save."}],
+ "isError": True,
+ }
+ path = arguments.get("path", "")
+ if not path:
+ return {
+ "content": [{"type": "text", "text": "Error: path is required."}],
+ "isError": True,
+ }
+ with open(path, "w") as f:
+ f.write(current_ir)
+ f.write("\n")
+ return {"content": [{"type": "text", "text": f"Saved to {path}"}]}
+
+ elif name == "verify":
+ if current_ir is None:
+ return {
+ "content": [{"type": "text", "text": "Error: no IR to verify."}],
+ "isError": True,
+ }
+ _, err = run_mlir_opt(current_ir, ["--verify-diagnostics"])
+ if err:
+ return {
+ "content": [{"type": "text", "text": f"Verification failed:\n{err}"}],
+ "isError": True,
+ }
+ return {"content": [{"type": "text", "text": "IR is valid."}]}
+
elif name == "history":
if not ir_history:
return {"content": [{"type": "text", "text": "(no history)"}]}
@@ -160,9 +246,11 @@ def handle_tool_call(name, arguments):
pretty = arguments.get("pretty", False)
width = arguments.get("width")
lines = []
+ bookmark_at = {v: k for k, v in bookmarks.items()}
for i, (desc, ir) in enumerate(ir_history):
marker = " <-- current" if i == len(ir_history) - 1 else ""
- lines.append(f"[{i}] {desc}{marker}")
+ bm = f" [{bookmark_at[i]}]" if i in bookmark_at else ""
+ lines.append(f"[{i}] {desc}{bm}{marker}")
if fmt == "side_by_side" and i > 0:
prev_ir = ir_history[i - 1][1]
sbs = render_side_by_side(
diff --git a/mlir/tools/mlir-opt-repl/src/mlir_opt_repl/mcp.py b/mlir/tools/mlir-opt-repl/src/mlir_opt_repl/mcp.py
index a25e2d0247b54..22e7c89b08c47 100644
--- a/mlir/tools/mlir-opt-repl/src/mlir_opt_repl/mcp.py
+++ b/mlir/tools/mlir-opt-repl/src/mlir_opt_repl/mcp.py
@@ -82,7 +82,7 @@
},
{
"name": "rewind",
- "description": "Rewind the IR state by N steps (default 1). Undoes the last N pass applications.",
+ "description": "Rewind the IR state by N steps or to a named bookmark.",
"inputSchema": {
"type": "object",
"properties": {
@@ -91,10 +91,50 @@
"description": "Number of steps to rewind (default 1)",
"default": 1,
"minimum": 1,
+ },
+ "target": {
+ "type": "string",
+ "description": "Bookmark name to rewind to (overrides steps)",
+ },
+ },
+ },
+ },
+ {
+ "name": "bookmark",
+ "description": "Bookmark the current history step with a name, or list all bookmarks (no name).",
+ "inputSchema": {
+ "type": "object",
+ "properties": {
+ "name": {
+ "type": "string",
+ "description": "Name for the bookmark. Omit to list existing bookmarks.",
+ "default": "",
}
},
},
},
+ {
+ "name": "save",
+ "description": "Save the current IR state to a file.",
+ "inputSchema": {
+ "type": "object",
+ "properties": {
+ "path": {
+ "type": "string",
+ "description": "File path to write the current IR to",
+ }
+ },
+ "required": ["path"],
+ },
+ },
+ {
+ "name": "verify",
+ "description": "Verify that the current IR is valid (runs mlir-opt --verify-diagnostics).",
+ "inputSchema": {
+ "type": "object",
+ "properties": {},
+ },
+ },
{
"name": "history",
"description": "Show the history of pass applications and their IR states.",
diff --git a/mlir/tools/mlir-opt-repl/src/mlir_opt_repl/repl.py b/mlir/tools/mlir-opt-repl/src/mlir_opt_repl/repl.py
index 6d5900b256345..50b712815f22a 100644
--- a/mlir/tools/mlir-opt-repl/src/mlir_opt_repl/repl.py
+++ b/mlir/tools/mlir-opt-repl/src/mlir_opt_repl/repl.py
@@ -1,18 +1,67 @@
+import readline
from textwrap import dedent
from mlir_opt_repl import engine
from mlir_opt_repl.diff import render_side_by_side, render_unified_diff
from mlir_opt_repl.render import BOLD, CYAN, DIM, GREEN, RED, RESET
+COMMANDS = [
+ "load",
+ "run",
+ "ir",
+ "history",
+ "diff",
+ "sbs",
+ "rewind",
+ "reset",
+ "passes",
+ "save",
+ "bookmark",
+ "verify",
+ "help",
+ "quit",
+ "exit",
+]
+
+_pass_names_cache = None
+bookmarks = {}
+
+
+def _get_pass_names():
+ global _pass_names_cache
+ if _pass_names_cache is None:
+ _pass_names_cache = [p["name"] for p in engine.list_passes()]
+ return _pass_names_cache
+
+
+def _completer(text, state):
+ line = readline.get_line_buffer()
+ parts = line.lstrip().split()
+
+ if len(parts) == 0 or (len(parts) == 1 and not line.endswith(" ")):
+ matches = [c + " " for c in COMMANDS if c.startswith(text)]
+ elif parts[0] == "run":
+ pass_names = _get_pass_names()
+ matches = [p + " " for p in pass_names if p.startswith(text)]
+ elif parts[0] == "rewind" and bookmarks:
+ matches = [b + " " for b in bookmarks if b.startswith(text)]
+ else:
+ matches = []
+
+ return matches[state] if state < len(matches) else None
-def interactive_main():
- import readline # noqa: F401 — enables line editing in input()
+def interactive_main():
engine.check_mlir_opt()
+ readline.set_completer(_completer)
+ readline.set_completer_delims(" ")
+ readline.parse_and_bind("tab: complete")
+
print(f"{BOLD}mlir-opt-repl{RESET} (using {engine.MLIR_OPT})")
print(
- "Commands: load <file>, run <passes...>, rewind [N], history, diff, sbs, ir, reset, quit"
+ "Commands: load, run, ir, history, diff, sbs, rewind, reset, "
+ "save, bookmark, verify, passes, help, quit"
)
print()
@@ -59,7 +108,6 @@ def interactive_main():
except FileNotFoundError:
print(f"{RED}File not found: {path}{RESET}")
continue
- # Validate the IR by running mlir-opt with no passes
_, err = engine.run_mlir_opt(mlir_text, [])
if err:
print(f"{RED}Invalid MLIR:{RESET}")
@@ -78,7 +126,11 @@ def interactive_main():
if engine.current_ir is None:
print(f"{RED}No IR loaded. Use 'load <file>' first.{RESET}")
continue
- passes = ["--" + p.lstrip("-") for p in parts[1:]]
+ pipeline_str = " ".join(parts[1:])
+ if "(" in pipeline_str:
+ passes = [f"--pass-pipeline={pipeline_str}"]
+ else:
+ passes = ["--" + p.lstrip("-") for p in parts[1:]]
output, err = engine.run_mlir_opt(engine.current_ir, passes)
if err:
print(f"{RED}{err}{RESET}")
@@ -92,17 +144,28 @@ def interactive_main():
if not engine.ir_history:
print(f"{RED}No history to rewind.{RESET}")
continue
- steps = int(parts[1]) if len(parts) > 1 else 1
- if steps >= len(engine.ir_history):
- desc, ir = engine.ir_history[0]
- engine.ir_history = [engine.ir_history[0]]
- engine.current_ir = ir
- print(f"{GREEN}Rewound to beginning ({desc}).{RESET}")
- else:
- engine.ir_history = engine.ir_history[:-steps]
+ target = parts[1] if len(parts) > 1 else "1"
+ if target in bookmarks:
+ idx = bookmarks[target]
+ if idx >= len(engine.ir_history):
+ print(f"{RED}Bookmark '{target}' points to invalid index.{RESET}")
+ continue
+ engine.ir_history = engine.ir_history[: idx + 1]
desc, ir = engine.ir_history[-1]
engine.current_ir = ir
- print(f"{GREEN}Rewound {steps} step(s). Now at: {desc}{RESET}")
+ print(f"{GREEN}Rewound to bookmark '{target}' ({desc}).{RESET}")
+ else:
+ steps = int(target)
+ if steps >= len(engine.ir_history):
+ desc, ir = engine.ir_history[0]
+ engine.ir_history = [engine.ir_history[0]]
+ engine.current_ir = ir
+ print(f"{GREEN}Rewound to beginning ({desc}).{RESET}")
+ else:
+ engine.ir_history = engine.ir_history[:-steps]
+ desc, ir = engine.ir_history[-1]
+ engine.current_ir = ir
+ print(f"{GREEN}Rewound {steps} step(s). Now at: {desc}{RESET}")
print()
print(engine.current_ir)
@@ -110,13 +173,15 @@ def interactive_main():
if not engine.ir_history:
print(f"{DIM}(no history){RESET}")
continue
+ bookmark_at = {v: k for k, v in bookmarks.items()}
for i, (desc, _) in enumerate(engine.ir_history):
marker = (
f" {GREEN}<-- current{RESET}"
if i == len(engine.ir_history) - 1
else ""
)
- print(f" {BOLD}[{i}]{RESET} {desc}{marker}")
+ bm = f" {CYAN}[{bookmark_at[i]}]{RESET}" if i in bookmark_at else ""
+ print(f" {BOLD}[{i}]{RESET} {desc}{bm}{marker}")
elif cmd == "diff":
if len(engine.ir_history) < 2:
@@ -169,8 +234,50 @@ def interactive_main():
elif cmd == "reset":
engine.current_ir = None
engine.ir_history = []
+ bookmarks.clear()
print(f"{GREEN}IR state cleared.{RESET}")
+ elif cmd == "save":
+ if len(parts) < 2:
+ print(f"{RED}Usage: save <file.mlir>{RESET}")
+ continue
+ if engine.current_ir is None:
+ print(f"{RED}No IR to save.{RESET}")
+ continue
+ path = " ".join(parts[1:])
+ with open(path, "w") as f:
+ f.write(engine.current_ir)
+ f.write("\n")
+ print(f"{GREEN}Saved to {path}{RESET}")
+
+ elif cmd == "bookmark":
+ if len(parts) < 2:
+ if not bookmarks:
+ print(f"{DIM}(no bookmarks){RESET}")
+ else:
+ for name, idx in sorted(bookmarks.items(), key=lambda x: x[1]):
+ desc = (
+ engine.ir_history[idx][0]
+ if idx < len(engine.ir_history)
+ else "?"
+ )
+ print(f" {BOLD}{name}{RESET} -> [{idx}] {desc}")
+ continue
+ name = parts[1]
+ idx = len(engine.ir_history) - 1
+ bookmarks[name] = idx
+ print(f"{GREEN}Bookmarked [{idx}] as '{name}'{RESET}")
+
+ elif cmd == "verify":
+ if engine.current_ir is None:
+ print(f"{RED}No IR loaded.{RESET}")
+ continue
+ _, err = engine.run_mlir_opt(engine.current_ir, ["--verify-diagnostics"])
+ if err:
+ print(f"{RED}{err}{RESET}")
+ else:
+ print(f"{GREEN}IR is valid.{RESET}")
+
elif cmd == "passes":
filt = parts[1] if len(parts) > 1 else ""
all_passes = engine.list_passes()
@@ -187,24 +294,24 @@ def interactive_main():
print(f"{DIM}(no passes matched){RESET}")
elif cmd == "help":
- print(
- dedent(
- f"""\
- {BOLD}Commands:{RESET}
- {CYAN}load <file.mlir>{RESET} Load MLIR from a file
- {CYAN}load -{RESET} Load MLIR from stdin (blank line to finish)
- {CYAN}run <passes...>{RESET} Apply passes to current IR
- {CYAN}ir{RESET} Show current IR
- {CYAN}history{RESET} Show pass application history
- {CYAN}diff [a b]{RESET} Unified diff (last step, or between indices a and b)
- {CYAN}sbs [a b]{RESET} Side-by-side diff (last step, or between indices a and b)
- {CYAN}rewind [N]{RESET} Undo last N steps (default 1)
- {CYAN}reset{RESET} Clear all state
- {CYAN}passes [filter]{RESET} List available passes
- {CYAN}quit{RESET} Exit
- """
- )
- )
+ print(dedent(f"""\
+ {BOLD}Commands:{RESET}
+ {CYAN}load <file.mlir>{RESET} Load MLIR from a file
+ {CYAN}load -{RESET} Load MLIR from stdin (blank line to finish)
+ {CYAN}run <passes...>{RESET} Apply passes to current IR
+ {CYAN}run <pipeline>{RESET} Apply a pass-pipeline string (e.g. builtin.module(...))
+ {CYAN}ir{RESET} Show current IR
+ {CYAN}history{RESET} Show pass application history
+ {CYAN}diff [a b]{RESET} Unified diff (last step, or between indices a and b)
+ {CYAN}sbs [a b]{RESET} Side-by-side diff (last step, or between indices a and b)
+ {CYAN}rewind [N|name]{RESET} Undo last N steps or rewind to a bookmark
+ {CYAN}bookmark [name]{RESET} Bookmark current step (no arg: list bookmarks)
+ {CYAN}save <file>{RESET} Save current IR to a file
+ {CYAN}verify{RESET} Verify current IR is valid
+ {CYAN}reset{RESET} Clear all state
+ {CYAN}passes [filter]{RESET} List available passes (tab-completable)
+ {CYAN}quit{RESET} Exit
+ """))
else:
print(f"{RED}Unknown command: {cmd}. Type 'help' for usage.{RESET}")
More information about the llvm-branch-commits
mailing list