add set_ir function to change each worker's ir

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)

python/mneme/futures.py

@@ -41,19 +41,35 @@ 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
         ----------
         job_id : int
             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

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']}"

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 = []