Repository navigation
perf: run large CPU matmuls on Accelerate on Apple Silicon #29
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,6 +1,7 @@ | ||
| #include "core/compute_workspace.h" | ||
|
|
||
| #include <algorithm> | ||
| #include <atomic> | ||
| #include <cstring> | ||
| #include <exception> | ||
| #include <map> | ||
|
|
@@ -14,10 +15,24 @@ | |
| #include "ggml-impl.h" | ||
|
|
||
| namespace sd { | ||
| ComputeWorkspace::ComputeWorkspace(ggml_backend_t backend) | ||
| : backend_(backend) { | ||
| if (sd_backend_is_cpu(backend_)) { | ||
| blas_backend_ = sd_backend_blas_init(); | ||
| } | ||
| } | ||
|
|
||
| ComputeWorkspace::~ComputeWorkspace() { | ||
| segment_end(); | ||
| release(); | ||
| ggml_backend_free(cpu_backend_); | ||
| ggml_backend_free(blas_backend_); | ||
| } | ||
|
|
||
| void ComputeWorkspace::set_accelerator_threads(int n_threads) const { | ||
| if (blas_backend_ != nullptr) { | ||
| sd_backend_set_n_threads_any(blas_backend_, n_threads); | ||
| } | ||
| } | ||
|
|
||
| void ComputeWorkspace::set_extra_backends(const std::vector<ggml_backend_t>& backends) { | ||
|
|
@@ -33,7 +48,11 @@ namespace sd { | |
| return true; | ||
| } | ||
| for (int i = 0; i < ggml_graph_n_nodes(graph); ++i) { | ||
| ggml_tensor* node = ggml_graph_node(graph, i); | ||
| if (!ggml_backend_supports_op(backend_, node)) { | ||
| return true; | ||
| } | ||
| if (uses_blas() && node->op == GGML_OP_MUL_MAT && ggml_backend_supports_op(blas_backend_, node)) { | ||
| return true; | ||
|
Comment on lines
+55
to
56
There was a problem hiding this comment.
When BLAS is available and a CPU graph contains a supported AGENTS.md reference: AGENTS.md:L69-L76 Useful? React with 👍 / 👎.
Member
Author
There was a problem hiding this comment. Fixed in d6736ba. When the workspace is CPU plus BLAS every tensor stays in host memory, so execute_segment now passes the eval callback to the scheduler through ggml_backend_sched_set_eval_callback; ggml's scheduler uses the same ask/collect protocol as the plain CPU path. Other scheduler setups keep the existing warning. Checked on SD1.5 at 256 px, 2 steps, with OpenBLAS: --imat-out with BLAS on used to save an empty imatrix (0 entries, the warning printed 3 times) and now saves 256 entries over 880 calls, the same tensor names and call counts as SD_CPU_BLAS=0. The rendered image is byte-identical to before. |
||
| } | ||
| } | ||
|
|
@@ -43,6 +62,10 @@ namespace sd { | |
| ggml_backend_sched_t ComputeWorkspace::make_scheduler(size_t graph_size) { | ||
| std::vector<ggml_backend_t> backends{backend_}; | ||
| backends.insert(backends.end(), extra_backends_.begin(), extra_backends_.end()); | ||
| if (uses_blas()) { | ||
| // Repacked weights are not host buffers, so their matmuls stay on the repack kernels. | ||
| backends.insert(backends.begin(), blas_backend_); | ||
| } | ||
| if (!sd_backend_is_cpu(backend_)) { | ||
| if (cpu_backend_ == nullptr) { | ||
| cpu_backend_ = sd_backend_cpu_init(); | ||
|
|
@@ -109,6 +132,17 @@ namespace sd { | |
| release(); | ||
| return false; | ||
| } | ||
| if (uses_blas()) { | ||
| static std::atomic<int> logged{0}; | ||
| if (logged.fetch_add(1) < 4) { | ||
| int on_blas = 0; | ||
| for (int i = 0; i < ggml_graph_n_nodes(graph); ++i) { | ||
| on_blas += ggml_backend_sched_get_tensor_backend(scheduler_, ggml_graph_node(graph, i)) == blas_backend_; | ||
| } | ||
| LOG_INFO("cpu blas: %d of %d graph nodes run on %s", on_blas, ggml_graph_n_nodes(graph), | ||
| ggml_backend_name(blas_backend_)); | ||
| } | ||
| } | ||
| } else { | ||
| if (scheduler_ != nullptr) { | ||
| release(); | ||
|
|
@@ -207,11 +241,13 @@ namespace sd { | |
| auto scheduler = make_scheduler(graph_size); | ||
| if (scheduler != nullptr) { | ||
| assign_nodes(scheduler, copy_graph); | ||
| std::vector<size_t> sizes(extra_backends_.size() + 3); | ||
| ggml_backend_sched_reserve_size(scheduler, copy_graph, sizes.data()); | ||
| // BLAS shares host memory: count it as the CPU backend's for residency. | ||
| const size_t first = uses_blas() ? 1 : 0; | ||
| result.buffers.push_back({backend_, sizes[first] + (first ? sizes[0] : 0)}); | ||
| for (size_t i = 0; i < extra_backends_.size(); ++i) { | ||
| result.buffers.push_back({extra_backends_[i], sizes[i + 1 + first]}); | ||
| } | ||
| ggml_backend_sched_free(scheduler); | ||
| } | ||
|
|
@@ -262,9 +298,22 @@ namespace sd { | |
| return true; | ||
| } | ||
|
|
||
| void ComputeWorkspace::release_blas_scratch() { | ||
| if (blas_backend_ == nullptr || active_ || scheduler_ != nullptr) { | ||
| return; | ||
| } | ||
| // ggml-blas keeps its F32 dequantization buffer until the backend is freed. | ||
| ggml_backend_free(blas_backend_); | ||
| blas_backend_ = sd_backend_blas_init(); | ||
| } | ||
|
|
||
| size_t ComputeWorkspace::bytes(ggml_backend_t backend) const { | ||
| if (scheduler_ != nullptr) { | ||
| size_t total = ggml_backend_sched_get_buffer_size(scheduler_, backend); | ||
| if (backend == backend_ && uses_blas()) { | ||
| total += ggml_backend_sched_get_buffer_size(scheduler_, blas_backend_); | ||
| } | ||
| return total; | ||
| } | ||
| return allocator_ != nullptr && backend == backend_ ? ggml_gallocr_get_buffer_size(allocator_, 0) : 0; | ||
| } | ||
|
|
||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
When BLAS processes an F16 or unrepacked quantized weight,
ggml_backend_blas_mul_matretains an F32 conversion buffer as large as the biggest weight seen. Because this backend is freed only by the workspace destructor,GGMLRunner::runner_end()releases the scheduler and model weights but leaves that scratch resident; running the conditioner, diffusion model, and VAE sequentially can therefore accumulate one large, untracked buffer per live runner and exhaust RAM on memory-constrained Macs. Reclaim or recreate the BLAS backend at runner-session boundaries rather than retaining it until model destruction.Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Fixed in 0fab806. ggml-blas only frees work_data when the backend is freed, so runner_end() now frees and recreates the idle BLAS backend right after releasing the workspace (thread count is set again on every segment). Within a run the buffer is kept: the diffusion runner does not end between sampling steps, so a txt2img run releases it three times (after the conditioner, after sampling, after the VAE). Peak RSS on SD1.5 at 256 px, 2 steps, CPU with OpenBLAS, 5 runs each: 2525 to 2564 MB before, 2475 to 2500 MB after. Rendered images are byte-identical with BLAS on and with SD_CPU_BLAS=0.