refactor: centralize runner weight staging and cleanup (#1644)

src/core/ggml_graph_cut.cpp

@@ -44,7 +44,9 @@ namespace sd::ggml_graph_cut {
         if (tensor == nullptr) {
             return false;
         }
-        return params_tensor_set.find(tensor) != params_tensor_set.end();
+        return params_tensor_set.find(tensor) != params_tensor_set.end() ||
+               (tensor->view_src != nullptr &&
+                params_tensor_set.find(tensor->view_src) != params_tensor_set.end());
     }
 
     static int graph_node_index_by_name(ggml_cgraph* gf, const char* name) {
@@ -135,6 +137,24 @@ namespace sd::ggml_graph_cut {
         return max_vram_bytes_to_gib(resolve_auto_max_vram_bytes(-max_vram, backend));
     }
 
+    static bool is_segment_output_needed_after(const Plan& plan,
+                                               size_t end_segment_index,
+                                               int output_node_index) {
+        if (end_segment_index + 1 >= plan.segments.size()) {
+            return false;
+        }
+        for (size_t seg_idx = end_segment_index + 1; seg_idx < plan.segments.size(); ++seg_idx) {
+            const auto& segment = plan.segments[seg_idx];
+            for (const auto& input_ref : segment.input_refs) {
+                if (input_ref.type == Segment::INPUT_PREVIOUS_CUT &&
+                    input_ref.node_index == output_node_index) {
+                    return true;
+                }
+            }
+        }
+        return false;
+    }
+
     static Segment make_segment_seed(const Plan& plan,
                                      size_t start_segment_index,
                                      size_t end_segment_index) {
@@ -147,8 +167,11 @@ namespace sd::ggml_graph_cut {
         const auto& target_segment = plan.segments[end_segment_index];
         std::unordered_set<int> seen_output_node_indices;
         for (size_t seg_idx = start_segment_index; seg_idx <= end_segment_index; ++seg_idx) {
+            const bool is_boundary_segment = seg_idx == end_segment_index;
             for (int output_node_index : plan.segments[seg_idx].output_node_indices) {
-                if (seen_output_node_indices.insert(output_node_index).second) {
+                if ((is_boundary_segment ||
+                     is_segment_output_needed_after(plan, end_segment_index, output_node_index)) &&
+                    seen_output_node_indices.insert(output_node_index).second) {
                     seed.output_node_indices.push_back(output_node_index);
                 }
             }
@@ -400,23 +423,6 @@ namespace sd::ggml_graph_cut {
         return tensors;
     }
 
-    std::vector<ggml_tensor*> runtime_param_tensors(ggml_cgraph* gf, const Segment& segment, const char* log_desc) {
-        std::vector<ggml_tensor*> tensors = param_tensors(gf, segment);
-        std::vector<ggml_tensor*> filtered_tensors;
-        filtered_tensors.reserve(tensors.size());
-        for (ggml_tensor* tensor : tensors) {
-            if (tensor_buffer(tensor) == nullptr) {
-                LOG_WARN("%s graph cut skipping param input without buffer: segment=%s tensor=%s",
-                         log_desc == nullptr ? "unknown" : log_desc,
-                         segment.group_name.c_str(),
-                         tensor->name);
-                continue;
-            }
-            filtered_tensors.push_back(tensor);
-        }
-        return filtered_tensors;
-    }
-
     std::unordered_set<std::string> collect_future_input_names(ggml_cgraph* gf,
                                                                const Plan& plan,
                                                                size_t current_segment_index) {
@@ -487,6 +493,44 @@ namespace sd::ggml_graph_cut {
             return 0;
         }
 
+        struct TensorRuntimeBinding {
+            ggml_backend_buffer_t buffer = nullptr;
+            void* data                   = nullptr;
+            void* extra                  = nullptr;
+        };
+        std::unordered_map<ggml_tensor*, TensorRuntimeBinding> saved_bindings;
+        auto mark_measurement_external = [&](ggml_tensor* tensor) {
+            if (tensor == nullptr) {
+                return;
+            }
+            auto save_tensor = [&](ggml_tensor* t) {
+                if (t == nullptr || saved_bindings.find(t) != saved_bindings.end()) {
+                    return;
+                }
+                saved_bindings[t] = {t->buffer, t->data, t->extra};
+                // During real execution params and previous-cut inputs already
+                // have backend/cache buffers, so gallocr must not reserve them.
+                t->data = reinterpret_cast<void*>(static_cast<uintptr_t>(1));
+            };
+            save_tensor(tensor);
+            save_tensor(tensor->view_src);
+        };
+        for (const auto& input : segment.input_refs) {
+            if (input.type != Segment::INPUT_PARAM &&
+                input.type != Segment::INPUT_PREVIOUS_CUT) {
+                continue;
+            }
+            mark_measurement_external(input_tensor(gf, input));
+        }
+
+        std::unordered_map<ggml_tensor*, int32_t> saved_output_flags;
+        for (int output_node_index : segment.output_node_indices) {
+            ggml_tensor* output = ggml_graph_node(gf, output_node_index);
+            if (output != nullptr && saved_output_flags.find(output) == saved_output_flags.end()) {
+                saved_output_flags[output] = output->flags;
+            }
+        }
+
         ggml_context* graph_ctx    = nullptr;
         ggml_cgraph* segment_graph = build_segment_graph(gf, segment, &graph_ctx);
         ggml_gallocr_t allocr      = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend));
@@ -502,6 +546,14 @@ namespace sd::ggml_graph_cut {
 
         ggml_gallocr_free(allocr);
         ggml_free(graph_ctx);
+        for (const auto& kv : saved_output_flags) {
+            kv.first->flags = kv.second;
+        }
+        for (const auto& kv : saved_bindings) {
+            kv.first->buffer = kv.second.buffer;
+            kv.first->data   = kv.second.data;
+            kv.first->extra  = kv.second.extra;
+        }
         return buffer_size;
     }
 
@@ -669,7 +721,8 @@ namespace sd::ggml_graph_cut {
                 GGML_ASSERT(!candidate_plan.segments.empty());
 
                 const auto& candidate_segment = candidate_plan.segments.back();
-                if (graph_cut_segment_vram_bytes(candidate_segment) > max_graph_vram_bytes) {
+                const size_t candidate_bytes  = graph_cut_segment_vram_bytes(candidate_segment);
+                if (candidate_bytes > max_graph_vram_bytes) {
                     break;
                 }
 

src/core/ggml_graph_cut.h

@@ -80,7 +80,6 @@ namespace sd::ggml_graph_cut {
     ggml_tensor* output_tensor(ggml_cgraph* gf, const Segment& segment, size_t output_index);
     ggml_tensor* input_tensor(ggml_cgraph* gf, const Segment::InputRef& input_ref);
     std::vector<ggml_tensor*> param_tensors(ggml_cgraph* gf, const Segment& segment);
-    std::vector<ggml_tensor*> runtime_param_tensors(ggml_cgraph* gf, const Segment& segment, const char* log_desc);
     std::unordered_set<std::string> collect_future_input_names(ggml_cgraph* gf,
                                                                const Plan& plan,
                                                                size_t current_segment_index);

src/core/layer_registry.cpp

@@ -1,132 +0,0 @@
-#include "core/layer_registry.h"
-
-#include <utility>
-
-#include "core/util.h"
-
-namespace sd::layer_registry {
-
-    void LayerRegistry::register_layer(const std::string& name, ggml_tensor* tensor) {
-        auto& info = layers_[name];
-        info.tensors.push_back(tensor);
-        info.bytes += ggml_nbytes(tensor);
-    }
-
-    bool LayerRegistry::move_layer_to_gpu(const std::string& name) {
-        auto it = layers_.find(name);
-        if (it == layers_.end())
-            return false;
-
-        LayerInfo& info = it->second;
-        if (info.on_gpu)
-            return true;
-        if (gpu_backend_ == nullptr || cpu_backend_ == nullptr) {
-            LOG_ERROR("layer_registry: backends not set; cannot move '%s' to GPU",
-                      name.c_str());
-            return false;
-        }
-        if (info.tensors.empty()) {
-            info.on_gpu = true;
-            return true;
-        }
-
-        // 1. Build a no_alloc context big enough to hold one twin tensor per CPU
-        //    tensor, plus a little overhead.
-        const size_t ctx_size = info.tensors.size() * ggml_tensor_overhead() + 1024;
-        ggml_init_params ctx_params{ctx_size, /*mem_buffer=*/nullptr, /*no_alloc=*/true};
-        ggml_context* twin_ctx = ggml_init(ctx_params);
-        if (twin_ctx == nullptr) {
-            LOG_ERROR("layer_registry: failed to allocate twin context for '%s'",
-                      name.c_str());
-            return false;
-        }
-
-        // 2. Create one GPU twin per CPU tensor. The twin shares the original
-        //    name so any name-based lookup keeps working.
-        std::vector<ggml_tensor*> gpu_twins;
-        gpu_twins.reserve(info.tensors.size());
-        for (ggml_tensor* cpu_t : info.tensors) {
-            ggml_tensor* twin = ggml_dup_tensor(twin_ctx, cpu_t);
-            if (cpu_t->name[0] != '\0') {
-                ggml_set_name(twin, cpu_t->name);
-            }
-            gpu_twins.push_back(twin);
-        }
-
-        // 3. Back the twins with a GPU buffer in one alloc call.
-        ggml_backend_buffer_t gpu_buffer = ggml_backend_alloc_ctx_tensors(twin_ctx, gpu_backend_);
-        if (gpu_buffer == nullptr) {
-            LOG_ERROR("layer_registry: failed to allocate GPU buffer for '%s'",
-                      name.c_str());
-            ggml_free(twin_ctx);
-            return false;
-        }
-
-        // 4. H2D copy + sync.
-        for (size_t i = 0; i < info.tensors.size(); ++i) {
-            ggml_backend_tensor_copy(info.tensors[i], gpu_twins[i]);
-        }
-        ggml_backend_synchronize(gpu_backend_);
-
-        // 5. Swap buffer/data/extra so the originals now point at GPU memory.
-        for (size_t i = 0; i < info.tensors.size(); ++i) {
-            std::swap(info.tensors[i]->buffer, gpu_twins[i]->buffer);
-            std::swap(info.tensors[i]->data, gpu_twins[i]->data);
-            std::swap(info.tensors[i]->extra, gpu_twins[i]->extra);
-        }
-
-        info.gpu_twins  = std::move(gpu_twins);
-        info.twin_ctx   = twin_ctx;
-        info.gpu_buffer = gpu_buffer;
-        info.on_gpu     = true;
-        return true;
-    }
-
-    bool LayerRegistry::move_layer_to_cpu(const std::string& name) {
-        auto it = layers_.find(name);
-        if (it == layers_.end())
-            return false;
-
-        LayerInfo& info = it->second;
-        if (!info.on_gpu)
-            return true;
-        if (info.tensors.size() != info.gpu_twins.size()) {
-            LOG_ERROR("layer_registry: twin/tensor count mismatch for '%s'",
-                      name.c_str());
-            return false;
-        }
-
-        // 1. Swap back: originals point at CPU memory again.
-        for (size_t i = 0; i < info.tensors.size(); ++i) {
-            if (info.gpu_twins[i] == nullptr)
-                continue;
-            std::swap(info.tensors[i]->buffer, info.gpu_twins[i]->buffer);
-            std::swap(info.tensors[i]->data, info.gpu_twins[i]->data);
-            std::swap(info.tensors[i]->extra, info.gpu_twins[i]->extra);
-        }
-
-        // 2. Free the GPU buffer + twin context.
-        if (info.gpu_buffer != nullptr) {
-            ggml_backend_buffer_free(info.gpu_buffer);
-            info.gpu_buffer = nullptr;
-        }
-        if (info.twin_ctx != nullptr) {
-            ggml_free(info.twin_ctx);
-            info.twin_ctx = nullptr;
-        }
-        info.gpu_twins.clear();
-        info.on_gpu = false;
-        return true;
-    }
-
-    bool LayerRegistry::is_layer_on_gpu(const std::string& name) const {
-        auto it = layers_.find(name);
-        return it != layers_.end() && it->second.on_gpu;
-    }
-
-    size_t LayerRegistry::get_layer_size(const std::string& name) const {
-        auto it = layers_.find(name);
-        return it != layers_.end() ? it->second.bytes : 0;
-    }
-
-}  // namespace sd::layer_registry

src/core/layer_registry.h

@@ -1,50 +0,0 @@
-#ifndef __SD_CORE_LAYER_REGISTRY_H__
-#define __SD_CORE_LAYER_REGISTRY_H__
-
-#include <map>
-#include <set>
-#include <string>
-#include <vector>
-
-#include "ggml-backend.h"
-#include "ggml.h"
-
-namespace sd::layer_registry {
-
-    struct LayerInfo {
-        std::vector<ggml_tensor*> tensors;
-        std::vector<ggml_tensor*> gpu_twins;
-        ggml_context* twin_ctx           = nullptr;
-        ggml_backend_buffer_t gpu_buffer = nullptr;
-        bool on_gpu                      = false;
-        size_t bytes                     = 0;
-    };
-
-    class LayerRegistry {
-    public:
-        LayerRegistry() = default;
-        LayerRegistry(ggml_backend_t gpu_backend, ggml_backend_t cpu_backend)
-            : gpu_backend_(gpu_backend), cpu_backend_(cpu_backend) {}
-
-        void set_backends(ggml_backend_t gpu_backend, ggml_backend_t cpu_backend) {
-            gpu_backend_ = gpu_backend;
-            cpu_backend_ = cpu_backend;
-        }
-        void register_layer(const std::string& name, ggml_tensor* tensor);
-        bool move_layer_to_gpu(const std::string& name);
-        bool move_layer_to_cpu(const std::string& name);
-        bool is_layer_on_gpu(const std::string& name) const;
-        size_t get_layer_size(const std::string& name) const;
-        size_t get_layer_count() const { return layers_.size(); }
-
-        const std::map<std::string, LayerInfo>& layers() const { return layers_; }
-
-    private:
-        ggml_backend_t gpu_backend_ = nullptr;
-        ggml_backend_t cpu_backend_ = nullptr;
-        std::map<std::string, LayerInfo> layers_;
-    };
-
-}  // namespace sd::layer_registry
-
-#endif  // __SD_CORE_LAYER_REGISTRY_H__

src/core/util.cpp

@@ -488,7 +488,7 @@ bool parse_strict_bool(const std::string& text, bool& value) {
     return false;
 }
 
-static std::string build_progress_bar(int step, int steps) {
+static std::string build_progress_bar(int step, int steps, char progress_char = '=', bool show_head = true) {
     std::string progress = "  |";
     int max_progress     = 50;
     int32_t current      = 0;
@@ -498,21 +498,21 @@ static std::string build_progress_bar(int step, int steps) {
     for (int i = 0; i < 50; i++) {
         if (i > current) {
             progress += " ";
-        } else if (i == current && i != max_progress - 1) {
+        } else if (show_head && i == current && i != max_progress - 1) {
             progress += ">";
         } else {
-            progress += "=";
+            progress += progress_char;
         }
     }
     progress += "|";
     return progress;
 }
 
-static void print_progress_line(int step, int steps, const std::string& speed_text) {
+static void print_progress_line(int step, int steps, const std::string& speed_text, char progress_char = '=', bool show_head = true) {
     if (step == 0) {
         return;
     }
-    std::string progress = build_progress_bar(step, steps);
+    std::string progress = build_progress_bar(step, steps, progress_char, show_head);
     const char* lf       = (step == steps ? "\n" : "");
     printf("\r%s %i/%i - %s\033[K%s", progress.c_str(), step, steps, speed_text.c_str(), lf);
     fflush(stdout);  // for linux
@@ -552,9 +552,9 @@ void pretty_bytes_progress(int step, int steps, uint64_t bytes_processed, float
 
     double speed_mb = bytes_per_second / (1024.0 * 1024.0);
     if (speed_mb >= 1024.0) {
-        print_progress_line(step, steps, sd_format("%.2fGB/s", speed_mb / 1024.0));
+        print_progress_line(step, steps, sd_format("%.2fGB/s", speed_mb / 1024.0), '#', false);
     } else {
-        print_progress_line(step, steps, sd_format("%.2fMB/s", speed_mb));
+        print_progress_line(step, steps, sd_format("%.2fMB/s", speed_mb), '#', false);
     }
 }