diff --git a/.gitlab-ci.yml b/.gitlab-ci.yml index bceb99bf..44773022 100644 --- a/.gitlab-ci.yml +++ b/.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' diff --git a/.gitlab/custom-jobs-and-variables.yml b/.gitlab/custom-jobs-and-variables.yml index 2a3f3156..678311ee 100644 --- a/.gitlab/custom-jobs-and-variables.yml +++ b/.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: diff --git a/.gitlab/subscribed-pipelines.yml b/.gitlab/subscribed-pipelines.yml index 5b752ea2..931e2c8b 100644 --- a/.gitlab/subscribed-pipelines.yml +++ b/.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] diff --git a/README.md b/README.md index a7eb2dc8..0fa93995 100644 --- a/README.md +++ b/README.md @@ -4,24 +4,19 @@ # 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 @@ -29,8 +24,7 @@ We welcome all kinds of contributions: new features, bug fixes, documentation ed 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} } ``` - diff --git a/cmake/MnemeFunctions.cmake b/cmake/MnemeFunctions.cmake index cea5b17d..90ef60d1 100644 --- a/cmake/MnemeFunctions.cmake +++ b/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") diff --git a/docs/examples.md b/docs/examples.md new file mode 100644 index 00000000..73ebfdf2 --- /dev/null +++ b/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 -- +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 --record-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 -- -f -n +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 --record-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 -- -nx -ny -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 --record-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. + diff --git a/docs/usage/getting-started.md b/docs/usage/getting-started.md index 1a760dc9..547a7153 100644 --- a/docs/usage/getting-started.md +++ b/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 diff --git a/docs/usage/install.md b/docs/usage/install.md index 8ec623d8..db725b59 100644 --- a/docs/usage/install.md +++ b/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. - diff --git a/examples/hecbench/bezier-surface/tune.py b/examples/hecbench/bezier-surface/tune.py index e8b20dc8..cc528dec 100644 --- a/examples/hecbench/bezier-surface/tune.py +++ b/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). """ diff --git a/examples/hecbench/miniFE/tune.py b/examples/hecbench/miniFE/tune.py index 8addf3ef..d2e6f315 100644 --- a/examples/hecbench/miniFE/tune.py +++ b/examples/hecbench/miniFE/tune.py @@ -3,18 +3,14 @@ Mneme tuning example (Optuna) This example demonstrates how to run a tuning session on a -previously recorded kernel execution. +previously recorded kernel execution using Optuna. Workflow: 1) Load a recorded execution (record-db) and select a kernel (record-id). - 2) Define a (Exhaustive) SearchSpace that exposes tunable parameters. + 2) Define a 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) Use Optuna to explore the search space and minimize execution time. + 5) Print the best configuration and its result. """ import argparse @@ -47,6 +43,12 @@ def dimensions(self): return self._search_space def derived(self, params) -> ExperimentConfiguration: + """ + Convert sampled parameters into a concrete ExperimentConfiguration. + + Maps the search space parameters (warp_fraction, grid_fraction, max_threads) + to the ExperimentConfiguration fields. + """ # Compute the number of active Warps and map that to the numThreads. maxWarpsInBlock = 1024 / 64 numActiveWarps = min( diff --git a/examples/hecbench/wsm5/tune.py b/examples/hecbench/wsm5/tune.py index 6b4d5632..d3d696d4 100644 --- a/examples/hecbench/wsm5/tune.py +++ b/examples/hecbench/wsm5/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. + 2) Define a 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 @@ -67,8 +63,8 @@ def derived(self, params) -> ExperimentConfiguration: """ Convert sampled parameters into a concrete ExperimentConfiguration. - If launch bounds are enabled, this example maps max_threads_fraction into a - proper max_threads integer in [block_dim_x, 1024] (rounded to multiples of 64). + Maps the search space parameters directly to the ExperimentConfiguration + fields. """ derived_config = { "block": { diff --git a/examples/hecbench/wsm5/tune_passes.py b/examples/hecbench/wsm5/tune_passes.py index acae9c41..e47eab89 100644 --- a/examples/hecbench/wsm5/tune_passes.py +++ b/examples/hecbench/wsm5/tune_passes.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. + 2) Define a 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 @@ -40,8 +36,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). """ diff --git a/include/mneme/DeviceTraits.hpp b/include/mneme/DeviceTraits.hpp index f6f5227c..aa6040df 100644 --- a/include/mneme/DeviceTraits.hpp +++ b/include/mneme/DeviceTraits.hpp @@ -187,7 +187,7 @@ template <> struct DeviceTraits { auto EC = DeviceErrorCheck( hipModuleGetFunction(&KernelFunc, HipModule, KernelName.c_str())); if (EC) - LOG_FATAL("Error with loading kernel from Module"); + LOG_FATAL("Error with loading kernel {} from Module {}", EC.value(), KernelName.c_str()); return KernelFunc; } diff --git a/include/mneme/MnemeJITProteus.hpp b/include/mneme/MnemeJITProteus.hpp index 1e7351fe..bbabe11c 100644 --- a/include/mneme/MnemeJITProteus.hpp +++ b/include/mneme/MnemeJITProteus.hpp @@ -1,9 +1,9 @@ -#include -#include +#include +#include #include #ifdef MNEME_ENABLE_HIP -#include +#include #elif defined(MNEME_ENABLE_CUDA) #else #error "Please define MNEME_ENABLE_HIP or MNEME_ENABLE_CUDA" diff --git a/include/mneme/MnemeRecord.hpp b/include/mneme/MnemeRecord.hpp index 700c30e6..e585a53a 100644 --- a/include/mneme/MnemeRecord.hpp +++ b/include/mneme/MnemeRecord.hpp @@ -22,7 +22,7 @@ #include #include -#include +#include #include "mneme/DeviceTraits.hpp" #include "mneme/MnemeKernelInfo.hpp" diff --git a/include/mneme/MnemeSnapshot.hpp b/include/mneme/MnemeSnapshot.hpp index 9c848dcc..588e2487 100644 --- a/include/mneme/MnemeSnapshot.hpp +++ b/include/mneme/MnemeSnapshot.hpp @@ -16,8 +16,8 @@ #include #include "proteus/CompilerInterfaceDevice.h" -#include "proteus/Hashing.hpp" -#include +#include "proteus/Hashing.h" +#include #include "mneme/DeviceTraits.hpp" #include "mneme/MnemeKernelInfo.hpp" diff --git a/mkdocs.yml b/mkdocs.yml index 35349533..0d0f02b0 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -36,7 +36,7 @@ nav: - Develop: - Python: dev/python.md - C++: dev/cpp.md - + - Examples: examples.md plugins: - search - mkdocstrings: diff --git a/python/mneme/alternate_ir.py b/python/mneme/alternate_ir.py new file mode 100644 index 00000000..9028867a --- /dev/null +++ b/python/mneme/alternate_ir.py @@ -0,0 +1,122 @@ +import json +import shutil +import tempfile +import warnings +import os +from pathlib import Path +from typing import Union, List, Optional +from mneme.llvm.module import parse_assembly + +class AlternateIRRecordDB: + """ + Context manager for creating a temporary record database with alternate LLVM IR. + Creates a *shallow copy* of the record DB and replaces the IR module with the new one. + + This class creates a temporary copy of a Mneme record database (JSON file) + and modifies it to point to a new LLVM IR module. This is useful for + replaying a recorded kernel with modified IR code. + + Usage: + with AlternateIRRecordDB(record_db="path/to/kernel.json", new_ir="path/to/new.ll") as rdb: + executor = AsyncReplayExecutor( + record_db=rdb.path, + record_id=rdb.rids[0], + ... + ) + """ + def __init__(self, record_db: Union[str, Path], new_ir: Union[str, Path], deep_copy: bool = False): + """ + Initialize the AlternateIRRecordDB context manager. + + Args: + record_db: Path to the existing record database JSON file (or directory containing one). + new_ir: Path to an LLVM IR file (.ll or .bc) or a string containing LLVM IR. + deep_copy: [Not Implemented] If True, create a deep copy of the record DB. Otherwise, create a shallow copy. + """ + self.record_db = Path(record_db).absolute() + self.new_ir = new_ir + self.temp_dir: Optional[str] = None + self.path: Optional[str] = None + self.rids: List[str] = [] + self.deep_copy = deep_copy + + def __post_init__(self): + if self.deep_copy: + raise NotImplementedError("Deep copy not implemented yet") + + def __enter__(self): + self.temp_dir = tempfile.mkdtemp(prefix="mneme_alternate_ir_") + + if not self.record_db.exists(): + raise FileNotFoundError(f"Record DB path not found: {self.record_db}") + + # A shortcut to pass DBs with only a single kernel in them + target_json = self.record_db + if self.record_db.is_dir(): + json_files = list(self.record_db.glob("*.json")) + if len(json_files) == 1: + target_json = json_files[0] + elif len(json_files) == 0: + raise ValueError(f"No JSON files found in {self.record_db}") + else: + raise ValueError(f"Multiple JSON files found in {self.record_db}. Please specify the JSON file directly.") + + with open(target_json, 'r') as f: + data = json.load(f) + + new_bc_path = Path(self.temp_dir) / "alternate_kernel.bc" + + # Handle new_ir -- either + # (1) a str containing LLVM IR + # (2) a Path to a .ll file + # (3) a Path to a .bc file + ir_is_content = False + if isinstance(self.new_ir, str): + if os.path.exists(self.new_ir): + ir_is_content = False + else: + # if a string and not a valid file, assume IR + ir_is_content = True + elif isinstance(self.new_ir, Path): + if self.new_ir.exists(): + ir_is_content = False + else: + raise FileNotFoundError(f"IR file not found: {self.new_ir}") + + if ir_is_content: + # parse assembly string and write bitcode + ir_str = str(self.new_ir) + mod = parse_assembly(ir_str) + mod.to_bitcode(str(new_bc_path)) + else: + ir_path = Path(self.new_ir) + if ir_path.suffix == '.bc': + shutil.copy(ir_path, new_bc_path) + else: + # assume .ll, read and convert + with open(ir_path, 'r') as f: + ir_content = f.read() + mod = parse_assembly(ir_content) + mod.to_bitcode(str(new_bc_path)) + + # update modules in existing record db + existing_modules = data.get("Modules", []) + if len(existing_modules) > 1: + warnings.warn(f"Original record DB had multiple modules: {existing_modules}. Replacing all with single alternate IR.") + + data["Modules"] = [str(new_bc_path)] + + # populate rids + self.rids = list(data.get("instances", {}).keys()) + + # write new JSON + new_json_path = Path(self.temp_dir) / target_json.name + with open(new_json_path, 'w') as f: + json.dump(data, f, indent=2) + + self.path = str(new_json_path) + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + if self.temp_dir and Path(self.temp_dir).exists(): + shutil.rmtree(self.temp_dir) diff --git a/python/mneme/async_executor.py b/python/mneme/async_executor.py index dd1ca864..7d2fe713 100644 --- a/python/mneme/async_executor.py +++ b/python/mneme/async_executor.py @@ -5,6 +5,7 @@ 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 Dict @@ -400,6 +401,26 @@ def __init__( for i in range(num_workers) ] + def set_ir(self, ir: str | Path): + """ Sets the LLVM IR for all workers to this IR. + Async, returns before the IR is set; But is guaranteed to set the IR + before any submit or evaluate calls after set_ir is called. + + 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 + + msg = {"payload": "set_ir", "data": ir_data} + for w in self.workers: + w._ipc_write_q.put(msg) + + # ------------------------------------------------------------------ # Submit new job (non-blocking) # ------------------------------------------------------------------ diff --git a/python/mneme/llvm/_lib_path_config.py b/python/mneme/llvm/_lib_path_config.py new file mode 100644 index 00000000..e39d062d --- /dev/null +++ b/python/mneme/llvm/_lib_path_config.py @@ -0,0 +1,16 @@ +# Auto-generated by setup.py +from __future__ import annotations +from pathlib import Path +import sys + +_PKG_DIR = Path(__file__).resolve().parents[1] +_NATIVE = _PKG_DIR / "native" +_LIB64 = _NATIVE / "lib64" + +def _lib(name_linux: str, name_darwin: str) -> str: + return str(_LIB64 / (name_darwin if sys.platform == "darwin" else name_linux)) + +MNEME_CORE_LIB = _lib("libmneme.so", "libmneme.dylib") +MNEME_PROFILE_LIB = _lib("libmneme_profile.so", "libmneme_profile.dylib") +MNEME_RECORD_LIB = _lib("librecord.so", "librecord.dylib") +MNEME_CONFIG_FILE = str(_NATIVE / "config.json") diff --git a/python/mneme/replay_executor.py b/python/mneme/replay_executor.py index 106d5689..bf524fce 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 Tuple from mneme.device import ( @@ -44,7 +45,8 @@ set_device, ) from mneme.llvm.buffer import MemBufferRef -from mneme.llvm.module import ModuleRef +from mneme.llvm.buffer import MemBufferRef +from mneme.llvm.module import ModuleRef, parse_assembly, parse_bitcode from mneme.logging import logger from mneme.mneme_types import ExperimentConfiguration, ExperimentResult from mneme.page_manager import PageManagerRef @@ -166,6 +168,23 @@ 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) + @cond_time("preprocess_ir_time") def _preprocess_ir( self, @@ -799,6 +818,13 @@ 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) + elif msg["payload"] == "process": logger.debug( f"Worker {worker.device_id} received processing request {msg['exp_id']}" diff --git a/scripts/gitlab/ci-build-test.sh b/scripts/gitlab/ci-build-test.sh index 6beb09da..85a2f7c5 100755 --- a/scripts/gitlab/ci-build-test.sh +++ b/scripts/gitlab/ci-build-test.sh @@ -1,26 +1,28 @@ #!/bin/bash -set -e +set -ex -temp_dir=$(pwd) #$(mktemp -d) +temp_dir=$(mktemp -d) echo "Temporary directory created at: $temp_dir" host=$(hostname) host=${host//[0-9]/} mkdir -p ${temp_dir}/build-${host}; build_dir=${temp_dir}/build-${host} -installDir="/dev/shm/install" +installDir=$(mktemp -d) +mneme_src=$(pwd) +pushd $build_dir build_proteus() { echo "Building PROTEUS" if [[ ! -d proteus ]]; then git clone --depth 1 git@github.com:Olympus-HPC/proteus.git - pushd proteus - git fetch --depth 1 origin 1d21c00008061704459a9b20300556e962c89043 - git checkout 1d21c00008061704459a9b20300556e962c89043 + pushd proteus + git fetch --depth 1 origin v2026.01.0 + git checkout -b v2026.01.0 FETCH_HEAD popd fi - pushd proteus + pushd proteus PROTEUS_ENABLE_HIP=$1 PROTEUS_ENABLE_CUDA=$2 PROTEUS_INSTALL_DIR=$3 @@ -68,7 +70,7 @@ build_spdlog() { -DCMAKE_C_COMPILER=${LLVM_INSTALL_DIR}/bin/clang \ -DCMAKE_CXX_COMPILER=${LLVM_INSTALL_DIR}/bin/clang++ \ -DCMAKE_INSTALL_PREFIX=${SPDLOG_INSTALL_DIR} \ - .. + .. make -j 10 make install -j 10 @@ -94,7 +96,7 @@ conda create -y -n mneme -c conda-forge \ else source ./${MINICONDA_DIR}/bin/activate fi -conda activate mneme +conda activate mneme LLVM_INSTALL_DIR=$(llvm-config --prefix) export LLVM_INSTALL_DIR=$(llvm-config --prefix) @@ -105,9 +107,6 @@ echo "Setting root dir to be ${LLVM_INSTALL_DIR}" build_proteus "OFF" "ON" $installDir ON echo "After proteus Current directory is $(pwd)" build_spdlog $installDir -mneme_src=$(pwd) -set -x -pushd $build_dir echo "Current dir is $(pwd)" cmake \ -DCMAKE_BUILD_TYPE=Debug \ @@ -136,8 +135,8 @@ build_proteus "ON" "OFF" $installDir OFF echo "After proteus Current directory is $(pwd)" build_spdlog $installDir echo "After spdlog Current directory is $(pwd)" -mneme_src=$(pwd) -pushd $build_dir + + cmake \ -DCMAKE_BUILD_TYPE=Relwithdebinfo \ -Dproteus_DIR=$installDir \ diff --git a/scripts/gitlab/ci-python-test.sh b/scripts/gitlab/ci-python-test.sh index 4771848a..69df33db 100755 --- a/scripts/gitlab/ci-python-test.sh +++ b/scripts/gitlab/ci-python-test.sh @@ -1,38 +1,75 @@ #!/bin/bash -set -e +set -euo pipefail +trap 'echo "[ERROR] line $LINENO" >&2' ERR + +log() { echo " #### [$(date +%T)] $* ###" >&2; } + +log "CI job $CI_JOB_ID starting" if [[ -n "$CODECOV_TOKEN" ]]; then - echo "CODECOV_TOKEN is set" + log "CODECOV_TOKEN is set" else - echo "CODECOV_TOKEN is not set" + log "CODECOV_TOKEN is not set" fi -mneme_src=$(pwd) +log "CI JOB ID IS ${CI_JOB_ID}" + +test_dir="$(mktemp -d)" +mkdir -p ${test_dir} + +# Make a copy of src to local FS to accelerate building etc. +log "Start copying src to temp" +mneme_orig_src=$(pwd) +mkdir -p $test_dir/mneme_src/ +rsync -a --delete "$mneme_orig_src" "$test_dir/mneme_src/" +mneme_src=$test_dir/mneme_src/Mneme +log "End copying mneme src to ${mneme_src} from ${mneme_orig_src}" + ml load python/${MNEME_CI_PYTHON_VERSION} ml load rocm/${MNEME_CI_ROCM_VERSION} export LLVM_INSTALL_DIR=${ROCM_PATH}/ -test_dir="$TMP/mneme-ci-${CI_JOB_ID}" -mkdir -p ${test_dir} -echo "Test dir is ${test_dir}" -VENV_NAME="/usr/workspace/LExperts/ci/gitlab/venv-${LCSCHEDCLUSTER}-${MNEME_CI_ROCM_VERSION}-${MNEME_CI_PYTHON_VERSION}/" -mkdir -p ${VENV_NAME} +log "Test dir is ${test_dir}" + + +# Make a copy of preinstalled deps to local FS for fast installation +#VENV_NAME="/usr/workspace/LExperts/ci/gitlab/venv-${LCSCHEDCLUSTER}-${MNEME_CI_ROCM_VERSION}-${MNEME_CI_PYTHON_VERSION}/" +#rm -rf ${VENV_NAME}/lib*/python*/site-packages/mneme + +log "Starting making environment" +LOCAL_VENV_NAME="${test_dir}/venv/" +mkdir -p ${LOCAL_VENV_NAME} +python -m venv "${LOCAL_VENV_NAME}" +#rsync -a "$VENV_NAME/" "$LOCAL_VENV_NAME" pushd ${test_dir} -python -m venv ${VENV_NAME} -source ${VENV_NAME}/bin/activate +source ${LOCAL_VENV_NAME}/bin/activate +echo "Environment is: ${LOCAL_VENV_NAME}" +echo "ENV will use python:" +which python + python -m pip uninstall -y mneme -rm -rf ${VENV_NAME}/lib*/python*/site-packages/mneme -python -m pip install ${mneme_src} +log "Environment made" + +log "Start installing mneme" +python -m pip -v install ${mneme_src} +log "Finalized with installation" + +log "Installing pytest" python -m pip install pytest pytest-cov + +log "Start running mneme tests" pytest -v -s ${mneme_src}/python/tests/ || exit $? +log "Done with testing" pushd ${mneme_src} pytest --cov-report=xml:coverage-${CI_JOB_ID}.xml --cov-config=.coveragerc python/tests/ +# we only need to upload reports once and we only need to test editable installs once. +if [[ "${MNEME_CI_PYTHON_VERSION}" == "3.10" && "${MNEME_CI_ROCM_VERSION}" == "6.4.2" ]]; then # Upload to Codecov (only if token is available) if [[ -n "$CODECOV_TOKEN" ]]; then - echo "Uploading coverage to Codecov..." - echo "SHA is $CI_COMMIT_SHA" - echo "Branch is $CI_COMMIT_BRANCH" + log "Uploading coverage to Codecov..." + log "SHA is $CI_COMMIT_SHA" + log "Branch is $CI_COMMIT_BRANCH" export NODE_TLS_REJECT_UNAUTHORIZED=0 curl -k -Os https://uploader.codecov.io/latest/linux/codecov @@ -50,4 +87,21 @@ else fi rm -f coverage-${CI_JOB_ID}.xml + + +python -m pip uninstall -y mneme +python -m pip install -U pip setuptools wheel + +# Try a editable install +log "Removed mneme" +log "Mneme src is ${mneme_src}" +log "Start editable install" +python -m pip install -e ${mneme_src} +# If everything is properly installed this command will not fail +log "End editable install" + +mneme config cxx || exit $? +fi + deactivate +rm -rf ${test_dir} diff --git a/setup.py b/setup.py index 6963efc1..01f4249b 100644 --- a/setup.py +++ b/setup.py @@ -1,15 +1,16 @@ import glob import json -import tempfile import os import subprocess import sys +import tempfile from pathlib import Path from setuptools import Extension, find_packages, setup from setuptools.command.build_ext import build_ext from setuptools.command.build_py import build_py from setuptools.command.develop import develop +from setuptools.command.egg_info import egg_info # Helper function to run shell commands @@ -58,54 +59,29 @@ class CMakeBuild(build_ext): PROTEUS_REPO = "https://github.com/Olympus-HPC/proteus.git" SPDLOG_REPO = "https://github.com/gabime/spdlog.git" - def write_python_config(self): - """Write a Python file with relocatable native library paths.""" - config = """\ -# Auto-generated by setup.py -from __future__ import annotations -from pathlib import Path -import sys - -_PKG_DIR = Path(__file__).resolve().parents[1] -_NATIVE = _PKG_DIR / "native" -_LIB64 = _NATIVE / "lib64" - -def _lib(name_linux: str, name_darwin: str) -> str: - return str(_LIB64 / (name_darwin if sys.platform == "darwin" else name_linux)) - -MNEME_CORE_LIB = _lib("libmneme.so", "libmneme.dylib") -MNEME_PROFILE_LIB = _lib("libmneme_profile.so", "libmneme_profile.dylib") -MNEME_RECORD_LIB = _lib("librecord.so", "librecord.dylib") -MNEME_CONFIG_FILE = str(_NATIVE / "config.json") -""" - - with open(self.config_py_path, "w") as fd: - fd.write(config) - def initialize_options(self): super().initialize_options() self.root_dir = Path(__file__).resolve().parent build_cmd = self.get_finalized_command("build") self.build_lib = Path(build_cmd.build_lib).resolve() - build_pkg_dir = (self.build_lib / "mneme") + build_pkg_dir = self.build_lib / "mneme" src_pkg_dir = self.root_dir / "python" / "mneme" inplace = bool(getattr(self, "inplace", False)) - self.editable = "editable_wheel" in sys.argv or inplace # include editable_wheel to be pep 660 compliant + self.editable = ( + "editable_wheel" in sys.argv or inplace + ) # include editable_wheel to be pep 660 compliant pkg_dir = src_pkg_dir if self.editable else build_pkg_dir self.install_dir = pkg_dir / "native" - self.config_py_path = pkg_dir / "llvm" / "_lib_path_config.py" - self.config_json = (self.install_dir / "config.json") + self.config_json = self.install_dir / "config.json" self.install_dir.mkdir(parents=True, exist_ok=True) (self.install_dir / "lib64").mkdir(parents=True, exist_ok=True) (self.install_dir / "include").mkdir(parents=True, exist_ok=True) (self.install_dir / "lib64" / "cmake").mkdir(parents=True, exist_ok=True) (self.install_dir / "llvm").mkdir(parents=True, exist_ok=True) - self.config_py_path.parent.mkdir(parents=True, exist_ok=True) - self.has_nvidia = "On" if has_nvidia_gpu() else "Off" self.has_amd = "On" if has_amd_gpu() else "Off" @@ -154,7 +130,6 @@ def run(self): spdlog_dir = self.clone_and_build_spdlog() self.build_mneme(proteus_dir, spdlog_dir) - self.write_python_config() def clone_and_build_proteus(self): if "PROTEUS_SRC" in os.environ: @@ -182,12 +157,12 @@ def clone_and_build_proteus(self): "--depth", "1", "origin", - "1d21c00008061704459a9b20300556e962c89043", + "v2026.01.0", ], cwd=str(Path(self.build_scratch) / "proteus"), ) run_command( - ["git", "checkout", "1d21c00008061704459a9b20300556e962c89043"], + ["git", "checkout", "-b", "v2026.01.0", "FETCH_HEAD"], cwd=str(Path(self.build_scratch) / "proteus"), ) @@ -283,8 +258,8 @@ def build_mneme(self, proteus_dir, spdlog_dir): "-DCMAKE_SKIP_INSTALL_RPATH=OFF", "-DCMAKE_INSTALL_RPATH_USE_LINK_PATH=ON", "-DMNEME_ENABLE_LOGGER=On", - #f"-Dproteus_DIR={proteus_dir}", - #f"-Dspdlog_DIR={spdlog_dir}", + # f"-Dproteus_DIR={proteus_dir}", + # f"-Dspdlog_DIR={spdlog_dir}", f"-DCMAKE_PREFIX_PATH={str(Path(self.install_dir).resolve())}", ] @@ -299,6 +274,21 @@ def run(self): self.run_command("build_ext") +class CustomEggInfo(egg_info): + def run(self): + try: + from pathlib import Path + + root = Path(__file__).resolve().parent + so = root / "python" / "mneme" / "native" / "lib64" / "libmneme.so" + if not so.exists(): + self.run_command("build_ext") + except Exception: + # if anything goes wrong, don't block egg_info + pass + super().run() + + class CustomBuildPy(build_py): def run(self): build_ext_cmd = self.get_finalized_command("build_ext") @@ -335,6 +325,7 @@ def run(self): "build_ext": CMakeBuild, "build_py": CustomBuildPy, "develop": CustomDevelop, + "egg_info": CustomEggInfo, }, classifiers=[ "Programming Language :: Python :: 3", diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 98bdac20..d44a0c1d 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -54,7 +54,9 @@ endif() target_include_directories(record SYSTEM PRIVATE ${LLVM_INCLUDE_DIRS} ${CMAKE_SOURCE_DIR}/include) -target_link_libraries(record PRIVATE proteus) +# Link the record preload shared library with proteus shared library to avoid +# LLVM static initialization issues. +target_link_libraries(record PRIVATE proteus_shared) if (MNEME_ENABLE_LOGGER) target_link_libraries(record PRIVATE spdlog::spdlog_header_only) diff --git a/src/MnemeReplay.cpp b/src/MnemeReplay.cpp index e3e7dc9b..954e2fc4 100644 --- a/src/MnemeReplay.cpp +++ b/src/MnemeReplay.cpp @@ -6,12 +6,12 @@ #include #include -#include -#include +#include +#include #include #ifdef MNEME_ENABLE_HIP -#include +#include #elif defined(MNEME_ENABLE_CUDA) #else #error "Please define MNEME_ENABLE_HIP or MNEME_ENABLE_CUDA" diff --git a/src/python/llvm/initfini.cpp b/src/python/llvm/initfini.cpp index 32f9479e..19fa3f26 100644 --- a/src/python/llvm/initfini.cpp +++ b/src/python/llvm/initfini.cpp @@ -1,7 +1,7 @@ #include "core.h" #include "llvm/Config/llvm-config.h" #include -#include +#include namespace { class LLVMInstance { diff --git a/src/python/proteus/jit.cpp b/src/python/proteus/jit.cpp index 093619f3..51d56bde 100644 --- a/src/python/proteus/jit.cpp +++ b/src/python/proteus/jit.cpp @@ -10,9 +10,9 @@ #include #include #include -#include -#include -#include +#include +#include +#include using namespace proteus; @@ -113,6 +113,8 @@ ProteusPY_linkModules(const char **LLVMIRFiles, int size, internalize(*Mod.get(), KernelSym); } + proteus::runCleanupPassPipeline(*Mod.get()); + return wrap(Mod.release()); }