From b1b15ef32322ed97c2f734f73af578fb051252d1 Mon Sep 17 00:00:00 2001 From: leejet Date: Sat, 19 Sep 2026 18:44:31 +0800 Subject: [PATCH] perf: parallelize host tensor elementwise and broadcast ops --- CMakeLists.txt | 2 + cmake/stable-diffusion-config.cmake.in | 3 +- cmake/stable-diffusion.pc.in | 2 +- src/core/parallel.cpp | 143 +++++++++++++++++ src/core/parallel.h | 77 +++++++++ src/core/tensor.hpp | 208 ++++++++++++++----------- src/pipeline/diffusion_engine.cpp | 6 +- src/pipeline/diffusion_engine.h | 8 +- src/upscaler.cpp | 4 + src/upscaler.h | 1 + 10 files changed, 360 insertions(+), 94 deletions(-) create mode 100644 src/core/parallel.cpp create mode 100644 src/core/parallel.h diff --git a/CMakeLists.txt b/CMakeLists.txt index 6991393ef..c6ebaded3 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -342,6 +342,8 @@ add_subdirectory(thirdparty) target_sources(${SD_LIB} PRIVATE $) target_link_libraries(${SD_LIB} PUBLIC ggml) +find_package(Threads REQUIRED) +target_link_libraries(${SD_LIB} PRIVATE Threads::Threads) target_link_libraries(${SD_LIB} PRIVATE onig sd-utf8proc) if (SD_CUDA) find_package(CUDAToolkit REQUIRED) diff --git a/cmake/stable-diffusion-config.cmake.in b/cmake/stable-diffusion-config.cmake.in index f3be7c031..0dfe5ef83 100644 --- a/cmake/stable-diffusion-config.cmake.in +++ b/cmake/stable-diffusion-config.cmake.in @@ -10,6 +10,7 @@ set(SD_BIN_DIR "@PACKAGE_SD_BIN_INSTALL_DIR@") include(CMakeFindDependencyMacro) find_dependency(ggml REQUIRED HINTS "${SD_LIB_DIR}/cmake") +find_dependency(Threads REQUIRED) if(@SD_CUDA@) find_dependency(CUDAToolkit REQUIRED) endif() @@ -25,7 +26,7 @@ if(NOT TARGET stable-diffusion) set_target_properties(stable-diffusion PROPERTIES INTERFACE_INCLUDE_DIRECTORIES "${SD_INCLUDE_DIR}" - INTERFACE_LINK_LIBRARIES "ggml::ggml" + INTERFACE_LINK_LIBRARIES "ggml::ggml;Threads::Threads" IMPORTED_LINK_INTERFACE_LANGUAGES "CXX" IMPORTED_LOCATION "${stable-diffusion_LIBRARY}" INTERFACE_COMPILE_FEATURES "c_std_11;cxx_std_17" diff --git a/cmake/stable-diffusion.pc.in b/cmake/stable-diffusion.pc.in index dad257b9f..ff357eee1 100644 --- a/cmake/stable-diffusion.pc.in +++ b/cmake/stable-diffusion.pc.in @@ -7,5 +7,5 @@ Name: stable-diffusion Description: Diffusion model(SD,Flux,Wan,Qwen Image,Z-Image,...) inference in pure C/C++ Version: @SDCPP_BUILD_VERSION@ Libs: -L${libdir} -lstable-diffusion -Libs.private: -lggml -lggml-base +Libs.private: -lggml -lggml-base @CMAKE_THREAD_LIBS_INIT@ Cflags: -I${includedir} diff --git a/src/core/parallel.cpp b/src/core/parallel.cpp new file mode 100644 index 000000000..a89c19755 --- /dev/null +++ b/src/core/parallel.cpp @@ -0,0 +1,143 @@ +#include "core/parallel.h" + +#include +#include +#include +#include +#include +#include + +namespace sd { + + namespace parallel_detail { + thread_local ParallelExecutor* executor = nullptr; + thread_local bool active = false; + } + + struct ParallelExecutor::Impl { + struct Worker { + std::condition_variable wake; + std::thread thread; + bool ready = false; + }; + + ParallelExecutor* owner; + std::mutex invocation_mutex; + std::mutex mutex; + std::condition_variable finished; + std::vector> workers; + bool stopping = false; + int pending = 0; + int participants = 1; + int64_t begin = 0; + int64_t count = 0; + const std::function* task = nullptr; + std::exception_ptr error; + + explicit Impl(ParallelExecutor* owner) + : owner(owner) {} + + ~Impl() { + { + std::lock_guard lock(mutex); + stopping = true; + } + for (auto& worker : workers) { + worker->wake.notify_one(); + } + for (auto& worker : workers) { + worker->thread.join(); + } + } + + void execute(int index) { + ParallelScope scope(owner); + parallel_detail::Region region; + const int64_t size = count / participants; + const int64_t extra = count % participants; + const int64_t first = begin + index * size + std::min(index, extra); + const int64_t last = first + size + (index < extra ? 1 : 0); + try { + (*task)(first, last); + } catch (...) { + std::lock_guard lock(mutex); + if (!error) { + error = std::current_exception(); + } + } + } + + void worker_loop(Worker* worker, int index) { + std::unique_lock lock(mutex); + for (;;) { + worker->wake.wait(lock, [&] { return stopping || worker->ready; }); + if (stopping) { + return; + } + worker->ready = false; + lock.unlock(); + execute(index); + lock.lock(); + if (--pending == 0) { + finished.notify_one(); + } + } + } + + void run(int64_t first, int64_t last, int n_tasks, const std::function& callback) { + std::lock_guard invocation_lock(invocation_mutex); + std::unique_lock lock(mutex); + while (static_cast(workers.size()) < n_tasks - 1) { + workers.push_back(std::make_unique()); + auto* worker = workers.back().get(); + const int index = static_cast(workers.size()); + try { + worker->thread = std::thread([this, worker, index] { worker_loop(worker, index); }); + } catch (...) { + workers.pop_back(); + throw; + } + } + begin = first; + count = last - first; + participants = n_tasks; + pending = n_tasks - 1; + task = &callback; + error = nullptr; + for (int i = 0; i < pending; ++i) { + workers[i]->ready = true; + workers[i]->wake.notify_one(); + } + lock.unlock(); + execute(0); + lock.lock(); + finished.wait(lock, [&] { return pending == 0; }); + task = nullptr; + if (error) { + std::rethrow_exception(error); + } + } + }; + + ParallelExecutor::ParallelExecutor(int n_threads) + : n_threads_(std::max(1, n_threads)), impl_(std::make_unique(this)) {} + + ParallelExecutor::~ParallelExecutor() = default; + + void ParallelExecutor::run(int64_t begin, int64_t end, int64_t grain_size, const std::function& task) { + if (begin < 0 || grain_size <= 0) { + throw std::invalid_argument("parallel_for requires begin >= 0 and grain_size > 0"); + } + if (end <= begin) { + return; + } + const int n_tasks = static_cast(std::min(n_threads_, (end - begin) / grain_size)); + if (parallel_detail::active || n_tasks <= 1) { + parallel_detail::Region region; + task(begin, end); + return; + } + impl_->run(begin, end, n_tasks, task); + } + +} diff --git a/src/core/parallel.h b/src/core/parallel.h new file mode 100644 index 000000000..3a2c98751 --- /dev/null +++ b/src/core/parallel.h @@ -0,0 +1,77 @@ +#ifndef __SD_CORE_PARALLEL_H__ +#define __SD_CORE_PARALLEL_H__ + +#include +#include +#include +#include +#include + +namespace sd { + + class ParallelExecutor { + struct Impl; + int n_threads_; + std::unique_ptr impl_; + + public: + explicit ParallelExecutor(int n_threads); + ~ParallelExecutor(); + ParallelExecutor(const ParallelExecutor&) = delete; + ParallelExecutor& operator=(const ParallelExecutor&) = delete; + + int num_threads() const { return n_threads_; } + void run(int64_t begin, int64_t end, int64_t grain_size, const std::function& task); + }; + + namespace parallel_detail { + extern thread_local ParallelExecutor* executor; + extern thread_local bool active; + + class Region { + bool previous_; + + public: + Region() + : previous_(active) { active = true; } + ~Region() { active = previous_; } + Region(const Region&) = delete; + Region& operator=(const Region&) = delete; + }; + } + + class ParallelScope { + ParallelExecutor* previous_; + + public: + explicit ParallelScope(ParallelExecutor* executor) + : previous_(parallel_detail::executor) { + parallel_detail::executor = executor; + } + ~ParallelScope() { parallel_detail::executor = previous_; } + ParallelScope(const ParallelScope&) = delete; + ParallelScope& operator=(const ParallelScope&) = delete; + }; + + // Ranges are non-negative. The callback may run concurrently and must own its writes. + template + inline void parallel_for(int64_t begin, int64_t end, int64_t grain_size, F&& task) { + if (begin < 0 || grain_size <= 0) { + throw std::invalid_argument("parallel_for requires begin >= 0 and grain_size > 0"); + } + if (end <= begin) { + return; + } + auto* executor = parallel_detail::executor; + if (parallel_detail::active || executor == nullptr || executor->num_threads() <= 1 || + (end - begin) / grain_size < 2) { + parallel_detail::Region region; + task(begin, end); + return; + } + executor->run(begin, end, grain_size, std::forward(task)); + } + +} + +#endif // __SD_CORE_PARALLEL_H__ diff --git a/src/core/tensor.hpp b/src/core/tensor.hpp index ba5dc137a..a925dbffb 100644 --- a/src/core/tensor.hpp +++ b/src/core/tensor.hpp @@ -16,6 +16,7 @@ #include #include +#include "core/parallel.h" #include "core/rng.hpp" namespace sd { @@ -59,6 +60,15 @@ namespace sd { return numel; } + template + inline void tensor_for_each(int64_t count, F&& fn, int64_t grain_size = 65536) { + parallel_for(0, count, grain_size, [&](int64_t begin, int64_t end) { + for (int64_t i = begin; i < end; ++i) { + fn(i); + } + }); + } + template class Tensor { public: @@ -230,7 +240,10 @@ namespace sd { } void fill_(const T& value) { - std::fill(data_.begin(), data_.end(), value); + const T fill_value = value; + parallel_for(0, numel(), 65536, [&](int64_t begin, int64_t end) { + std::fill_n(data_.data() + begin, end - begin, fill_value); + }); } Tensor& masked_fill_(const Tensor& mask, const T& value); @@ -390,7 +403,7 @@ namespace sd { tensor_shape_to_string(lhs) + ", rhs_shape=" + tensor_shape_to_string(rhs)); } - shape[i] = std::max(lhs_dim, rhs_dim); + shape[i] = lhs_dim == 1 ? rhs_dim : lhs_dim; } return shape; } @@ -425,39 +438,55 @@ namespace sd { const std::vector& rhs_shape_raw, const std::vector& rhs_strides_raw, F&& fn) { - const size_t ndim = out_shape.size(); - std::vector out_strides = tensor_compute_strides(out_shape); - std::vector lhs_shape(ndim, 1); - std::vector lhs_strides(ndim, 0); - std::vector rhs_shape(ndim, 1); - std::vector rhs_strides(ndim, 0); - - for (size_t i = 0; i < lhs_shape_raw.size(); ++i) { - lhs_shape[i] = lhs_shape_raw[i]; - lhs_strides[i] = lhs_strides_raw[i]; - } - for (size_t i = 0; i < rhs_shape_raw.size(); ++i) { - rhs_shape[i] = rhs_shape_raw[i]; - rhs_strides[i] = rhs_strides_raw[i]; + const int64_t numel = tensor_numel(out_shape); + const size_t ndim = out_shape.size(); + auto broadcast_strides = [&](const std::vector& shape, + const std::vector& strides) { + if ((numel != 0 && tensor_numel(shape) == 0) || strides.size() != shape.size()) { + tensor_throw_invalid_argument("Tensor broadcast requires non-empty inputs and matching strides"); + } + std::vector result(ndim, 0); + for (size_t i = 0; i < std::max(ndim, shape.size()); ++i) { + const int64_t input_dim = i < shape.size() ? shape[i] : 1; + const int64_t output_dim = i < ndim ? out_shape[i] : 1; + if (input_dim != 1 && input_dim != output_dim) { + tensor_throw_invalid_argument("Tensor broadcast cannot expand the destination: input_shape=" + + tensor_shape_to_string(shape) + ", output_shape=" + + tensor_shape_to_string(out_shape)); + } + if (i < ndim && input_dim != 1) { + result[i] = strides[i]; + } + } + return result; + }; + const auto lhs_strides = broadcast_strides(lhs_shape_raw, lhs_strides_raw); + const auto rhs_strides = broadcast_strides(rhs_shape_raw, rhs_strides_raw); + if (numel == 0) { + return; } - - const int64_t numel = tensor_numel(out_shape); - for (int64_t flat = 0; flat < numel; ++flat) { - int64_t remaining = flat; + parallel_for(0, numel, 16384, [&](int64_t begin, int64_t end) { + auto coord = tensor_unravel_index(begin, out_shape); int64_t lhs_offset = 0; int64_t rhs_offset = 0; - for (size_t i = ndim; i-- > 0;) { - int64_t coord = remaining / out_strides[i]; - remaining %= out_strides[i]; - if (lhs_shape[i] != 1) { - lhs_offset += coord * lhs_strides[i]; - } - if (rhs_shape[i] != 1) { - rhs_offset += coord * rhs_strides[i]; + for (size_t i = 0; i < ndim; ++i) { + lhs_offset += coord[i] * lhs_strides[i]; + rhs_offset += coord[i] * rhs_strides[i]; + } + for (int64_t flat = begin; flat < end; ++flat) { + fn(flat, lhs_offset, rhs_offset); + for (size_t i = 0; i < ndim; ++i) { + lhs_offset += lhs_strides[i]; + rhs_offset += rhs_strides[i]; + if (++coord[i] < out_shape[i]) { + break; + } + coord[i] = 0; + lhs_offset -= out_shape[i] * lhs_strides[i]; + rhs_offset -= out_shape[i] * rhs_strides[i]; } } - fn(flat, lhs_offset, rhs_offset); - } + }); } template @@ -469,6 +498,7 @@ namespace sd { const std::vector data_strides = tensor_compute_strides(shape_); const std::vector mask_strides = tensor_compute_strides(mask.shape()); const uint8_t* mask_data = mask.data(); + const T fill_value = value; tensor_for_each_broadcast_offset(shape_, shape_, data_strides, @@ -476,7 +506,7 @@ namespace sd { mask_strides, [&](int64_t, int64_t data_offset, int64_t mask_offset) { if (mask_data[mask_offset] != 0) { - data_[static_cast(data_offset)] = value; + data_[static_cast(data_offset)] = fill_value; } }); return *this; @@ -486,9 +516,9 @@ namespace sd { inline Tensor operator<(const Tensor& lhs, Scalar rhs) { Tensor result(lhs.shape()); const T value = static_cast(rhs); - for (int64_t i = 0; i < lhs.numel(); ++i) { - result[i] = lhs[i] < value ? 1 : 0; - } + tensor_for_each(lhs.numel(), [&](int64_t i) { + result.data()[i] = lhs.data()[i] < value ? 1 : 0; + }); return result; } @@ -496,9 +526,9 @@ namespace sd { inline Tensor operator<(Scalar lhs, const Tensor& rhs) { Tensor result(rhs.shape()); const T value = static_cast(lhs); - for (int64_t i = 0; i < rhs.numel(); ++i) { - result[i] = value < rhs[i] ? 1 : 0; - } + tensor_for_each(rhs.numel(), [&](int64_t i) { + result.data()[i] = value < rhs.data()[i] ? 1 : 0; + }); return result; } @@ -516,7 +546,7 @@ namespace sd { rhs.shape(), rhs_strides, [&](int64_t flat, int64_t lhs_offset, int64_t rhs_offset) { - result[flat] = lhs_data[lhs_offset] < rhs_data[rhs_offset] ? 1 : 0; + result.data()[flat] = lhs_data[lhs_offset] < rhs_data[rhs_offset] ? 1 : 0; }); return result; } @@ -524,9 +554,9 @@ namespace sd { template inline Tensor& operator+=(Tensor& lhs, const Tensor& rhs) { if (lhs.shape() == rhs.shape()) { - for (int64_t i = 0; i < lhs.numel(); ++i) { - lhs[i] += rhs[i]; - } + tensor_for_each(lhs.numel(), [&](int64_t i) { + lhs.data()[i] += rhs.data()[i]; + }); return lhs; } tensor_broadcast_shape(lhs.shape(), rhs.shape()); @@ -539,7 +569,7 @@ namespace sd { rhs.shape(), rhs_strides, [&](int64_t, int64_t lhs_offset, int64_t rhs_offset) { - lhs[static_cast(lhs_offset)] += rhs_data[rhs_offset]; + lhs.data()[lhs_offset] += rhs_data[rhs_offset]; }); return lhs; } @@ -547,18 +577,18 @@ namespace sd { template ::value>> inline Tensor& operator+=(Tensor& lhs, Scalar rhs) { const T value = static_cast(rhs); - for (int64_t i = 0; i < lhs.numel(); ++i) { - lhs[i] += value; - } + tensor_for_each(lhs.numel(), [&](int64_t i) { + lhs.data()[i] += value; + }); return lhs; } template inline Tensor& operator-=(Tensor& lhs, const Tensor& rhs) { if (lhs.shape() == rhs.shape()) { - for (int64_t i = 0; i < lhs.numel(); ++i) { - lhs[i] -= rhs[i]; - } + tensor_for_each(lhs.numel(), [&](int64_t i) { + lhs.data()[i] -= rhs.data()[i]; + }); return lhs; } tensor_broadcast_shape(lhs.shape(), rhs.shape()); @@ -571,7 +601,7 @@ namespace sd { rhs.shape(), rhs_strides, [&](int64_t, int64_t lhs_offset, int64_t rhs_offset) { - lhs[static_cast(lhs_offset)] -= rhs_data[rhs_offset]; + lhs.data()[lhs_offset] -= rhs_data[rhs_offset]; }); return lhs; } @@ -579,18 +609,18 @@ namespace sd { template ::value>> inline Tensor& operator-=(Tensor& lhs, Scalar rhs) { const T value = static_cast(rhs); - for (int64_t i = 0; i < lhs.numel(); ++i) { - lhs[i] -= value; - } + tensor_for_each(lhs.numel(), [&](int64_t i) { + lhs.data()[i] -= value; + }); return lhs; } template inline Tensor& operator*=(Tensor& lhs, const Tensor& rhs) { if (lhs.shape() == rhs.shape()) { - for (int64_t i = 0; i < lhs.numel(); ++i) { - lhs[i] *= rhs[i]; - } + tensor_for_each(lhs.numel(), [&](int64_t i) { + lhs.data()[i] *= rhs.data()[i]; + }); return lhs; } tensor_broadcast_shape(lhs.shape(), rhs.shape()); @@ -603,7 +633,7 @@ namespace sd { rhs.shape(), rhs_strides, [&](int64_t, int64_t lhs_offset, int64_t rhs_offset) { - lhs[static_cast(lhs_offset)] *= rhs_data[rhs_offset]; + lhs.data()[lhs_offset] *= rhs_data[rhs_offset]; }); return lhs; } @@ -611,18 +641,18 @@ namespace sd { template ::value>> inline Tensor& operator*=(Tensor& lhs, Scalar rhs) { const T value = static_cast(rhs); - for (int64_t i = 0; i < lhs.numel(); ++i) { - lhs[i] *= value; - } + tensor_for_each(lhs.numel(), [&](int64_t i) { + lhs.data()[i] *= value; + }); return lhs; } template inline Tensor& operator/=(Tensor& lhs, const Tensor& rhs) { if (lhs.shape() == rhs.shape()) { - for (int64_t i = 0; i < lhs.numel(); ++i) { - lhs[i] /= rhs[i]; - } + tensor_for_each(lhs.numel(), [&](int64_t i) { + lhs.data()[i] /= rhs.data()[i]; + }); return lhs; } tensor_broadcast_shape(lhs.shape(), rhs.shape()); @@ -635,7 +665,7 @@ namespace sd { rhs.shape(), rhs_strides, [&](int64_t, int64_t lhs_offset, int64_t rhs_offset) { - lhs[static_cast(lhs_offset)] /= rhs_data[rhs_offset]; + lhs.data()[lhs_offset] /= rhs_data[rhs_offset]; }); return lhs; } @@ -643,9 +673,9 @@ namespace sd { template ::value>> inline Tensor& operator/=(Tensor& lhs, Scalar rhs) { const T value = static_cast(rhs); - for (int64_t i = 0; i < lhs.numel(); ++i) { - lhs[i] /= value; - } + tensor_for_each(lhs.numel(), [&](int64_t i) { + lhs.data()[i] /= value; + }); return lhs; } @@ -664,7 +694,7 @@ namespace sd { rhs.shape(), rhs_strides, [&](int64_t flat, int64_t lhs_offset, int64_t rhs_offset) { - result[flat] = lhs_data[lhs_offset] + rhs_data[rhs_offset]; + result.data()[flat] = lhs_data[lhs_offset] + rhs_data[rhs_offset]; }); return result; } @@ -699,7 +729,7 @@ namespace sd { rhs.shape(), rhs_strides, [&](int64_t flat, int64_t lhs_offset, int64_t rhs_offset) { - result[flat] = lhs_data[lhs_offset] - rhs_data[rhs_offset]; + result.data()[flat] = lhs_data[lhs_offset] - rhs_data[rhs_offset]; }); return result; } @@ -717,9 +747,9 @@ namespace sd { inline Tensor operator-(Scalar lhs, const Tensor& rhs) { Tensor result = rhs; const T value = static_cast(lhs); - for (int64_t i = 0; i < result.numel(); ++i) { - result[i] = value - result[i]; - } + tensor_for_each(result.numel(), [&](int64_t i) { + result.data()[i] = value - result.data()[i]; + }); return result; } @@ -738,7 +768,7 @@ namespace sd { rhs.shape(), rhs_strides, [&](int64_t flat, int64_t lhs_offset, int64_t rhs_offset) { - result[flat] = lhs_data[lhs_offset] * rhs_data[rhs_offset]; + result.data()[flat] = lhs_data[lhs_offset] * rhs_data[rhs_offset]; }); return result; } @@ -773,7 +803,7 @@ namespace sd { rhs.shape(), rhs_strides, [&](int64_t flat, int64_t lhs_offset, int64_t rhs_offset) { - result[flat] = lhs_data[lhs_offset] / rhs_data[rhs_offset]; + result.data()[flat] = lhs_data[lhs_offset] / rhs_data[rhs_offset]; }); return result; } @@ -791,18 +821,18 @@ namespace sd { inline Tensor operator/(Scalar lhs, const Tensor& rhs) { Tensor result = rhs; const T value = static_cast(lhs); - for (int64_t i = 0; i < result.numel(); ++i) { - result[i] = value / result[i]; - } + tensor_for_each(result.numel(), [&](int64_t i) { + result.data()[i] = value / result.data()[i]; + }); return result; } template inline Tensor operator-(const Tensor& tensor) { Tensor result = tensor; - for (int64_t i = 0; i < result.numel(); ++i) { - result[i] = -result[i]; - } + tensor_for_each(result.numel(), [&](int64_t i) { + result.data()[i] = -result.data()[i]; + }); return result; } @@ -1067,9 +1097,11 @@ namespace sd { template inline Tensor exp(const Tensor& input) { Tensor output(input.shape()); - for (int64_t i = 0; i < input.numel(); ++i) { - output[i] = static_cast(std::exp(static_cast(input[i]))); - } + tensor_for_each( + input.numel(), [&](int64_t i) { + output.data()[i] = static_cast(std::exp(static_cast(input.data()[i]))); + }, + 4096); return output; } @@ -1079,18 +1111,18 @@ namespace sd { tensor_throw_invalid_argument("Tensor clamp requires min_value <= max_value"); } Tensor output(input.shape()); - for (int64_t i = 0; i < input.numel(); ++i) { - output[i] = std::clamp(input[i], min_value, max_value); - } + tensor_for_each(input.numel(), [&](int64_t i) { + output.data()[i] = std::clamp(input.data()[i], min_value, max_value); + }); return output; } template inline Tensor round(const Tensor& input) { Tensor output(input.shape()); - for (int64_t i = 0; i < input.numel(); ++i) { - output[i] = static_cast(std::round(static_cast(input[i]))); - } + tensor_for_each(input.numel(), [&](int64_t i) { + output.data()[i] = static_cast(std::round(static_cast(input.data()[i]))); + }); return output; } diff --git a/src/pipeline/diffusion_engine.cpp b/src/pipeline/diffusion_engine.cpp index 4fb90807e..dabd795cc 100644 --- a/src/pipeline/diffusion_engine.cpp +++ b/src/pipeline/diffusion_engine.cpp @@ -863,8 +863,10 @@ bool StableDiffusionGGML::init(const sd_ctx_params_t* sd_ctx_params) { return false; } } - auto configuration = std::make_unique(*sd_ctx_params); - n_threads = sd_ctx_params->n_threads; + auto configuration = std::make_unique(*sd_ctx_params); + n_threads = sd_ctx_params->n_threads; + tensor_executor = std::make_unique(n_threads > 0 ? n_threads : sd_get_num_physical_cores()); + sd::ParallelScope tensor_scope(tensor_executor.get()); enable_mmap = sd_ctx_params->enable_mmap; disable_prefetch = sd_ctx_params->disable_prefetch; disable_segmented_compute = sd_ctx_params->disable_segmented_compute; diff --git a/src/pipeline/diffusion_engine.h b/src/pipeline/diffusion_engine.h index a3c530b21..682825fc3 100644 --- a/src/pipeline/diffusion_engine.h +++ b/src/pipeline/diffusion_engine.h @@ -8,6 +8,7 @@ #include #include #include +#include #include #include #include @@ -56,8 +57,9 @@ class StableDiffusionGGML { std::shared_ptr rng; std::shared_ptr sampler_rng = nullptr; int n_threads = -1; - float default_flow_shift = INFINITY; - float active_flow_shift = INFINITY; + std::unique_ptr tensor_executor; + float default_flow_shift = INFINITY; + float active_flow_shift = INFINITY; std::shared_ptr cond_stage_model; std::shared_ptr clip_vision; // for svd or wan2.1 i2v @@ -206,6 +208,7 @@ class StableDiffusionGGML { StableDiffusionGGML& sd; std::unique_lock lock; bool acquired = false; + std::optional tensor_scope; explicit ContextOperation(StableDiffusionGGML& sd) : sd(sd), lock(sd.execution_mutex, std::try_to_lock) { @@ -215,6 +218,7 @@ class StableDiffusionGGML { } sd.executing_ = true; acquired = true; + tensor_scope.emplace(sd.tensor_executor.get()); } ~ContextOperation() { diff --git a/src/upscaler.cpp b/src/upscaler.cpp index 6b7b0efb0..341c76ecf 100644 --- a/src/upscaler.cpp +++ b/src/upscaler.cpp @@ -14,6 +14,7 @@ UpscalerGGML::UpscalerGGML(int n_threads, std::string backend_spec, std::string params_backend_spec) : n_threads(n_threads), + tensor_executor(n_threads > 0 ? n_threads : sd_get_num_physical_cores()), direct(direct), tile_size(tile_size), backend_spec(std::move(backend_spec)), @@ -35,6 +36,7 @@ void UpscalerGGML::set_max_graph_vram_bytes(size_t max_vram_bytes) { bool UpscalerGGML::load_from_file(const std::string& esrgan_path, int n_threads) { + sd::ParallelScope tensor_scope(&tensor_executor); ggml_log_set(sd_ggml_log_callback, nullptr); std::string error; @@ -108,6 +110,7 @@ bool UpscalerGGML::load_from_file(const std::string& esrgan_path, } sd::Tensor UpscalerGGML::upscale_tensor(const sd::Tensor& input_tensor) { + sd::ParallelScope tensor_scope(&tensor_executor); sd::Tensor upscaled; const int scale = esrgan_upscaler->config.scale; if (tile_size <= 0 || (input_tensor.shape()[0] <= tile_size && input_tensor.shape()[1] <= tile_size)) { @@ -142,6 +145,7 @@ sd::Tensor UpscalerGGML::upscale_tensor(const sd::Tensor& input_te } sd_image_t UpscalerGGML::upscale(sd_image_t input_image, uint32_t upscale_factor) { + sd::ParallelScope tensor_scope(&tensor_executor); // upscale_factor, unused for RealESRGAN_x4plus_anime_6B.pth sd_image_t upscaled_image = {0, 0, 0, nullptr}; const int scale = esrgan_upscaler->config.scale; diff --git a/src/upscaler.h b/src/upscaler.h index 867f64440..c23e9d0f2 100644 --- a/src/upscaler.h +++ b/src/upscaler.h @@ -17,6 +17,7 @@ struct UpscalerGGML { std::shared_ptr esrgan_upscaler; std::string esrgan_path; int n_threads; + sd::ParallelExecutor tensor_executor; bool direct = false; int tile_size = 128; size_t max_graph_vram_bytes = 0;