From 3bf39f592f6f0c59607bf83917aec87d55783e1f Mon Sep 17 00:00:00 2001 From: Jhin Lee Date: Sat, 14 Nov 2026 18:30:08 -0500 Subject: [PATCH] Carry loader_scratch_vm_allocate.patch: Apple loader scratch from vm_allocate (#45) --- .github/workflows/validate.yml | 6 + CMakeLists.txt | 1 + README.md | 6 +- patches/README.md | 2 + patches/loader_scratch_vm_allocate.patch | 92 +++++++++++++ tests/native/loader_scratch_runtime_test.cpp | 110 +++++++++++++++ tests/test_loader_scratch_runtime.py | 134 +++++++++++++++++++ 7 files changed, 349 insertions(+), 2 deletions(-) create mode 100644 patches/loader_scratch_vm_allocate.patch create mode 100644 tests/native/loader_scratch_runtime_test.cpp create mode 100644 tests/test_loader_scratch_runtime.py diff --git a/.github/workflows/validate.yml b/.github/workflows/validate.yml index 7dbd1c1..3b01953 100644 --- a/.github/workflows/validate.yml +++ b/.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 diff --git a/CMakeLists.txt b/CMakeLists.txt index 91cbde5..35fc8bc 100644 --- a/CMakeLists.txt +++ b/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) diff --git a/README.md b/README.md index bcc6644..33f5d81 100644 --- a/README.md +++ b/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 diff --git a/patches/README.md b/patches/README.md index 4b3ae7d..ba4e49d 100644 --- a/patches/README.md +++ b/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. diff --git a/patches/loader_scratch_vm_allocate.patch b/patches/loader_scratch_vm_allocate.patch new file mode 100644 index 0000000..bae1431 --- /dev/null +++ b/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 + #include + #include ++#include + #include + #include + #include +@@ -33,6 +34,55 @@ + #include "zip.h" + + #include "name_conversion.h" ++ ++#if defined(__APPLE__) ++#include ++ ++// 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; ++#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 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 read_buffer; +- std::vector convert_buffer; +- std::vector zip_entry_buffer; ++ TensorScratch read_buffer; ++ TensorScratch convert_buffer; ++ TensorScratch zip_entry_buffer; + + while (true) { + int64_t t0, t1; diff --git a/tests/native/loader_scratch_runtime_test.cpp b/tests/native/loader_scratch_runtime_test.cpp new file mode 100644 index 0000000..3f55b77 --- /dev/null +++ b/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 +// +// 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 `. + +#include "sd_dart_wrapper.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +static std::atomic 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 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 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(¶ms); + 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, ¶ms, &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 \n", argv[0]); + return 2; +} diff --git a/tests/test_loader_scratch_runtime.py b/tests/test_loader_scratch_runtime.py new file mode 100644 index 0000000..6926645 --- /dev/null +++ b/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()