Mneme and Proteus

[koparasy]

Move Mneme Recording into the Proteus Repository

This proposal suggests migrating the LD_PRELOAD_LIBRARY (librecord.so) and the standalone C++ replay tool (non-Python) from Mneme into the Proteus repository.

Motivation

Over the past year, both projects have developed overlapping functionality — particularly in their LLVM instrumentation layers. Both currently embed LLVM IR into binaries and provide runtime mechanisms to retrieve and associate kernel invocations with their respective LLVM IR modules.

Consolidating these efforts into a shared implementation within Proteus would:

• Reduce code duplication and maintenance overhead, especially around LLVM pass integration.
• Eliminate compatibility issues such as static/dynamic hashing mismatches that currently require intermediate JSON-based Mneme → Proteus translation.
• Simplify installation and deployment by allowing users to integrate and autotune directly through Proteus, with only the lightweight Python interface (mneme) remaining as a separate, pip-installable component.
• Enable broader compiler integration. In particular, this move will allow extending the original Flang OpenMP-offload implementation to emit output files compatible with those generated by Mneme today. As a result, the remaining Python-based Mneme tool will be able to tune these offload cases as well, further expanding its applicability.

Implementation Outline

Only a focused subset of Mneme functionality would be migrated, specifically the core recording and replay components:

Recording

1. Track memory allocations and accesses.
2. Intercept JITDeviceImplT::instance() to extract the associated LLVM IR.
3. Store memory and metadata for replay.

Replay

1. Initialize and load recorded memory.
2. JIT-compile the corresponding kernel.
3. Execute the kernel with recorded inputs.

These components (librecord and the replay tool) would be defined as independent CMake targets within the Proteus build system. JITProteus would remain independent, with the recording/replay modules depending on it rather than the other way around. We can further guard these tools with ENABLE_MNEME cmake functionality.

[ZwFink]

Is it feasible to extract shared functionality into a library both Proteus/Mneme use, rather than having have one subsume the other?

[koparasy]

Technically yes, there is extra implementation cost on this.

[ggeorgakoudis]

I voted "no" at this time for these reasons:

1. It creates extra maintenance costs while Proteus is an fledgling state and I want to focus the code on the core JIT functionality
2. Mneme can already use Proteus as an optional dependency, including the LLVM pass, esp. since Mneme has larger functionality on record-replay, memory management, and tuning for both Proteus and non-Proteus instrumented codes
3. The librecord.so LD_PRELOAD approach is useful for retrofitting Mneme on non-Proteus instrumented apps but out-of-scope for Proteus

That said, I think we can integrate better Proteus with Mneme especially regarding recording as the Proteus runtime already extracts LLVM IR so it. For replay, the JSON configuration interface is a good start and we can evolve it.

I don't see any compelling factors for merging the codebases since we can interface them successfully and the philosophy of each tool doing "one task well" makes more sense to me.

[koparasy]

Thanks for the feedback — I completely understand the desire to keep Proteus focused while it’s still maturing. I want to clarify a few points on the maintenance and integration aspects to make sure we’re aligned.

1. Maintenance Cost:
   Unless we plan to remove LLVM IR embedding from Proteus’ core, I don’t see this adding significant ongoing maintenance overhead. The main addition would be a few extra tests to validate memory state consistency. Beyond that, the shared LLVM IR handling logic already exists on both sides.

2. Dependency Direction:
   Proteus is not actually an optional dependency for Mneme — Mneme always depends on it today. So, the proposal wouldn’t introduce a new dependency relationship, just consolidate existing shared functionality in one place.

3. LD_PRELOAD Layer:
   The librecord.so component would continue to exist as a standalone mechanism for managing memory state and retrofitting non-Proteus applications. The overlap with Proteus is primarily around LLVM IR access, not the full record/replay logic. Keep in mind, adding the proteus pass without annotations creates a very thin shim layer with negligible overheads.

My main concern is that maintaining separate paths for Proteus and Mneme effectively doubles the integration and testing effort, both for us and for end users. Starting from the same codebase for kernel-level recording and replay would simplify things considerably and reduce friction long-term.

That said, I fully agree that Mneme and Proteus should continue to do “one task well.” The intent here is to share the lowest common layer (LLVM IR handling and basic runtime hooks), not to merge higher-level logic or tuning functionality. I am not suggesting for Proteus to become an auto-tuning library, and I am not suggesting Mneme to become a JIT library.