diff --git a/python/mneme/async_executor.py b/python/mneme/async_executor.py index 19c94174..24052265 100644 --- a/python/mneme/async_executor.py +++ b/python/mneme/async_executor.py @@ -5,9 +5,10 @@ from multiprocessing import Event as ProcessEvent from multiprocessing import Process from multiprocessing import Queue as ProcessQueue +from pathlib import Path from queue import Queue as ThreadQueue from threading import Event as ThreadEvent -from typing import Callable, Dict, Optional +from typing import Callable, Dict, Optional, Union from mneme.futures import EvalFuture from mneme.mneme_logging import logger @@ -193,6 +194,7 @@ def _spawn_process(self): ) self._process.start() self._action = self.StateMachine.SUBMIT + self._ir_revision = 0 def _startup_failure_error(self): return ( @@ -275,6 +277,10 @@ def _submit(self): return self.current = future + if future.ir_revision != self._ir_revision: + self._ipc_write_q.put({"payload": "set_ir", "data": future.ir_data}) + self._ir_revision = future.ir_revision + msg = { "payload": "process", "data": future.config.to_dict(), @@ -427,6 +433,8 @@ def __init__( self._futures: Dict[int, EvalFuture] = {} self._next_id = 0 self._lock = threading.Lock() + self._ir_revision = 0 + self._ir_data = None self.iterations = iterations self.warmup = warmup self.max_startup_failures = max_startup_failures @@ -467,6 +475,23 @@ def _fail_pending_futures(self, error: str): return future.set_error(error) + def set_ir(self, ir: Union[str, Path]): + """Use this LLVM IR for evaluations submitted after this call. + + Parameters + ---------- + ir : str | Path + Path to the LLVM IR file (.bc or .ll) or the IR as a string. + """ + if isinstance(ir, Path): + ir_data = str(ir.absolute()) + else: + ir_data = ir + + with self._lock: + self._ir_revision += 1 + self._ir_data = ir_data + # ------------------------------------------------------------------ # Submit new job (non-blocking) # ------------------------------------------------------------------ @@ -492,7 +517,7 @@ def submit(self, config: ExperimentConfiguration) -> EvalFuture: job_id = self._next_id self._next_id += 1 logger.debug(f"[{self.__class__.__name__}] Submitting job {job_id}") - future = EvalFuture(job_id, config) + future = EvalFuture(job_id, config, self._ir_revision, self._ir_data) self._futures[job_id] = future if self._broken_error is not None: future.set_error(self._broken_error) diff --git a/python/mneme/futures.py b/python/mneme/futures.py index 44cbd723..669849c4 100644 --- a/python/mneme/futures.py +++ b/python/mneme/futures.py @@ -41,9 +41,19 @@ class EvalFuture: Stable identifier assigned by the submitting executor. config : ExperimentConfiguration Configuration associated with this evaluation request. + ir_revision : int + IR revision captured when this evaluation was submitted. + ir_data : str | None + Replacement IR associated with ``ir_revision``. """ - def __init__(self, job_id: int, config: ExperimentConfiguration): + def __init__( + self, + job_id: int, + config: ExperimentConfiguration, + ir_revision: int = 0, + ir_data: Optional[str] = None, + ): """ Parameters ---------- @@ -51,9 +61,15 @@ def __init__(self, job_id: int, config: ExperimentConfiguration): Unique identifier for this evaluation. config : ExperimentConfiguration Configuration to be evaluated. + ir_revision : int, optional + IR revision captured for this evaluation. + ir_data : str | None, optional + Replacement IR associated with the captured revision. """ self.job_id = job_id self.config = config # small dict of input params + self.ir_revision = ir_revision + self.ir_data = ir_data self._cond = threading.Condition() self._done = False diff --git a/python/mneme/replay_executor.py b/python/mneme/replay_executor.py index 1084ccb9..e597d0ae 100644 --- a/python/mneme/replay_executor.py +++ b/python/mneme/replay_executor.py @@ -34,6 +34,7 @@ import os from datetime import datetime, timezone from multiprocessing import Event, Queue +from pathlib import Path from typing import Optional, Tuple from mneme.device import ( @@ -44,7 +45,7 @@ set_device, ) from mneme.llvm.buffer import MemBufferRef -from mneme.llvm.module import ModuleRef +from mneme.llvm.module import ModuleRef, parse_assembly, parse_bitcode from mneme.mneme_logging import logger from mneme.mneme_types import ExperimentConfiguration, ExperimentResult from mneme.page_manager import PageManagerRef @@ -178,6 +179,25 @@ def __exit__(self, exc_type, exc_val, exc_tb): def link_ir(self): return self.records.link_llvm_modules(prune=True, internalize=True) + def set_new_ir(self, ir_path_or_asm: str): + if isinstance(ir_path_or_asm, str) and (ir_path_or_asm.endswith('.ll') or ir_path_or_asm.endswith('.bc')): + ir_path = Path(ir_path_or_asm) + if ir_path.suffix == '.bc': + with open(ir_path, 'rb') as f: + new_ir = parse_bitcode(f.read()) + else: # .ll + with open(ir_path, 'r') as f: + new_ir = parse_assembly(f.read()) + else: + # assume str is text IR + new_ir = parse_assembly(ir_path_or_asm) + + # apply internalization and pruning + jit.internalize(new_ir, self.kernel_descr.kernel_name) + jit.pruneIR(new_ir) + + return new_ir + @cond_time("preprocess_ir_time") def _preprocess_ir( self, @@ -817,6 +837,16 @@ def run( f"Worker {worker.device_id} received terminate request, exiting ..." ) break + elif msg["payload"] == "set_ir": + # update the root_ir for subsequent requests + logger.debug(f"Worker {worker.device_id} received set_ir request") + new_ir_data = msg["data"] + + root_ir = worker.set_new_ir(new_ir_data) + + if root_ir is None: + logger.error(f"Worker {worker.device_id} failed to set new IR") + elif msg["payload"] == "process": logger.debug( f"Worker {worker.device_id} received processing request {msg['exp_id']}" diff --git a/python/tests/test_async_executor.py b/python/tests/test_async_executor.py index d84ae8ab..a9b7a5ff 100644 --- a/python/tests/test_async_executor.py +++ b/python/tests/test_async_executor.py @@ -144,6 +144,23 @@ def test_submit_moves_state_and_sets_current(handle): assert handle.current is future assert handle._action == handle.StateMachine.RECEIVE + assert handle._ipc_write_q.get_nowait()["payload"] == "process" + + +def test_submit_installs_future_ir_revision(handle): + future = EvalFuture( + 3, ExperimentConfiguration(), ir_revision=1, ir_data="replacement ir" + ) + handle.global_q.put(future) + + handle._submit() + + assert handle._ipc_write_q.get_nowait() == { + "payload": "set_ir", + "data": "replacement ir", + } + assert handle._ipc_write_q.get_nowait()["payload"] == "process" + assert handle._ir_revision == 1 def test_try_receive_no_message(handle): @@ -393,6 +410,45 @@ def test_async_executor_evaluate(monkeypatch): assert out.executed is True +def test_async_executor_set_ir_only_affects_subsequent_work(monkeypatch): + monkeypatch.setattr("mneme.async_executor.TuneWorkerHandle", lambda *a, **k: None) + + exe = AsyncReplayExecutor( + record_db="db", + record_id="rid", + iterations=3, + results_db_dir="/tmp", + num_workers=0, + ) + + old_future = exe.submit(ExperimentConfiguration()) + exe.set_ir("define void @kernel() { ret void }") + new_future = exe.submit(ExperimentConfiguration()) + + assert old_future.ir_revision == 0 + assert old_future.ir_data is None + assert new_future.ir_revision == 1 + assert new_future.ir_data == "define void @kernel() { ret void }" + + +def test_async_executor_set_ir_normalizes_path(monkeypatch, tmp_path): + monkeypatch.setattr("mneme.async_executor.TuneWorkerHandle", lambda *a, **k: None) + + exe = AsyncReplayExecutor( + record_db="db", + record_id="rid", + iterations=3, + results_db_dir="/tmp", + num_workers=0, + ) + ir_path = tmp_path / "kernel.ll" + + exe.set_ir(ir_path) + future = exe.submit(ExperimentConfiguration()) + + assert future.ir_data == str(ir_path.absolute()) + + def test_async_executor_shutdown(monkeypatch): joined = [] diff --git a/python/tests/test_replay_executor.py b/python/tests/test_replay_executor.py index 519dd5da..0d09f828 100644 --- a/python/tests/test_replay_executor.py +++ b/python/tests/test_replay_executor.py @@ -360,6 +360,85 @@ def test_link_ir_forwards_prune_and_internalize(monkeypatch): assert rec.link_calls == [(True, True)] +@pytest.mark.parametrize( + ("ir_input", "file_mode", "file_contents", "expected_parser", "expected_data"), + [ + ( + "define void @K() { ret void }", + None, + None, + "asm", + "define void @K() { ret void }", + ), + ( + "kernel.ll", + "w", + "define void @K() { ret void }", + "asm", + "define void @K() { ret void }", + ), + ("kernel.bc", "wb", b"bitcode", "bitcode", b"bitcode"), + ], +) +def test_set_new_ir_parses_and_prunes_replacement_ir( + monkeypatch, + tmp_path, + ir_input, + file_mode, + file_contents, + expected_parser, + expected_data, +): + mod = _reload_with_identity_decorators(monkeypatch) + + kernel = FakeKernelDescr(kernel_name="K") + rec = FakeRecordedExecution(kernel) + + monkeypatch.setattr( + mod.RecordedExecution, "from_json", staticmethod(lambda _: rec), raising=True + ) + monkeypatch.setattr(mod, "set_device", lambda _: None, raising=True) + monkeypatch.setattr(mod, "get_device_arch", lambda: "sm", raising=True) + monkeypatch.setattr(mod, "get_device_count", lambda: 1, raising=True) + + parsed_module = FakeModule("parsed") + calls = [] + + def fake_parse_assembly(data): + calls.append(("asm", data)) + return parsed_module + + def fake_parse_bitcode(data): + calls.append(("bitcode", data)) + return parsed_module + + def fake_internalize(ir, kernel_name): + calls.append(("internalize", ir, kernel_name)) + + def fake_prune(ir): + calls.append(("prune", ir)) + + monkeypatch.setattr(mod, "parse_assembly", fake_parse_assembly, raising=True) + monkeypatch.setattr(mod, "parse_bitcode", fake_parse_bitcode, raising=True) + monkeypatch.setattr(mod.jit, "internalize", fake_internalize, raising=True) + monkeypatch.setattr(mod.jit, "pruneIR", fake_prune, raising=True) + + if file_mode is not None: + ir_path = tmp_path / ir_input + with open(ir_path, file_mode) as f: + f.write(file_contents) + ir_input = str(ir_path) + + ex = mod.BaseExecutor(record_db="x", record_id="rid") + + assert ex.set_new_ir(ir_input) is parsed_module + assert calls == [ + (expected_parser, expected_data), + ("internalize", parsed_module, "K"), + ("prune", parsed_module), + ] + + def test_preprocess_ir_calls_jit_hooks_based_on_config(monkeypatch): mod = _reload_with_identity_decorators(monkeypatch) @@ -633,6 +712,9 @@ def test_tuneworker_run_process_and_terminate(monkeypatch, tmp_path): monkeypatch.setattr(mod.os, "open", lambda *a, **k: 999, raising=True) monkeypatch.setattr(mod.os, "dup2", lambda *a, **k: None, raising=True) + set_ir_calls = [] + process_ir_names = [] + class FakeWorker: def __init__(self, record_db, record_id, device_id, iterations, warmup): self.record_db = record_db @@ -644,6 +726,10 @@ def __init__(self, record_db, record_id, device_id, iterations, warmup): def link_ir(self): return FakeModule("root_ir") + def set_new_ir(self, ir_data): + set_ir_calls.append(ir_data) + return FakeModule("replacement_ir") + def __enter__(self): return self @@ -651,6 +737,7 @@ def __exit__(self, exc_type, exc, tb): return False def process_payload(self, ir_module, config): + process_ir_names.append(ir_module.name) res = mod.ExperimentResult(executed=True, verified=True) return res, FakeModule("ir_out") @@ -671,6 +758,7 @@ def set(self): state = FakeEvent() + req_q.put({"payload": "set_ir", "data": "replacement asm"}) req_q.put( { "payload": "process", @@ -693,8 +781,74 @@ def set(self): ) assert state.set_called == 1 + assert set_ir_calls == ["replacement asm"] + assert process_ir_names == ["replacement_ir_clone"] msg = resp_q.get_nowait() assert msg["payload"] == "result" assert msg["exp_id"] == 7 assert isinstance(msg["data"], dict) assert msg["llvm_ir"] == "" + + +def test_tuneworker_run_logs_failed_set_ir(monkeypatch, tmp_path): + mod = _reload_with_identity_decorators(monkeypatch) + + real_run = mod.TuneWorker.run + + monkeypatch.setattr(mod.os, "open", lambda *a, **k: 999, raising=True) + monkeypatch.setattr(mod.os, "dup2", lambda *a, **k: None, raising=True) + + errors = [] + monkeypatch.setattr( + mod.logger, "error", lambda msg: errors.append(msg), raising=True + ) + + class FakeWorker: + def __init__(self, record_db, record_id, device_id, iterations, warmup): + self.device_id = device_id + + def link_ir(self): + return FakeModule("root_ir") + + def set_new_ir(self, ir_data): + return None + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc, tb): + return False + + monkeypatch.setattr( + mod, "TuneWorker", lambda *a, **k: FakeWorker(*a, **k), raising=True + ) + + req_q = queue.Queue() + resp_q = queue.Queue() + + class FakeEvent: + def __init__(self): + self.set_called = 0 + + def set(self): + self.set_called += 1 + + state = FakeEvent() + + req_q.put({"payload": "set_ir", "data": "bad asm"}) + req_q.put({"payload": "terminate"}) + + real_run( + request_q=req_q, + response_q=resp_q, + record_db="db.json", + record_id="rid", + device_id=3, + iterations=3, + results_db_dir=str(tmp_path), + state=state, + ) + + assert state.set_called == 1 + assert errors == ["Worker 3 failed to set new IR"] + assert resp_q.empty()