This is an automated email from the ASF dual-hosted git repository.

tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm-ffi.git


The following commit(s) were added to refs/heads/main by this push:
     new 51341adc [FIX][ORCJIT] Discard CRT callbacks after failed ARM64 links 
(#880)
51341adc is described below

commit 51341adc9a0c4c0de329852e24d57b1587b0e0c6
Author: Yaxing Cai <[email protected]>
AuthorDate: Mon Oct 26 23:38:40 2026 +0800

    [FIX][ORCJIT] Discard CRT callbacks after failed ARM64 links (#880)
    
    An ARM64 COFF object with CRT callback sections and an unresolved symbol
    can crash during module teardown. RuntimeDyld frees the failed object's
    memory, but the addon previously retained indirect
    constructor/destructor slots into that memory.
    
    Stage CRT entries per object until `NotifyEmitted` confirms successful
    linking. The object's memory manager discards staged entries on failure,
    before releasing their memory. Keep staging records separate when
    object-buffer addresses are reused.
    
    Add native Windows ARM64 regressions for clang-cl and MSVC. Each
    repeatedly fails materialization, unloads the failed module without
    callbacks, and then verifies constructor/destructor ordering on a valid
    module using the same session.
    
    Validation:
    - Reproduced the original ARM64 COFF crash in the C++ harness against
    `mlc-llvm-static==23.1.1.post0`; the fix leaves zero stale callbacks and
    exits successfully.
    - The same harness still executes the native ARM64 COFF fixtures,
    imported host calls, arithmetic, CRT ordering, null/sentinel slots, and
    teardown successfully.
    - macOS ARM64: built and repaired the wheel against the exact
    `23.1.1.post0` SDK, installed it with the published
    `apache-tvm-ffi==0.1.15.post1` core in an isolated environment without
    the LLVM SDK, and passed **136 tests, 57 skipped**, plus both C and C++
    quick-start examples.
    - Pre-commit checks pass.
    - Native Windows ARM64: **107 passed, 51 skipped**, including the new
    failed-materialization regressions for both clang-cl and MSVC; the C
    quick-start passes. [Validation
    run](https://github.com/cyx-6/tvm-ffi/actions/runs/40341249761).
    - Linux ARM64: **178 passed, 15 skipped**; Linux x86-64: **180 passed,
    13 skipped**. C and C++ quick-start examples pass on both. [Remaining
    platform validation
    run](https://github.com/cyx-6/tvm-ffi/actions/runs/40341882212) also
    builds Windows AMD64 and macOS ARM64 wheels.
    - Windows AMD64: **82 passed, 42 skipped**; the C quick-start passes.
    All four remote platform jobs have passed; the macOS CI job is still
    queued, with local native wheel validation completed above.
    
    The fork validation commits change only the CI workflow; addon and core
    sources are identical to this PR.
    
    All five candidate wheels are built and saved with SHA-256 checksums.
    Wheel metadata declares only the TVM FFI core as a runtime dependency;
    LLVM remains statically linked.
    
    The ORCJIT 0.1.3 publishing run was cancelled before PyPI upload so this
    fix can be included in the release. No LLVM SDK rebuild is needed.
---
 .../src/ffi/llvm_patches/win_arm64_rtdyld.cc       | 73 +++++++++++++++++++---
 .../src/ffi/llvm_patches/win_arm64_rtdyld.h        |  5 +-
 addons/tvm_ffi_orcjit/tests/test_basic.py          | 64 +++++++++++++++++++
 3 files changed, 132 insertions(+), 10 deletions(-)

diff --git a/addons/tvm_ffi_orcjit/src/ffi/llvm_patches/win_arm64_rtdyld.cc 
b/addons/tvm_ffi_orcjit/src/ffi/llvm_patches/win_arm64_rtdyld.cc
index 6d7d0f8a..01664c75 100644
--- a/addons/tvm_ffi_orcjit/src/ffi/llvm_patches/win_arm64_rtdyld.cc
+++ b/addons/tvm_ffi_orcjit/src/ffi/llvm_patches/win_arm64_rtdyld.cc
@@ -30,22 +30,64 @@
 
 #include <cstdint>
 #include <memory>
+#include <mutex>
+#include <unordered_map>
+#include <vector>
 
 namespace tvm {
 namespace ffi {
 namespace orcjit {
+namespace {
+
+struct PendingCrtEntries {
+  std::vector<ORCJITExecutionSessionObj::InitFiniEntry> initializers;
+  std::vector<ORCJITExecutionSessionObj::InitFiniEntry> deinitializers;
+};
+
+struct PendingCrtObjects {
+  std::mutex mutex;
+  std::unordered_map<const char*, std::shared_ptr<PendingCrtEntries>> objects;
+};
+
+// RuntimeDyld destroys the memory manager on failed materialization. Discard
+// the object's staged callbacks then, before their slot memory is released.
+class CrtSectionMemoryManager : public llvm::SectionMemoryManager {
+ public:
+  CrtSectionMemoryManager(std::shared_ptr<PendingCrtObjects> pending, const 
char* key)
+      : pending_(std::move(pending)), key_(key), 
entries_(std::make_shared<PendingCrtEntries>()) {
+    std::lock_guard<std::mutex> lock(pending_->mutex);
+    pending_->objects[key_] = entries_;
+  }
+
+  ~CrtSectionMemoryManager() override {
+    std::lock_guard<std::mutex> lock(pending_->mutex);
+    auto it = pending_->objects.find(key_);
+    // A successful object's buffer can be reused while its JIT memory lives.
+    if (it != pending_->objects.end() && it->second == entries_) 
pending_->objects.erase(it);
+  }
+
+ private:
+  std::shared_ptr<PendingCrtObjects> pending_;
+  const char* key_;
+  std::shared_ptr<PendingCrtEntries> entries_;
+};
+
+}  // namespace
 
 std::unique_ptr<llvm::orc::ObjectLayer> CreateWindowsArm64ObjectLinkingLayer(
     llvm::orc::ExecutionSession& execution_session, ORCJITExecutionSessionObj* 
session) {
+  auto pending = std::make_shared<PendingCrtObjects>();
   auto layer = std::make_unique<llvm::orc::RTDyldObjectLinkingLayer>(
-      execution_session,
-      [](const llvm::MemoryBuffer&) { return 
std::make_unique<llvm::SectionMemoryManager>(); });
+      execution_session, [pending](const llvm::MemoryBuffer& buffer) {
+        return std::make_unique<CrtSectionMemoryManager>(pending, 
buffer.getBufferStart());
+      });
   layer->setOverrideObjectFlagsWithResponsibilityFlags(true);
   layer->setAutoClaimResponsibilityForObjectSymbols(true);
-  layer->setNotifyLoaded([session](llvm::orc::MaterializationResponsibility& 
responsibility,
+  layer->setNotifyLoaded([pending](llvm::orc::MaterializationResponsibility&,
                                    const llvm::object::ObjectFile& object,
                                    const llvm::RuntimeDyld::LoadedObjectInfo& 
loaded) {
-    auto& dylib = responsibility.getTargetJITDylib();
+    std::lock_guard<std::mutex> lock(pending->mutex);
+    auto& entries = 
*pending->objects.at(object.getMemoryBufferRef().getBufferStart());
     for (const auto& section : object.sections()) {
       auto name = section.getName();
       if (!name) {
@@ -66,21 +108,36 @@ std::unique_ptr<llvm::orc::ObjectLayer> 
CreateWindowsArm64ObjectLinkingLayer(
           priority += number;
         }
       }
-      // NotifyLoaded precedes relocation. Keep addresses of the pointer
-      // slots; RunInitFiniEntries dereferences them after materialization.
+      // NotifyLoaded precedes relocation. Stage pointer slots here and publish
+      // them only after NotifyEmitted confirms successful materialization.
       constexpr uint64_t kPointerSize = 8;
       for (uint64_t offset = 0; offset + kPointerSize <= section.getSize();
            offset += kPointerSize) {
         ORCJITExecutionSessionObj::InitFiniEntry 
entry{llvm::orc::ExecutorAddr(address + offset),
                                                        priority, true};
         if (is_init) {
-          session->AddPendingInitializer(&dylib, entry);
+          entries.initializers.push_back(entry);
         } else {
-          session->AddPendingDeinitializer(&dylib, entry);
+          entries.deinitializers.push_back(entry);
         }
       }
     }
   });
+  layer->setNotifyEmitted(
+      [pending, session](llvm::orc::MaterializationResponsibility& 
responsibility,
+                         std::unique_ptr<llvm::MemoryBuffer> buffer) {
+        PendingCrtEntries entries;
+        {
+          std::lock_guard<std::mutex> lock(pending->mutex);
+          auto it = pending->objects.find(buffer->getBufferStart());
+          entries = std::move(*it->second);
+          pending->objects.erase(it);
+        }
+        auto* dylib = &responsibility.getTargetJITDylib();
+        for (const auto& entry : entries.initializers) 
session->AddPendingInitializer(dylib, entry);
+        for (const auto& entry : entries.deinitializers)
+          session->AddPendingDeinitializer(dylib, entry);
+      });
   return layer;
 }
 
diff --git a/addons/tvm_ffi_orcjit/src/ffi/llvm_patches/win_arm64_rtdyld.h 
b/addons/tvm_ffi_orcjit/src/ffi/llvm_patches/win_arm64_rtdyld.h
index 6917f475..30de6b99 100644
--- a/addons/tvm_ffi_orcjit/src/ffi/llvm_patches/win_arm64_rtdyld.h
+++ b/addons/tvm_ffi_orcjit/src/ffi/llvm_patches/win_arm64_rtdyld.h
@@ -23,8 +23,9 @@
  *
  * LLVM 23.1.1 JITLink rejects COFF/AArch64 objects. RuntimeDyld supports
  * their relocations, long branches and DLL-import pointer slots, so use
- * RTDyldObjectLinkingLayer on this platform. A loaded-object callback records
- * .CRT$XC* / .CRT$XT* entries for the session's normal init/fini drain.
+ * RTDyldObjectLinkingLayer on this platform. A loaded-object callback stages
+ * .CRT$XC* / .CRT$XT* slots; successful emission publishes them for the
+ * session's normal init/fini drain. Failed objects discard their staged slots.
  *
  * Trigger: loading any native Windows ARM64 COFF object.
  * Symptom without the patch: "Unsupported target machine architecture in
diff --git a/addons/tvm_ffi_orcjit/tests/test_basic.py 
b/addons/tvm_ffi_orcjit/tests/test_basic.py
index a1902c00..33869b6c 100644
--- a/addons/tvm_ffi_orcjit/tests/test_basic.py
+++ b/addons/tvm_ffi_orcjit/tests/test_basic.py
@@ -26,6 +26,7 @@ from __future__ import annotations
 
 import gc
 import platform
+import shutil
 import sys
 import tempfile
 import threading
@@ -59,6 +60,69 @@ def test_windows_arm64_native_compilers() -> None:
         assert mod.test_add(10, 20) == 30
 
 
[email protected](
+    sys.platform != "win32" or platform.machine().lower() not in ("arm64", 
"aarch64"),
+    reason="requires native Windows ARM64 COFF lifecycle handling",
+)
[email protected]("compiler", ["clang-cl", "cl"])
+def test_failed_materialization_discards_crt_callbacks(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch, compiler: str
+) -> None:
+    """Failed links release CRT slots without running them or poisoning later 
loads."""
+    assert shutil.which(compiler), f"{compiler} is required for native Windows 
ARM64 tests"
+    monkeypatch.setenv("CC", compiler)
+    source = Path(__file__).parent / "sources" / "c" / "test_ctor_dtor.c"
+    failing_source = tmp_path / "failed_ctor.c"
+    failing_source.write_text(
+        source.read_text().replace(
+            "static void puts_log(const char* msg) {",
+            "extern void tvm_ffi_test_unresolved_symbol(void);\n"
+            "static void puts_log(const char* msg) {\n"
+            "  tvm_ffi_test_unresolved_symbol();",
+        )
+    )
+    flags = ["/GS-"]
+    if compiler == "clang-cl":
+        flags.append("--target=aarch64-pc-windows-msvc")
+    failed_obj = tvm_ffi.cpp.build(
+        name="failed_ctor",
+        sources=[str(failing_source)],
+        output="failed_ctor.o",
+        extra_cflags=flags,
+        build_directory=str(tmp_path / "build"),
+    )
+    log: list[str] = []
+
+    @tvm_ffi.register_global_func("append_log", override=True)
+    def append_ctor_log(value: str) -> None:
+        log.append(value)
+
+    session = ExecutionSession()
+    for _ in range(3):
+        failed = None
+        with pytest.raises(Exception, match=r"Failed to materialize|Symbols 
not found"):
+            failed = session.load_module(failed_obj)
+            failed.get_function("main")
+        del failed
+        gc.collect()
+        assert log == []
+
+        mod = session.load_module(str(OBJ_DIR / "c" / "test_ctor_dtor.o"))
+        mod.get_function("main")()
+        del mod
+        gc.collect()
+        assert log == [
+            "<crt.XCA>",
+            "<crt.XCB>",
+            "<crt.XCC>",
+            "<crt.XCU>",
+            "<main>",
+            "<crt.XTA>",
+            "<crt.XTZ>",
+        ]
+        log.clear()
+
+
 _KNOWN_SUBDIRS = [
     "cc",
     "c",  # default (LLVM clang)

Reply via email to