mirror of
https://github.com/leejet/stable-diffusion.cpp.git
synced 2026-07-25 20:30:51 -05:00
222 lines
9.8 KiB
C++
222 lines
9.8 KiB
C++
#include "core/layer_split_partition.h"
|
|
|
|
#include <algorithm>
|
|
#include <cstdint>
|
|
#include <cstdlib>
|
|
#include <cstring>
|
|
|
|
#include "core/util.h"
|
|
|
|
namespace sd {
|
|
|
|
static bool layer_split_path_segment_starts_at(const std::string& name, size_t pos) {
|
|
return pos == 0 || name[pos - 1] == '.';
|
|
}
|
|
|
|
static bool layer_split_has_path_segment(const std::string& name, const char* segment) {
|
|
size_t pos = name.find(segment);
|
|
while (pos != std::string::npos) {
|
|
if (layer_split_path_segment_starts_at(name, pos)) {
|
|
return true;
|
|
}
|
|
pos = name.find(segment, pos + 1);
|
|
}
|
|
return false;
|
|
}
|
|
|
|
int layer_split_tensor_block_index(const std::string& name) {
|
|
static const char* unet_block_segments[] = {"input_blocks.", "output_blocks.", "middle_block.",
|
|
"down_blocks.", "up_blocks.", "mid_block."};
|
|
for (const char* segment : unet_block_segments) {
|
|
if (layer_split_has_path_segment(name, segment)) {
|
|
return -1;
|
|
}
|
|
}
|
|
|
|
static const char* block_keywords[] = {"transformer_blocks.", "joint_blocks.", "double_blocks.",
|
|
"single_blocks.", "blocks.", "block.", "layers."};
|
|
for (const char* keyword : block_keywords) {
|
|
size_t pos = name.find(keyword);
|
|
while (pos != std::string::npos) {
|
|
if (!layer_split_path_segment_starts_at(name, pos)) {
|
|
pos = name.find(keyword, pos + 1);
|
|
continue;
|
|
}
|
|
pos += std::strlen(keyword);
|
|
size_t end = pos;
|
|
while (end < name.size() && name[end] >= '0' && name[end] <= '9') {
|
|
end++;
|
|
}
|
|
if (end > pos && (end == name.size() || name[end] == '.')) {
|
|
return std::atoi(name.substr(pos, end - pos).c_str());
|
|
}
|
|
break;
|
|
}
|
|
}
|
|
return -1;
|
|
}
|
|
|
|
std::string layer_split_backend_device_display_name(ggml_backend_t backend) {
|
|
ggml_backend_dev_t dev = ggml_backend_get_device(backend);
|
|
const char* name = dev != nullptr ? ggml_backend_dev_name(dev) : ggml_backend_name(backend);
|
|
return name != nullptr ? name : "unknown";
|
|
}
|
|
|
|
static bool layer_split_backend_supports_tensor(ggml_backend_t backend, const ggml_tensor* tensor) {
|
|
return backend != nullptr && tensor != nullptr && ggml_backend_supports_op(backend, tensor);
|
|
}
|
|
|
|
static size_t layer_split_supported_target(const std::string& desc,
|
|
const std::string& tensor_name,
|
|
const ggml_tensor* tensor,
|
|
const std::vector<ggml_backend_t>& backends,
|
|
size_t preferred) {
|
|
if (tensor == nullptr || backends.empty()) {
|
|
return preferred;
|
|
}
|
|
size_t preferred_safe = std::min(preferred, backends.size() - 1);
|
|
if (layer_split_backend_supports_tensor(backends[preferred_safe], tensor)) {
|
|
return preferred_safe;
|
|
}
|
|
for (size_t i = 0; i < backends.size(); i++) {
|
|
if (layer_split_backend_supports_tensor(backends[i], tensor)) {
|
|
LOG_WARN("%s layer split: moving tensor '%s' from %s to %s because the preferred backend cannot run op=%s type=%s nbytes=%.2f MB",
|
|
desc.c_str(),
|
|
tensor_name.c_str(),
|
|
layer_split_backend_device_display_name(backends[preferred_safe]).c_str(),
|
|
layer_split_backend_device_display_name(backends[i]).c_str(),
|
|
ggml_op_name(tensor->op),
|
|
ggml_type_name(tensor->type),
|
|
ggml_nbytes(tensor) / (1024.0 * 1024.0));
|
|
return i;
|
|
}
|
|
}
|
|
LOG_WARN("%s layer split: tensor '%s' is not supported by any split backend: op=%s type=%s nbytes=%.2f MB",
|
|
desc.c_str(),
|
|
tensor_name.c_str(),
|
|
ggml_op_name(tensor->op),
|
|
ggml_type_name(tensor->type),
|
|
ggml_nbytes(tensor) / (1024.0 * 1024.0));
|
|
return preferred_safe;
|
|
}
|
|
|
|
std::vector<std::map<std::string, ggml_tensor*>> partition_layer_split_tensors(
|
|
const std::string& desc,
|
|
const std::map<std::string, ggml_tensor*>& tensors,
|
|
const std::map<std::string, ggml_tensor*>& split_tensors,
|
|
const std::vector<ggml_backend_t>& backends) {
|
|
std::vector<std::map<std::string, ggml_tensor*>> partitions(backends.size());
|
|
if (backends.empty()) {
|
|
LOG_WARN("%s: no backend available for a layer split", desc.c_str());
|
|
return partitions;
|
|
}
|
|
|
|
std::map<int, int64_t> block_bytes;
|
|
std::map<std::string, size_t> non_block_targets;
|
|
std::vector<int64_t> other_bytes_by_backend(backends.size(), 0);
|
|
int64_t total_block_bytes = 0;
|
|
int64_t total_other_bytes = 0;
|
|
int n_blocks = 0;
|
|
for (const auto& kv : tensors) {
|
|
int64_t bytes = (int64_t)ggml_nbytes(kv.second);
|
|
int idx = split_tensors.count(kv.first) != 0 ? layer_split_tensor_block_index(kv.first) : -1;
|
|
if (idx >= 0) {
|
|
block_bytes[idx] += bytes;
|
|
total_block_bytes += bytes;
|
|
n_blocks = std::max(n_blocks, idx + 1);
|
|
} else {
|
|
size_t target = layer_split_supported_target(desc, kv.first, kv.second, backends, 0);
|
|
non_block_targets[kv.first] = target;
|
|
other_bytes_by_backend[target] += bytes;
|
|
total_other_bytes += bytes;
|
|
}
|
|
}
|
|
if (n_blocks == 0) {
|
|
LOG_WARN("%s: no transformer blocks found for a layer split; keeping tensors on compatible backends starting from %s",
|
|
desc.c_str(),
|
|
layer_split_backend_device_display_name(backends[0]).c_str());
|
|
for (const auto& kv : tensors) {
|
|
size_t target = 0;
|
|
auto target_it = non_block_targets.find(kv.first);
|
|
if (target_it != non_block_targets.end()) {
|
|
target = target_it->second;
|
|
}
|
|
partitions[target][kv.first] = kv.second;
|
|
}
|
|
return partitions;
|
|
}
|
|
|
|
// Reserve compute headroom and subtract each device's actual non-block
|
|
// bytes from its block budget.
|
|
constexpr int64_t compute_headroom_bytes = 2ll * 1024 * 1024 * 1024;
|
|
std::vector<double> device_weights(backends.size(), 1.0);
|
|
double weight_sum = 0.0;
|
|
for (size_t i = 0; i < backends.size(); i++) {
|
|
ggml_backend_dev_t dev = ggml_backend_get_device(backends[i]);
|
|
size_t free_bytes = 0, total_bytes = 0;
|
|
if (dev != nullptr) {
|
|
ggml_backend_dev_memory(dev, &free_bytes, &total_bytes);
|
|
}
|
|
// Keep a small share even for tight devices instead of dropping them.
|
|
int64_t usable_bytes = std::max<int64_t>((int64_t)free_bytes - compute_headroom_bytes,
|
|
(int64_t)free_bytes / 8);
|
|
device_weights[i] = usable_bytes > 0 ? (double)usable_bytes : 1.0;
|
|
weight_sum += device_weights[i];
|
|
}
|
|
|
|
std::vector<int64_t> block_budgets(backends.size(), 0);
|
|
const int64_t total_bytes = total_block_bytes + total_other_bytes;
|
|
for (size_t i = 0; i < backends.size(); i++) {
|
|
int64_t budget = (int64_t)((double)total_bytes * device_weights[i] / weight_sum);
|
|
budget = std::max<int64_t>(budget - other_bytes_by_backend[i], 0);
|
|
block_budgets[i] = budget;
|
|
}
|
|
|
|
std::vector<int> boundaries(backends.size(), n_blocks);
|
|
size_t current = 0;
|
|
int64_t used = 0;
|
|
for (int b = 0; b < n_blocks; b++) {
|
|
int64_t bytes = block_bytes.count(b) != 0 ? block_bytes[b] : 0;
|
|
if (current + 1 < backends.size() && used > 0 && used + bytes > block_budgets[current]) {
|
|
boundaries[current] = b;
|
|
current++;
|
|
used = 0;
|
|
}
|
|
used += bytes;
|
|
}
|
|
|
|
for (const auto& kv : tensors) {
|
|
size_t target = 0;
|
|
int idx = split_tensors.count(kv.first) != 0 ? layer_split_tensor_block_index(kv.first) : -1;
|
|
if (idx >= 0) {
|
|
while (target < boundaries.size() && idx >= boundaries[target]) {
|
|
target++;
|
|
}
|
|
target = std::min(target, backends.size() - 1);
|
|
target = layer_split_supported_target(desc, kv.first, kv.second, backends, target);
|
|
} else {
|
|
auto target_it = non_block_targets.find(kv.first);
|
|
if (target_it != non_block_targets.end()) {
|
|
target = target_it->second;
|
|
}
|
|
}
|
|
partitions[target][kv.first] = kv.second;
|
|
}
|
|
|
|
int range_start = 0;
|
|
for (size_t i = 0; i < backends.size(); i++) {
|
|
int range_end = boundaries[i];
|
|
const char* non_block_suffix = other_bytes_by_backend[i] > 0 ? " + non-block tensors" : "";
|
|
LOG_INFO("%s layer split: %s <- blocks [%d, %d)%s",
|
|
desc.c_str(),
|
|
layer_split_backend_device_display_name(backends[i]).c_str(),
|
|
range_start,
|
|
range_end,
|
|
non_block_suffix);
|
|
range_start = range_end;
|
|
}
|
|
return partitions;
|
|
}
|
|
|
|
} // namespace sd
|