Add set_ir interface to AsyncReplayExecutor

.gitlab-ci.yml

@@ -74,7 +74,7 @@ stages:
     include:
       - local: '.gitlab/custom-jobs-and-variables.yml'
       - project: 'radiuss/radiuss-shared-ci'
-        ref: 'v2025.06.0'
+        ref: 'v2025.12.1'
         file: 'pipelines/${CI_MACHINE}.yml'
       # Add your jobs
       # you can use a local file
@@ -93,7 +93,7 @@ stages:
     include:
       - local: '.gitlab/custom-jobs-and-variables.yml'
       - project: 'radiuss/radiuss-shared-ci'
-        ref: 'v2025.06.0'
+        ref: 'v2025.12.1'
         file: 'pipelines/${CI_MACHINE}.yml'
       - local: '.gitlab/jobs/${CI_MACHINE}-python-cov.yml'
     strategy: depend
@@ -106,7 +106,7 @@ include:
     file: 'id_tokens.yml'
   # [Optional] checks preliminary to running the actual CI test
   - project: 'radiuss/radiuss-shared-ci'
-    ref: 'v2025.06.0'
+    ref: 'v2025.12.1'
     file: 'utilities/preliminary-ignore-draft-pr.yml'
   # pipelines subscribed by the project
   - local: '.gitlab/subscribed-pipelines.yml'

.gitlab/custom-jobs-and-variables.yml

@@ -40,7 +40,7 @@ variables:
 # Tioga
 # Arguments for top level allocation
 # OPTIONAL: "-o per-resource.count=2" allows to get 2 jobs running on each node.
-  TIOGA_SHARED_ALLOC: "--queue=pci --exclusive --time-limit=1h --nodes=1"
+  TIOGA_SHARED_ALLOC: "--queue=pci --exclusive --time-limit=1h --nodes=1 -o per-resource.count=4"
 # Arguments for job level allocation
   TIOGA_JOB_ALLOC: "--nodes=1 --begin-time=+5s"
 # Add variables that should apply to all the jobs on a machine:
@@ -55,7 +55,7 @@ variables:
 # Tuo
 # Arguments for top level allocation
 # OPTIONAL: "-o per-resource.count=2" allows to get 2 jobs running on each node.
-  TUOLUMNE_SHARED_ALLOC: "--queue=pci --exclusive --time-limit=1h --nodes=1"
+  TUOLUMNE_SHARED_ALLOC: "--queue=pci --exclusive --time-limit=1h --nodes=1 -o per-resource.count=4"
 # Arguments for job level allocation
   TUOLUMNE_JOB_ALLOC: "--nodes=1 --begin-time=+5s"
 # Add variables that should apply to all the jobs on a machine:

.gitlab/subscribed-pipelines.yml

@@ -81,7 +81,7 @@ tuolumne-test-python-cov:
     JOB_CMD:
       value: "scripts/gitlab/ci-python-test.sh"
       expand: false
-  needs: [tioga-up-check]
+  needs: [tuolumne-up-check]
   extends: [.test-python-cov]
 
 

README.md

@@ -4,33 +4,27 @@
 
 # <img src="docs/assets/images/MnemeLogoNoText.png" width="128" align="middle" /> Mneme (Μνήμη)
 
-*Named after the Greek goddess of memory, preserves and replays the essence of your application's execution, allowing developers to revisit, analyze, and refine specific moments in code with precision.* 
+*Named after the Greek goddess of memory, preserves and replays the essence of your application's execution, allowing developers to revisit, analyze, and refine specific moments in code with precision.*
 
-## Description
 [Mneme](https://en.wikipedia.org/wiki/Mneme) is a tool allowing recording the execution of a GPU (CUDA/HIP) kernel and replaying that kernel as an independent executable.
 
-Mneme operates in 3 phases. First, during compile time, the user needs to apply a provided LLVM pass to instrument the code. This pass detects the global variables
-and functions on the GPU device, and stores this information with the respective LLVM-IR in the global device memory. The compilation generates a _recordable_ executable. 
+## Documentation
 
-The second phase involves running the _recordable_ executable with a desired input and using `LD_PRELOAD` to enable recording. When recording, before invoking a device kernel, 
-the pre-loaded library stores device memory in persistent storage and associates the memory with the device kernel and an LLVM IR file. At the end of the recorded execution,
-the pre-loaded library generates a database in the form of a collection of `json` files, each containing information regarding the LLVM-IR files and the snapshots of device memory for a single GPU kernel.
+For full usage instructions, tutorials, and API reference, please visit the **[Documentation](https://olympus-hpc.github.io/Mneme/)**.
 
-During the third and last phase, the user can replay the execution of a kernel as a separate independent executable. In addition to executing the kernel, the user can also modify the LLVM IR file and
-auto-tune parameters such as kernel launch-bounds or kernel runtime execution parameters (e.g. Kernel Block and Grid Dimensions).
-
-This documentation contains the user guide and  developers' manual for
-[Mneme](https://github.com/Olympus-HPC/Mneme).
+## Key Features
 
+*   **Record**: Capture GPU kernels from large applications into isolated replayable units.
+*   **Replay**: Execute captured kernels independently without the original application context.
+*   **Tune**: Optimize kernel parameters (block size, grid size) and compiler passes using Python tools like Optuna.
 
 ## Contributions
 
 We welcome all kinds of contributions: new features, bug fixes, documentation edits; it's all great!
 
 To contribute, make a pull request, with `develop` as the destination branch.
 
-
-# Release
+## Release
 
 Mneme is released under Apache License (Version 2.0) with LLVM exceptions. For more details, please see the [LICENSE](./LICENSE).
 
@@ -49,4 +43,3 @@ If you use this software, please cite it as below:
   year={2023}
 }
 ```
-

cmake/MnemeFunctions.cmake

@@ -1,5 +1,7 @@
 function(add_mneme target)
-    add_proteus(${target} FORCE_JIT_ANNOTATE_ALL)
+    # Link proteus as a shared library for preloading.
+    # TODO: Change to static linking and use linker wrapper flags for interposing.
+    add_proteus(${target} FORCE_JIT_ANNOTATE_ALL LINK_SHARED)
     string(TOLOWER "${CMAKE_CXX_COMPILER}" _cxx)
     set(_looks_like_clang FALSE)
     if(_cxx MATCHES "clang|amdclang|hipcc")

docs/examples.md

@@ -0,0 +1,175 @@
+# Examples
+
+The examples directory contains benchmarks derived from the [HecBench](https://github.com/zjin-lcf/HecBench) suite, adapted to demonstrate Mneme's **kernel-level record, replay, and tuning** capabilities.
+
+Mneme enables users to **isolate** GPU kernels from larger applications by capturing the exact code (LLVM IR) and device memory state required for execution. This isolation allows for rapid experimentation, search space exploration, and auto-tuning without the overhead of running the full application or managing complex host-side dependencies.
+
+The examples provided here serve as a reference for:
+
+*   **Integration**: How to link Mneme with existing C/C++ build systems (CMake).
+*   **Recording**: Capturing kernel executions at runtime.
+*   **Tuning**: Using Python-based tools like **Optuna** to explore replay configurations (block size, grid size, etc.) and optimize performance.
+
+For instance, the `miniFE` example explicitly demonstrates how to define a search space over kernel parameters and drive the replay engine to find optimal configurations.
+
+## Building
+
+All examples are built using CMake. For each benchmark, the build system generates three distinct binaries:
+
+1.  **Original**: The baseline application without any instrumentation.
+2.  **Proteus**: The application instrumented with the base Proteus runtime.
+3.  **Mneme**: The application fully instrumented with Mneme, enabled for **recording** and **replay**.
+
+### Requirements
+
+To build the examples, you need to provide the installation paths for **Mneme**, **Proteus**, and your **LLVM** installation (if not in standard paths).
+
+### Build Command
+
+You can build the examples by configuring with CMake. You must enable HIP backend support using the`WITH_MNEME_EXAMPLE_HIP` flags.
+
+```bash
+mkdir build && cd build
+cmake -DCMAKE_C_COMPILER=$(mneme config cc) \
+      -DCMAKE_CXX_COMPILER=$(mneme config cxx) \
+      -DCMAKE_PREFIX_PATH=$(mneme config cmakedir)
+      -DWITH_MNEME_EXAMPLE_HIP=On \
+      ../examples/hecbench
+make -j
+```
+
+## WSM5
+
+**Location:** `examples/hecbench/wsm5`
+
+WSM5 (WRF Single Moment 5-class Microphysics) is a kernel extracted from the Weather Research and Forecasting (WRF) model. It simulates microphysics processes including vapor, rain, snow, cloud ice, and cloud water.
+
+The implementation relies on CUDA kernels to perform the heavy lifting. The `main.cpp` orchestrates the data movement and kernel launches.
+
+- **Key Files**:
+    - `main.cpp`: Setup and execution of the WSM5 kernel.
+    - `kernel.h`: formatting and logic for the CUDA kernel.
+    - `tune.py`: Tuning script exploring kernel launch parameters.
+    - `tune_passes.py`: Tuning script exploring compiler optimizations.
+
+### Recording
+
+To record the execution of the WSM5 kernel, you use the `mneme record` command. The application takes a single argument indicating the number of repetitions.
+
+```bash
+# syntax: mneme record -rdb <database_dir> -- <application> <args>
+mneme record -rdb wsm5_db -- ./build/examples/hecbench/wsm5/wsm5-mneme 1
+```
+
+This will create a `wsm5_db` directory containing the recording database and artifacts (memory snapshots, LLVM IR).
+
+### Tuning
+
+Once recorded, you can tune the kernel using the provided Python scripts. These scripts demonstrate how to use Mneme's Python API to drive replay with different configurations.
+
+**1. Kernel Parameter Tuning (`tune.py`)**
+
+This script focuses on tuning **kernel launch parameters** and **specialization**. It explores a search space defined by:
+
+- Block dimensions (`block_dim_x`, `block_dim_y`, `block_dim_z`)
+- Grid dimensions (`grid_dim_x`, `grid_dim_y`, `grid_dim_z`)
+- Specialization of kernel arguments
+- Launch bounds
+
+```bash
+# syntax: ./tune.py --record-db <database_dir> --record-id <kernel_id>
+./examples/hecbench/wsm5/tune.py --record-db wsm5_db/123456789.json --record-id 987654321
+```
+
+**2. Compiler Pass Tuning (`tune_passes.py`)**
+
+This script focuses on tuning the **code generation** process itself. Instead of just changing launch parameters, it modifies how the kernel is compiled from the recorded LLVM IR. It explores:
+
+- Optimization pipelines (e.g., `-O1`, `-O2`, `-O3`, custom pass lists)
+- Backend code generation options
+- Specialization (as it interacts with optimization)
+
+```bash
+./examples/hecbench/wsm5/tune_passes.py --record-db wsm5_db/123456789.json --record-id 987654321
+```
+
+**Conceptual Difference**:
+*   `tune.py` optimizes **how the kernel is run** (threads per block, grid size).
+*   `tune_passes.py` optimizes **how the kernel is built** (compiler optimizations, register allocation strategies).
+
+Both approaches rely on Mneme's unique ability to **recompile** the recorded LLVM IR on-the-fly during replay.
+
+## Bezier Surface
+
+**Location:** `examples/hecbench/bezier-surface`
+
+This example computes a Bezier surface. It demonstrates a basic integration where the application can run on GPU and contrasts the computer output with a CPU implementation.
+
+- **Key Files**:
+    - `main.cpp`: Contains the host and device code for Bezier surface calculation.
+    - `tune.py`: Python script for auto-tuning the kernel parameters and compiler passes.
+
+### Recording
+
+To record the execution of the Bezier Surface kernel, you use the `mneme record` command. The application requires an input file and an output size.
+
+```bash
+# syntax: mneme record -rdb <database_dir> -- <application> -f <input_file> -n <size>
+mneme record -rdb bezier_db -- ./build/examples/hecbench/bezier-surface/bezier-mneme -f examples/hecbench/bezier-surface/input/control.txt -n 8192
+```
+
+### Tuning
+
+The provided `tune.py` script for this example combines both parameter tuning and compiler pass selection into a single search space. It uses `PipelineManager` to generate a list of potential optimization pipelines.
+
+```bash
+# syntax: ./tune.py --record-db <database_dir> --record-id <kernel_id>
+./examples/hecbench/bezier-surface/tune.py --record-db bezier_db/123456789.json --record-id 987654321
+```
+
+This script exhaustively searches through combinations of:
+
+- **Specialization**: Toggling kernel argument specialization.
+- **Launch Bounds**: Enabling/disabling launch bounds.
+- **Optimization Pipelines**: Iterating through standard levels (`-O1`, `-O2`, `-O3`) and a procedurally generated set of custom pass sequences.
+
+## MiniFE
+
+**Location:** `examples/hecbench/miniFE`
+
+MiniFE is a Finite Element mini-application which implements a couple of kernels representative of implicit finite-element applications.
+
+This example is particularly notable for its integration with **Optuna** for auto-tuning. The `tune.py` script provided in this directory demonstrates a workflow where:
+
+1. A kernel execution is recorded using Mneme.
+2. An `EntireSpace` class defines a search space for tuning parameters like `warp_fraction`, `grid_fraction`, and `max_threads`.
+3. An `AsyncReplayExecutor` is used to replay the recorded kernel with different configurations proposed by Optuna to find the optimal speedup.
+
+- **Key Files**:
+    - `tune.py`: An advanced example of defining a search space and using Optuna to tune a recorded kernel execution.
+    - `src/`: Source code for the miniFE application.
+
+### Recording
+
+To record the execution of the MiniFE kernel, you use the `mneme record` command. The application requires dimensions for the problem size (x, y, z).
+
+```bash
+# syntax: mneme record -rdb <database_dir> -- <application> -nx <nx> -ny <ny> -nz <nz>
+mneme record -rdb miniFE_db -- ./build/examples/hecbench/miniFE/miniFE-mneme -nx 220 -ny 200 -nz 190
+```
+
+### Tuning
+
+The `tune.py` script for MiniFE uses the **Optuna** library to drive the search for optimal kernel parameters. Unlike the other examples that use a simple exhaustive search, this demonstrates how to integrate Mneme with external optimization frameworks.
+
+```bash
+# syntax: ./tune.py --record-db <database_dir> --record-id <kernel_id>
+./examples/hecbench/miniFE/tune.py --record-db miniFE_db/123456789.json --record-id 987654321
+```
+
+Keys parameters tuned in this example include:
+
+- **Warp Fraction**: Adjusting the number of active warps.
+- **Grid Fraction**: Scaling the grid size.
+- **Max Threads**: Limiting the maximum threads per block.
+

docs/usage/getting-started.md

@@ -16,7 +16,8 @@ For full installation details, see **Usage → Install**.
 ```bash
 git clone https://github.com/Olympus-HPC/Mneme.git
 cd Mneme
-LLVM_INSTALL_DIR=${ROCM_PATH} pip install -e .
+export LLVM_INSTALL_DIR=${ROCM_PATH}
+pip install -e .
 ```
 
 ## Execute Example Code

docs/usage/install.md

@@ -21,8 +21,8 @@ on internal test systems.
 | ROCm version | Python 3.9 | Python 3.10 | Python 3.11 | Python 3.12 |
 |-------------|------------|-------------|-------------|-------------|
 | **6.3**     | ✅         | ✅          | ✅          | ✅          |
-| **6.4**     | ✅         | ✅          | ✅          | ✅          | 
-| **7.0**     | ⏳         | ⏳          | ⏳          | ⏳          | 
+| **6.4**     | ✅         | ✅          | ✅          | ✅          |
+| **7.0**     | ⏳         | ⏳          | ⏳          | ⏳          |
 
 #### Notes
 
@@ -58,7 +58,7 @@ use the corresponding Proteus commit to avoid incompatibilities.
 #### Tested Proteus commit
 
 - Repository: https://github.com/Olympus-HPC/Proteus
-- Commit: `1d21c00008061704459a9b20300556e962c89043`  
+- Commit: `v2026.01.0`
 - Tested with: Mneme `develop`
 
 !!! note
@@ -107,7 +107,8 @@ to record and replay kernels.
 ```bash
 git clone https://github.com/Olympus-HPC/Mneme.git
 cd Mneme
-LLVM_INSTALL_DIR=${ROCM_PATH} pip install .
+export LLVM_INSTALL_DIR=${ROCM_PATH}
+pip install .
 ```
 
 This installs the Mneme CLI (mneme) and Python bindings along with all
@@ -197,4 +198,3 @@ Python bindings are correctly installed.
 Once Mneme is installed and the test suite completes successfully,
 proceed to **Getting Started** for a guided, end-to-end example of
 building, recording, and replaying a GPU kernel with Mneme.
-

examples/hecbench/bezier-surface/tune.py

@@ -1,20 +1,16 @@
 #!/usr/bin/env python3
 """
-Mneme tuning example (Optuna)
+Mneme tuning example (Exhaustive Search)
 
 This example demonstrates how to run a tuning session on a
-previously recorded kernel execution.
+previously recorded kernel execution using an exhaustive search strategy.
 
 Workflow:
   1) Load a recorded execution (record-db) and select a kernel (record-id).
   2) Define a (Exhaustive) SearchSpace that exposes tunable parameters.
   3) Run a baseline configuration to verify replay correctness and measure baseline time.
-  4) Print the best configuration and its result.
-
-Notes:
-  - This example intentionally keeps the API usage explicit and minimal.
-  - The Optuna objective is configured as direction="minimize" and the script
-    reports a speedup value to Optuna, exactly as shown in the original example.
+  4) Exhaustively explore the defined search space.
+  5) Print the best configuration and its result.
 """
 
 import argparse
@@ -42,8 +38,8 @@ class EntireSpace(SearchSpace):
 
     This SearchSpace:
       - Uses the recorded grid/block dims as fixed reference values where needed.
-      - Exposes a set of tunable parameters (block_dim_x, specialization toggles,
-        min_blocks_per_sm, launch bounds, codegen choices, and pass pipeline).
+      - Exposes a set of tunable parameters (specialization toggles,
+        launch bounds, and optimization pass pipeline).
       - Produces an ExperimentConfiguration via derived(params).
     """