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 <[email protected]> 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) <[email protected]> --- 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}") _______________________________________________ llvm-branch-commits mailing list [email protected] https://lists.llvm.org/cgi-bin/mailman/listinfo/llvm-branch-commits
