Convert purge-autoinit to real pass

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}

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)

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):
     """

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
 

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<O3>"), (True, "purge-autoinit,default<O3>")]
+
+    # The returned module is a clone; the caller's module is left untouched.
+    assert out_ir is not ir
 
     ex.close()
 

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()

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
+  "$<$<PLATFORM_ID:Darwin>:LINKER:SHELL:-undefined dynamic_lookup>"
+  "$<$<AND:$<NOT:$<PLATFORM_ID:Darwin>>,$<BOOL:${UNIX}>>:-Wl,--version-script=${CMAKE_CURRENT_SOURCE_DIR}/MnemePassPlugin.exports.map>")
+set_target_properties(MnemePassPlugin PROPERTIES
+  LINK_DEPENDS "${CMAKE_CURRENT_SOURCE_DIR}/MnemePassPlugin.exports.map")

src/pass/MnemePassPlugin.exports.map

@@ -0,0 +1,5 @@
+{
+  global:
+    llvmGetPassPluginInfo;
+  local: *;
+};

src/pass/MnemePurgeAutoInitPass.cpp

@@ -0,0 +1,90 @@
+#include <llvm/IR/Instructions.h>
+#include <llvm/IR/IntrinsicInst.h>
+#include <llvm/IR/LLVMContext.h>
+#include <llvm/IR/Metadata.h>
+#include <llvm/IR/Module.h>
+#include <llvm/IR/PassManager.h>
+#include <llvm/Passes/PassBuilder.h>
+#if __has_include(<llvm/Plugins/PassPlugin.h>)
+#include <llvm/Plugins/PassPlugin.h>
+#elif __has_include(<llvm/Passes/PassPlugin.h>)
+#include <llvm/Passes/PassPlugin.h>
+#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<MDString>(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<MDNode>(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<PurgeAutoInitPass> {
+public:
+  PreservedAnalyses run(Module &M, ModuleAnalysisManager &) {
+    SmallVector<CallInst *, 32> ToErase;
+
+    for (Function &F : M) {
+      if (F.isDeclaration())
+        continue;
+
+      for (BasicBlock &BB : F)
+        for (Instruction &I : BB)
+          if (auto *II = dyn_cast<IntrinsicInst>(&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<PassBuilder::PipelineElement>) {
+                  if (Name != "purge-autoinit")
+                    return false;
+                  MPM.addPass(PurgeAutoInitPass());
+                  return true;
+                });
+          }};
+}

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)

src/python/proteus/jit.cpp

@@ -11,6 +11,7 @@
 #include <mneme/MnemeLogger.hpp>
 #include <optional>
 #include <proteus/CompilerInterfaceTypes.h>
+#include <proteus/Init.h>
 #include <proteus/impl/CompilerInterfaceRuntimeConstantInfo.h>
 #include <proteus/impl/CoreLLVM.h>
 #include <proteus/impl/CoreLLVMDevice.h>
@@ -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);
+}
 }

tests/CMakeLists.txt

@@ -48,3 +48,4 @@ endif()
 
 add_subdirectory(unit_tests)
 add_subdirectory(record)
+add_subdirectory(pass)