DarthAffe/stable-diffusion.cpp
1
0
Fork 0
mirror of https://github.com/leejet/stable-diffusion.cpp.git synced 2026-09-24 20:20:37 +00:00
Code Issues Packages Projects Releases Wiki Activity

perf: parallelize host tensor elementwise and broadcast ops (#1998)

Browse Source
This commit is contained in:
leejet 2026-09-19 21:46:16 +08:00 committed by GitHub
parent 275ab58e01
commit 17860c0e45
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
10 changed files with 361 additions and 95 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
CMakeLists.txt
View File

@ -342,6 +342,8 @@ add_subdirectory(thirdparty)
target_sources(${SD_LIB} PRIVATE $<TARGET_OBJECTS:zip>)
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)

3
cmake/stable-diffusion-config.cmake.in
View File

@ -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"

2
cmake/stable-diffusion.pc.in
View File

@ -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}

143
src/core/parallel.cpp Normal file
View File

@ -0,0 +1,143 @@
#include "core/parallel.h"
#include <algorithm>
#include <condition_variable>
#include <exception>
#include <mutex>
#include <thread>
#include <vector>
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<std::unique_ptr<Worker>> workers;
bool stopping = false;
int pending = 0;
int participants = 1;
int64_t begin = 0;
int64_t count = 0;
const std::function<void(int64_t, int64_t)>* task = nullptr;
std::exception_ptr error;
explicit Impl(ParallelExecutor* owner)
: owner(owner) {}
~Impl() {
{
std::lock_guard<std::mutex> 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<int64_t>(index, extra);
const int64_t last = first + size + (index < extra ? 1 : 0);
try {
(*task)(first, last);
} catch (...) {
std::lock_guard<std::mutex> lock(mutex);
if (!error) {
error = std::current_exception();
}
}
}
void worker_loop(Worker* worker, int index) {
std::unique_lock<std::mutex> 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<void(int64_t, int64_t)>& callback) {
std::lock_guard<std::mutex> invocation_lock(invocation_mutex);
std::unique_lock<std::mutex> lock(mutex);
while (static_cast<int>(workers.size()) < n_tasks - 1) {
workers.push_back(std::make_unique<Worker>());
auto* worker = workers.back().get();
const int index = static_cast<int>(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<Impl>(this)) {}
ParallelExecutor::~ParallelExecutor() = default;
void ParallelExecutor::run(int64_t begin, int64_t end, int64_t grain_size, const std::function<void(int64_t, int64_t)>& 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<int>(std::min<int64_t>(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);
}
}

77
src/core/parallel.h Normal file
View File

@ -0,0 +1,77 @@
#ifndef __SD_CORE_PARALLEL_H__
#define __SD_CORE_PARALLEL_H__
#include <cstdint>
#include <functional>
#include <memory>
#include <stdexcept>
#include <utility>
namespace sd {
class ParallelExecutor {
struct Impl;
int n_threads_;
std::unique_ptr<Impl> 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<void(int64_t, int64_t)>& 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 <typename F>
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<F>(task));
}
}
#endif // __SD_CORE_PARALLEL_H__

210
src/core/tensor.hpp
View File

@ -16,6 +16,7 @@
#include <utility>
#include <vector>
#include "core/parallel.h"
#include "core/rng.hpp"
namespace sd {
@ -59,6 +60,15 @@ namespace sd {
return numel;
}
template <typename F>
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 <typename T>
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<uint8_t>& 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<int64_t>& rhs_shape_raw,
const std::vector<int64_t>& rhs_strides_raw,
F&& fn) {
const size_t ndim = out_shape.size();
std::vector<int64_t> out_strides = tensor_compute_strides(out_shape);
std::vector<int64_t> lhs_shape(ndim, 1);
std::vector<int64_t> lhs_strides(ndim, 0);
std::vector<int64_t> rhs_shape(ndim, 1);
std::vector<int64_t> 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);
for (int64_t flat = 0; flat < numel; ++flat) {
int64_t remaining = flat;
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];
const int64_t numel = tensor_numel(out_shape);
const size_t ndim = out_shape.size();
auto broadcast_strides = [&](const std::vector<int64_t>& shape,
const std::vector<int64_t>& 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<int64_t> 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 (rhs_shape[i] != 1) {
rhs_offset += coord * rhs_strides[i];
if (i < ndim && input_dim != 1) {
result[i] = strides[i];
}
}
fn(flat, lhs_offset, rhs_offset);
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;
}
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 = 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];
}
}
});
}
template <typename T>
@ -469,6 +498,7 @@ namespace sd {
const std::vector<int64_t> data_strides = tensor_compute_strides(shape_);
const std::vector<int64_t> 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<size_t>(data_offset)] = value;
data_[static_cast<size_t>(data_offset)] = fill_value;
}
});
return *this;
@ -486,9 +516,9 @@ namespace sd {
inline Tensor<uint8_t> operator<(const Tensor<T>& lhs, Scalar rhs) {
Tensor<uint8_t> result(lhs.shape());
const T value = static_cast<T>(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<uint8_t> operator<(Scalar lhs, const Tensor<T>& rhs) {
Tensor<uint8_t> result(rhs.shape());
const T value = static_cast<T>(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 <typename T>
inline Tensor<T>& operator+=(Tensor<T>& lhs, const Tensor<T>& 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<int64_t>(lhs_offset)] += rhs_data[rhs_offset];
lhs.data()[lhs_offset] += rhs_data[rhs_offset];
});
return lhs;
}
@ -547,18 +577,18 @@ namespace sd {
template <typename T, typename Scalar, typename = std::enable_if_t<std::is_arithmetic<Scalar>::value>>
inline Tensor<T>& operator+=(Tensor<T>& lhs, Scalar rhs) {
const T value = static_cast<T>(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 <typename T>
inline Tensor<T>& operator-=(Tensor<T>& lhs, const Tensor<T>& 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<int64_t>(lhs_offset)] -= rhs_data[rhs_offset];
lhs.data()[lhs_offset] -= rhs_data[rhs_offset];
});
return lhs;
}
@ -579,18 +609,18 @@ namespace sd {
template <typename T, typename Scalar, typename = std::enable_if_t<std::is_arithmetic<Scalar>::value>>
inline Tensor<T>& operator-=(Tensor<T>& lhs, Scalar rhs) {
const T value = static_cast<T>(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 <typename T>
inline Tensor<T>& operator*=(Tensor<T>& lhs, const Tensor<T>& 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<int64_t>(lhs_offset)] *= rhs_data[rhs_offset];
lhs.data()[lhs_offset] *= rhs_data[rhs_offset];
});
return lhs;
}
@ -611,18 +641,18 @@ namespace sd {
template <typename T, typename Scalar, typename = std::enable_if_t<std::is_arithmetic<Scalar>::value>>
inline Tensor<T>& operator*=(Tensor<T>& lhs, Scalar rhs) {
const T value = static_cast<T>(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 <typename T>
inline Tensor<T>& operator/=(Tensor<T>& lhs, const Tensor<T>& 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<int64_t>(lhs_offset)] /= rhs_data[rhs_offset];
lhs.data()[lhs_offset] /= rhs_data[rhs_offset];
});
return lhs;
}
@ -643,9 +673,9 @@ namespace sd {
template <typename T, typename Scalar, typename = std::enable_if_t<std::is_arithmetic<Scalar>::value>>
inline Tensor<T>& operator/=(Tensor<T>& lhs, Scalar rhs) {
const T value = static_cast<T>(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<T> operator-(Scalar lhs, const Tensor<T>& rhs) {
Tensor<T> result = rhs;
const T value = static_cast<T>(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<T> operator/(Scalar lhs, const Tensor<T>& rhs) {
Tensor<T> result = rhs;
const T value = static_cast<T>(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 <typename T>
inline Tensor<T> operator-(const Tensor<T>& tensor) {
Tensor<T> 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 <typename T>
inline Tensor<T> exp(const Tensor<T>& input) {
Tensor<T> output(input.shape());
for (int64_t i = 0; i < input.numel(); ++i) {
output[i] = static_cast<T>(std::exp(static_cast<double>(input[i])));
}
tensor_for_each(
input.numel(), [&](int64_t i) {
output.data()[i] = static_cast<T>(std::exp(static_cast<double>(input.data()[i])));
},
4096);
return output;
}
@ -1079,18 +1111,18 @@ namespace sd {
tensor_throw_invalid_argument("Tensor clamp requires min_value <= max_value");
}
Tensor<T> 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 <typename T>
inline Tensor<T> round(const Tensor<T>& input) {
Tensor<T> output(input.shape());
for (int64_t i = 0; i < input.numel(); ++i) {
output[i] = static_cast<T>(std::round(static_cast<double>(input[i])));
}
tensor_for_each(input.numel(), [&](int64_t i) {
output.data()[i] = static_cast<T>(std::round(static_cast<double>(input.data()[i])));
});
return output;
}

6
src/pipeline/diffusion_engine.cpp
View File

@ -863,8 +863,10 @@ bool StableDiffusionGGML::init(const sd_ctx_params_t* sd_ctx_params) {
return false;
}
}
auto configuration = std::make_unique<ModelConfig>(*sd_ctx_params);
n_threads = sd_ctx_params->n_threads;
auto configuration = std::make_unique<ModelConfig>(*sd_ctx_params);
n_threads = sd_ctx_params->n_threads;
tensor_executor = std::make_unique<sd::ParallelExecutor>(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;

8
src/pipeline/diffusion_engine.h
View File

@ -8,6 +8,7 @@
#include <map>
#include <memory>
#include <mutex>
#include <optional>
#include <set>
#include <string>
#include <vector>
@ -56,8 +57,9 @@ public:
std::shared_ptr<RNG> rng;
std::shared_ptr<RNG> sampler_rng = nullptr;
int n_threads = -1;
float default_flow_shift = INFINITY;
float active_flow_shift = INFINITY;
std::unique_ptr<sd::ParallelExecutor> tensor_executor;
float default_flow_shift = INFINITY;
float active_flow_shift = INFINITY;
std::shared_ptr<Conditioner> cond_stage_model;
std::shared_ptr<FrozenCLIPVisionEmbedder> clip_vision; // for svd or wan2.1 i2v
@ -206,6 +208,7 @@ public:
StableDiffusionGGML& sd;
std::unique_lock<std::recursive_mutex> lock;
bool acquired = false;
std::optional<sd::ParallelScope> tensor_scope;
explicit ContextOperation(StableDiffusionGGML& sd)
: sd(sd), lock(sd.execution_mutex, std::try_to_lock) {
@ -215,6 +218,7 @@ public:
}
sd.executing_ = true;
acquired = true;
tensor_scope.emplace(sd.tensor_executor.get());
}
~ContextOperation() {

4
src/upscaler.cpp
View File

@ -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<float> UpscalerGGML::upscale_tensor(const sd::Tensor<float>& input_tensor) {
sd::ParallelScope tensor_scope(&tensor_executor);
sd::Tensor<float> 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<float> UpscalerGGML::upscale_tensor(const sd::Tensor<float>& 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;

1
src/upscaler.h
View File

@ -17,6 +17,7 @@ struct UpscalerGGML {
std::shared_ptr<ESRGAN> 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;