Carry loader_scratch_vm_allocate.patch: Apple loader scratch from vm_allocate

.github/workflows/validate.yml

@@ -144,6 +144,12 @@ jobs:
           SD_GRAPH_SPLIT_TARGET: ${{ matrix.target }}
           SD_REQUIRE_RUNTIME: "1"
           GGML_VK_VISIBLE_DEVICES: "0"
+      - name: Check loader scratch memory
+        # Metal needs a Mac: GitHub's macOS runners list no Metal device.
+        if: runner.os == 'macOS' && matrix.smoke
+        run: python3 -m unittest discover -s tests -p test_loader_scratch_runtime.py -v
+        env:
+          SD_REQUIRE_RUNTIME: "1"
       - uses: dart-lang/setup-dart@v1
         if: runner.os == 'Linux' && matrix.smoke
       - name: Verify cooperative host shutdown

CMakeLists.txt

@@ -105,6 +105,7 @@ sd_native_carry_patch(torch_zip_entry_lookup.patch src/model_io/torch_zip_io.cpp
 sd_native_carry_patch(fp8_metal_f16_load.patch src/model_manager.cpp src/core/backend_fit.cpp)
 sd_native_carry_patch(lora_immediate_fallback.patch src/pipeline/diffusion_engine.cpp)
 sd_native_carry_patch(intrinsic_tiny_decoder_policy.patch src/stable-diffusion.cpp src/pipeline/diffusion_engine.cpp)
+sd_native_carry_patch(loader_scratch_vm_allocate.patch src/model_loader.cpp)
 sd_native_carry_patch(ggml_apple_vm_allocate.patch TARGET ggml-base ggml/src/ggml.c)
 if(TARGET ggml-metal)
     sd_native_carry_patch(ggml_metal_apple_vm_allocate.patch TARGET ggml-metal ggml/src/ggml-metal/ggml-metal-device.m)

README.md

@@ -657,7 +657,7 @@ the macOS probe and builds the iOS device and simulator slices.
 Upstream 1017 has its
 [ABI audit, wrapper review, FP8 evidence and outstanding gates](docs/upstream_1017_qualification.md);
 the carry inventory is in the [upstream 982 record](docs/upstream_982_qualification.md).
-It carries five fixes the pin lacks: unbatched CFG in Metal builds (below),
+It carries six fixes the pin lacks: unbatched CFG in Metal builds (below),
 the [`pickle_line_span.patch`](patches/README.md) for legacy checkpoints above
 2 GiB, the [`fp8_metal_f16_load.patch`](patches/README.md), which loads FP8
 weights as F16 on Metal instead of aborting, so they take twice their file size
@@ -667,7 +667,9 @@ cannot merge it: into FP8 weights, and on Metal, which cannot add into the F16
 convolution weights every model has, and the
 [`ggml_apple_vm_allocate.patch`](patches/README.md), which backports
 [llama.cpp#31437](https://github.com/ggml-org/llama.cpp/pull/31437) so iOS
-host buffers go back to the system when freed.
+host buffers go back to the system when freed, and the
+[`loader_scratch_vm_allocate.patch`](patches/README.md), which does the same for
+the loader's tensor scratch buffers.
 Upstream's `generate_image` now returns an owned result with per-image
 generation metadata; `sd_dart_generate_image` keeps returning an image array
 that `free_sd_images` frees and drops that metadata. Since upstream 982 a null

patches/README.md

@@ -26,5 +26,7 @@ reports.
 | `intrinsic_tiny_decoder_policy.patch` | `src/stable-diffusion.cpp`, `src/pipeline/diffusion_engine.cpp` | Owner-only `sd_dart_new_sd_ctx_intrinsic_tiny_v1` context contract for [llamadart #1047](https://github.com/leehack/llamadart/issues/1047), not an upstream defect: automatic model refresh is valid upstream behavior. The public C API cannot preserve a decoder invariant, so minimal context creation/free and candidate/execution insertion points reject a full-VAE replacement before runner probes or decode. Real unchanged/refreshed SDXS and unrestricted reload controls exercise the boundary. Retire when upstream exposes an equivalent persistent decoder policy. | `tests/test_decoder_policy_runtime.py` |
 | `ggml_apple_vm_allocate.patch`, `ggml_metal_apple_vm_allocate.patch` | `ggml/src/ggml.c`, `ggml/src/ggml-metal/ggml-metal-device.m` | On iOS, ggml's CPU and Metal host buffers come from `posix_memalign`/`free`, and libmalloc keeps the freed pages in the app's memory footprint, so an unloaded model still counts toward the jetsam limit ([llamadart #2004](https://github.com/leehack/llamadart/issues/2004)). Every Apple platform now uses `vm_allocate`/`vm_deallocate`, as macOS already did. A backport of llama.cpp `ea24b952c` ([ggml-org/llama.cpp#31437](https://github.com/ggml-org/llama.cpp/pull/31437)); retire when the pinned ggml includes it. Two patches because the files belong to the `ggml-base` and `ggml-metal` targets. | `tools/validate_artifacts.py` (an iOS library must not import `_posix_memalign`), `tests/test_tools.py` |
 
+| `loader_scratch_vm_allocate.patch` | `src/model_loader.cpp` | On iOS 18 the loader's per-thread read, convert and zip scratch vectors and `convert_tensor`'s buffer leave about 240 MiB in libmalloc's cache after an unload, which still counts toward the jetsam limit ([llamadart #2004](https://github.com/leehack/llamadart/issues/2004)). On Apple they come from `vm_allocate`/`vm_deallocate`; other platforms keep `std::vector`. **Patch-specific maintainer exception:** no upstream report is filed; [owner #45](https://github.com/leehack/stable-diffusion-native/issues/45) tracks retirement until one is authorized. Retire when upstream's loader no longer takes tensor-sized scratch from the heap on Apple. | `tests/test_loader_scratch_runtime.py` |
+
 Patches touching the same source compose in declaration order. Each next
 patch applies to the preceding build-tree copy; the submodule stays pristine.

patches/loader_scratch_vm_allocate.patch

@@ -0,0 +1,92 @@
+diff --git a/src/model_loader.cpp b/src/model_loader.cpp
+--- a/src/model_loader.cpp
++++ b/src/model_loader.cpp
+@@ -8,6 +8,7 @@
+ #include <fstream>
+ #include <functional>
+ #include <mutex>
++#include <new>
+ #include <regex>
+ #include <set>
+ #include <stdexcept>
+@@ -33,6 +34,55 @@
+ #include "zip.h"
+
+ #include "name_conversion.h"
++
++#if defined(__APPLE__)
++#include <mach/mach.h>
++
++// Tensor-sized scratch memory. libmalloc can keep a freed large block in the
++// process footprint (leehack/llamadart#2004); vm_deallocate returns it at once.
++// Unlike a vector, growing keeps no contents.
++class TensorScratch {
++public:
++    TensorScratch()                                = default;
++    TensorScratch(const TensorScratch&)            = delete;
++    TensorScratch& operator=(const TensorScratch&) = delete;
++    ~TensorScratch() {
++        release();
++    }
++
++    void resize(size_t size) {
++        if (size <= capacity_) {
++            return;
++        }
++        const size_t capacity = std::max(size, 2 * capacity_);
++        release();
++        vm_address_t address = 0;
++        if (vm_allocate(mach_task_self(), &address, capacity, VM_FLAGS_ANYWHERE) != KERN_SUCCESS) {
++            throw std::bad_alloc();
++        }
++        data_     = (uint8_t*)address;
++        capacity_ = capacity;
++    }
++
++    uint8_t* data() {
++        return data_;
++    }
++
++private:
++    void release() {
++        if (data_ != nullptr) {
++            vm_deallocate(mach_task_self(), (vm_address_t)data_, capacity_);
++        }
++        data_     = nullptr;
++        capacity_ = 0;
++    }
++
++    uint8_t* data_   = nullptr;
++    size_t capacity_ = 0;
++};
++#else
++using TensorScratch = std::vector<uint8_t>;
++#endif
+
+ /*================================================= Preprocess ==================================================*/
+
+@@ -148,9 +198,9 @@
+             throw std::runtime_error(sd_format("type %s unsupported for integer quantization: no dequantization available",
+                                                ggml_type_name(src_type)));
+         }
+-        std::vector<char> buf;
++        TensorScratch buf;
+         buf.resize(sizeof(float) * n);
+-        char* src_data_f32 = buf.data();
++        char* src_data_f32 = (char*)buf.data();
+         qtype->to_float(src, (float*)src_data_f32, n);
+         if (dst_type == GGML_TYPE_F16) {
+             ggml_fp32_to_fp16_row((float*)src_data_f32, (ggml_fp16_t*)dst, n);
+@@ -1168,9 +1218,9 @@
+                     }
+                 }
+
+-                std::vector<uint8_t> read_buffer;
+-                std::vector<uint8_t> convert_buffer;
+-                std::vector<uint8_t> zip_entry_buffer;
++                TensorScratch read_buffer;
++                TensorScratch convert_buffer;
++                TensorScratch zip_entry_buffer;
+
+                 while (true) {
+                     int64_t t0, t1;

tests/native/loader_scratch_runtime_test.cpp

@@ -0,0 +1,110 @@
+// Loads a model, generates once and reports the largest `operator new` request
+// a thread other than the caller's made meanwhile.
+//
+// Usage: loader_scratch_runtime_test devices
+//        loader_scratch_runtime_test run <model> <backend>
+//
+// upstream's loader reads and converts tensors on its own threads. This
+// program replaces the process's `operator new`, which the library's vectors
+// use too, so a loader that takes its tensor-sized scratch from the heap shows
+// up as `WORKER_NEW <bytes>`.
+
+#include "sd_dart_wrapper.h"
+
+#include <atomic>
+#include <cstdio>
+#include <cstdlib>
+#include <cstring>
+#include <dlfcn.h>
+#include <new>
+#include <pthread.h>
+#include <vector>
+
+static std::atomic<size_t> largest_worker_new{0};
+
+// Metal's own threads allocate too; only the library's requests count.
+static bool from_library(void* caller) {
+    Dl_info info;
+    return dladdr(caller, &info) != 0 && info.dli_fname != nullptr &&
+           std::strstr(info.dli_fname, "libstable-diffusion") != nullptr;
+}
+
+void* operator new(size_t size) {
+    size_t seen = largest_worker_new.load();
+    if (size > seen && pthread_main_np() == 0 && from_library(__builtin_return_address(0))) {
+        while (size > seen && !largest_worker_new.compare_exchange_weak(seen, size)) {
+        }
+    }
+    void* memory = std::malloc(size == 0 ? 1 : size);
+    if (memory == nullptr) {
+        throw std::bad_alloc();
+    }
+    return memory;
+}
+
+void operator delete(void* memory) noexcept {
+    std::free(memory);
+}
+
+namespace {
+
+int list_devices() {
+    std::vector<char> devices(sd_list_devices(nullptr, 0) + 1);
+    sd_list_devices(devices.data(), devices.size());
+    std::printf("%s", devices.data());
+    return 0;
+}
+
+int fail(const char* what) {
+    std::vector<char> error(4096);
+    sd_dart_last_error(error.data(), error.size());
+    std::fprintf(stderr, "%s: %s\n", what, error.data());
+    return 3;
+}
+
+int run(const char* model, const char* backend) {
+    sd_dart_log_enable();
+    sd_ctx_params_t context_params;
+    sd_ctx_params_init(&context_params);
+    context_params.model_path = model;
+    context_params.taesd_path = model;
+    context_params.backend    = backend;
+    sd_ctx_t* context         = sd_dart_new_sd_ctx(&context_params);
+    if (context == nullptr) {
+        return fail("load failed");
+    }
+    sd_img_gen_params_t params;
+    sd_img_gen_params_init(&params);
+    params.prompt                         = "a lighthouse";
+    params.negative_prompt                = "";
+    params.width                          = 64;
+    params.height                         = 64;
+    params.seed                           = 7;
+    params.sample_params.sample_steps     = 1;
+    params.sample_params.guidance.txt_cfg = 1.0f;
+    sd_image_t* images                    = nullptr;
+    int generated                         = 0;
+    const bool ok                         = sd_dart_generate_image(context, &params, &images, &generated);
+    if (ok) {
+        free_sd_images(images, generated);
+    }
+    sd_dart_exit_free(context);
+    if (!ok) {
+        return fail("generation failed");
+    }
+    std::printf("WORKER_NEW %zu\n", largest_worker_new.load());
+    return 0;
+}
+
+}  // namespace
+
+int main(int argc, char** argv) {
+    if (argc == 2 && std::strcmp(argv[1], "devices") == 0) {
+        return list_devices();
+    }
+    if (argc == 4 && std::strcmp(argv[1], "run") == 0) {
+        return run(argv[2], argv[3]);
+    }
+    std::fprintf(stderr, "usage: %s devices | run <model> <backend>\n", argv[0]);
+    return 2;
+}

tests/test_loader_scratch_runtime.py

@@ -0,0 +1,134 @@
+"""On Apple, the model loader takes no tensor-sized scratch memory from the heap.
+
+Upstream's loader reads and converts tensors in scratch vectors on its own
+threads. libmalloc can keep such a freed block in the process footprint, about
+200 MiB of an iPhone app after an SD-Turbo unload (leehack/llamadart#2004), so
+`patches/loader_scratch_vm_allocate.patch` takes them from `vm_allocate`.
+
+The probe replaces `operator new` and reports the largest request the library
+made off the calling thread while it loaded the 12 MB PixArt model
+`exit_teardown_runtime_test make-model` writes and generated once. Without the
+patch that is a whole tensor: a convolution weight read for conversion on the
+CPU, every tensor on Metal. A copy with BF16 convolution weights reaches the
+conversion's own scratch. Nothing is downloaded. GitHub's macOS runners list
+no Metal device; run this on a Mac with `SD_REQUIRE_METAL_GENERATION=1`.
+
+    SD_REQUIRE_RUNTIME=1            fail instead of skipping
+    SD_REQUIRE_METAL_GENERATION=1   fail when no Metal device generates
+"""
+
+from __future__ import annotations
+
+import os
+import platform
+import re
+import shutil
+import subprocess
+import sys
+import tempfile
+import unittest
+from pathlib import Path
+
+REPO_ROOT = Path(__file__).resolve().parent.parent
+sys.path.insert(0, str(REPO_ROOT / "tools"))
+sys.path.insert(0, str(REPO_ROOT / "tests"))
+
+import build  # noqa: E402
+from test_fp8_runtime import read_safetensors, write_safetensors, write_weight_copies  # noqa: E402
+
+COMPILER = os.environ.get("CXX") or shutil.which("c++")
+MODEL_WRITER = REPO_ROOT / "tests" / "native" / "exit_teardown_runtime_test.cpp"
+PROBE = REPO_ROOT / "tests" / "native" / "loader_scratch_runtime_test.cpp"
+# The loader's threads ask for a few hundred bytes; the model's smallest
+# scratch block without the patch is a 73728-byte convolution weight.
+MAX_WORKER_NEW = 16 * 1024
+WORKER_NEW = re.compile(r"^WORKER_NEW (\d+)$", re.MULTILINE)
+
+
+def write_bf16_convolutions(source: Path, copy: Path) -> None:
+    """Stores the 4-D weights as BF16, which the loader converts to F16."""
+    header, data = read_safetensors(source)
+    header.pop("__metadata__", None)
+    blobs = {}
+    for name, info in header.items():
+        begin, end = info["data_offsets"]
+        blobs[name] = data[begin:end]
+        if len(info["shape"]) == 4 and info["dtype"] == "F32":
+            blob = blobs[name]
+            blobs[name] = bytes(b for pair in zip(blob[2::4], blob[3::4]) for b in pair)
+            info["dtype"] = "BF16"
+    write_safetensors(copy, header, blobs)
+
+
+@unittest.skipUnless(sys.platform == "darwin", "the carried patch changes Apple builds only")
+class LoaderScratchRuntimeTest(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls) -> None:
+        cls.skip_reason = ""
+        target = build.TARGETS["macos-arm64" if platform.machine() == "arm64" else "macos-x64"]
+        library = build.BIN_ROOT / target.name / "lib" / target.library
+        if not library.is_file() or COMPILER is None:
+            cls.skip_reason = f"no built runtime for {target.name} or no C++ compiler"
+            return
+        cls.directory = tempfile.TemporaryDirectory()
+        work = Path(cls.directory.name)
+        binaries = {}
+        for source in (MODEL_WRITER, PROBE):
+            binary = work / source.stem
+            subprocess.run(
+                [COMPILER, "-std=c++17", "-O1", "-pthread", f"-I{REPO_ROOT / 'src'}",
+                 f"-I{build.HEADER.parent}", str(source), f"-L{library.parent}",
+                 "-lstable-diffusion", f"-Wl,-rpath,{library.parent}", "-o", str(binary)],
+                check=True)
+            binaries[source] = binary
+        cls.probe = binaries[PROBE]
+        source_model = work / "source.safetensors"
+        subprocess.run([str(binaries[MODEL_WRITER]), "make-model", str(source_model)],
+                       check=True, timeout=600, capture_output=True)
+        cls.models = {"F32": work / "f32.safetensors", "BF16 convolutions": work / "bf16.safetensors"}
+        write_weight_copies(source_model, {"F32": cls.models["F32"]})
+        write_bf16_convolutions(cls.models["F32"], cls.models["BF16 convolutions"])
+        listing = subprocess.run([str(cls.probe), "devices"], check=True, capture_output=True,
+                                 text=True, timeout=120).stdout
+        cls.devices = [line.split("\t")[0] for line in listing.splitlines() if line.strip()]
+
+    @classmethod
+    def tearDownClass(cls) -> None:
+        if hasattr(cls, "directory"):
+            cls.directory.cleanup()
+
+    def setUp(self) -> None:
+        if self.skip_reason:
+            if os.environ.get("SD_REQUIRE_RUNTIME"):
+                self.fail(self.skip_reason)
+            self.skipTest(self.skip_reason)
+
+    def check_device(self, device: str) -> None:
+        for name, model in self.models.items():
+            with self.subTest(device=device, model=name):
+                result = subprocess.run([str(self.probe), "run", str(model), device],
+                                        capture_output=True, text=True, timeout=600)
+                if result.returncode != 0 and device.startswith("MTL") and \
+                        not os.environ.get("SD_REQUIRE_METAL_GENERATION"):
+                    self.skipTest(f"{device} cannot generate:\n{result.stderr[-2000:]}")
+                self.assertEqual(0, result.returncode, f"{device}:\n{result.stderr[-4000:]}")
+                requests = WORKER_NEW.findall(result.stdout)
+                self.assertEqual(1, len(requests), result.stdout[-4000:])
+                self.assertLess(int(requests[0]), MAX_WORKER_NEW,
+                                "a loader thread took tensor scratch from the heap")
+
+    def test_cpu_load_takes_no_tensor_scratch_from_the_heap(self) -> None:
+        self.check_device("CPU")
+
+    def test_metal_load_takes_no_tensor_scratch_from_the_heap(self) -> None:
+        metal = [device for device in self.devices if device.startswith("MTL")]
+        if not metal:
+            if os.environ.get("SD_REQUIRE_METAL_GENERATION"):
+                self.fail(f"no Metal device in {self.devices}")
+            self.skipTest(f"no Metal device in {self.devices}")
+        for device in metal:
+            self.check_device(device)
+
+
+if __name__ == "__main__":
+    unittest.main()