From 1c0cca9c4343a7b0d53ab9164c143d2bf1fa0fd3 Mon Sep 17 00:00:00 2001 From: Zane Fink Date: Wed, 5 Aug 2026 12:51:26 -0700 Subject: [PATCH] Convert to real pass --- CMakeLists.txt | 2 +- python/mneme/llvm/_lib_path_config.py | 1 + python/mneme/proteus/jit.py | 19 +++++ python/mneme/replay_executor.py | 26 +++---- python/mneme/transforms/__init__.py | 0 python/mneme/transforms/transform.py | 12 ---- python/tests/test_replay_executor.py | 28 +++----- src/CMakeLists.txt | 2 + src/pass/CMakeLists.txt | 14 ++++ src/pass/MnemePassPlugin.exports.map | 5 ++ src/pass/MnemePurgeAutoInitPass.cpp | 90 ++++++++++++++++++++++++ src/python/CMakeLists.txt | 3 +- src/python/proteus/jit.cpp | 6 ++ src/python/transforms/purge_autoinit.cpp | 57 --------------- tests/CMakeLists.txt | 1 + tests/pass/CMakeLists.txt | 40 +++++++++++ tests/pass/purge-autoinit.ll | 24 +++++++ 17 files changed, 225 insertions(+), 105 deletions(-) delete mode 100644 python/mneme/transforms/__init__.py delete mode 100644 python/mneme/transforms/transform.py create mode 100644 src/pass/CMakeLists.txt create mode 100644 src/pass/MnemePassPlugin.exports.map create mode 100644 src/pass/MnemePurgeAutoInitPass.cpp delete mode 100644 src/python/transforms/purge_autoinit.cpp create mode 100644 tests/pass/CMakeLists.txt create mode 100644 tests/pass/purge-autoinit.ll diff --git a/CMakeLists.txt b/CMakeLists.txt index 54f5180a..67275619 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -164,7 +164,7 @@ write_basic_package_version_file( set(_mneme_targets mnemert record) if(MNEME_ENABLE_PYTHON) - list(APPEND _mneme_targets mneme_profile mneme) + list(APPEND _mneme_targets mneme_profile mneme MnemePassPlugin) endif() install(TARGETS ${_mneme_targets} diff --git a/python/mneme/llvm/_lib_path_config.py b/python/mneme/llvm/_lib_path_config.py index b40c28c1..b52a0996 100644 --- a/python/mneme/llvm/_lib_path_config.py +++ b/python/mneme/llvm/_lib_path_config.py @@ -33,4 +33,5 @@ def _lib(name_linux: str, name_darwin: str) -> str: 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_PASS_PLUGIN_LIB = _lib("libMnemePassPlugin.so", "libMnemePassPlugin.dylib") MNEME_CONFIG_FILE = str(_CONFIG_FILE) diff --git a/python/mneme/proteus/jit.py b/python/mneme/proteus/jit.py index 32f4b1c3..a54a0697 100644 --- a/python/mneme/proteus/jit.py +++ b/python/mneme/proteus/jit.py @@ -90,6 +90,25 @@ ffi.lib.ProteusPY_getCodegenMethod.argtypes = [] ffi.lib.ProteusPY_getCodegenMethod.restype = c_char_p +ffi.lib.ProteusPY_registerJITPassPlugin.argtypes = [c_char_p] + + +def register_pass_plugin(path: str): + """ + Register an LLVM pass plugin shared library with Proteus. + + Uses the load-only registration mode: the plugin's pass names become + parseable in pipeline strings passed to :func:`optimize`, but nothing is + added to the pipeline automatically. Registration is idempotent; Proteus + deduplicates by resolved path. + + Parameters + ---------- + path : str + Filesystem path to the plugin shared library. + """ + ffi.lib.ProteusPY_registerJITPassPlugin(_encode_string(path)) + def pruneIR(mod: ModuleRef): """ diff --git a/python/mneme/replay_executor.py b/python/mneme/replay_executor.py index 1084ccb9..3161c887 100644 --- a/python/mneme/replay_executor.py +++ b/python/mneme/replay_executor.py @@ -32,6 +32,7 @@ """ import os +from dataclasses import replace from datetime import datetime, timezone from multiprocessing import Event, Queue from typing import Optional, Tuple @@ -43,6 +44,7 @@ get_device_count, set_device, ) +from mneme.llvm._lib_path_config import MNEME_PASS_PLUGIN_LIB from mneme.llvm.buffer import MemBufferRef from mneme.llvm.module import ModuleRef from mneme.mneme_logging import logger @@ -51,7 +53,6 @@ from mneme.profile import init_profiler from mneme.proteus import jit from mneme.recorded_execution import RecordedExecution, MemStateRef -from mneme.transforms import transform from mneme.utils import cond_gpu_time, cond_time @@ -139,6 +140,7 @@ def __init__( logger.debug( f"GPU Affinity of process was set to device:{self.device_id} out of {self.num_devices}" ) + jit.register_pass_plugin(MNEME_PASS_PLUGIN_LIB) def open(self): # Note the 'executor' allocates all resources and picks address space. @@ -568,9 +570,10 @@ def _execute( match, allowing the system to validate kernel determinism and correctness. 2. **IR sanitization** - A custom transformation is applied to remove automatically inserted Clang - initialization code. Only IR regions explicitly marked by Clang are - removed to avoid disturbing user code. + The tracked build prepends the ``purge-autoinit`` pass to the optimization + pipeline, removing automatically inserted Clang initialization code. Only + IR regions explicitly marked by Clang are removed to avoid disturbing user + code. 3. **Instrumented execution** The cleaned up version of the kernel is built with tracking enabled. @@ -613,15 +616,12 @@ def _execute( mem_buffer = self._build(result, config, ver_mod, False) self._run(result, config, mem_buffer, self.prologue, self.epilogue, True, False, 1) - # NOTE: 2. We apply a custom pass to delete all clang insered code. - # It is hard to identify these cases, So we delete only things - # that have been attributed by clang - ir_module = transform.remove_auto_initialize(ir_module.clone()) - # Done with verification. Moving to next stage - - # NOTE: 3. We build and run. We set tracking on and execute warmups plus iterations, - # to enalbe later computation of statistical metrics etc. - mem_buffer = self._build(result, config, ir_module, True) + # NOTE: 2-3. Build and run with tracking on. The tracked build prepends the + # purge-autoinit pass (registered via Proteus' JIT pass plugin API) to strip + # the Clang-inserted auto-init memsets; verification above ran without it. + ir_module = ir_module.clone() + tracked_config = replace(config, passes=f"purge-autoinit,{config.passes}") + mem_buffer = self._build(result, tracked_config, ir_module, True) self._run(result, config, mem_buffer, self.prologue, self.epilogue, False, True, self._iterations + self._warmup) result.executed = True diff --git a/python/mneme/transforms/__init__.py b/python/mneme/transforms/__init__.py deleted file mode 100644 index e69de29b..00000000 diff --git a/python/mneme/transforms/transform.py b/python/mneme/transforms/transform.py deleted file mode 100644 index dd5bb99e..00000000 --- a/python/mneme/transforms/transform.py +++ /dev/null @@ -1,12 +0,0 @@ -from ctypes import c_char_p - -from ..llvm import ffi as ffi -from ..llvm.common import _encode_string -from ..llvm.module import ModuleRef - -ffi.lib.TransformPy_RemoveAutoInitMemset.argtypes = [ffi.LLVMModuleRef] - - -def remove_auto_initialize(mod: ModuleRef) -> ModuleRef: - ffi.lib.TransformPy_RemoveAutoInitMemset(mod) - return mod diff --git a/python/tests/test_replay_executor.py b/python/tests/test_replay_executor.py index 519dd5da..9c3478c7 100644 --- a/python/tests/test_replay_executor.py +++ b/python/tests/test_replay_executor.py @@ -71,7 +71,6 @@ class FakeModule: def __init__(self, name="m", clone_counter=None): self.name = name self._clone_counter = clone_counter if clone_counter is not None else {"n": 0} - self.removed_auto_init = False def clone(self): self._clone_counter["n"] += 1 @@ -186,6 +185,7 @@ def wrapper(*args, **kwargs): # Now reload target module so the decorators wrap with the identity versions. mod = importlib.import_module(MODULE_PATH) mod = importlib.reload(mod) + monkeypatch.setattr(mod.jit, "register_pass_plugin", lambda path: None, raising=True) return mod @@ -534,27 +534,12 @@ def test_execute_orchestrates_verification_and_tracked_run(monkeypatch): monkeypatch.setattr(mod, "get_device_count", lambda: 1, raising=True) monkeypatch.setattr(mod, "PageManagerRef", FakePageManager, raising=True) - # transform.remove_auto_initialize should be called on ir.clone() - transform_calls = [] - - def fake_remove_auto_initialize(ir_mod): - transform_calls.append(ir_mod.name) - ir_mod.removed_auto_init = True - return ir_mod - - monkeypatch.setattr( - mod.transform, - "remove_auto_initialize", - fake_remove_auto_initialize, - raising=True, - ) - # Spy on _build and _run build_calls = [] run_calls = [] def fake_build(result, cfg, ir_mod, track): - build_calls.append((ir_mod.name, track)) + build_calls.append((track, cfg.passes)) return FakeMemBuffer() def fake_run(result, cfg, mem_buf, prologue, epilogue, verify, track, iters): @@ -580,9 +565,12 @@ def fake_run(result, cfg, mem_buf, prologue, epilogue, verify, track, iters): assert res.executed is True assert res.verified is True - # Remove-auto-init called once - assert len(transform_calls) == 1 - assert out_ir.removed_auto_init is True + # The verification build uses the caller's pipeline; the tracked build + # prepends the purge pass. + assert build_calls == [(False, "default"), (True, "purge-autoinit,default")] + + # The returned module is a clone; the caller's module is left untouched. + assert out_ir is not ir ex.close() diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index bdf8cd73..1a618f02 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -54,6 +54,8 @@ endif() set_target_properties(record PROPERTIES INSTALL_RPATH_USE_LINK_PATH TRUE) +add_subdirectory(pass) + if (MNEME_ENABLE_PYTHON) add_subdirectory(python) endif() diff --git a/src/pass/CMakeLists.txt b/src/pass/CMakeLists.txt new file mode 100644 index 00000000..82d94fc5 --- /dev/null +++ b/src/pass/CMakeLists.txt @@ -0,0 +1,14 @@ +add_library(MnemePassPlugin MODULE MnemePurgeAutoInitPass.cpp) + +target_include_directories(MnemePassPlugin SYSTEM PRIVATE ${LLVM_INCLUDE_DIRS}) + +if(MNEME_LINK_SHARED_LLVM AND LLVM_LINK_LLVM_DYLIB) + llvm_config(MnemePassPlugin USE_SHARED) +endif() +target_link_libraries(MnemePassPlugin PRIVATE ${MNEME_LLVM_LIBS}) + +target_link_options(MnemePassPlugin PRIVATE + "$<$:LINKER:SHELL:-undefined dynamic_lookup>" + "$<$>,$>:-Wl,--version-script=${CMAKE_CURRENT_SOURCE_DIR}/MnemePassPlugin.exports.map>") +set_target_properties(MnemePassPlugin PROPERTIES + LINK_DEPENDS "${CMAKE_CURRENT_SOURCE_DIR}/MnemePassPlugin.exports.map") diff --git a/src/pass/MnemePassPlugin.exports.map b/src/pass/MnemePassPlugin.exports.map new file mode 100644 index 00000000..d352de5c --- /dev/null +++ b/src/pass/MnemePassPlugin.exports.map @@ -0,0 +1,5 @@ +{ + global: + llvmGetPassPluginInfo; + local: *; +}; diff --git a/src/pass/MnemePurgeAutoInitPass.cpp b/src/pass/MnemePurgeAutoInitPass.cpp new file mode 100644 index 00000000..5bbe5950 --- /dev/null +++ b/src/pass/MnemePurgeAutoInitPass.cpp @@ -0,0 +1,90 @@ +#include +#include +#include +#include +#include +#include +#include +#if __has_include() +#include +#elif __has_include() +#include +#else +#error "Cannot find LLVM PassPlugin.h" +#endif + +using namespace llvm; + +namespace { + +// Clang tags the initialization it emits for -ftrivial-auto-var-init with +// !annotation !{!"auto-init"}. Annotation metadata holds the strings directly +// or, when annotations are grouped, tuples of strings; both are accepted by +// the IR verifier, so check for either shape. +bool hasAutoInitAnnotation(const Instruction &I) { + const MDNode *Annotations = I.getMetadata(LLVMContext::MD_annotation); + if (!Annotations) + return false; + + auto IsAutoInit = [](const Metadata *MD) { + const auto *Str = dyn_cast_or_null(MD); + return Str && Str->getString() == "auto-init"; + }; + + for (const MDOperand &Op : Annotations->operands()) { + if (IsAutoInit(Op.get())) + return true; + + if (const auto *Group = dyn_cast_or_null(Op.get())) + for (const MDOperand &GroupOp : Group->operands()) + if (IsAutoInit(GroupOp.get())) + return true; + } + + return false; +} + +// Removes the llvm.memset calls Clang emits for -ftrivial-auto-var-init. +// Mneme records with forced zero-initialization; replay strips it to match +// the original application. +class PurgeAutoInitPass : public PassInfoMixin { +public: + PreservedAnalyses run(Module &M, ModuleAnalysisManager &) { + SmallVector ToErase; + + for (Function &F : M) { + if (F.isDeclaration()) + continue; + + for (BasicBlock &BB : F) + for (Instruction &I : BB) + if (auto *II = dyn_cast(&I)) + if (II->getIntrinsicID() == Intrinsic::memset && + hasAutoInitAnnotation(*II)) + ToErase.push_back(II); + } + + for (CallInst *CI : ToErase) + CI->eraseFromParent(); + + return ToErase.empty() ? PreservedAnalyses::all() + : PreservedAnalyses::none(); + } +}; + +} // namespace + +extern "C" LLVM_ATTRIBUTE_WEAK ::llvm::PassPluginLibraryInfo +llvmGetPassPluginInfo() { + return {LLVM_PLUGIN_API_VERSION, "MnemePurgeAutoInit", "0.1", + [](PassBuilder &PB) { + PB.registerPipelineParsingCallback( + [](StringRef Name, ModulePassManager &MPM, + ArrayRef) { + if (Name != "purge-autoinit") + return false; + MPM.addPass(PurgeAutoInitPass()); + return true; + }); + }}; +} diff --git a/src/python/CMakeLists.txt b/src/python/CMakeLists.txt index 6fd88447..ced835de 100644 --- a/src/python/CMakeLists.txt +++ b/src/python/CMakeLists.txt @@ -3,10 +3,9 @@ add_subdirectory(profile) file(GLOB LLVM_SRC "llvm/*.cpp") file(GLOB PROTEUS_SRC "proteus/*.cpp") file(GLOB MNEME_SRC "./*.cpp") -file(GLOB TRANSFORMS "transforms/*.cpp") -add_library(mneme SHARED ${LLVM_SRC} ${PROTEUS_SRC} ${MNEME_SRC} ${TRANSFORMS} ${MnemeCoreSrcDir}/MnemePageManager.cpp ${MnemeCoreSrcDir}/MnemeDeviceCode.cpp ${MnemeCoreSrcDir}/MnemeAnnotationInternal.cpp) +add_library(mneme SHARED ${LLVM_SRC} ${PROTEUS_SRC} ${MNEME_SRC} ${MnemeCoreSrcDir}/MnemePageManager.cpp ${MnemeCoreSrcDir}/MnemeDeviceCode.cpp ${MnemeCoreSrcDir}/MnemeAnnotationInternal.cpp) if (MNEME_ENABLE_HIP) set_source_files_properties(${MnemeCoreSrcDir}/MnemeDeviceCode.cpp PROPERTIES LANGUAGE HIP) set_target_properties(mneme PROPERTIES LINKER_LANGUAGE HIP) diff --git a/src/python/proteus/jit.cpp b/src/python/proteus/jit.cpp index d6010984..dae35a7c 100644 --- a/src/python/proteus/jit.cpp +++ b/src/python/proteus/jit.cpp @@ -11,6 +11,7 @@ #include #include #include +#include #include #include #include @@ -218,4 +219,9 @@ ProteusPY_setLaunchBounds(LLVMModuleRef Mod, uint64_t CurrentHash, API_EXPORT(const char*) ProteusPY_getCodegenMethod(){ return getRTCMethod(); } + +API_EXPORT(void) +ProteusPY_registerJITPassPlugin(const char *Path) { + proteus::registerJITPassPlugin(Path); +} } diff --git a/src/python/transforms/purge_autoinit.cpp b/src/python/transforms/purge_autoinit.cpp deleted file mode 100644 index f473f675..00000000 --- a/src/python/transforms/purge_autoinit.cpp +++ /dev/null @@ -1,57 +0,0 @@ -// ZeroInitAllocas.cpp -#include "../llvm/core.h" - -#include - -#include "llvm/IR/Attributes.h" -#include "llvm/IR/IntrinsicInst.h" -#include "llvm/IR/Instructions.h" -#include "llvm/IR/Module.h" - -using namespace llvm; - -namespace { - -static bool hasAutoInitAttr(const CallInst *CI) { - // Match string attribute "auto-init" (value optional, e.g., "zero") - return CI->getAttributes() - .hasAttributeAtIndex(AttributeList::FunctionIndex, "auto-init"); -} - -struct RemoveAutoInitMemsetsTransform { - static bool run(Module &M) { - SmallVector ToErase; - - for (Function &F : M) { - if (F.isDeclaration()) continue; - - for (BasicBlock &BB : F) - for (Instruction &I : BB) - if (auto *CI = dyn_cast(&I)) { - if (auto *II = dyn_cast(CI)) { - if (II->getIntrinsicID() == Intrinsic::memset && - hasAutoInitAttr(CI)) { - ToErase.push_back(CI); - } - } - } - } - - for (CallInst *CI : ToErase) - CI->eraseFromParent(); - - return !ToErase.empty(); - } -}; - -} // namespace - -extern "C" { -API_EXPORT(void) -TransformPy_RemoveAutoInitMemset(LLVMModuleRef Mod) { - auto *M = unwrap(Mod); - bool Modified = RemoveAutoInitMemsetsTransform::run(*M); - LOG_DEBUG("RemoveAutoInitMemsets {} modify LLVM IR", - Modified ? "did" : "did not"); -} -} diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index c2e5bbb6..64811388 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -48,3 +48,4 @@ endif() add_subdirectory(unit_tests) add_subdirectory(record) +add_subdirectory(pass) diff --git a/tests/pass/CMakeLists.txt b/tests/pass/CMakeLists.txt new file mode 100644 index 00000000..eeb50bf1 --- /dev/null +++ b/tests/pass/CMakeLists.txt @@ -0,0 +1,40 @@ +# The plugin is built against LLVM_INSTALL_DIR, so an opt from a different LLVM +# cannot load it. NO_DEFAULT_PATH keeps the search off the system paths. +find_program(OPT opt + PATHS ${LLVM_INSTALL_DIR}/bin ${LLVM_INSTALL_DIR}/libexec/llvm + NO_DEFAULT_PATH) +if(NOT OPT) + message(FATAL_ERROR "Pass tests require opt from the LLVM install") +endif() +message(STATUS "Found opt at ${OPT}") + +file(WRITE ${CMAKE_CURRENT_BINARY_DIR}/lit.cfg.py " +import lit.formats +import os +import tempfile +import atexit +import shutil + +config.name = 'Mneme pass tests' +config.test_format = lit.formats.ShTest(True) +config.environment = os.environ.copy() + +config.suffixes = ['.ll'] +config.test_source_root = '${CMAKE_CURRENT_SOURCE_DIR}' +exec_root = tempfile.mkdtemp(prefix='lit.tmp.', dir='${CMAKE_CURRENT_BINARY_DIR}') +config.test_exec_root = exec_root +atexit.register(lambda: shutil.rmtree(exec_root, ignore_errors=False)) + +config.substitutions.append(('%opt', lit_config.params['OPT'])) +config.substitutions.append(('%plugin', lit_config.params['PLUGIN'])) +config.substitutions.append(('%FILECHECK', lit_config.params['FILECHECK'])) +" +) + +add_test(NAME purge_autoinit_pass + COMMAND ${LIT} -a -vv + -DFILECHECK=${FILECHECK} + -DOPT=${OPT} + -DPLUGIN=$ + purge-autoinit.ll) +set_tests_properties(purge_autoinit_pass PROPERTIES LABELS "MnemePass") diff --git a/tests/pass/purge-autoinit.ll b/tests/pass/purge-autoinit.ll new file mode 100644 index 00000000..3ee0caaa --- /dev/null +++ b/tests/pass/purge-autoinit.ll @@ -0,0 +1,24 @@ +; RUN: %opt -load-pass-plugin=%plugin -passes=purge-autoinit -S %s | %FILECHECK %s + +define void @f(ptr %p) { +entry: + call void @llvm.memset.p0.i64(ptr %p, i8 0, i64 64, i1 false), !annotation !0 + call void @llvm.memset.p0.i64(ptr %p, i8 1, i64 64, i1 false), !annotation !1 + call void @llvm.memset.p0.i64(ptr %p, i8 2, i64 64, i1 false), !annotation !2 + call void @llvm.memset.p0.i64(ptr %p, i8 3, i64 64, i1 false) + ret void +} + +declare void @llvm.memset.p0.i64(ptr writeonly, i8, i64, i1 immarg) + +; !0 and !1 are the two annotation shapes the IR verifier accepts, a flat string +; and a tuple of strings. !2 is an unrelated annotation. +!0 = !{!"auto-init"} +!1 = !{!{!"auto-init"}} +!2 = !{!"other"} + +; CHECK-LABEL: @f +; CHECK-NEXT: entry: +; CHECK-NEXT: call void @llvm.memset.p0.i64(ptr %p, i8 2 +; CHECK-NEXT: call void @llvm.memset.p0.i64(ptr %p, i8 3 +; CHECK-NEXT: ret void