diff --git a/Makefile b/Makefile index 1139cfe55..0743fe8c3 100644 --- a/Makefile +++ b/Makefile @@ -701,7 +701,7 @@ budget.o: common/reasoning-budget.cpp common/reasoning-budget.h chat.o: common/chat.cpp common/chat.h $(CXX) $(CXXFLAGS) -c $< -o $@ -SDCPP_COMMON_BASENAMES := include/stable-diffusion.h src/conditioning/conditioner.hpp src/core/backend_fit.cpp src/core/backend_fit.h src/core/ggml_extend_backend.cpp src/core/ggml_extend_backend.h src/core/ggml_extend.hpp src/core/ggml_graph_cut.cpp src/core/ggml_graph_cut.h src/core/layer_split_partition.cpp src/core/layer_split_partition.h src/core/ordered_map.hpp src/core/rng.hpp src/core/rng_mt19937.hpp src/core/rng_philox.hpp src/core/tensor_ggml.hpp src/core/tensor.hpp src/core/util.cpp src/core/util.h src/detailer.cpp src/detailer.h src/extensions/generation_extension.h src/extensions/photomaker_extension.cpp src/extensions/pulid_extension.cpp src/kcpp_sd_extensions.h src/model/adapter/ip_adapter.hpp src/model/adapter/lora.hpp src/model/adapter/pmid.hpp src/model/adapter/pulid.hpp src/model/common/block.hpp src/model/common/rope.hpp src/model/detector/yolov8.h src/model/diffusion/anima.hpp src/model/diffusion/animatediff.hpp src/model/diffusion/boogu.hpp src/model/diffusion/control.hpp src/model/diffusion/dit.hpp src/model/diffusion/ernie_image.hpp src/model/diffusion/flux.hpp src/model/diffusion/hidream_o1.hpp src/model/diffusion/hunyuan.hpp src/model/diffusion/ideogram4.hpp src/model/diffusion/krea2.hpp src/model/diffusion/lens.hpp src/model/diffusion/lingbot_video.hpp src/model/diffusion/ltxv.hpp src/model/diffusion/mage_flow.hpp src/model/diffusion/minimax_h3.hpp src/model/diffusion/minit2i.hpp src/model/diffusion/mmdit.hpp src/model/diffusion/model.hpp src/model/diffusion/pid.hpp src/model/diffusion/qwen_image.hpp src/model/diffusion/sefi_image.hpp src/model/diffusion/unet.hpp src/model/diffusion/wan.hpp src/model/diffusion/z_image.hpp src/model.h src/model_io/binary_io.h src/model_io/gguf_io.cpp src/model_io/gguf_io.h src/model_io/gguf_reader_ext.h src/model_io/kcpp_sdcpp_quantized_safetensors.hpp src/model_io/pickle_io.cpp src/model_io/pickle_io.h src/model_io/safetensors_io.cpp src/model_io/safetensors_io.h src/model_io/streaming_writer.h src/model_io/tensor_storage.h src/model_io/torch_legacy_io.cpp src/model_io/torch_legacy_io.h src/model_io/torch_zip_io.cpp src/model_io/torch_zip_io.h src/model_loader.cpp src/model_loader.h src/model_manager.cpp src/model_manager.h src/model/te/clip.hpp src/model/te/llm.hpp src/model/te/t5.hpp src/model/upscaler/esrgan.hpp src/model/upscaler/ltx_latent_upscaler.hpp src/model/vae/audio_vae.hpp src/model/vae/auto_encoder_kl.hpp src/model/vae/hunyuan_vae.hpp src/model/vae/ltx_audio_vae.hpp src/model/vae/ltx_vae.hpp src/model/vae/mage_vae.hpp src/model/vae/minimax_h3_audio_vae.hpp src/model/vae/minimax_h3_vae.hpp src/model/vae/tae.hpp src/model/vae/vae.hpp src/model/vae/wan_vae.hpp src/name_conversion.cpp src/name_conversion.h src/runtime/cache_dit.hpp src/runtime/condition_cache_utils.hpp src/runtime/denoiser.hpp src/runtime/easycache.hpp src/runtime/gits_noise.h src/runtime/guidance.cpp src/runtime/guidance.h src/runtime/imatrix.cpp src/runtime/imatrix.h src/runtime/latent-preview.h src/runtime/preprocessing.hpp src/runtime/sample-cache.cpp src/runtime/sample-cache.h src/runtime/spectrum.hpp src/runtime/ucache.hpp src/stable-diffusion.cpp src/tokenizers/bpe_tokenizer.cpp src/tokenizers/bpe_tokenizer.h src/tokenizers/clip_tokenizer.cpp src/tokenizers/clip_tokenizer.h src/tokenizers/gemma_tokenizer.cpp src/tokenizers/gemma_tokenizer.h src/tokenizers/gpt_oss_tokenizer.cpp src/tokenizers/gpt_oss_tokenizer.h src/tokenizers/mistral_tokenizer.cpp src/tokenizers/mistral_tokenizer.h src/tokenizers/qwen2_tokenizer.cpp src/tokenizers/qwen2_tokenizer.h src/tokenizers/t5_unigram_tokenizer.cpp src/tokenizers/t5_unigram_tokenizer.h src/tokenizers/tokenizer.cpp src/tokenizers/tokenizer.h src/tokenizers/tokenize_util.cpp src/tokenizers/tokenize_util.h src/tokenizers/vocab/vocab.h src/upscaler.cpp src/upscaler.h src/weight_manager.h +SDCPP_COMMON_BASENAMES := include/stable-diffusion.h src/conditioning/conditioner.hpp src/core/backend_fit.cpp src/core/backend_fit.h src/core/compute_workspace.cpp src/core/compute_workspace.h src/core/ggml_extend_backend.cpp src/core/ggml_extend_backend.h src/core/ggml_extend.cpp src/core/ggml_extend.h src/core/ggml_graph_cut.cpp src/core/ggml_graph_cut.h src/core/ggml_runner.cpp src/core/ggml_runner.h src/core/ggml_tensor_utils.cpp src/core/ggml_tensor_utils.h src/core/layer_split_partition.cpp src/core/layer_split_partition.h src/core/ordered_map.hpp src/core/rng.hpp src/core/rng_mt19937.hpp src/core/rng_philox.hpp src/core/runner_cache.cpp src/core/runner_cache.h src/core/segment_graph_bindings.cpp src/core/segment_graph_bindings.h src/core/segment_weight_pipeline.cpp src/core/segment_weight_pipeline.h src/core/tensor_ggml.hpp src/core/tensor.hpp src/core/util.cpp src/core/util.h src/detailer.cpp src/detailer.h src/device_residency_manager.h src/extensions/generation_extension.h src/extensions/photomaker_extension.cpp src/extensions/pulid_extension.cpp src/model/adapter/ip_adapter.hpp src/model/adapter/lora.hpp src/model/adapter/lora_ops.cpp src/model/adapter/lora_ops.h src/model/adapter/pmid.hpp src/model/adapter/pulid.hpp src/model/common/block.hpp src/model/common/ggml_block.hpp src/model/common/rope.hpp src/model/detector/yolov8.h src/model/diffusion/anima.hpp src/model/diffusion/animatediff.hpp src/model/diffusion/boogu.hpp src/model/diffusion/control.hpp src/model/diffusion/dit.hpp src/model/diffusion/ernie_image.hpp src/model/diffusion/flux.hpp src/model/diffusion/hidream_o1.hpp src/model/diffusion/hunyuan.hpp src/model/diffusion/ideogram4.hpp src/model/diffusion/krea2.hpp src/model/diffusion/lens.hpp src/model/diffusion/lingbot_video.hpp src/model/diffusion/ltxv.hpp src/model/diffusion/mage_flow.hpp src/model/diffusion/minimax_h3.hpp src/model/diffusion/minit2i.hpp src/model/diffusion/mmdit.hpp src/model/diffusion/model.hpp src/model/diffusion/pid.hpp src/model/diffusion/qwen_image.hpp src/model/diffusion/sefi_image.hpp src/model/diffusion/unet.hpp src/model/diffusion/wan.hpp src/model/diffusion/z_image.hpp src/model.h src/model_io/binary_io.h src/model_io/gguf_io.cpp src/model_io/gguf_io.h src/model_io/gguf_reader_ext.h src/model_io/pickle_io.cpp src/model_io/pickle_io.h src/model_io/safetensors_io.cpp src/model_io/safetensors_io.h src/model_io/streaming_writer.h src/model_io/tensor_storage.h src/model_io/torch_legacy_io.cpp src/model_io/torch_legacy_io.h src/model_io/torch_zip_io.cpp src/model_io/torch_zip_io.h src/model_loader.cpp src/model_loader.h src/model_manager.cpp src/model_manager.h src/model_manager_prefetch.cpp src/model/te/clip.hpp src/model/te/llm.hpp src/model/te/t5.hpp src/model/upscaler/esrgan.hpp src/model/upscaler/ltx_latent_upscaler.hpp src/model/vae/audio_vae.hpp src/model/vae/auto_encoder_kl.hpp src/model/vae/hunyuan_vae.hpp src/model/vae/ltx_audio_vae.hpp src/model/vae/ltx_vae.hpp src/model/vae/mage_vae.hpp src/model/vae/minimax_h3_audio_vae.hpp src/model/vae/minimax_h3_vae.hpp src/model/vae/tae.hpp src/model/vae/vae.hpp src/model/vae/vae_tiling.hpp src/model/vae/wan_vae.hpp src/name_conversion.cpp src/name_conversion.h src/runtime/cache_dit.hpp src/runtime/condition_cache_utils.hpp src/runtime/denoiser.hpp src/runtime/easycache.hpp src/runtime/gits_noise.h src/runtime/guidance.cpp src/runtime/guidance.h src/runtime/imatrix.cpp src/runtime/imatrix.h src/runtime/latent-preview.h src/runtime/preprocessing.hpp src/runtime/preview_interval.h src/runtime/sample-cache.cpp src/runtime/sample-cache.h src/runtime/spectrum.hpp src/runtime/tiling.cpp src/runtime/tiling.h src/runtime/ucache.hpp src/stable-diffusion.cpp src/tokenizers/bpe_tokenizer.cpp src/tokenizers/bpe_tokenizer.h src/tokenizers/clip_tokenizer.cpp src/tokenizers/clip_tokenizer.h src/tokenizers/gemma_tokenizer.cpp src/tokenizers/gemma_tokenizer.h src/tokenizers/gpt_oss_tokenizer.cpp src/tokenizers/gpt_oss_tokenizer.h src/tokenizers/mistral_tokenizer.cpp src/tokenizers/mistral_tokenizer.h src/tokenizers/qwen2_tokenizer.cpp src/tokenizers/qwen2_tokenizer.h src/tokenizers/t5_unigram_tokenizer.cpp src/tokenizers/t5_unigram_tokenizer.h src/tokenizers/tokenizer.cpp src/tokenizers/tokenizer.h src/tokenizers/tokenize_util.cpp src/tokenizers/tokenize_util.h src/tokenizers/vocab/vocab.h src/upscaler.cpp src/upscaler.h SDCPP_MAIN_BASENAMES := examples/cli/image_metadata.cpp examples/cli/image_metadata.h examples/cli/main.cpp examples/cli/msf_gif.h examples/common/common.cpp examples/common/common.h examples/common/log.cpp examples/common/log.h examples/common/media_io.cpp examples/common/media_io.h examples/common/resource_owners.hpp src/tokenizers/vocab/clip_merges.hpp src/tokenizers/vocab/gemma2_merges.hpp src/tokenizers/vocab/gemma2_vocab.hpp src/tokenizers/vocab/gemma_merges.hpp src/tokenizers/vocab/gemma_vocab.hpp src/tokenizers/vocab/gpt_oss_merges.hpp src/tokenizers/vocab/gpt_oss_vocab.hpp src/tokenizers/vocab/mistral_merges.hpp src/tokenizers/vocab/mistral_vocab.hpp src/tokenizers/vocab/qwen_merges.hpp src/tokenizers/vocab/t5.hpp src/tokenizers/vocab/umt5.hpp src/tokenizers/vocab/vocab.cpp src/convert.cpp src/version.cpp diff --git a/otherarch/sdcpp/examples/cli/main.cpp b/otherarch/sdcpp/examples/cli/main.cpp index d77356b12..22b26da79 100644 --- a/otherarch/sdcpp/examples/cli/main.cpp +++ b/otherarch/sdcpp/examples/cli/main.cpp @@ -40,9 +40,9 @@ struct SDCliParams { std::string image_path; std::string metadata_format = "text"; - bool verbose = false; - bool canny_preprocess = false; - bool convert_name = false; + sd_log_level_t log_level = SD_LOG_INFO; + bool canny_preprocess = false; + bool convert_name = false; preview_t preview_method = PREVIEW_NONE; int preview_interval = 1; @@ -115,10 +115,6 @@ struct SDCliParams { "--convert-name", "convert tensor name (for convert mode)", true, &convert_name}, - {"-v", - "--verbose", - "print extra info", - true, &verbose}, {"", "--color", "colors the logging tags according to level", @@ -220,6 +216,7 @@ struct SDCliParams { on_imatrix_in_arg}, }; + add_log_options(options, log_level); return options; }; @@ -269,7 +266,7 @@ struct SDCliParams { << " output_path: \"" << output_path << "\",\n" << " image_path: \"" << image_path << "\",\n" << " metadata_format: \"" << metadata_format << "\",\n" - << " verbose: " << (verbose ? "true" : "false") << ",\n" + << " log_level: " << log_level_name(log_level) << ",\n" << " color: " << (color ? "true" : "false") << ",\n" << " canny_preprocess: " << (canny_preprocess ? "true" : "false") << ",\n" << " convert_name: " << (convert_name ? "true" : "false") << ",\n" @@ -307,6 +304,9 @@ void parse_args(int argc, const char** argv, SDCliParams& cli_params, SDContextP exit(cli_params.normal_exit ? 0 : 1); } + log_level = cli_params.log_level; + log_color = cli_params.color; + bool valid = cli_params.resolve_and_validate(); if (valid && cli_params.mode != METADATA) { valid = ctx_params.resolve_and_validate(cli_params.mode) && @@ -323,15 +323,14 @@ void parse_args(int argc, const char** argv, SDCliParams& cli_params, SDContextP void sd_log_cb(enum sd_log_level_t level, const char* log, void* data) { SDCliParams* cli_params = (SDCliParams*)data; - log_print(level, log, cli_params->verbose, cli_params->color); + log_print(level, log, cli_params->log_level, cli_params->color); } bool load_images_from_dir(const std::string dir, std::vector& images, int expected_width = 0, int expected_height = 0, - int max_image_num = 0, - bool verbose = false) { + int max_image_num = 0) { if (!fs::exists(dir) || !fs::is_directory(dir)) { LOG_ERROR("'%s' is not a valid directory\n", dir.c_str()); return false; @@ -355,7 +354,7 @@ bool load_images_from_dir(const std::string dir, std::transform(ext.begin(), ext.end(), ext.begin(), ::tolower); if (ext == ".jpg" || ext == ".jpeg" || ext == ".png" || ext == ".bmp" || ext == ".webp") { - LOG_DEBUG("load image %zu from '%s'", images.size(), path.c_str()); + LOG_VERBOSE("load image %zu from '%s'", images.size(), path.c_str()); int width = 0; int height = 0; uint8_t* image_buffer = load_image_from_file(path.c_str(), width, height, expected_width, expected_height); @@ -651,8 +650,6 @@ int main(int argc, const char* argv[]) { parse_args(argc, argv, cli_params, ctx_params, gen_params); sd_set_log_callback(sd_log_cb, (void*)&cli_params); - log_verbose = cli_params.verbose; - log_color = cli_params.color; if (cli_params.mode == METADATA) { MetadataReadOptions options; @@ -700,11 +697,11 @@ int main(int argc, const char* argv[]) { cli_params.preview_noisy, (void*)&cli_params); - LOG_DEBUG("version: %s", version_string().c_str()); - LOG_DEBUG("%s", sd_get_system_info()); - LOG_DEBUG("%s", cli_params.to_string().c_str()); - LOG_DEBUG("%s", ctx_params.to_string().c_str()); - LOG_DEBUG("%s", gen_params.to_string().c_str()); + LOG_VERBOSE("version: %s", version_string().c_str()); + LOG_VERBOSE("%s", sd_get_system_info()); + LOG_VERBOSE("%s", cli_params.to_string().c_str()); + LOG_VERBOSE("%s", ctx_params.to_string().c_str()); + LOG_VERBOSE("%s", gen_params.to_string().c_str()); if (!cli_params.imatrix_out.empty()) { if (fs::exists(cli_params.imatrix_out) && @@ -808,7 +805,7 @@ int main(int argc, const char* argv[]) { gen_params.ref_videos.reserve(gen_params.ref_video_paths.size()); for (const auto& path : gen_params.ref_video_paths) { std::vector frames; - if (!load_images_from_dir(path, frames, 0, 0, 0, cli_params.verbose) || frames.empty()) { + if (!load_images_from_dir(path, frames) || frames.empty()) { LOG_ERROR("load reference video frames from '%s' failed", path.c_str()); return 1; } @@ -890,8 +887,7 @@ int main(int argc, const char* argv[]) { gen_params.control_frames, gen_params.get_resolved_width(), gen_params.get_resolved_height(), - gen_params.video_frames, - cli_params.verbose)) { + gen_params.video_frames)) { return 1; } } @@ -902,8 +898,7 @@ int main(int argc, const char* argv[]) { gen_params.pm_id_images, 0, 0, - 0, - cli_params.verbose)) { + 0)) { return 1; } } diff --git a/otherarch/sdcpp/examples/common/common.cpp b/otherarch/sdcpp/examples/common/common.cpp index d46be91a4..335a6226f 100644 --- a/otherarch/sdcpp/examples/common/common.cpp +++ b/otherarch/sdcpp/examples/common/common.cpp @@ -239,6 +239,26 @@ void ArgOptions::print() const { } } +void add_log_options(ArgOptions& options, sd_log_level_t& level) { + options.manual_options.push_back({"", "--log-level", + "minimum log level, one of [debug, verbose, info, warn, error] (default: info)", + [&level](int argc, const char** argv, int index) { + if (++index >= argc) { + return -1; + } + if (!parse_log_level(argv[index], level)) { + LOG_ERROR("invalid log level %s, must be one of [debug, verbose, info, warn, error]", argv[index]); + return -1; + } + return 1; + }}); + options.manual_options.push_back({"-v", "--verbose", "equivalent to --log-level verbose", + [&level](int, const char**, int) { + level = SD_LOG_VERBOSE; + return 0; + }}); +} + bool parse_options(int argc, const char** argv, const std::vector& options_list) { bool invalid_arg = false; std::string arg; @@ -502,7 +522,7 @@ ArgOptions SDContextParams::get_options() { &rpc_servers}, {"", "--max-vram", - "maximum VRAM budget in GiB for graph-cut segmented execution. Accepts a single value or assignments by backend/device, e.g. 6 or cuda0=6,vulkan0=4. 0 disables graph splitting; a negative value auto-detects free VRAM, sparing the specified value", + "optional per-device budget in GiB for managed weights and runner buffers during automatic graph-cut execution. Accepts a single value or assignments by backend/device, e.g. 6 or cuda0=6,vulkan0=4. 0 uses live free VRAM without an explicit budget; a negative value reserves that much free VRAM", 0, &max_vram}, }; @@ -517,19 +537,17 @@ ArgOptions SDContextParams::get_options() { options.bool_options = { {"", - "--stream-layers", - "enable residency+prefetch streaming on top of --max-vram (no effect without --max-vram; defaults to false)", - true, &stream_layers}, + "--disable-prefetch", + "disable asynchronous next-segment weight prefetch (defaults to false)", + true, &disable_prefetch}, + {"", + "--disable-segmented-compute", + "force monolithic graph execution even when automatic graph cutting is needed (defaults to false)", + true, &disable_segmented_compute}, {"", "--eager-load", "load all params into the params backend at model-load time instead of lazily on first use (defaults to false)", true, &eager_load}, - {"", - "--auto-fit", - "pick the diffusion/te/vae device placements automatically from the model size and the per-device " - "memory budgets (--max-vram; defaults to free memory minus a small margin). Overrides --backend and " - "--params-backend; may split modules across GPUs (--split-mode still selects layer or row)", - true, &auto_fit}, {"", "--force-sdxl-vae-conv-scale", "force use of conv scale on sdxl vae", @@ -572,6 +590,23 @@ ArgOptions SDContextParams::get_options() { true, &vae_conv_direct}, }; + auto on_auto_fit_arg = [&](int argc, const char** argv, int index) { + if (++index >= argc) { + LOG_ERROR("--auto-fit requires 'on' or 'off'"); + return -1; + } + const std::string arg = argv[index]; + if (arg == "on") { + auto_fit = true; + } else if (arg == "off") { + auto_fit = false; + } else { + LOG_ERROR("invalid --auto-fit value '%s'; expected 'on' or 'off'", argv[index]); + return -1; + } + return 1; + }; + auto on_type_arg = [&](int argc, const char** argv, int index) { if (++index >= argc) { return -1; @@ -643,6 +678,12 @@ ArgOptions SDContextParams::get_options() { }; options.manual_options = { + {"", + "--auto-fit", + "on|off (default: on). Use one GPU for diffusion/te/vae computation and place weights on that GPU, " + "RAM, another GPU, or disk in that order, according to available memory (--max-vram limits GPU budgets). " + "Disabled by explicit --backend or --params-backend; uses automatic graph segmentation when needed", + on_auto_fit_arg}, {"", "--type", "weight type (examples: f32, f16, q4_0, q4_1, q5_0, q5_1, q8_0, q2_K, q3_K, q4_K). " @@ -831,7 +872,8 @@ std::string SDContextParams::to_string() const { << " sampler_rng_type: " << sd_rng_type_name(sampler_rng_type) << ",\n" << " offload_params_to_cpu: " << (offload_params_to_cpu ? "true" : "false") << ",\n" << " max_vram: \"" << max_vram << "\",\n" - << " stream_layers: " << (stream_layers ? "true" : "false") << ",\n" + << " disable_prefetch: " << (disable_prefetch ? "true" : "false") << ",\n" + << " disable_segmented_compute: " << (disable_segmented_compute ? "true" : "false") << ",\n" << " eager_load: " << (eager_load ? "true" : "false") << ",\n" << " backend: \"" << backend << "\",\n" << " params_backend: \"" << params_backend << "\",\n" @@ -903,7 +945,8 @@ sd_ctx_params_t SDContextParams::to_sd_ctx_params_t(bool taesd_preview) { sd_ctx_params.force_sdxl_vae_conv_scale = force_sdxl_vae_conv_scale; sd_ctx_params.vae_format = str_to_vae_format(vae_format); sd_ctx_params.max_vram = max_vram.c_str(); - sd_ctx_params.stream_layers = stream_layers; + sd_ctx_params.disable_prefetch = disable_prefetch; + sd_ctx_params.disable_segmented_compute = disable_segmented_compute; sd_ctx_params.eager_load = eager_load; sd_ctx_params.backend = effective_backend.c_str(); sd_ctx_params.params_backend = effective_params_backend.c_str(); diff --git a/otherarch/sdcpp/examples/common/common.h b/otherarch/sdcpp/examples/common/common.h index 34b4a013b..ef084dfd4 100644 --- a/otherarch/sdcpp/examples/common/common.h +++ b/otherarch/sdcpp/examples/common/common.h @@ -107,6 +107,7 @@ struct ArgOptions { void print() const; }; +void add_log_options(ArgOptions& options, sd_log_level_t& level); bool parse_options(int argc, const char** argv, const std::vector& options_list); bool decode_base64_image(const std::string& encoded_input, int target_channels, @@ -146,17 +147,18 @@ struct SDContextParams { std::map embedding_map; std::vector embedding_vec; - rng_type_t rng_type = CUDA_RNG; - rng_type_t sampler_rng_type = RNG_TYPE_COUNT; - bool offload_params_to_cpu = false; - std::string max_vram = "0"; - bool stream_layers = false; - bool eager_load = false; + rng_type_t rng_type = CUDA_RNG; + rng_type_t sampler_rng_type = RNG_TYPE_COUNT; + bool offload_params_to_cpu = false; + std::string max_vram = "0"; + bool disable_prefetch = false; + bool disable_segmented_compute = false; + bool eager_load = false; std::string backend; std::string params_backend; std::string split_mode; std::string model_args; - bool auto_fit = false; + bool auto_fit = true; std::string rpc_servers; std::string effective_backend; std::string effective_params_backend; diff --git a/otherarch/sdcpp/examples/common/log.cpp b/otherarch/sdcpp/examples/common/log.cpp index 2c4343912..c272c2376 100644 --- a/otherarch/sdcpp/examples/common/log.cpp +++ b/otherarch/sdcpp/examples/common/log.cpp @@ -2,8 +2,8 @@ #include -bool log_verbose = false; -bool log_color = false; +sd_log_level_t log_level = SD_LOG_INFO; +bool log_color = false; std::string sd_basename(const std::string& path) { size_t pos = path.find_last_of('/'); @@ -51,12 +51,40 @@ void print_utf8(FILE* stream, const char* utf8) { #endif } -void log_print(enum sd_log_level_t level, const char* log, bool verbose, bool color) { +const char* log_level_name(sd_log_level_t level) { + switch (level) { + case SD_LOG_DEBUG: + return "debug"; + case SD_LOG_VERBOSE: + return "verbose"; + case SD_LOG_INFO: + return "info"; + case SD_LOG_WARN: + return "warn"; + case SD_LOG_ERROR: + return "error"; + default: + return "unknown"; + } +} + +bool parse_log_level(const std::string& name, sd_log_level_t& level) { + const sd_log_level_t levels[] = {SD_LOG_DEBUG, SD_LOG_VERBOSE, SD_LOG_INFO, SD_LOG_WARN, SD_LOG_ERROR}; + for (sd_log_level_t candidate : levels) { + if (name == log_level_name(candidate)) { + level = candidate; + return true; + } + } + return false; +} + +void log_print(enum sd_log_level_t level, const char* log, sd_log_level_t min_level, bool color) { int tag_color; const char* level_str; FILE* out_stream = (level == SD_LOG_ERROR) ? stderr : stdout; - if (!log || (!verbose && level <= SD_LOG_DEBUG)) { + if (!log || level < min_level) { return; } @@ -65,6 +93,10 @@ void log_print(enum sd_log_level_t level, const char* log, bool verbose, bool co tag_color = 37; level_str = "DEBUG"; break; + case SD_LOG_VERBOSE: + tag_color = 37; + level_str = "VERBOSE"; + break; case SD_LOG_INFO: tag_color = 34; level_str = "INFO"; @@ -84,10 +116,11 @@ void log_print(enum sd_log_level_t level, const char* log, bool verbose, bool co } if (color) { - fprintf(out_stream, "\033[%d;1m[%-5s]\033[0m ", tag_color, level_str); + fprintf(out_stream, "\033[%d;1m[%-7s]\033[0m ", tag_color, level_str); } else { - fprintf(out_stream, "[%-5s] ", level_str); + fprintf(out_stream, "[%-7s] ", level_str); } + fflush(out_stream); print_utf8(out_stream, log); fflush(out_stream); } @@ -109,7 +142,7 @@ void example_log_printf(sd_log_level_t level, const char* file, int line, const strncat(log_buffer, "\n", LOG_BUFFER_SIZE - len); } - log_print(level, log_buffer, log_verbose, log_color); + log_print(level, log_buffer, log_level, log_color); va_end(args); } diff --git a/otherarch/sdcpp/examples/common/log.h b/otherarch/sdcpp/examples/common/log.h index f28b4b4ea..f623b0ead 100644 --- a/otherarch/sdcpp/examples/common/log.h +++ b/otherarch/sdcpp/examples/common/log.h @@ -16,15 +16,18 @@ #include "stable-diffusion.h" -extern bool log_verbose; +extern sd_log_level_t log_level; extern bool log_color; std::string sd_basename(const std::string& path); void print_utf8(FILE* stream, const char* utf8); -void log_print(sd_log_level_t level, const char* log, bool verbose, bool color); +const char* log_level_name(sd_log_level_t level); +bool parse_log_level(const std::string& name, sd_log_level_t& level); +void log_print(sd_log_level_t level, const char* log, sd_log_level_t min_level, bool color); void example_log_printf(sd_log_level_t level, const char* file, int line, const char* format, ...); #define LOG_DEBUG(format, ...) example_log_printf(SD_LOG_DEBUG, __FILE__, __LINE__, format, ##__VA_ARGS__) +#define LOG_VERBOSE(format, ...) example_log_printf(SD_LOG_VERBOSE, __FILE__, __LINE__, format, ##__VA_ARGS__) #define LOG_INFO(format, ...) example_log_printf(SD_LOG_INFO, __FILE__, __LINE__, format, ##__VA_ARGS__) #define LOG_WARN(format, ...) example_log_printf(SD_LOG_WARN, __FILE__, __LINE__, format, ##__VA_ARGS__) #define LOG_ERROR(format, ...) example_log_printf(SD_LOG_ERROR, __FILE__, __LINE__, format, ##__VA_ARGS__) diff --git a/otherarch/sdcpp/examples/common/media_io.cpp b/otherarch/sdcpp/examples/common/media_io.cpp index 4fb26e980..3dedb8e71 100644 --- a/otherarch/sdcpp/examples/common/media_io.cpp +++ b/otherarch/sdcpp/examples/common/media_io.cpp @@ -850,7 +850,7 @@ std::vector create_mjpg_avi_from_sd_images_to_vector(sd_image_t* images const uint32_t audio_data_size = has_audio ? static_cast(audio_pcm.size()) : 0; if (mjpg_quality != quality) - LOG_DEBUG("create_mjpg_avi...(): compression quality was limited from %i to %i", quality, mjpg_quality); + LOG_VERBOSE("create_mjpg_avi...(): compression quality was limited from %i to %i", quality, mjpg_quality); std::vector avi_data; avi_data.reserve(static_cast(num_images) * 1024); diff --git a/otherarch/sdcpp/include/stable-diffusion.h b/otherarch/sdcpp/include/stable-diffusion.h index 60e06fa13..4100cf615 100644 --- a/otherarch/sdcpp/include/stable-diffusion.h +++ b/otherarch/sdcpp/include/stable-diffusion.h @@ -147,6 +147,7 @@ enum sd_type_t { enum sd_log_level_t { SD_LOG_DEBUG, + SD_LOG_VERBOSE, SD_LOG_INFO, SD_LOG_WARN, SD_LOG_ERROR @@ -229,8 +230,8 @@ typedef struct { bool vae_conv_direct; bool force_sdxl_vae_conv_scale; enum sd_vae_format_t vae_format; - const char* max_vram; // GiB budget or backend assignment spec for graph-cut segmented param offload (0 = disabled, -1 = auto) - bool stream_layers; // Enable residency+prefetch streaming on top of --max-vram (no effect without --max-vram) + const char* max_vram; // Optional per-device GiB budget for managed weights and runner buffers; 0 uses live free VRAM without an explicit budget + bool disable_prefetch; // Disable asynchronous next-segment weight prefetch bool eager_load; // Load all params into the params backend at model-load time instead of lazily on first use const char* backend; const char* params_backend; @@ -238,6 +239,7 @@ typedef struct { bool auto_fit; const char* rpc_servers; const char* model_args; + bool disable_segmented_compute; // Force monolithic graph execution even when automatic graph cutting would fit memory better } sd_ctx_params_t; typedef struct { diff --git a/otherarch/sdcpp/sdtype_adapter.cpp b/otherarch/sdcpp/sdtype_adapter.cpp index a702e0c2f..fe754f52e 100644 --- a/otherarch/sdcpp/sdtype_adapter.cpp +++ b/otherarch/sdcpp/sdtype_adapter.cpp @@ -558,7 +558,7 @@ bool sdtype_load_model(const sd_load_model_inputs inputs) { params.vae_conv_direct = sd_params->vae_conv_direct; params.model_args = "chroma_use_dit_mask=true"; params.max_vram = max_vram.c_str(); - params.stream_layers = inputs.stream_layers; + //params.stream_layers = inputs.stream_layers; // removed in master-843 params.eager_load = true; //kcpp should preload everything params.enable_mmap = inputs.use_mmap; params.backend = backend.c_str(); diff --git a/otherarch/sdcpp/src/conditioning/conditioner.hpp b/otherarch/sdcpp/src/conditioning/conditioner.hpp index 8968676b3..23cfbddd2 100644 --- a/otherarch/sdcpp/src/conditioning/conditioner.hpp +++ b/otherarch/sdcpp/src/conditioning/conditioner.hpp @@ -1,11 +1,13 @@ #ifndef __SD_CONDITIONING_CONDITIONER_HPP__ #define __SD_CONDITIONING_CONDITIONER_HPP__ +#include #include #include #include #include #include +#include "core/ggml_tensor_utils.h" #include "core/tensor_ggml.hpp" #include "core/util.h" @@ -142,14 +144,13 @@ public: virtual void get_param_tensors(std::map& tensors) = 0; virtual void get_param_tensor_ops(std::map& tensor_ops) {} virtual void set_max_graph_vram_bytes(size_t max_vram_bytes) {} - virtual void set_stream_layers_enabled(bool enabled) {} virtual void set_runtime_backends(const std::vector& backends) {} virtual void set_graph_cut_layer_split_enabled(bool enabled) {} virtual void set_graph_cut_layer_split_backend_vram_limits(const std::vector& limits) {} virtual void get_layer_split_param_tensors(std::map& tensors) {} virtual void set_flash_attention_enabled(bool enabled) = 0; virtual void set_weight_adapter(const std::shared_ptr& adapter) {} - virtual void runner_done() {} + virtual void runner_end() {} }; // ldm.modules.encoders.modules.FrozenCLIPEmbedder @@ -161,9 +162,9 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner { std::shared_ptr text_model2; std::map embedding_map; - int32_t num_custom_embeddings = 0; - int32_t num_custom_embeddings_2 = 0; + int32_t num_custom_embeddings = 0; std::vector token_embed_custom; + std::vector token_embed_custom2; std::map> embedding_pos_map; FrozenCLIPEmbedderWithCustomWords(ggml_backend_t backend, @@ -202,13 +203,6 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner { } } - void set_stream_layers_enabled(bool enabled) override { - text_model->set_stream_layers_enabled(enabled); - if (sd_version_is_sdxl(version)) { - text_model2->set_stream_layers_enabled(enabled); - } - } - void set_runtime_backends(const std::vector& backends) override { text_model->set_runtime_backends(backends); if (sd_version_is_sdxl(version)) { @@ -244,10 +238,10 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner { } } - void runner_done() override { - text_model->runner_done(); + void runner_end() override { + text_model->runner_end(); if (sd_version_is_sdxl(version)) { - text_model2->runner_done(); + text_model2->runner_end(); } } @@ -257,74 +251,93 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner { LOG_ERROR("embedding '%s' failed", embd_name.c_str()); return false; } - auto iter = embedding_pos_map.find(embd_name); - if (iter != embedding_pos_map.end()) { - LOG_DEBUG("embedding already read in: %s", embd_name.c_str()); - for (int i = iter->second.first; i < iter->second.second; i++) { + auto push_ids = [&](int pos_start, int pos_end, bool cached) { + for (int i = pos_start; i < pos_end; i++) { bpe_tokens.push_back(text_model->model.vocab_size + i); } + if (!cached) { + LOG_VERBOSE("embedding '%s' applied: %i token(s), custom embeddings: %i", embd_name.c_str(), pos_end - pos_start, num_custom_embeddings); + } + }; + auto iter = embedding_pos_map.find(embd_name); + if (iter != embedding_pos_map.end()) { + LOG_VERBOSE("embedding already read in: %s", embd_name.c_str()); + push_ids(iter->second.first, iter->second.second, true); return true; } ggml_init_params params; - params.mem_size = 100 * 1024 * 1024; // max for custom embeddings 100 MB - params.mem_buffer = nullptr; - params.no_alloc = false; - ggml_context* embd_ctx = ggml_init(params); - ggml_tensor* embd = nullptr; - ggml_tensor* embd2 = nullptr; - auto on_load = [&](const TensorStorage& tensor_storage, ggml_tensor** dst_tensor) { - if (tensor_storage.ne[0] != text_model->model.hidden_size) { - if (text_model2) { - if (tensor_storage.ne[0] == text_model2->model.hidden_size) { - embd2 = ggml_new_tensor_2d(embd_ctx, tensor_storage.type, text_model2->model.hidden_size, tensor_storage.n_dims > 1 ? tensor_storage.ne[1] : 1); - *dst_tensor = embd2; - } else { - LOG_DEBUG("embedding wrong hidden size, got %i, expected %i or %i", tensor_storage.ne[0], text_model->model.hidden_size, text_model2->model.hidden_size); - return false; - } - } else { - LOG_DEBUG("embedding wrong hidden size, got %i, expected %i", tensor_storage.ne[0], text_model->model.hidden_size); + params.mem_size = 100 * 1024 * 1024; // max for custom embeddings 100 MB + params.mem_buffer = nullptr; + params.no_alloc = false; + auto ggml_ctx_deleter = [](ggml_context* ctx) { ggml_free(ctx); }; + auto embd_ctx = std::unique_ptr(ggml_init(params), ggml_ctx_deleter); + if (!embd_ctx.get()) { + LOG_ERROR("ggml_init failed when loading embeddings file"); + return false; + } + ggml_tensor* embd = nullptr; + ggml_tensor* embd2 = nullptr; + ggml_type embd_type = text_model->model.get_token_embed_weight()->type; + ggml_type embd2_type = text_model2 ? text_model2->model.get_token_embed_weight()->type : embd_type; + int64_t hidden_size = text_model->model.hidden_size; + int64_t hidden_size2 = text_model2 ? text_model2->model.hidden_size : 0; + auto on_load = [&](const TensorStorage& tensor_storage, ggml_tensor** dst_tensor) { + if (tensor_storage.ne[0] == hidden_size) { + embd = ggml_new_tensor_2d(embd_ctx.get(), embd_type, hidden_size, tensor_storage.n_dims > 1 ? tensor_storage.ne[1] : 1); + if (embd == nullptr) { return false; } - } else { - embd = ggml_new_tensor_2d(embd_ctx, tensor_storage.type, text_model->model.hidden_size, tensor_storage.n_dims > 1 ? tensor_storage.ne[1] : 1); *dst_tensor = embd; + } else if (text_model2) { + if (tensor_storage.ne[0] == hidden_size2) { + embd2 = ggml_new_tensor_2d(embd_ctx.get(), embd2_type, hidden_size2, tensor_storage.n_dims > 1 ? tensor_storage.ne[1] : 1); + if (embd2 == nullptr) { + return false; + } + *dst_tensor = embd2; + } else { + LOG_VERBOSE("embedding skipped, wrong hidden size, got %i, expected %i or %i", tensor_storage.ne[0], hidden_size, hidden_size2); + } + } else { + LOG_VERBOSE("embedding skipped, wrong hidden size, got %i, expected %i", tensor_storage.ne[0], hidden_size); } return true; }; model_loader.set_n_threads(1); - model_loader.load_tensors(on_load); - int pos_start = num_custom_embeddings; - if (embd) { - int64_t hidden_size = text_model->model.hidden_size; - token_embed_custom.resize(token_embed_custom.size() + ggml_nbytes(embd)); - memcpy((void*)(token_embed_custom.data() + num_custom_embeddings * hidden_size * ggml_type_size(embd->type)), - embd->data, - ggml_nbytes(embd)); - for (int i = 0; i < embd->ne[1]; i++) { - bpe_tokens.push_back(text_model->model.vocab_size + num_custom_embeddings); - // LOG_DEBUG("new custom token: %i", text_model.vocab_size + num_custom_embeddings); - num_custom_embeddings++; - } - LOG_DEBUG("embedding '%s' applied, custom embeddings: %i", embd_name.c_str(), num_custom_embeddings); - } - if (embd2) { - int64_t hidden_size = text_model2->model.hidden_size; - token_embed_custom.resize(token_embed_custom.size() + ggml_nbytes(embd2)); - memcpy((void*)(token_embed_custom.data() + num_custom_embeddings_2 * hidden_size * ggml_type_size(embd2->type)), - embd2->data, - ggml_nbytes(embd2)); - for (int i = 0; i < embd2->ne[1]; i++) { - bpe_tokens.push_back(text_model2->model.vocab_size + num_custom_embeddings_2); - // LOG_DEBUG("new custom token: %i", text_model.vocab_size + num_custom_embeddings); - num_custom_embeddings_2++; - } - LOG_DEBUG("embedding '%s' applied, custom embeddings: %i (text model 2)", embd_name.c_str(), num_custom_embeddings_2); - } - int pos_end = num_custom_embeddings; - if (pos_end == pos_start) { + if (!model_loader.load_tensors(on_load)) { + LOG_ERROR("embedding '%s' failed", embd_name.c_str()); return false; } + if (!embd && !embd2) { + LOG_WARN("embedding '%s' has no usable tensor", embd_name.c_str()); + return false; + } + int pos_start = num_custom_embeddings; + int64_t embd_rows = embd ? embd->ne[1] : 0; + int64_t embd2_rows = embd2 ? embd2->ne[1] : 0; + if (embd_rows < embd2_rows) { + LOG_WARN("embedding '%s' has fewer rows for text model 1, zero-padding", embd_name.c_str()); + } else if (text_model2 && embd2_rows < embd_rows) { + LOG_WARN("embedding '%s' has fewer rows for text model 2, zero-padding", embd_name.c_str()); + } + int64_t rows = std::max(embd_rows, embd2_rows); + size_t embd_bytes = hidden_size * ggml_type_size(embd_type); + token_embed_custom.resize(token_embed_custom.size() + embd_bytes * rows); + if (embd) { + memcpy((void*)(token_embed_custom.data() + embd_bytes * num_custom_embeddings), + embd->data, embd_bytes * embd_rows); + } + if (text_model2) { + size_t embd2_bytes = hidden_size2 * ggml_type_size(embd2_type); + token_embed_custom2.resize(token_embed_custom2.size() + embd2_bytes * rows); + if (embd2) { + memcpy((void*)(token_embed_custom2.data() + embd2_bytes * num_custom_embeddings), + embd2->data, embd2_bytes * embd2_rows); + } + } + num_custom_embeddings += (int)rows; + int pos_end = num_custom_embeddings; + push_ids(pos_start, pos_end, false); embedding_pos_map[embd_name] = std::pair{pos_start, pos_end}; return true; } @@ -368,7 +381,7 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner { ss << "['" << item.first << "', " << item.second << "], "; } ss << "]"; - LOG_DEBUG("parse '%s' to %s", text.c_str(), ss.str().c_str()); + LOG_VERBOSE("parse '%s' to %s", text.c_str(), ss.str().c_str()); } auto on_new_token_cb = [&](std::string& str, std::vector& bpe_tokens) -> bool { @@ -389,7 +402,7 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner { size_t padding_size = (75 - (current_size % 75)) % 75; // Ensure no negative padding if (padding_size > 0) { - LOG_DEBUG("BREAK token encountered, padding current chunk by %zu tokens.", padding_size); + LOG_VERBOSE("BREAK token encountered, padding current chunk by %zu tokens.", padding_size); tokens.insert(tokens.end(), padding_size, tokenizer.EOS_TOKEN_ID); weights.insert(weights.end(), padding_size, 1.0f); } @@ -461,21 +474,17 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner { max_token_idx, false, clip_skip, - false, - true, - true); + false); GGML_ASSERT(!chunk_hidden_states.empty()); if (sd_version_is_sdxl(version)) { auto chunk_hidden_states2 = text_model2->compute(n_threads, input_ids2, num_custom_embeddings, - token_embed_custom.data(), + token_embed_custom2.data(), max_token_idx, false, clip_skip, - false, - true, - true); + false); GGML_ASSERT(!chunk_hidden_states2.empty()); chunk_hidden_states = sd::ops::concat(chunk_hidden_states, chunk_hidden_states2, 0); @@ -483,18 +492,16 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner { pooled = text_model2->compute(n_threads, input_ids2, num_custom_embeddings, - token_embed_custom.data(), + token_embed_custom2.data(), max_token_idx, true, clip_skip, - false, - true, - true); + false); GGML_ASSERT(!pooled.empty()); } } int64_t t1 = ggml_time_ms(); - LOG_DEBUG("computing condition graph completed, taking %" PRId64 " ms", t1 - t0); + LOG_VERBOSE("computing condition graph completed, taking %" PRId64 " ms", t1 - t0); chunk_hidden_states = apply_token_weights(std::move(chunk_hidden_states), chunk_weights); @@ -608,7 +615,7 @@ struct FrozenCLIPVisionEmbedder : public GGMLRunner { auto get_graph = [&]() -> ggml_cgraph* { return build_graph(pixel_values, return_pooled, clip_skip); }; - return take_or_empty(GGMLRunner::compute(get_graph, n_threads, true, true, true)); + return take_or_empty(GGMLRunner::compute(get_graph, n_threads, true)); } }; @@ -675,18 +682,6 @@ struct SD3CLIPEmbedder : public Conditioner { } } - void set_stream_layers_enabled(bool enabled) override { - if (clip_l) { - clip_l->set_stream_layers_enabled(enabled); - } - if (clip_g) { - clip_g->set_stream_layers_enabled(enabled); - } - if (t5) { - t5->set_stream_layers_enabled(enabled); - } - } - void set_runtime_backends(const std::vector& backends) override { if (clip_l) { clip_l->set_runtime_backends(backends); @@ -753,15 +748,15 @@ struct SD3CLIPEmbedder : public Conditioner { } } - void runner_done() override { + void runner_end() override { if (clip_l) { - clip_l->runner_done(); + clip_l->runner_end(); } if (clip_g) { - clip_g->runner_done(); + clip_g->runner_end(); } if (t5) { - t5->runner_done(); + t5->runner_end(); } } @@ -778,7 +773,7 @@ struct SD3CLIPEmbedder : public Conditioner { ss << "['" << item.first << "', " << item.second << "], "; } ss << "]"; - LOG_DEBUG("parse '%s' to %s", text.c_str(), ss.str().c_str()); + LOG_VERBOSE("parse '%s' to %s", text.c_str(), ss.str().c_str()); } auto on_new_token_cb = [&](std::string& str, std::vector& bpe_tokens) -> bool { @@ -881,9 +876,7 @@ struct SD3CLIPEmbedder : public Conditioner { max_token_idx, false, clip_skip, - false, - true, - true); + false); GGML_ASSERT(!chunk_hidden_states_l.empty()); chunk_hidden_states_l = ::apply_token_weights(std::move(chunk_hidden_states_l), chunk_weights); @@ -897,9 +890,7 @@ struct SD3CLIPEmbedder : public Conditioner { max_token_idx, true, clip_skip, - false, - true, - true); + false); GGML_ASSERT(!pooled_l.empty()); } } else { @@ -928,9 +919,7 @@ struct SD3CLIPEmbedder : public Conditioner { max_token_idx, false, clip_skip, - false, - true, - true); + false); GGML_ASSERT(!chunk_hidden_states_g.empty()); chunk_hidden_states_g = ::apply_token_weights(std::move(chunk_hidden_states_g), chunk_weights); @@ -944,9 +933,7 @@ struct SD3CLIPEmbedder : public Conditioner { max_token_idx, true, clip_skip, - false, - true, - true); + false); GGML_ASSERT(!pooled_g.empty()); } } else { @@ -969,9 +956,7 @@ struct SD3CLIPEmbedder : public Conditioner { chunk_hidden_states_t5 = t5->compute(n_threads, input_ids, sd::Tensor(), - false, - true, - true); + false); GGML_ASSERT(!chunk_hidden_states_t5.empty()); chunk_hidden_states_t5 = ::apply_token_weights(std::move(chunk_hidden_states_t5), chunk_weights); } else { @@ -996,7 +981,7 @@ struct SD3CLIPEmbedder : public Conditioner { } int64_t t1 = ggml_time_ms(); - LOG_DEBUG("computing condition graph completed, taking %" PRId64 " ms", t1 - t0); + LOG_VERBOSE("computing condition graph completed, taking %" PRId64 " ms", t1 - t0); if (zero_out_masked) { chunk_hidden_states.fill_(0.0f); } @@ -1079,15 +1064,6 @@ struct FluxCLIPEmbedder : public Conditioner { } } - void set_stream_layers_enabled(bool enabled) override { - if (clip_l) { - clip_l->set_stream_layers_enabled(enabled); - } - if (t5) { - t5->set_stream_layers_enabled(enabled); - } - } - void set_runtime_backends(const std::vector& backends) override { if (clip_l) { clip_l->set_runtime_backends(backends); @@ -1139,12 +1115,12 @@ struct FluxCLIPEmbedder : public Conditioner { } } - void runner_done() override { + void runner_end() override { if (clip_l) { - clip_l->runner_done(); + clip_l->runner_end(); } if (t5) { - t5->runner_done(); + t5->runner_end(); } } @@ -1160,7 +1136,7 @@ struct FluxCLIPEmbedder : public Conditioner { ss << "['" << item.first << "', " << item.second << "], "; } ss << "]"; - LOG_DEBUG("parse '%s' to %s", text.c_str(), ss.str().c_str()); + LOG_VERBOSE("parse '%s' to %s", text.c_str(), ss.str().c_str()); } auto on_new_token_cb = [&](std::string& str, std::vector& bpe_tokens) -> bool { @@ -1247,9 +1223,7 @@ struct FluxCLIPEmbedder : public Conditioner { max_token_idx, true, clip_skip, - false, - true, - true); + false); GGML_ASSERT(!pooled.empty()); } else { pooled = sd::Tensor::zeros({768}); @@ -1268,9 +1242,7 @@ struct FluxCLIPEmbedder : public Conditioner { chunk_hidden_states = t5->compute(n_threads, input_ids, sd::Tensor(), - false, - true, - true); + false); GGML_ASSERT(!chunk_hidden_states.empty()); chunk_hidden_states = ::apply_token_weights(std::move(chunk_hidden_states), chunk_weights); if (zero_out_masked) { @@ -1281,7 +1253,7 @@ struct FluxCLIPEmbedder : public Conditioner { } int64_t t1 = ggml_time_ms(); - LOG_DEBUG("computing condition graph completed, taking %" PRId64 " ms", t1 - t0); + LOG_VERBOSE("computing condition graph completed, taking %" PRId64 " ms", t1 - t0); if (!hidden_states.empty()) { hidden_states = sd::ops::concat(hidden_states, chunk_hidden_states, 1); } else { @@ -1366,12 +1338,6 @@ struct T5CLIPEmbedder : public Conditioner { } } - void set_stream_layers_enabled(bool enabled) override { - if (t5) { - t5->set_stream_layers_enabled(enabled); - } - } - void set_runtime_backends(const std::vector& backends) override { if (t5) { t5->set_runtime_backends(backends); @@ -1408,9 +1374,9 @@ struct T5CLIPEmbedder : public Conditioner { } } - void runner_done() override { + void runner_end() override { if (t5) { - t5->runner_done(); + t5->runner_end(); } } @@ -1426,7 +1392,7 @@ struct T5CLIPEmbedder : public Conditioner { ss << "['" << item.first << "', " << item.second << "], "; } ss << "]"; - LOG_DEBUG("parse '%s' to %s", text.c_str(), ss.str().c_str()); + LOG_VERBOSE("parse '%s' to %s", text.c_str(), ss.str().c_str()); } auto on_new_token_cb = [&](std::string& str, std::vector& bpe_tokens) -> bool { @@ -1467,7 +1433,7 @@ struct T5CLIPEmbedder : public Conditioner { ++num_pad; } } - // LOG_DEBUG("PAD: %d", num_pad); + // LOG_VERBOSE("PAD: %d", num_pad); } SDCondition get_learned_condition_common(int n_threads, @@ -1508,9 +1474,7 @@ struct T5CLIPEmbedder : public Conditioner { auto chunk_hidden_states = t5->compute(n_threads, input_ids, t5_attn_mask_chunk, - false, - true, - true); + false); GGML_ASSERT(!chunk_hidden_states.empty()); chunk_hidden_states = apply_token_weights(std::move(chunk_hidden_states), chunk_weights); @@ -1521,7 +1485,7 @@ struct T5CLIPEmbedder : public Conditioner { } int64_t t1 = ggml_time_ms(); - LOG_DEBUG("computing condition graph completed, taking %" PRId64 " ms", t1 - t0); + LOG_VERBOSE("computing condition graph completed, taking %" PRId64 " ms", t1 - t0); if (!hidden_states.empty()) { hidden_states = sd::ops::concat(hidden_states, chunk_hidden_states, 1); @@ -1582,12 +1546,6 @@ struct MiniT2IConditioner : public Conditioner { } } - void set_stream_layers_enabled(bool enabled) override { - if (t5) { - t5->set_stream_layers_enabled(enabled); - } - } - void set_runtime_backends(const std::vector& backends) override { if (t5) { t5->set_runtime_backends(backends); @@ -1624,9 +1582,9 @@ struct MiniT2IConditioner : public Conditioner { } } - void runner_done() override { + void runner_end() override { if (t5) { - t5->runner_done(); + t5->runner_end(); } } @@ -1657,9 +1615,7 @@ struct MiniT2IConditioner : public Conditioner { sd::Tensor hidden_states = t5->compute(n_threads, input_ids, sd::Tensor::from_vector(t5_mask), - false, - true, - true); + false); GGML_ASSERT(!hidden_states.empty()); result.c_crossattn = std::move(hidden_states); result.c_vector = sd::Tensor::from_vector(mask); @@ -1696,10 +1652,6 @@ struct AnimaConditioner : public Conditioner { llm->set_max_graph_vram_bytes(max_vram_bytes); } - void set_stream_layers_enabled(bool enabled) override { - llm->set_stream_layers_enabled(enabled); - } - void set_runtime_backends(const std::vector& backends) override { llm->set_runtime_backends(backends); } @@ -1724,8 +1676,8 @@ struct AnimaConditioner : public Conditioner { llm->set_weight_adapter(adapter); } - void runner_done() override { - llm->runner_done(); + void runner_end() override { + llm->runner_end(); } std::tuple, std::vector, std::vector, std::vector> tokenize(std::string text) { @@ -1738,7 +1690,7 @@ struct AnimaConditioner : public Conditioner { ss << "['" << item.first << "', " << item.second << "], "; } ss << "]"; - LOG_DEBUG("parse '%s' to %s", text.c_str(), ss.str().c_str()); + LOG_VERBOSE("parse '%s' to %s", text.c_str(), ss.str().c_str()); } std::vector qwen_tokens; @@ -1787,16 +1739,14 @@ struct AnimaConditioner : public Conditioner { {}, {}, false, - false, - true, - true); + false); GGML_ASSERT(!hidden_states.empty()); hidden_states = apply_token_weights(std::move(hidden_states), qwen_weights); auto t5_ids_tensor = sd::Tensor::from_vector(t5_tokens); auto t5_weight_tensor = sd::Tensor::from_vector(t5_weights); int64_t t1 = ggml_time_ms(); - LOG_DEBUG("computing condition graph completed, taking %" PRId64 " ms", t1 - t0); + LOG_VERBOSE("computing condition graph completed, taking %" PRId64 " ms", t1 - t0); SDCondition result; result.c_crossattn = std::move(hidden_states); @@ -1887,13 +1837,6 @@ struct LLMEmbedder : public Conditioner { } } - void set_stream_layers_enabled(bool enabled) override { - llm->set_stream_layers_enabled(enabled); - if (byt5) { - byt5->set_stream_layers_enabled(enabled); - } - } - void set_runtime_backends(const std::vector& backends) override { llm->set_runtime_backends(backends); if (byt5) { @@ -1942,12 +1885,12 @@ struct LLMEmbedder : public Conditioner { } } - void runner_done() override { + void runner_end() override { if (llm) { - llm->runner_done(); + llm->runner_end(); } if (byt5) { - byt5->runner_done(); + byt5->runner_end(); } } @@ -1984,7 +1927,7 @@ struct LLMEmbedder : public Conditioner { ss << "['" << item.first << "', " << item.second << "], "; } ss << "]"; - LOG_DEBUG("parse '%s' to %s", text.c_str(), ss.str().c_str()); + LOG_VERBOSE("parse '%s' to %s", text.c_str(), ss.str().c_str()); } std::vector tokens; @@ -2051,8 +1994,6 @@ struct LLMEmbedder : public Conditioner { out_layers, false, false, - true, - true, deepstack_image_embeds, image_grids); GGML_ASSERT(!hidden_states.empty()); @@ -2220,9 +2161,7 @@ struct LLMEmbedder : public Conditioner { prompt += ": "; add_vision_outputs(llm->encode_image_outputs(n_threads, resized, - false, - true, - true), + false), static_cast(resized.shape()[1]) / patch_size, static_cast(resized.shape()[0]) / patch_size); continue; @@ -2250,9 +2189,7 @@ struct LLMEmbedder : public Conditioner { auto pair = sd::ops::concat(first.unsqueeze(2), second.unsqueeze(2), 2); add_vision_outputs(llm->encode_video_block_outputs(n_threads, pair, - false, - true, - true), + false), static_cast(first.shape()[1]) / patch_size, static_cast(first.shape()[0]) / patch_size); } @@ -2263,9 +2200,7 @@ struct LLMEmbedder : public Conditioner { prompt += ": "; add_vision_outputs(llm->encode_image_outputs(n_threads, resized, - false, - true, - true), + false), static_cast(resized.shape()[1]) / patch_size, static_cast(resized.shape()[0]) / patch_size); } @@ -2310,7 +2245,7 @@ struct LLMEmbedder : public Conditioner { prompt_template_encode_start_idx++; } } - LOG_DEBUG("prompt_template_encode_start_idx %d", prompt_template_encode_start_idx); + LOG_VERBOSE("prompt_template_encode_start_idx %d", prompt_template_encode_start_idx); prompt = prompt_prefix; if (llm->enable_vision && conditioner_params.ref_images != nullptr && !conditioner_params.ref_images->empty()) { @@ -2350,9 +2285,9 @@ struct LLMEmbedder : public Conditioner { resize_image_dims(height, width, h_bar, w_bar, factor, min_pixels, max_pixels, resize_mode); - LOG_DEBUG("resize LingBotVideo ref image %d from %dx%d to %dx%d", i, height, width, h_bar, w_bar); + LOG_VERBOSE("resize LingBotVideo ref image %d from %dx%d to %dx%d", i, height, width, h_bar, w_bar); auto resized_image = clip_preprocess(image, w_bar, h_bar); - auto image_embed = llm->encode_image(n_threads, resized_image, false, true, true); + auto image_embed = llm->encode_image(n_threads, resized_image, false); GGML_ASSERT(!image_embed.empty()); std::string image_prefix = prompt + img_prompt + "<|vision_start|>"; @@ -2407,11 +2342,11 @@ struct LLMEmbedder : public Conditioner { resize_image_dims(height, width, h_bar, w_bar, factor, min_pixels, max_pixels, resize_mode); - LOG_DEBUG("resize conditioner ref image %d from %dx%d to %dx%d", i, height, width, h_bar, w_bar); + LOG_VERBOSE("resize conditioner ref image %d from %dx%d to %dx%d", i, height, width, h_bar, w_bar); auto resized_image = clip_preprocess(image, w_bar, h_bar); - auto image_embed = llm->encode_image(n_threads, resized_image, false, true, true); + auto image_embed = llm->encode_image(n_threads, resized_image, false); GGML_ASSERT(!image_embed.empty()); image_embeds.emplace_back(image_embed_idx, image_embed); image_embed_idx += 1 + static_cast(image_embed.shape()[1]) + 6; @@ -2491,10 +2426,10 @@ struct LLMEmbedder : public Conditioner { resize_image_dims(height, width, h_bar, w_bar, factor, min_pixels, max_pixels, resize_mode); - LOG_DEBUG("resize conditioner ref image %d from %dx%d to %dx%d", i, height, width, h_bar, w_bar); + LOG_VERBOSE("resize conditioner ref image %d from %dx%d to %dx%d", i, height, width, h_bar, w_bar); auto resized_image = clip_preprocess(image, w_bar, h_bar); - auto image_embed = llm->encode_image(n_threads, resized_image, false, true, true); + auto image_embed = llm->encode_image(n_threads, resized_image, false); GGML_ASSERT(!image_embed.empty()); std::string image_prefix = prompt_prefix + img_prompt + "<|vision_start|>"; @@ -2559,10 +2494,10 @@ struct LLMEmbedder : public Conditioner { resize_image_dims(height, width, h_bar, w_bar, factor, min_pixels, max_pixels, resize_mode); - LOG_DEBUG("resize conditioner ref image %d from %dx%d to %dx%d", i, height, width, h_bar, w_bar); + LOG_VERBOSE("resize conditioner ref image %d from %dx%d to %dx%d", i, height, width, h_bar, w_bar); auto resized_image = clip_preprocess(image, w_bar, h_bar); - auto image_embed = llm->encode_image(n_threads, resized_image, false, true, true); + auto image_embed = llm->encode_image(n_threads, resized_image, false); GGML_ASSERT(!image_embed.empty()); std::string image_prefix = prompt + img_prompt + "Picture " + std::to_string(i + 1) + ": <|vision_start|>"; @@ -2622,10 +2557,10 @@ struct LLMEmbedder : public Conditioner { resize_image_dims(height, width, h_bar, w_bar, factor, min_pixels, max_pixels, resize_mode); - LOG_DEBUG("resize conditioner ref image %d from %dx%d to %dx%d", i, height, width, h_bar, w_bar); + LOG_VERBOSE("resize conditioner ref image %d from %dx%d to %dx%d", i, height, width, h_bar, w_bar); auto resized_image = clip_preprocess(image, w_bar, h_bar); - auto image_embed = llm->encode_image(n_threads, resized_image, false, true, true); + auto image_embed = llm->encode_image(n_threads, resized_image, false); GGML_ASSERT(!image_embed.empty()); image_embeds.emplace_back(image_embed_idx, image_embed); image_embed_idx += 1 + static_cast(image_embed.shape()[1]) + 6; @@ -2802,7 +2737,7 @@ struct LLMEmbedder : public Conditioner { } int64_t t1 = ggml_time_ms(); - LOG_DEBUG("computing condition graph completed, taking %" PRId64 " ms", t1 - t0); + LOG_VERBOSE("computing condition graph completed, taking %" PRId64 " ms", t1 - t0); SDCondition result; result.c_crossattn = std::move(hidden_states); @@ -2857,9 +2792,7 @@ struct LLMEmbedder : public Conditioner { auto byt5_hidden_states = byt5->compute(n_threads, input_ids, sd::Tensor(), - false, - true, - true); + false); GGML_ASSERT(!byt5_hidden_states.empty()); extra_hidden_states_vec.push_back(std::move(byt5_hidden_states)); } @@ -2879,7 +2812,7 @@ struct LLMEmbedder : public Conditioner { } int64_t t1 = ggml_time_ms(); - LOG_DEBUG("computing condition graph completed, taking %" PRId64 " ms", t1 - t0); + LOG_VERBOSE("computing condition graph completed, taking %" PRId64 " ms", t1 - t0); SDCondition result; result.c_crossattn = std::move(hidden_states); result.extra_c_crossattns = std::move(extra_hidden_states_vec); @@ -2960,13 +2893,11 @@ struct LTXAVTextProjectionRunner : public GGMLRunner { sd::Tensor compute(int n_threads, const sd::Tensor& x, - bool auto_free = true, - bool free_compute_buffer = true, - bool free_compute_params = true) { + bool auto_runner_end = true) { auto get_graph = [&]() -> ggml_cgraph* { return build_graph(x); }; - return take_or_empty(GGMLRunner::compute(get_graph, n_threads, auto_free, free_compute_buffer, free_compute_params)); + return take_or_empty(GGMLRunner::compute(get_graph, n_threads, auto_runner_end)); } }; @@ -3060,9 +2991,9 @@ struct LTXAVEmbedder : public Conditioner { projector->set_weight_adapter(adapter); } - void runner_done() override { - llm->runner_done(); - projector->runner_done(); + void runner_end() override { + llm->runner_end(); + projector->runner_end(); } std::tuple, std::vector, std::vector> tokenize(std::string text, @@ -3129,9 +3060,7 @@ struct LTXAVEmbedder : public Conditioner { {}, {}, true, - false, - true, - true); + false); GGML_ASSERT(!hidden_states.empty()); hidden_states = apply_token_weights(std::move(hidden_states), weights); @@ -3190,7 +3119,7 @@ struct LTXAVEmbedder : public Conditioner { } hidden_states.reshape_({kNumStates * kHiddenSize, valid_tokens}); - return projector->compute(n_threads, hidden_states, false, true, true); + return projector->compute(n_threads, hidden_states, false); } SDCondition get_learned_condition(int n_threads, @@ -3207,7 +3136,7 @@ struct LTXAVEmbedder : public Conditioner { GGML_ASSERT(!hidden_states.empty()); int64_t t1 = ggml_time_ms(); - LOG_DEBUG("computing LTXAV condition graph completed, taking %" PRId64 " ms", t1 - t0); + LOG_VERBOSE("computing LTXAV condition graph completed, taking %" PRId64 " ms", t1 - t0); SDCondition result; result.c_crossattn = std::move(hidden_states); diff --git a/otherarch/sdcpp/src/core/backend_fit.cpp b/otherarch/sdcpp/src/core/backend_fit.cpp index 0ba2df14f..0f2dd67cc 100644 --- a/otherarch/sdcpp/src/core/backend_fit.cpp +++ b/otherarch/sdcpp/src/core/backend_fit.cpp @@ -2,364 +2,384 @@ #include #include +#include #include +#include #include #include +#if defined(_WIN32) +#ifndef NOMINMAX +#define NOMINMAX +#endif +#include +#elif defined(__APPLE__) +#include +#endif + #include "core/ggml_extend_backend.h" #include "core/util.h" #include "ggml-backend.h" namespace sd::backend_fit { - namespace { - constexpr int64_t MiB = 1024ll * 1024; + static constexpr int64_t MiB = 1024ll * 1024; - enum class ComponentKind { - DIT = 0, - VAE = 1, - CONDITIONER = 2, - }; + enum class ComponentKind { + DIT, + CONDITIONER, + VAE, + }; - struct Component { + struct Component { + ComponentKind kind; + const char* name; + int64_t params_bytes = 0; + int64_t reserve_bytes = 0; + int64_t staging_bytes = 0; + }; + + struct Device { + std::string name; + std::string description; + int64_t free_bytes = 0; + int64_t budget_bytes = 0; + }; + + enum class ParamsLocation { + MAIN_GPU, + CPU, + OTHER_GPU, + DISK, + }; + + struct Decision { + ParamsLocation params_location = ParamsLocation::DISK; + size_t params_device = SIZE_MAX; + }; + + struct Plan { + bool valid = false; + size_t main_device = SIZE_MAX; + std::vector decisions; + }; + + static bool classify_tensor(const std::string& name, ComponentKind& out) { + auto contains = [&](const char* s) { return name.find(s) != std::string::npos; }; + + if (contains("model.diffusion_model.") || contains("unet.")) { + out = ComponentKind::DIT; + return true; + } + if (contains("first_stage_model.") || + name.rfind("vae.", 0) == 0 || + name.rfind("tae.", 0) == 0) { + out = ComponentKind::VAE; + return true; + } + if (contains("text_encoders") || + contains("cond_stage_model") || + contains("te.text_model.") || + contains("conditioner") || + name.rfind("text_encoder.", 0) == 0 || + name.rfind("text_embedding_projection.", 0) == 0 || + contains(".aggregate_embed.")) { + out = ComponentKind::CONDITIONER; + return true; + } + return false; + } + + static std::vector estimate_components(ModelLoader& loader, ggml_type override_wtype) { + int64_t bytes[3] = {0, 0, 0}; + int64_t largest_tensor[3] = {0, 0, 0}; + for (const auto& [name, stored_tensor] : loader.get_tensor_storage_map()) { + TensorStorage ts = stored_tensor; ComponentKind kind; - const char* name; - int64_t params_bytes = 0; - int64_t reserve_bytes = 0; - bool splittable = false; - }; - - struct Device { - ggml_backend_dev_t dev = nullptr; - std::string name; - std::string description; - int64_t free_bytes = 0; - int64_t total_bytes = 0; - int64_t budget_bytes = 0; - }; - - struct Decision { - ComponentKind kind; - bool on_cpu = false; - std::vector device_idxs; - }; - - struct Plan { - bool valid = false; - bool time_share = false; - std::vector decisions; - }; - - bool classify_tensor(const std::string& name, ComponentKind& out) { - auto contains = [&](const char* s) { return name.find(s) != std::string::npos; }; - - if (contains("model.diffusion_model.") || contains("unet.")) { - out = ComponentKind::DIT; - return true; + if (is_unused_tensor(ts.name) || !classify_tensor(ts.name, kind)) { + continue; } - if (contains("first_stage_model.") || - name.rfind("vae.", 0) == 0 || - name.rfind("tae.", 0) == 0) { - out = ComponentKind::VAE; - return true; + if (ts.expected_type != GGML_TYPE_COUNT) { + ts.type = ts.expected_type; + } else if (override_wtype != GGML_TYPE_COUNT && loader.tensor_should_be_converted(ts, override_wtype)) { + ts.type = override_wtype; } - if (contains("text_encoders") || - contains("cond_stage_model") || - contains("te.text_model.") || - contains("conditioner") || - name.rfind("text_encoder.", 0) == 0 || - name.rfind("text_embedding_projection.", 0) == 0 || - contains(".aggregate_embed.")) { - out = ComponentKind::CONDITIONER; - return true; - } - return false; + const int64_t tensor_bytes = (int64_t)ts.nbytes() + 64; + bytes[int(kind)] += tensor_bytes; + largest_tensor[int(kind)] = std::max(largest_tensor[int(kind)], tensor_bytes); } - std::vector estimate_components(ModelLoader& loader, ggml_type override_wtype) { - const auto& storage = loader.get_tensor_storage_map(); + return { + {ComponentKind::DIT, "DiT", bytes[int(ComponentKind::DIT)], 2048 * MiB, largest_tensor[int(ComponentKind::DIT)]}, + {ComponentKind::CONDITIONER, "Conditioner", bytes[int(ComponentKind::CONDITIONER)], 2048 * MiB, largest_tensor[int(ComponentKind::CONDITIONER)]}, + {ComponentKind::VAE, "VAE", bytes[int(ComponentKind::VAE)], 1024 * MiB, largest_tensor[int(ComponentKind::VAE)]}, + }; + } - int64_t bytes[3] = {0, 0, 0}; - for (const auto& [name, ts_const] : storage) { - TensorStorage ts = ts_const; - if (is_unused_tensor(ts.name)) { - continue; - } - ComponentKind kind; - if (!classify_tensor(ts.name, kind)) { - continue; - } - if (override_wtype != GGML_TYPE_COUNT && - loader.tensor_should_be_converted(ts, override_wtype)) { - ts.type = override_wtype; - } else if (ts.expected_type != GGML_TYPE_COUNT && ts.expected_type != ts.type) { - ts.type = ts.expected_type; - } - bytes[int(kind)] += (int64_t)ts.nbytes() + 64; + static std::string budget_key(std::string name) { + std::transform(name.begin(), name.end(), name.begin(), [](unsigned char c) { return (char)std::tolower(c); }); + return name; + } + + static std::vector enumerate_gpu_devices(const sd::ggml_graph_cut::MaxVramAssignment& budgets) { + std::vector out; + for (size_t i = 0; i < ggml_backend_dev_count(); ++i) { + ggml_backend_dev_t dev = ggml_backend_dev_get(i); + if (ggml_backend_dev_type(dev) != GGML_BACKEND_DEVICE_TYPE_GPU) { + continue; } + Device device; + device.name = ggml_backend_dev_name(dev); + device.description = ggml_backend_dev_description(dev); + size_t free_bytes = 0, total_bytes = 0; + ggml_backend_dev_memory(dev, &free_bytes, &total_bytes); + device.free_bytes = (int64_t)free_bytes; - std::vector out; - out.push_back({ComponentKind::DIT, "DiT", bytes[int(ComponentKind::DIT)], 2048 * MiB, true}); - out.push_back({ComponentKind::VAE, "VAE", bytes[int(ComponentKind::VAE)], 1024 * MiB, false}); - out.push_back({ComponentKind::CONDITIONER, "Conditioner", bytes[int(ComponentKind::CONDITIONER)], 2048 * MiB, true}); - return out; + float gib = budgets.default_gib; + auto it = budgets.backend_gib.find(budget_key(device.name)); + if (it != budgets.backend_gib.end()) { + gib = it->second; + } + if (gib > 0.f) { + device.budget_bytes = (int64_t)std::min(gib * 1024.0 * MiB, (double)device.free_bytes); + } else if (gib < 0.f) { + device.budget_bytes = (int64_t)std::max(device.free_bytes + gib * 1024.0 * MiB, 0); + } else { + device.budget_bytes = std::max(device.free_bytes - 512 * MiB, 0); + } + out.push_back(std::move(device)); } + return out; + } - std::vector enumerate_gpu_devices(const sd::ggml_graph_cut::MaxVramAssignment& budgets) { - std::vector out; - for (size_t i = 0; i < ggml_backend_dev_count(); i++) { - ggml_backend_dev_t dev = ggml_backend_dev_get(i); - if (ggml_backend_dev_type(dev) != GGML_BACKEND_DEVICE_TYPE_GPU) { - continue; - } - Device d; - d.dev = dev; - d.name = ggml_backend_dev_name(dev); - d.description = ggml_backend_dev_description(dev); - size_t free_bytes = 0, total_bytes = 0; - ggml_backend_dev_memory(dev, &free_bytes, &total_bytes); - d.free_bytes = (int64_t)free_bytes; - d.total_bytes = (int64_t)total_bytes; - - std::string budget_key = d.name; - std::transform(budget_key.begin(), budget_key.end(), budget_key.begin(), - [](unsigned char c) { return (char)std::tolower(c); }); - float gib = budgets.default_gib; - auto it = budgets.backend_gib.find(budget_key); - if (it != budgets.backend_gib.end()) { - gib = it->second; - } - if (gib > 0.f) { - d.budget_bytes = std::min((int64_t)(gib * 1024.0 * 1024.0 * 1024.0), d.free_bytes); - } else if (gib < 0.f) { - d.budget_bytes = d.free_bytes + (int64_t)(gib * 1024.0 * 1024.0 * 1024.0); - } else { - d.budget_bytes = d.free_bytes - 512 * MiB; - } - d.budget_bytes = std::max(d.budget_bytes, 0); - out.push_back(d); - } - return out; + static int64_t available_ram_bytes() { +#if defined(_WIN32) + MEMORYSTATUSEX status{}; + status.dwLength = sizeof(status); + if (GlobalMemoryStatusEx(&status)) { + return (int64_t)status.ullAvailPhys; } - - Plan compute_plan(const std::vector& components, const std::vector& devices) { - Plan plan; - if (devices.empty()) { - return plan; +#elif defined(__linux__) + std::ifstream meminfo("/proc/meminfo"); + std::string key, unit; + int64_t kib = 0; + while (meminfo >> key >> kib >> unit) { + if (key == "MemAvailable:" && unit == "kB" && kib >= 0) { + return kib * 1024; } + } +#elif defined(__APPLE__) + const mach_port_t host = mach_host_self(); + vm_size_t page_size = 0; + vm_statistics64_data_t stats{}; + mach_msg_type_number_t count = HOST_VM_INFO64_COUNT; + const bool ok = host_page_size(host, &page_size) == KERN_SUCCESS && + host_statistics64(host, HOST_VM_INFO64, (host_info64_t)&stats, &count) == KERN_SUCCESS; + mach_port_deallocate(mach_task_self(), host); + if (ok) { + return ((int64_t)stats.free_count + stats.inactive_count) * page_size; + } +#endif + return -1; + } - std::vector order(components.size()); - for (size_t i = 0; i < order.size(); i++) { - order[i] = i; + static Plan compute_plan(const std::vector& components, + const std::vector& devices, + int64_t ram_budget_bytes) { + Plan plan; + for (size_t di = 0; di < devices.size(); ++di) { + if (devices[di].budget_bytes > 0 && + (plan.main_device == SIZE_MAX || devices[di].budget_bytes > devices[plan.main_device].budget_bytes)) { + plan.main_device = di; } - std::sort(order.begin(), order.end(), [&](size_t a, size_t b) { - return components[a].params_bytes > components[b].params_bytes; - }); - - { - std::vector params_sum(devices.size(), 0); - std::vector max_reserve(devices.size(), 0); - std::vector decisions(components.size()); - bool ok = true; - for (size_t ci : order) { - const Component& comp = components[ci]; - decisions[ci].kind = comp.kind; - if (comp.params_bytes == 0) { - continue; - } - int best = -1; - for (size_t di = 0; di < devices.size(); di++) { - int64_t need = params_sum[di] + comp.params_bytes + std::max(max_reserve[di], comp.reserve_bytes); - if (need <= devices[di].budget_bytes && - (best < 0 || devices[di].budget_bytes - params_sum[di] > devices[best].budget_bytes - params_sum[best])) { - best = (int)di; - } - } - if (best < 0) { - ok = false; - break; - } - params_sum[best] += comp.params_bytes; - max_reserve[best] = std::max(max_reserve[best], comp.reserve_bytes); - decisions[ci].device_idxs.push_back((size_t)best); - } - if (ok) { - plan.valid = true; - plan.time_share = false; - plan.decisions = std::move(decisions); - return plan; - } - } - - plan.decisions.assign(components.size(), {}); - for (size_t ci : order) { - const Component& comp = components[ci]; - Decision& decision = plan.decisions[ci]; - decision.kind = comp.kind; - if (comp.params_bytes == 0) { - continue; - } - int best = -1; - for (size_t di = 0; di < devices.size(); di++) { - if (comp.params_bytes + comp.reserve_bytes <= devices[di].budget_bytes && - (best < 0 || devices[di].budget_bytes > devices[best].budget_bytes)) { - best = (int)di; - } - } - if (best >= 0) { - decision.device_idxs.push_back((size_t)best); - continue; - } - if (comp.splittable && devices.size() > 1) { - int64_t capacity = 0; - for (const Device& d : devices) { - capacity += std::max(d.budget_bytes - comp.reserve_bytes, 0); - } - if (comp.params_bytes <= capacity) { - std::vector idxs(devices.size()); - for (size_t i = 0; i < idxs.size(); i++) { - idxs[i] = i; - } - std::sort(idxs.begin(), idxs.end(), [&](size_t a, size_t b) { - return devices[a].budget_bytes > devices[b].budget_bytes; - }); - decision.device_idxs = std::move(idxs); - continue; - } - } - decision.on_cpu = true; - } - plan.valid = true; - plan.time_share = true; + } + if (plan.main_device == SIZE_MAX) { return plan; } - void print_plan(const Plan& plan, - const std::vector& components, - const std::vector& devices) { - LOG_INFO("auto-fit plan%s:", plan.time_share ? " (time-share: params load per phase and free after)" : ""); - LOG_INFO(" devices:"); - for (const Device& d : devices) { - LOG_INFO(" %-12s %-32s free %6lld MiB, budget %6lld MiB", - d.name.c_str(), d.description.c_str(), - (long long)(d.free_bytes / MiB), (long long)(d.budget_bytes / MiB)); - } - LOG_INFO(" components:"); - for (size_t ci = 0; ci < components.size(); ci++) { - const Component& comp = components[ci]; - const Decision& decision = plan.decisions[ci]; - std::string target; - if (comp.params_bytes == 0) { - target = "(not present)"; - } else if (decision.on_cpu) { - target = "CPU"; - } else { - for (size_t k = 0; k < decision.device_idxs.size(); k++) { - if (k > 0) { - target += " & "; - } - target += devices[decision.device_idxs[k]].name; - } - if (decision.device_idxs.size() > 1) { - target += " (split)"; - } - } - LOG_INFO(" %-12s params %6lld MiB, compute reserve %5lld MiB -> %s", - comp.name, - (long long)(comp.params_bytes / MiB), - (long long)(comp.reserve_bytes / MiB), - target.c_str()); - } + std::vector order(components.size()); + for (size_t ci = 0; ci < components.size(); ++ci) { + order[ci] = ci; } + std::stable_sort(order.begin(), order.end(), [&](size_t a, size_t b) { + return components[a].kind < components[b].kind; + }); - void append_assignment(std::string& spec, const char* key, const std::string& value) { - if (!spec.empty()) { - spec += ","; - } - spec += key; - spec += "="; - spec += value; + std::vector remaining; + for (const Device& device : devices) { + remaining.push_back(std::max(device.budget_bytes, 0)); } + ram_budget_bytes = std::max(ram_budget_bytes, 0); + plan.decisions.resize(components.size()); - void append_component_decision(const std::vector& components, - const std::vector& devices, - const Plan& plan, - ComponentKind kind, - const char* module_key, - std::string& runtime_spec, - std::string& params_spec) { - for (size_t ci = 0; ci < components.size(); ci++) { - if (components[ci].kind != kind || components[ci].params_bytes == 0) { + for (size_t ci : order) { + const Component& comp = components[ci]; + Decision& decision = plan.decisions[ci]; + if (comp.params_bytes == 0) { + continue; + } + + // Higher-priority offloaded weights need GPU cache space across graph runs. + int64_t headroom = 0; + for (size_t other = 0; other < components.size(); ++other) { + if (components[other].params_bytes == 0) { continue; } - const Decision& decision = plan.decisions[ci]; - if (decision.on_cpu) { - append_assignment(runtime_spec, module_key, "cpu"); - return; + const bool resident = other == ci || plan.decisions[other].params_location == ParamsLocation::MAIN_GPU; + const int64_t cached_weights = components[other].kind < comp.kind + ? components[other].params_bytes + : components[other].staging_bytes; + headroom = std::max(headroom, components[other].reserve_bytes + + (resident ? 0 : cached_weights)); + } + int64_t& main_remaining = remaining[plan.main_device]; + if (headroom <= main_remaining && comp.params_bytes <= main_remaining - headroom) { + decision.params_location = ParamsLocation::MAIN_GPU; + decision.params_device = plan.main_device; + main_remaining -= comp.params_bytes; + continue; + } + if (comp.params_bytes <= ram_budget_bytes) { + decision.params_location = ParamsLocation::CPU; + ram_budget_bytes -= comp.params_bytes; + continue; + } + + size_t best = SIZE_MAX; + for (size_t di = 0; di < devices.size(); ++di) { + if (di != plan.main_device && comp.params_bytes <= remaining[di] && + (best == SIZE_MAX || remaining[di] > remaining[best])) { + best = di; } - if (decision.device_idxs.empty()) { - return; - } - std::string device_list; - for (size_t k = 0; k < decision.device_idxs.size(); k++) { - if (k > 0) { - device_list += "&"; - } - device_list += devices[decision.device_idxs[k]].name; - } - append_assignment(runtime_spec, module_key, device_list); - if (plan.time_share) { - append_assignment(params_spec, module_key, "disk"); - } - return; + } + if (best != SIZE_MAX) { + decision.params_location = ParamsLocation::OTHER_GPU; + decision.params_device = best; + remaining[best] -= comp.params_bytes; } } + plan.valid = true; + return plan; + } - } // namespace + static std::string params_backend_name(const Decision& decision, const std::vector& devices) { + switch (decision.params_location) { + case ParamsLocation::MAIN_GPU: + case ParamsLocation::OTHER_GPU: + return devices[decision.params_device].name; + case ParamsLocation::CPU: + return "cpu"; + case ParamsLocation::DISK: + return "disk"; + } + return "disk"; + } + + static void print_plan(const Plan& plan, + const std::vector& components, + const std::vector& devices, + int64_t free_ram, + int64_t ram_budget) { + LOG_INFO("auto-fit plan (single-GPU compute on %s):", devices[plan.main_device].name.c_str()); + LOG_INFO(" devices:"); + for (const Device& device : devices) { + LOG_INFO(" %-12s %-32s free %6lld MiB, budget %6lld MiB", + device.name.c_str(), device.description.c_str(), + (long long)(device.free_bytes / MiB), (long long)(device.budget_bytes / MiB)); + } + if (free_ram < 0) { + LOG_WARN("auto-fit: available RAM is unknown; skipping CPU parameter residency"); + } else { + LOG_INFO(" RAM free %6lld MiB, params budget %6lld MiB", + (long long)(free_ram / MiB), (long long)(ram_budget / MiB)); + } + LOG_INFO(" main-GPU weight cache priority: diffusion > te > vae"); + LOG_INFO(" components (params: main GPU -> RAM -> other GPU -> disk):"); + for (size_t ci = 0; ci < components.size(); ++ci) { + const Component& comp = components[ci]; + if (comp.params_bytes == 0) { + continue; + } + const std::string params = params_backend_name(plan.decisions[ci], devices); + LOG_INFO(" %-12s params %6lld MiB, compute reserve %5lld MiB -> compute %s, params %s", + comp.name, (long long)(comp.params_bytes / MiB), (long long)(comp.reserve_bytes / MiB), + devices[plan.main_device].name.c_str(), params.c_str()); + } + } + + static void append_assignment(std::string& spec, const char* key, const std::string& value) { + if (!spec.empty()) { + spec += ","; + } + spec += key; + spec += "="; + spec += value; + } + + static const char* module_key(ComponentKind kind) { + switch (kind) { + case ComponentKind::DIT: + return "diffusion"; + case ComponentKind::CONDITIONER: + return "te"; + case ComponentKind::VAE: + return "vae"; + } + return ""; + } bool derive_backend_specs(ModelLoader& loader, ggml_type override_wtype, sd::ggml_graph_cut::MaxVramAssignment& budgets, std::string& runtime_spec, std::string& params_spec) { - if (!runtime_spec.empty() || !params_spec.empty()) { - LOG_WARN("--auto-fit is enabled; ignoring --backend / --params-backend"); + std::string error; + if (!budgets.canonicalize_backend_keys(&error)) { + LOG_ERROR("%s", error.c_str()); + return false; } - { - std::string error; - if (!budgets.canonicalize_backend_keys(&error)) { - LOG_ERROR("%s", error.c_str()); - return false; - } - } - - auto components = estimate_components(loader, override_wtype); - auto devices = enumerate_gpu_devices(budgets); - auto plan = compute_plan(components, devices); + const auto components = estimate_components(loader, override_wtype); + const auto devices = enumerate_gpu_devices(budgets); + const int64_t free_ram = available_ram_bytes(); + const int64_t ram_budget = std::max(free_ram - std::max(2048 * MiB, free_ram / 10), 0); + const auto plan = compute_plan(components, devices, ram_budget); + runtime_spec.clear(); + params_spec.clear(); if (!plan.valid) { - LOG_WARN("auto-fit: no usable GPU devices; using the default backend"); - runtime_spec.clear(); - params_spec.clear(); + if (devices.empty()) { + LOG_WARN("auto-fit: no GPU devices; using the default backend"); + } else { + LOG_WARN("auto-fit: no GPU memory budget available; using CPU"); + runtime_spec = "cpu"; + } return true; } - print_plan(plan, components, devices); + print_plan(plan, components, devices, free_ram, ram_budget); + for (size_t ci = 0; ci < components.size(); ++ci) { + if (components[ci].params_bytes == 0) { + continue; + } + const char* key = module_key(components[ci].kind); + append_assignment(runtime_spec, key, devices[plan.main_device].name); + if (plan.decisions[ci].params_location != ParamsLocation::MAIN_GPU) { + append_assignment(params_spec, key, params_backend_name(plan.decisions[ci], devices)); + } + } - std::string derived_runtime_spec; - std::string derived_params_spec; - append_component_decision(components, devices, plan, ComponentKind::DIT, "diffusion", derived_runtime_spec, derived_params_spec); - append_component_decision(components, devices, plan, ComponentKind::CONDITIONER, "te", derived_runtime_spec, derived_params_spec); - append_component_decision(components, devices, plan, ComponentKind::VAE, "vae", derived_runtime_spec, derived_params_spec); - - runtime_spec = std::move(derived_runtime_spec); - params_spec = std::move(derived_params_spec); + // Keep the planner's safety margin when the runner resolves its device limits. + for (const Device& device : devices) { + if (device.budget_bytes > 0) { + budgets.backend_gib[budget_key(device.name)] = (float)(device.budget_bytes / (1024.0 * MiB)); + } + } + budgets.resolved_backend_bytes.clear(); LOG_INFO("auto-fit: --backend \"%s\"%s%s%s", runtime_spec.empty() ? "(default)" : runtime_spec.c_str(), params_spec.empty() ? "" : " --params-backend \"", - params_spec.c_str(), - params_spec.empty() ? "" : "\""); + params_spec.c_str(), params_spec.empty() ? "" : "\""); return true; } diff --git a/otherarch/sdcpp/src/core/compute_workspace.cpp b/otherarch/sdcpp/src/core/compute_workspace.cpp new file mode 100644 index 000000000..a40e2dbd8 --- /dev/null +++ b/otherarch/sdcpp/src/core/compute_workspace.cpp @@ -0,0 +1,257 @@ +#include "core/compute_workspace.h" + +#include +#include +#include +#include +#include + +#include "core/ggml_extend_backend.h" +#include "core/ggml_graph_cut.h" +#include "ggml-cpu.h" +#include "ggml/src/ggml-impl.h" + +namespace sd { + ComputeWorkspace::~ComputeWorkspace() { + segment_end(); + release(); + ggml_backend_free(cpu_backend_); + } + + void ComputeWorkspace::set_extra_backends(const std::vector& backends) { + if (extra_backends_ != backends) { + GGML_ASSERT(!active_); + release(); + extra_backends_ = backends; + } + } + + bool ComputeWorkspace::needs_scheduler(ggml_cgraph* graph) const { + if (!extra_backends_.empty()) { + return true; + } + for (int i = 0; i < ggml_graph_n_nodes(graph); ++i) { + if (!ggml_backend_supports_op(backend_, ggml_graph_node(graph, i))) { + return true; + } + } + return false; + } + + ggml_backend_sched_t ComputeWorkspace::make_scheduler(size_t graph_size) { + std::vector backends{backend_}; + backends.insert(backends.end(), extra_backends_.begin(), extra_backends_.end()); + if (!sd_backend_is_cpu(backend_)) { + if (cpu_backend_ == nullptr) { + cpu_backend_ = sd_backend_cpu_init(); + } + if (cpu_backend_ == nullptr) { + return nullptr; + } + backends.push_back(cpu_backend_); + } + std::vector bufts; + for (auto backend : backends) { + auto buft = backend == cpu_backend_ + ? ggml_backend_dev_host_buffer_type(ggml_backend_get_device(backend_)) + : nullptr; + bufts.push_back(buft != nullptr ? buft : ggml_backend_get_default_buffer_type(backend)); + } + return ggml_backend_sched_new(backends.data(), bufts.data(), static_cast(backends.size()), + graph_size, false, false); + } + + bool ComputeWorkspace::measurement_matches(ggml_cgraph* graph, const Measurement& measurement) const { + return measurement.scheduler == needs_scheduler(graph); + } + + bool ComputeWorkspace::prepare(const Measurement& measurement) { + GGML_ASSERT(!active_); + if (measurement.buffers.empty()) { + return false; + } + const bool grows = std::any_of(measurement.buffers.begin(), measurement.buffers.end(), + [&](const BackendBufferSize& size) { return size.bytes > bytes(size.backend); }); + if (measurement.scheduler != (scheduler_ != nullptr) || grows) { + release(); + } + return true; + } + + bool ComputeWorkspace::release_excess(const Measurement& measurement) { + return std::any_of(measurement.buffers.begin(), measurement.buffers.end(), + [&](const BackendBufferSize& size) { return bytes(size.backend) > size.bytes; }) && + release(); + } + + bool ComputeWorkspace::allocate(ggml_cgraph* graph, const AssignNodes& assign_nodes) { + GGML_ASSERT(!active_); + const bool use_scheduler = needs_scheduler(graph); + if (use_scheduler) { + if (allocator_ != nullptr) { + release(); + } + const size_t capacity = static_cast(graph->n_nodes + graph->n_leafs) + 8; + if (scheduler_ == nullptr || capacity > scheduler_capacity_) { + release(); + scheduler_ = make_scheduler(capacity); + scheduler_capacity_ = capacity; + } + if (scheduler_ == nullptr) { + return false; + } + ggml_backend_sched_reset(scheduler_); + assign_nodes(scheduler_, graph); + // Scheduler allocation rewrites sources. Split the execution graph only once. + if (!ggml_backend_sched_alloc_graph(scheduler_, graph)) { + release(); + return false; + } + } else { + if (scheduler_ != nullptr) { + release(); + } + if (allocator_ == nullptr) { + allocator_ = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend_)); + } + auto signature = ggml_graph_cut::graph_layout(graph, true); + if (signature != reservation_) { + if (!ggml_gallocr_reserve(allocator_, graph)) { + release(); + return false; + } + reservation_ = std::move(signature); + ++reservations_; + } + if (!ggml_gallocr_alloc_graph(allocator_, graph)) { + release(); + return false; + } + } + active_ = true; + return true; + } + + ComputeWorkspace::Measurement ComputeWorkspace::measure( + ggml_cgraph* graph, + size_t direct_bytes, + const std::function& external_backend, + const AssignNodes& assign_nodes) { + if (!needs_scheduler(graph)) { + return {{{backend_, direct_bytes}}, false}; + } + std::vector tensors; + std::unordered_set seen; + auto visit = [&](const ggml_tensor* tensor) { + if (tensor != nullptr && seen.insert(tensor).second) { + tensors.push_back(tensor); + } + }; + for (int i = 0; i < graph->n_nodes; ++i) { + visit(graph->nodes[i]); + } + for (int i = 0; i < graph->n_leafs; ++i) { + visit(graph->leafs[i]); + } + for (size_t i = 0; i < tensors.size(); ++i) { + visit(tensors[i]->view_src); + for (auto source : tensors[i]->src) { + visit(source); + } + } + const size_t graph_size = tensors.size() + 8; + auto context = ggml_init({tensors.size() * ggml_tensor_overhead() + ggml_graph_overhead_custom(graph_size, false), nullptr, true}); + if (context == nullptr) { + return {}; + } + std::unordered_map copies; + std::map external_buffers; + for (auto tensor : tensors) { + auto copy = ggml_dup_tensor(context, tensor); + *copy = *tensor; + copies[tensor] = copy; + } + for (const auto& entry : copies) { + auto source = entry.first; + auto copy = entry.second; + copy->view_src = source->view_src == nullptr ? nullptr : copies.at(source->view_src); + for (int i = 0; i < GGML_MAX_SRC; ++i) { + copy->src[i] = source->src[i] == nullptr ? nullptr : copies.at(source->src[i]); + } + auto external = external_backend(source); + if (external != nullptr && source->view_src == nullptr) { + auto& buffer = external_buffers[external]; + if (buffer == nullptr) { + buffer = ggml_backend_alloc_buffer(external, 0); + GGML_ASSERT(buffer != nullptr); + ggml_backend_buffer_set_usage(buffer, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); + } + copy->buffer = buffer; + copy->data = reinterpret_cast(static_cast(1)); + copy->extra = nullptr; + } + } + auto copy_graph = ggml_new_graph_custom(context, graph_size, false); + copy_graph->n_nodes = graph->n_nodes; + copy_graph->n_leafs = graph->n_leafs; + for (int i = 0; i < graph->n_nodes; ++i) { + copy_graph->nodes[i] = copies.at(graph->nodes[i]); + } + for (int i = 0; i < graph->n_leafs; ++i) { + copy_graph->leafs[i] = copies.at(graph->leafs[i]); + } + Measurement result; + result.scheduler = true; + auto scheduler = make_scheduler(graph_size); + if (scheduler != nullptr) { + assign_nodes(scheduler, copy_graph); + std::vector sizes(extra_backends_.size() + 2); + ggml_backend_sched_reserve_size(scheduler, copy_graph, sizes.data()); + result.buffers.push_back({backend_, sizes[0]}); + for (size_t i = 0; i < extra_backends_.size(); ++i) { + result.buffers.push_back({extra_backends_[i], sizes[i + 1]}); + } + ggml_backend_sched_free(scheduler); + } + for (const auto& entry : external_buffers) { + ggml_backend_buffer_free(entry.second); + } + ggml_free(context); + return result; + } + + void ComputeWorkspace::synchronize() const { + if (scheduler_ != nullptr) { + ggml_backend_sched_synchronize(scheduler_); + } else { + ggml_backend_synchronize(backend_); + } + } + + void ComputeWorkspace::segment_end() { + if (active_) { + synchronize(); + active_ = false; + } + } + + bool ComputeWorkspace::release() { + if (active_) { + return false; + } + ggml_gallocr_free(allocator_); + allocator_ = nullptr; + ggml_backend_sched_free(scheduler_); + scheduler_ = nullptr; + scheduler_capacity_ = 0; + reservation_.clear(); + return true; + } + + size_t ComputeWorkspace::bytes(ggml_backend_t backend) const { + if (scheduler_ != nullptr) { + return ggml_backend_sched_get_buffer_size(scheduler_, backend); + } + return allocator_ != nullptr && backend == backend_ ? ggml_gallocr_get_buffer_size(allocator_, 0) : 0; + } +} diff --git a/otherarch/sdcpp/src/core/compute_workspace.h b/otherarch/sdcpp/src/core/compute_workspace.h new file mode 100644 index 000000000..7e73c23e3 --- /dev/null +++ b/otherarch/sdcpp/src/core/compute_workspace.h @@ -0,0 +1,64 @@ +#ifndef __SD_CORE_COMPUTE_WORKSPACE_H__ +#define __SD_CORE_COMPUTE_WORKSPACE_H__ + +#include +#include + +#include "ggml-alloc.h" +#include "ggml-backend.h" + +namespace sd { + struct BackendBufferSize { + ggml_backend_t backend = nullptr; + size_t bytes = 0; + }; + + class ComputeWorkspace { + ggml_backend_t backend_; + std::vector extra_backends_; + ggml_backend_t cpu_backend_ = nullptr; + ggml_gallocr_t allocator_ = nullptr; + ggml_backend_sched_t scheduler_ = nullptr; + size_t scheduler_capacity_ = 0; + std::vector reservation_; + bool active_ = false; + size_t reservations_ = 0; + + ggml_backend_sched_t make_scheduler(size_t graph_size); + bool needs_scheduler(ggml_cgraph* graph) const; + + public: + struct Measurement { + std::vector buffers; + bool scheduler = false; + }; + using AssignNodes = std::function; + + explicit ComputeWorkspace(ggml_backend_t backend) + : backend_(backend) {} + ~ComputeWorkspace(); + ComputeWorkspace(const ComputeWorkspace&) = delete; + ComputeWorkspace& operator=(const ComputeWorkspace&) = delete; + + void set_extra_backends(const std::vector& backends); + bool measurement_matches(ggml_cgraph* graph, const Measurement& measurement) const; + bool prepare(const Measurement& measurement); + bool release_excess(const Measurement& measurement); + bool allocate(ggml_cgraph* graph, const AssignNodes& assign_nodes); + Measurement measure( + ggml_cgraph* graph, + size_t direct_bytes, + const std::function& external_backend, + const AssignNodes& assign_nodes); + void synchronize() const; + void segment_end(); + bool release(); + bool active() const { return active_; } + ggml_backend_sched_t scheduler() const { return scheduler_; } + ggml_backend_t cpu_backend() const { return cpu_backend_; } + size_t bytes(ggml_backend_t backend) const; + size_t reservation_count() const { return reservations_; } + }; +} + +#endif // __SD_CORE_COMPUTE_WORKSPACE_H__ diff --git a/otherarch/sdcpp/src/core/ggml_extend.cpp b/otherarch/sdcpp/src/core/ggml_extend.cpp new file mode 100644 index 000000000..afa30a26d --- /dev/null +++ b/otherarch/sdcpp/src/core/ggml_extend.cpp @@ -0,0 +1,729 @@ +#include "core/ggml_extend.h" + +#include +#include + +#include "core/ggml_extend_backend.h" + +ggml_tensor* ggml_ext_mul_n_mode(ggml_context* ctx, ggml_tensor* a, ggml_tensor* b, int mode) { + // reshape A + // swap 0th and nth axis + a = ggml_cont(ctx, ggml_permute(ctx, a, mode, mode != 1 ? 1 : 0, mode != 2 ? 2 : 0, mode != 3 ? 3 : 0)); + int64_t ne1 = a->ne[1]; + int64_t ne2 = a->ne[2]; + int64_t ne3 = a->ne[3]; + // make 2D + a = ggml_cont(ctx, ggml_reshape_2d(ctx, a, a->ne[0], (ne3 * ne2 * ne1))); + + ggml_tensor* result = ggml_cont(ctx, ggml_transpose(ctx, ggml_mul_mat(ctx, a, b))); + + // reshape output (same shape as a after permutation except first dim) + result = ggml_reshape_4d(ctx, result, result->ne[0], ne1, ne2, ne3); + // swap back 0th and nth axis + result = ggml_permute(ctx, result, mode, mode != 1 ? 1 : 0, mode != 2 ? 2 : 0, mode != 3 ? 3 : 0); + return result; +} + +ggml_tensor* ggml_ext_kronecker(ggml_context* ctx, ggml_tensor* a, ggml_tensor* b) { + return ggml_mul(ctx, + ggml_interpolate(ctx, + a, + a->ne[0] * b->ne[0], + a->ne[1] * b->ne[1], + a->ne[2] * b->ne[2], + a->ne[3] * b->ne[3], + GGML_SCALE_MODE_NEAREST), + b); +} + +ggml_tensor* ggml_ext_cont(ggml_context* ctx, + ggml_tensor* x) { + if (ggml_is_contiguous(x)) { + return x; + } + return ggml_cont(ctx, x); +} + +ggml_tensor* ggml_ext_torch_permute(ggml_context* ctx, + ggml_tensor* x, + int axis0, + int axis1, + int axis2, + int axis3) { + int torch_axes[4] = {axis0, axis1, axis2, axis3}; + + int ggml_axes[4] = {0}; + for (int i = 0; i < 4; ++i) { + int found = 0; + for (int j = 0; j < 4; ++j) { + if (torch_axes[j] == i) { + ggml_axes[i] = j; + found = 1; + break; + } + } + GGML_ASSERT(found && "Invalid permute input: must be a permutation of 0-3"); + } + + return ggml_permute(ctx, x, ggml_axes[0], ggml_axes[1], ggml_axes[2], ggml_axes[3]); +} + +ggml_tensor* ggml_ext_slice(ggml_context* ctx, + ggml_tensor* x, + int dim, + int64_t start, + int64_t end, + bool cont) { + GGML_ASSERT(dim >= 0 && dim < 4); + if (x->ne[dim] == 1) { + return x; + } + while (start < 0) { + start = x->ne[dim] + start; + } + while (end < 0) { + end = x->ne[dim] + end; + } + GGML_ASSERT(end > start); + GGML_ASSERT(start >= 0 && start < x->ne[dim]); + GGML_ASSERT(end > start && end <= x->ne[dim]); + + int64_t slice_size = end - start; + int64_t slice_ne[4] = {x->ne[0], x->ne[1], x->ne[2], x->ne[3]}; + slice_ne[dim] = slice_size; + + x = ggml_view_4d(ctx, x, + slice_ne[0], slice_ne[1], slice_ne[2], slice_ne[3], + x->nb[1], x->nb[2], x->nb[3], start * x->nb[dim]); + + if (cont) { + x = ggml_cont(ctx, x); + } + + return x; +} + +std::vector ggml_ext_chunk(ggml_context* ctx, + ggml_tensor* x, + int num, + int64_t dim, + bool cont) { + GGML_ASSERT(dim >= 0 && dim < 4); + GGML_ASSERT(x->ne[dim] % num == 0); + + std::vector chunks; + int64_t chunk_size = x->ne[dim] / num; + int64_t stride = chunk_size * x->nb[dim]; + int64_t chunk_ne[4] = {x->ne[0], x->ne[1], x->ne[2], x->ne[3]}; + chunk_ne[dim] = chunk_size; + for (int i = 0; i < num; i++) { + auto chunk = ggml_view_4d( + ctx, x, + chunk_ne[0], chunk_ne[1], chunk_ne[2], chunk_ne[3], + x->nb[1], x->nb[2], x->nb[3], stride * i); + if (cont) { + chunk = ggml_cont(ctx, chunk); + } + chunks.push_back(chunk); + } + + return chunks; +} + +ggml_tensor* ggml_ext_silu_act(ggml_context* ctx, ggml_tensor* x, bool gate_first) { + // x: [ne3, ne2, ne1, ne0] + // return: [ne3, ne2, ne1, ne0/2] + + auto x_vec = ggml_ext_chunk(ctx, x, 2, 0, false); + ggml_tensor* gate; + if (gate_first) { + gate = x_vec[0]; + x = x_vec[1]; + } else { + x = x_vec[0]; + gate = x_vec[1]; + } + gate = ggml_cont(ctx, gate); + gate = ggml_silu_inplace(ctx, gate); + + x = ggml_mul(ctx, x, gate); // [ne3, ne2, ne1, ne0/2] + + return x; +} + +ggml_tensor* ggml_ext_group_norm_32(ggml_context* ctx, + ggml_tensor* a) { + const float eps = 1e-6f; // default eps parameter + return ggml_group_norm(ctx, a, 32, eps); +} + +static bool ggml_ext_is_padded_1d(const ggml_tensor* x) { + return x->nb[0] == ggml_type_size(x->type) && + x->nb[2] == x->nb[1] * x->ne[1] && + x->nb[3] == x->nb[2] * x->ne[2]; +} + +ggml_tensor* ggml_ext_scale(ggml_context* ctx, + ggml_tensor* x, + float factor, + bool inplace) { + if (!ggml_ext_is_padded_1d(x)) { + x = ggml_cont(ctx, x); + } + if (inplace) { + x = ggml_scale_inplace(ctx, x, factor); + } else { + x = ggml_scale(ctx, x, factor); + } + return x; +} + +ggml_tensor* ggml_ext_gelu(ggml_context* ctx, + ggml_tensor* x, + bool inplace) { + if (!ggml_is_contiguous(x)) { + x = ggml_cont(ctx, x); + } + if (inplace) { + x = ggml_gelu_inplace(ctx, x); + } else { + x = ggml_gelu(ctx, x); + } + return x; +} + +ggml_tensor* ggml_ext_gelu_quick(ggml_context* ctx, + ggml_tensor* x, + bool inplace) { + if (!ggml_is_contiguous(x)) { + x = ggml_cont(ctx, x); + } + if (inplace) { + x = ggml_gelu_quick_inplace(ctx, x); + } else { + x = ggml_gelu_quick(ctx, x); + } + return x; +} + +ggml_tensor* ggml_ext_linear(ggml_context* ctx, + ggml_tensor* x, + ggml_tensor* w, + ggml_tensor* b, + bool force_prec_f32, + float scale) { + if (scale != 1.f) { + x = ggml_ext_scale(ctx, x, scale); + } + if (x->ne[2] * x->ne[3] > 1024) { + // workaround: avoid ggml cuda error + int64_t ne2 = x->ne[2]; + int64_t ne3 = x->ne[3]; + x = ggml_reshape_2d(ctx, x, x->ne[0], x->ne[1] * x->ne[2] * x->ne[3]); + x = ggml_mul_mat(ctx, w, x); + if (force_prec_f32) { + ggml_mul_mat_set_prec(x, GGML_PREC_F32); + } + x = ggml_reshape_4d(ctx, x, x->ne[0], x->ne[1] / ne2 / ne3, ne2, ne3); + } else { + x = ggml_mul_mat(ctx, w, x); + if (force_prec_f32) { + ggml_mul_mat_set_prec(x, GGML_PREC_F32); + } + } + if (scale != 1.f) { + x = ggml_ext_scale(ctx, x, 1.f / scale); + } + if (b != nullptr) { + x = ggml_add_inplace(ctx, x, b); + } + return x; +} + +#if KCPP_MAINLINE_INT8_CONVROT +ggml_tensor* ggml_ext_linear_i8_tensorwise(ggml_context* ctx, + ggml_tensor* x, + ggml_tensor* w, + ggml_tensor* weight_scale, + ggml_tensor* b, + int convrot_group_size, + float scale) { + GGML_ASSERT(x->type == GGML_TYPE_F32 || (x->type == GGML_TYPE_I8 && scale == 1.f)); + if (scale != 1.f) { + x = ggml_ext_scale(ctx, x, scale); + } + + ggml_tensor* fused_bias = scale == 1.f ? b : nullptr; + if (x->ne[2] * x->ne[3] > 1024) { + int64_t ne2 = x->ne[2]; + int64_t ne3 = x->ne[3]; + x = ggml_reshape_2d(ctx, x, x->ne[0], x->ne[1] * x->ne[2] * x->ne[3]); + x = ggml_mul_mat_i8_tensorwise(ctx, w, x, weight_scale, fused_bias, convrot_group_size); + x = ggml_reshape_4d(ctx, x, x->ne[0], x->ne[1] / ne2 / ne3, ne2, ne3); + } else { + x = ggml_mul_mat_i8_tensorwise(ctx, w, x, weight_scale, fused_bias, convrot_group_size); + } + + if (scale != 1.f) { + x = ggml_ext_scale(ctx, x, 1.f / scale); + if (b != nullptr) { + x = ggml_add_inplace(ctx, x, b); + } + } + return x; +} +#endif //kcpp + +ggml_tensor* ggml_ext_pad_ext(ggml_context* ctx, + ggml_backend_t backend, + ggml_tensor* x, + int lp0, + int rp0, + int lp1, + int rp1, + int lp2, + int rp2, + int lp3, + int rp3, + bool circular_x, + bool circular_y) { + if (circular_x && circular_y) { + return ggml_pad_ext_circular(ctx, x, lp0, rp0, lp1, rp1, lp2, rp2, lp3, rp3); + } + + if (circular_x && (lp0 != 0 || rp0 != 0)) { + x = ggml_pad_ext_circular(ctx, x, lp0, rp0, 0, 0, 0, 0, 0, 0); + lp0 = rp0 = 0; + } + if (circular_y && (lp1 != 0 || rp1 != 0)) { + x = ggml_pad_ext_circular(ctx, x, 0, 0, lp1, rp1, 0, 0, 0, 0); + lp1 = rp1 = 0; + } + + if (lp0 != 0 || rp0 != 0 || lp1 != 0 || rp1 != 0 || lp2 != 0 || rp2 != 0 || lp3 != 0 || rp3 != 0) { + ggml_tensor* padded = ggml_pad_ext(ctx, x, lp0, rp0, lp1, rp1, lp2, rp2, lp3, rp3); + if (backend == nullptr || ggml_backend_supports_op(backend, padded)) { + x = padded; + } else { + // Some backends (e.g. Metal) only implement right-padding for + // GGML_OP_PAD (see #850): pad right by lp+rp instead, then roll + // the padding around to the left. shift < ne always holds because + // ne grew by lp+rp. + x = ggml_pad_ext(ctx, x, 0, lp0 + rp0, 0, lp1 + rp1, 0, lp2 + rp2, 0, lp3 + rp3); + x = ggml_roll(ctx, x, lp0, lp1, lp2, lp3); + } + } + return x; +} + +ggml_tensor* ggml_ext_pad(ggml_context* ctx, + ggml_tensor* x, + int p0, + int p1, + int p2, + int p3, + bool circular_x, + bool circular_y) { + return ggml_ext_pad_ext(ctx, nullptr, x, 0, p0, 0, p1, 0, p2, 0, p3, circular_x, circular_y); +} + +ggml_tensor* ggml_ext_conv_2d(ggml_context* ctx, + ggml_tensor* x, + ggml_tensor* w, + ggml_tensor* b, + int s0, + int s1, + int p0, + int p1, + int d0, + int d1, + bool direct, + bool circular_x, + bool circular_y, + float scale) { + if (scale != 1.f) { + x = ggml_ext_scale(ctx, x, scale); + } + if (w->ne[2] != x->ne[2] && ggml_n_dims(w) == 2) { + w = ggml_reshape_4d(ctx, w, 1, 1, w->ne[0], w->ne[1]); + } + + if ((p0 != 0 || p1 != 0) && (circular_x || circular_y)) { + x = ggml_ext_pad_ext(ctx, nullptr, x, p0, p0, p1, p1, 0, 0, 0, 0, circular_x, circular_y); + p0 = 0; + p1 = 0; + } + + if (direct) { + x = ggml_conv_2d_direct(ctx, w, x, s0, s1, p0, p1, d0, d1); + } else { + x = ggml_conv_2d(ctx, w, x, s0, s1, p0, p1, d0, d1); + } + if (scale != 1.f) { + x = ggml_ext_scale(ctx, x, 1.f / scale); + } + if (b != nullptr) { + b = ggml_reshape_4d(ctx, b, 1, 1, b->ne[0], 1); + x = ggml_add_inplace(ctx, x, b); + } + return x; +} + +ggml_tensor* ggml_ext_conv_3d(ggml_context* ctx, + ggml_backend_t backend, + ggml_tensor* x, + ggml_tensor* w, + ggml_tensor* b, + int64_t IC, + int s0, + int s1, + int s2, + int p0, + int p1, + int p2, + int d0, + int d1, + int d2, + bool force_prec_f32) { + if (force_prec_f32) { + ggml_tensor* im2col = ggml_im2col_3d(ctx, w, x, IC, s0, s1, s2, p0, p1, p2, d0, d1, d2, w->type); + + int64_t OC = w->ne[3] / IC; + int64_t N = x->ne[3] / IC; + x = ggml_mul_mat(ctx, + ggml_reshape_2d(ctx, im2col, im2col->ne[0], im2col->ne[3] * im2col->ne[2] * im2col->ne[1]), + ggml_reshape_2d(ctx, w, w->ne[0] * w->ne[1] * w->ne[2] * IC, OC)); + ggml_mul_mat_set_prec(x, GGML_PREC_F32); + + int64_t OD = im2col->ne[3] / N; + x = ggml_reshape_4d(ctx, x, im2col->ne[1] * im2col->ne[2], OD, N, OC); + x = ggml_cont(ctx, ggml_permute(ctx, x, 0, 1, 3, 2)); + x = ggml_reshape_4d(ctx, x, im2col->ne[1], im2col->ne[2], OD, OC * N); + } else { + // ggml_conv_3d decomposes into GGML_OP_IM2COL_3D, which some backends + // (e.g. Metal, see #850) do not implement. Fall back to + // GGML_OP_CONV_3D on those backends. + bool im2col_3d_supported = true; + if (backend != nullptr) { + ggml_tensor* im2col = ggml_im2col_3d(ctx, w, x, IC, s0, s1, s2, p0, p1, p2, d0, d1, d2, w->type); + im2col_3d_supported = ggml_backend_supports_op(backend, im2col); + } + if (im2col_3d_supported) { + x = ggml_conv_3d(ctx, w, x, IC, s0, s1, s2, p0, p1, p2, d0, d1, d2); + } else { + int64_t OC = w->ne[3] / IC; + int64_t N = x->ne[3] / IC; + x = ggml_conv_3d_direct(ctx, w, x, s0, s1, s2, p0, p1, p2, d0, d1, d2, (int)IC, (int)N, (int)OC); + } + } + + if (b != nullptr) { + b = ggml_reshape_4d(ctx, b, 1, 1, 1, b->ne[0]); // [OC, 1, 1, 1] + x = ggml_add_inplace(ctx, x, b); + } + return x; +} + +ggml_tensor* ggml_ext_conv_3d_nx1x1(ggml_context* ctx, + ggml_tensor* x, + ggml_tensor* w, + ggml_tensor* b, + int s2, + int p2, + int d2) { + x = ggml_conv_2d(ctx, w, x, 1, s2, 0, p2, 1, d2); // [N, OC, T, OH * OW] + if (b != nullptr) { + b = ggml_reshape_4d(ctx, b, 1, 1, b->ne[0], 1); + x = ggml_add(ctx, x, b); + } + return x; // [N, OC, T, OH * OW] +} + +std::vector split_qkv(ggml_context* ctx, + ggml_tensor* qkv) { + qkv = ggml_reshape_4d(ctx, qkv, qkv->ne[0] / 3, 3, qkv->ne[1], qkv->ne[2]); // [N, L, 3, C] + qkv = ggml_cont(ctx, ggml_permute(ctx, qkv, 0, 3, 1, 2)); // [3, N, L, C] + + int64_t offset = qkv->nb[2] * qkv->ne[2]; + auto q = ggml_view_3d(ctx, qkv, qkv->ne[0], qkv->ne[1], qkv->ne[2], qkv->nb[1], qkv->nb[2], offset * 0); // [N, L, C] + auto k = ggml_view_3d(ctx, qkv, qkv->ne[0], qkv->ne[1], qkv->ne[2], qkv->nb[1], qkv->nb[2], offset * 1); // [N, L, C] + auto v = ggml_view_3d(ctx, qkv, qkv->ne[0], qkv->ne[1], qkv->ne[2], qkv->nb[1], qkv->nb[2], offset * 2); // [N, L, C] + return {q, k, v}; +} + +std::vector split_image_qkv(ggml_context* ctx, + ggml_tensor* qkv) { + int64_t W = qkv->ne[0]; + int64_t H = qkv->ne[1]; + int64_t C = qkv->ne[2] / 3; + int64_t N = qkv->ne[3]; + int64_t nb1 = qkv->nb[1]; + int64_t nb2 = qkv->nb[2]; + qkv = ggml_reshape_4d(ctx, qkv, W * H, C, 3, N); // [N, 3, C, H*W] + qkv = ggml_cont(ctx, ggml_ext_torch_permute(ctx, qkv, 0, 1, 3, 2)); // [3, N, C, H*W] + + int64_t offset = qkv->nb[2] * qkv->ne[2]; + auto q = ggml_view_4d(ctx, qkv, W, H, C, N, nb1, nb2, qkv->nb[3], offset * 0); // [N, C, H, W] + auto k = ggml_view_4d(ctx, qkv, W, H, C, N, nb1, nb2, qkv->nb[3], offset * 1); // [N, C, H, W] + auto v = ggml_view_4d(ctx, qkv, W, H, C, N, nb1, nb2, qkv->nb[3], offset * 2); // [N, C, H, W] + return {q, k, v}; +} + +ggml_tensor* ggml_ext_full(ggml_context* ctx, + float value, + int64_t ne0, + int64_t ne1, + int64_t ne2, + int64_t ne3) { + auto one = ggml_get_tensor(ctx, "ggml_runner_build_in_tensor:one"); + auto t = ggml_ext_scale(ctx, one, value); // [1,] + t = ggml_repeat_4d(ctx, t, ne0, ne1, ne2, ne3); // [ne0, ne1, ne2, ne3] + return t; +} + +ggml_tensor* ggml_ext_zeros(ggml_context* ctx, + int64_t ne0, + int64_t ne1, + int64_t ne2, + int64_t ne3) { + return ggml_ext_full(ctx, 0.f, ne0, ne1, ne2, ne3); +} + +ggml_tensor* ggml_ext_zeros_like(ggml_context* ctx, + ggml_tensor* x) { + return ggml_ext_zeros(ctx, x->ne[0], x->ne[1], x->ne[2], x->ne[3]); +} + +ggml_tensor* ggml_ext_ones(ggml_context* ctx, + int64_t ne0, + int64_t ne1, + int64_t ne2, + int64_t ne3) { + return ggml_ext_full(ctx, 1.f, ne0, ne1, ne2, ne3); +} + +ggml_tensor* ggml_ext_ones_like(ggml_context* ctx, + ggml_tensor* x) { + return ggml_ext_ones(ctx, x->ne[0], x->ne[1], x->ne[2], x->ne[3]); +} + +ggml_tensor* ggml_ext_cast_f32(ggml_context* ctx, ggml_backend_t backend, ggml_tensor* a) { + if (sd_backend_is(backend, "Vulkan")) { + auto zero_index = ggml_get_tensor(ctx, "ggml_runner_build_in_tensor:zero_int"); + auto out = ggml_reshape_1d(ctx, a, ggml_nelements(a)); + out = ggml_get_rows(ctx, out, zero_index); + out = ggml_reshape(ctx, out, a); + // auto out = ggml_cast(ctx, a, GGML_TYPE_F32); + return out; + } else { + auto out = ggml_reshape_2d(ctx, a, 1, ggml_nelements(a)); + ggml_tensor* one = ggml_ext_ones(ctx, 1, 1, 1, 1); // [1,] + if (ggml_is_transposed(out)) { + out = ggml_mul_mat(ctx, one, out); + } else { + out = ggml_mul_mat(ctx, out, one); + } + out = ggml_reshape(ctx, out, a); + return out; + } +} + +ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx, + ggml_backend_t backend, + ggml_tensor* q, + ggml_tensor* k, + ggml_tensor* v, + int64_t n_head, + ggml_tensor* mask, + bool skip_reshape, + bool flash_attn, + float kv_scale) { // avoid overflow + int64_t L_q; + int64_t L_k; + int64_t C; + int64_t N; + int64_t d_head; + int64_t n_kv_head; + if (!skip_reshape) { + L_q = q->ne[1]; + L_k = k->ne[1]; + C = q->ne[0]; + N = q->ne[2]; + d_head = C / n_head; + n_kv_head = k->ne[0] / d_head; + + q = ggml_reshape_4d(ctx, q, d_head, n_head, L_q, N); // [N, L_q, n_head, d_head] + q = ggml_ext_cont(ctx, ggml_permute(ctx, q, 0, 2, 1, 3)); // [N, n_head, L_q, d_head] + q = ggml_reshape_3d(ctx, q, d_head, L_q, n_head * N); // [N * n_head, L_q, d_head] + + k = ggml_reshape_4d(ctx, k, d_head, n_kv_head, L_k, N); // [N, L_k, n_kv_head, d_head] + k = ggml_ext_cont(ctx, ggml_permute(ctx, k, 0, 2, 1, 3)); // [N, n_kv_head, L_k, d_head] + k = ggml_reshape_3d(ctx, k, d_head, L_k, n_kv_head * N); // [N * n_kv_head, L_k, d_head] + + v = ggml_reshape_4d(ctx, v, d_head, n_kv_head, L_k, N); // [N, L_k, n_kv_head, d_head] + } else { + L_q = q->ne[1]; + L_k = k->ne[1]; + d_head = v->ne[0]; + N = v->ne[3]; + n_kv_head = k->ne[2] / N; + C = d_head * n_head; + } + + float scale = (1.0f / sqrt((float)d_head)); + + ggml_tensor* kqv = nullptr; + + auto build_kqv = [&](ggml_tensor* q_in, ggml_tensor* k_in, ggml_tensor* v_in, ggml_tensor* mask_in) -> ggml_tensor* { + if (kv_scale != 1.0f) { + k_in = ggml_ext_scale(ctx, k_in, kv_scale); + } + k_in = ggml_cast(ctx, k_in, GGML_TYPE_F16); + + v_in = ggml_ext_cont(ctx, ggml_permute(ctx, v_in, 0, 2, 1, 3)); + v_in = ggml_reshape_3d(ctx, v_in, d_head, L_k, n_kv_head * N); + if (kv_scale != 1.0f) { + v_in = ggml_ext_scale(ctx, v_in, kv_scale); + } + v_in = ggml_cast(ctx, v_in, GGML_TYPE_F16); + + if (mask_in != nullptr) { + // ggml_flash_attn_ext expects the mask as a contiguous F16 tensor shaped + // [n_kv, n_q, (heads), (batch)] (ne0 = key length, ne1 = query length) and, + // unlike the manual-attention path, does not broadcast the query dimension. + // Some callers (e.g. Chroma/T5) pass a per-key padding mask broadcast over + // queries ([n_kv, 1, ...]); materialize the query dimension to L_q so the + // kernel indexes it correctly. (A bare ggml_transpose here produced a + // [1, n_kv, ...] mask that the kernel silently misreads, yielding NaN/blank + // output for masked flash attention.) + if (mask_in->ne[1] != L_q) { + mask_in = ggml_repeat(ctx, mask_in, + ggml_new_tensor_4d(ctx, mask_in->type, mask_in->ne[0], L_q, mask_in->ne[2], mask_in->ne[3])); + } + mask_in = ggml_cast(ctx, mask_in, GGML_TYPE_F16); + } + + auto out = ggml_flash_attn_ext(ctx, q_in, k_in, v_in, mask_in, scale / kv_scale, 0, 0); + if (!ggml_backend_supports_op(backend, out)) { + return nullptr; + } + ggml_flash_attn_ext_set_prec(out, GGML_PREC_F32); + if (kv_scale != 1.0f) { + out = ggml_ext_scale(ctx, out, 1.0f / kv_scale); + } + return out; + }; + + if (flash_attn) { + // LOG_VERBOSE("attention_ext L_q:%d L_k:%d n_head:%d C:%d d_head:%d N:%d", L_q, L_k, n_head, C, d_head, N); + bool can_use_flash_attn = true; + if (mask != nullptr) { + // TODO: figure out if we can bend t5 to work too + can_use_flash_attn = can_use_flash_attn && mask->ne[3] == 1; + } + + if (can_use_flash_attn) { + kqv = build_kqv(q, k, v, mask); + if (kqv != nullptr) { + kqv = ggml_view_4d(ctx, + kqv, + d_head, + n_head, + L_q, + N, + kqv->nb[1], + kqv->nb[2], + kqv->nb[1] * n_head, + 0); + } + } + } + + if (kqv == nullptr) { + // if (flash_attn) { + // LOG_VERBOSE("fallback to default attention, L_q:%d L_k:%d n_head:%d C:%d d_head:%d N:%d", L_q, L_k, n_head, C, d_head, N); + // } + v = ggml_ext_cont(ctx, ggml_permute(ctx, v, 1, 2, 0, 3)); // [N, n_kv_head, d_head, L_k] + v = ggml_reshape_3d(ctx, v, L_k, d_head, n_kv_head * N); // [N * n_kv_head, d_head, L_k] + + auto kq = ggml_mul_mat(ctx, k, q); // [N * n_head, L_q, L_k] + ggml_mul_mat_set_prec(kq, GGML_PREC_F32); + kq = ggml_scale_inplace(ctx, kq, scale); + if (mask) { + kq = ggml_add_inplace(ctx, kq, mask); + } + kq = ggml_soft_max_inplace(ctx, kq); + + kqv = ggml_mul_mat(ctx, v, kq); // [N * n_head, L_q, d_head] + + kqv = ggml_reshape_4d(ctx, kqv, d_head, L_q, n_head, N); // [N, n_head, L_q, d_head] + kqv = ggml_permute(ctx, kqv, 0, 2, 1, 3); // [N, L_q, n_head, d_head] + } + + kqv = ggml_ext_cont(ctx, kqv); + kqv = ggml_reshape_3d(ctx, kqv, d_head * n_head, L_q, N); // [N, L_q, C] + + return kqv; +} + +ggml_tensor* ggml_ext_layer_norm(ggml_context* ctx, + ggml_tensor* x, + ggml_tensor* w, + ggml_tensor* b, + float eps) { + x = ggml_norm(ctx, x, eps); + if (w != nullptr) { + x = ggml_mul_inplace(ctx, x, w); + if (b != nullptr) { + x = ggml_add_inplace(ctx, x, b); + } + } + return x; +} + +ggml_tensor* ggml_ext_group_norm(ggml_context* ctx, + ggml_tensor* x, + ggml_tensor* w, + ggml_tensor* b, + int num_groups) { + if (ggml_n_dims(x) >= 3 && w != nullptr && b != nullptr) { + w = ggml_reshape_4d(ctx, w, 1, 1, w->ne[0], 1); + b = ggml_reshape_4d(ctx, b, 1, 1, b->ne[0], 1); + } + + const float eps = 1e-6f; // default eps parameter + x = ggml_group_norm(ctx, x, num_groups, eps); + if (w != nullptr && b != nullptr) { + x = ggml_mul_inplace(ctx, x, w); + // b = ggml_repeat(ctx, b, x); + x = ggml_add_inplace(ctx, x, b); + } + return x; +} + +ggml_tensor* ggml_ext_timestep_embedding( + ggml_context* ctx, + ggml_tensor* timesteps, + int dim, + int max_period, + float time_factor) { + timesteps = ggml_ext_scale(ctx, timesteps, time_factor); + return ggml_timestep_embedding(ctx, timesteps, dim, max_period); +} + +ggml_tensor* ggml_ext_vec_concat(ggml_context* ctx, + std::vector& tensors, + int dim) { + while (tensors.size() > 1) { + std::vector next_level; + for (size_t i = 0; i < tensors.size(); i += 2) { + if (i + 1 < tensors.size()) { + next_level.push_back(ggml_concat(ctx, tensors[i], tensors[i + 1], dim)); + } else { + next_level.push_back(tensors[i]); + } + } + tensors = std::move(next_level); + } + return tensors[0]; +} diff --git a/otherarch/sdcpp/src/core/ggml_extend.h b/otherarch/sdcpp/src/core/ggml_extend.h new file mode 100644 index 000000000..4d5b1c897 --- /dev/null +++ b/otherarch/sdcpp/src/core/ggml_extend.h @@ -0,0 +1,235 @@ +#ifndef __SD_CORE_GGML_EXTEND_H__ +#define __SD_CORE_GGML_EXTEND_H__ + +#include +#include + +#include "ggml-backend.h" +#include "ggml.h" + +#define EPS 1e-05f + +static_assert(GGML_MAX_NAME >= 128, "GGML_MAX_NAME must be at least 128"); + +// n-mode tensor-matrix product +// example: 2-mode product +// A: [ne03, k, ne01, ne00] +// B: k rows, m columns => [k, m] +// result is [ne03, m, ne01, ne00] +ggml_tensor* ggml_ext_mul_n_mode(ggml_context* ctx, ggml_tensor* a, ggml_tensor* b, int mode = 0); + +// Kronecker product +// [ne03,ne02,ne01,ne00] x [ne13,ne12,ne11,ne10] => [ne03*ne13,ne02*ne12,ne01*ne11,ne00*ne10] +ggml_tensor* ggml_ext_kronecker(ggml_context* ctx, ggml_tensor* a, ggml_tensor* b); + +ggml_tensor* ggml_ext_cont(ggml_context* ctx, + ggml_tensor* x); + +// torch like permute +ggml_tensor* ggml_ext_torch_permute(ggml_context* ctx, + ggml_tensor* x, + int axis0, + int axis1, + int axis2, + int axis3); + +ggml_tensor* ggml_ext_slice(ggml_context* ctx, + ggml_tensor* x, + int dim, + int64_t start, + int64_t end, + bool cont = true); + +// example: [N, 3*C, H, W] => ([N, C, H, W], [N, C, H, W], [N, C, H, W]) +std::vector ggml_ext_chunk(ggml_context* ctx, + ggml_tensor* x, + int num, + int64_t dim, + bool cont = true); + +ggml_tensor* ggml_ext_silu_act(ggml_context* ctx, ggml_tensor* x, bool gate_first = true); + +ggml_tensor* ggml_ext_group_norm_32(ggml_context* ctx, + ggml_tensor* a); + +ggml_tensor* ggml_ext_scale(ggml_context* ctx, + ggml_tensor* x, + float factor, + bool inplace = false); + +ggml_tensor* ggml_ext_gelu(ggml_context* ctx, + ggml_tensor* x, + bool inplace = false); + +ggml_tensor* ggml_ext_gelu_quick(ggml_context* ctx, + ggml_tensor* x, + bool inplace = false); + +ggml_tensor* ggml_ext_linear(ggml_context* ctx, + ggml_tensor* x, + ggml_tensor* w, + ggml_tensor* b, + bool force_prec_f32 = false, + float scale = 1.f); + +ggml_tensor* ggml_ext_linear_i8_tensorwise(ggml_context* ctx, + ggml_tensor* x, + ggml_tensor* w, + ggml_tensor* weight_scale, + ggml_tensor* b, + int convrot_group_size, + float scale = 1.f); + +ggml_tensor* ggml_ext_pad_ext(ggml_context* ctx, + ggml_backend_t backend, + ggml_tensor* x, + int lp0, + int rp0, + int lp1, + int rp1, + int lp2, + int rp2, + int lp3, + int rp3, + bool circular_x = false, + bool circular_y = false); + +ggml_tensor* ggml_ext_pad(ggml_context* ctx, + ggml_tensor* x, + int p0, + int p1, + int p2 = 0, + int p3 = 0, + bool circular_x = false, + bool circular_y = false); + +// w: [OC，IC, KH, KW] +// x: [N, IC, IH, IW] +// b: [OC,] +// result: [N, OC, OH, OW] +ggml_tensor* ggml_ext_conv_2d(ggml_context* ctx, + ggml_tensor* x, + ggml_tensor* w, + ggml_tensor* b, + int s0 = 1, + int s1 = 1, + int p0 = 0, + int p1 = 0, + int d0 = 1, + int d1 = 1, + bool direct = false, + bool circular_x = false, + bool circular_y = false, + float scale = 1.f); + +// w: [OC，IC, KD, 1 * 1] +// x: [N, IC, IH, IW] +// b: [OC,] +// result: [N*OC, OD, OH, OW] +ggml_tensor* ggml_ext_conv_3d(ggml_context* ctx, + ggml_backend_t backend, + ggml_tensor* x, + ggml_tensor* w, + ggml_tensor* b, + int64_t IC, + int s0 = 1, + int s1 = 1, + int s2 = 1, + int p0 = 0, + int p1 = 0, + int p2 = 0, + int d0 = 1, + int d1 = 1, + int d2 = 1, + bool force_prec_f32 = false); + +// w: [OC，IC, KD, 1 * 1] +// x: [N, IC, ID, IH*IW] +// b: [OC,] +// result: [N, OC, OD, OH*OW] +ggml_tensor* ggml_ext_conv_3d_nx1x1(ggml_context* ctx, + ggml_tensor* x, + ggml_tensor* w, + ggml_tensor* b, + int s2 = 1, + int p2 = 1, + int d2 = 1); + +// qkv: [N, L, 3*C] +// return: ([N, L, C], [N, L, C], [N, L, C]) +std::vector split_qkv(ggml_context* ctx, + ggml_tensor* qkv); + +// qkv: [N, 3*C, H, W] +// return: ([N, C, H, W], [N, C, H, W], [N, C, H, W]) +std::vector split_image_qkv(ggml_context* ctx, + ggml_tensor* qkv); + +// Constant and cast helpers require the built-in tensors initialized by GGMLRunner. +ggml_tensor* ggml_ext_full(ggml_context* ctx, + float value, + int64_t ne0, + int64_t ne1, + int64_t ne2, + int64_t ne3); + +ggml_tensor* ggml_ext_zeros(ggml_context* ctx, + int64_t ne0, + int64_t ne1, + int64_t ne2, + int64_t ne3); + +ggml_tensor* ggml_ext_zeros_like(ggml_context* ctx, + ggml_tensor* x); + +ggml_tensor* ggml_ext_ones(ggml_context* ctx, + int64_t ne0, + int64_t ne1, + int64_t ne2, + int64_t ne3); + +ggml_tensor* ggml_ext_ones_like(ggml_context* ctx, + ggml_tensor* x); + +ggml_tensor* ggml_ext_cast_f32(ggml_context* ctx, ggml_backend_t backend, ggml_tensor* a); + +// q: [N, L_q, C(n_head*d_head)] or [N*n_head, L_q, d_head] +// k: [N, L_k, n_kv_head*d_head] or [N*n_kv_head, L_k, d_head] +// v: [N, L_k, n_kv_head*d_head] or [N, L_k, n_kv_head, d_head] +// mask: [N, L_q, L_k] +// return: [N, L_q, C] +ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx, + ggml_backend_t backend, + ggml_tensor* q, + ggml_tensor* k, + ggml_tensor* v, + int64_t n_head, + ggml_tensor* mask = nullptr, + bool skip_reshape = false, + bool flash_attn = false, + float kv_scale = 1.0f); + +ggml_tensor* ggml_ext_layer_norm(ggml_context* ctx, + ggml_tensor* x, + ggml_tensor* w, + ggml_tensor* b, + float eps = EPS); + +ggml_tensor* ggml_ext_group_norm(ggml_context* ctx, + ggml_tensor* x, + ggml_tensor* w, + ggml_tensor* b, + int num_groups = 32); + +ggml_tensor* ggml_ext_timestep_embedding( + ggml_context* ctx, + ggml_tensor* timesteps, + int dim, + int max_period = 10000, + float time_factor = 1.0f); + +ggml_tensor* ggml_ext_vec_concat(ggml_context* ctx, + std::vector& tensors, + int dim); + +#endif // __SD_CORE_GGML_EXTEND_H__ diff --git a/otherarch/sdcpp/src/core/ggml_extend.hpp b/otherarch/sdcpp/src/core/ggml_extend.hpp deleted file mode 100644 index e9227986e..000000000 --- a/otherarch/sdcpp/src/core/ggml_extend.hpp +++ /dev/null @@ -1,4316 +0,0 @@ -#ifndef __SD_CORE_GGML_EXTEND_HPP__ -#define __SD_CORE_GGML_EXTEND_HPP__ - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#include "core/ggml_extend_backend.h" -#include "core/ggml_graph_cut.h" -#include "core/layer_split_partition.h" -#include "ggml-alloc.h" -#include "ggml-backend.h" -#include "ggml.h" - -// kcpp sidestep int8 convrot support -#ifndef KCPP_MAINLINE_INT8_CONVROT -#define KCPP_MAINLINE_INT8_CONVROT 0 -#endif - -// kcpp sidestep fp8 scaled support -#ifndef KCPP_MAINLINE_FP8_SCALED -#define KCPP_MAINLINE_FP8_SCALED 0 -#endif - -#include "core/tensor.hpp" -#include "model.h" - -#include "core/rng.hpp" -#include "core/tensor_ggml.hpp" -#include "core/util.h" -#include "weight_manager.h" - -#define EPS 1e-05f - -#ifndef __STATIC_INLINE__ -#define __STATIC_INLINE__ static inline -#endif - -#ifndef SD_UNUSED -#define SD_UNUSED(x) (void)(x) -#endif - -__STATIC_INLINE__ int align_up_offset(int n, int multiple) { - return (multiple - n % multiple) % multiple; -} - -__STATIC_INLINE__ int align_up(int n, int multiple) { - return n + align_up_offset(n, multiple); -} - -__STATIC_INLINE__ void ggml_log_callback_default(ggml_log_level level, const char* text, void*) { - switch (level) { - case GGML_LOG_LEVEL_DEBUG: - LOG_DEBUG(text); - break; - case GGML_LOG_LEVEL_INFO: - LOG_INFO(text); - break; - case GGML_LOG_LEVEL_WARN: - LOG_WARN(text); - break; - case GGML_LOG_LEVEL_ERROR: - LOG_ERROR(text); - break; - default: - LOG_DEBUG(text); - } -} - -static_assert(GGML_MAX_NAME >= 128, "GGML_MAX_NAME must be at least 128"); - -// n-mode tensor-matrix product -// example: 2-mode product -// A: [ne03, k, ne01, ne00] -// B: k rows, m columns => [k, m] -// result is [ne03, m, ne01, ne00] -__STATIC_INLINE__ ggml_tensor* ggml_ext_mul_n_mode(ggml_context* ctx, ggml_tensor* a, ggml_tensor* b, int mode = 0) { - // reshape A - // swap 0th and nth axis - a = ggml_cont(ctx, ggml_permute(ctx, a, mode, mode != 1 ? 1 : 0, mode != 2 ? 2 : 0, mode != 3 ? 3 : 0)); - int64_t ne1 = a->ne[1]; - int64_t ne2 = a->ne[2]; - int64_t ne3 = a->ne[3]; - // make 2D - a = ggml_cont(ctx, ggml_reshape_2d(ctx, a, a->ne[0], (ne3 * ne2 * ne1))); - - ggml_tensor* result = ggml_cont(ctx, ggml_transpose(ctx, ggml_mul_mat(ctx, a, b))); - - // reshape output (same shape as a after permutation except first dim) - result = ggml_reshape_4d(ctx, result, result->ne[0], ne1, ne2, ne3); - // swap back 0th and nth axis - result = ggml_permute(ctx, result, mode, mode != 1 ? 1 : 0, mode != 2 ? 2 : 0, mode != 3 ? 3 : 0); - return result; -} - -__STATIC_INLINE__ ggml_tensor* ggml_ext_merge_lora(ggml_context* ctx, - ggml_tensor* lora_down, - ggml_tensor* lora_up, - ggml_tensor* lora_mid = nullptr) { - ggml_tensor* updown; - // flat lora tensors to multiply it - int64_t lora_up_rows = lora_up->ne[ggml_n_dims(lora_up) - 1]; - lora_up = ggml_reshape_2d(ctx, lora_up, ggml_nelements(lora_up) / lora_up_rows, lora_up_rows); - auto lora_down_n_dims = ggml_n_dims(lora_down); - // assume n_dims should always be a multiple of 2 (otherwise rank 1 doesn't work) - lora_down_n_dims = (lora_down_n_dims + lora_down_n_dims % 2); - int64_t lora_down_rows = lora_down->ne[lora_down_n_dims - 1]; - lora_down = ggml_reshape_2d(ctx, lora_down, ggml_nelements(lora_down) / lora_down_rows, lora_down_rows); - - // ggml_mul_mat requires tensor b transposed - lora_down = ggml_cont(ctx, ggml_transpose(ctx, lora_down)); - if (lora_mid == nullptr) { - updown = ggml_mul_mat(ctx, lora_up, lora_down); - updown = ggml_cont(ctx, ggml_transpose(ctx, updown)); - } else { - // undoing tucker decomposition for conv layers. - // lora_mid has shape (3, 3, Rank, Rank) - // lora_down has shape (Rank, In, 1, 1) - // lora_up has shape (Rank, Out, 1, 1) - // conv layer shape is (3, 3, Out, In) - updown = ggml_ext_mul_n_mode(ctx, ggml_ext_mul_n_mode(ctx, lora_mid, lora_down, 3), lora_up, 2); - updown = ggml_cont(ctx, updown); - } - return updown; -} - -// Kronecker product -// [ne03,ne02,ne01,ne00] x [ne13,ne12,ne11,ne10] => [ne03*ne13,ne02*ne12,ne01*ne11,ne00*ne10] -__STATIC_INLINE__ ggml_tensor* ggml_ext_kronecker(ggml_context* ctx, ggml_tensor* a, ggml_tensor* b) { - return ggml_mul(ctx, - ggml_interpolate(ctx, - a, - a->ne[0] * b->ne[0], - a->ne[1] * b->ne[1], - a->ne[2] * b->ne[2], - a->ne[3] * b->ne[3], - GGML_SCALE_MODE_NEAREST), - b); -} - -__STATIC_INLINE__ void ggml_ext_im_set_randn_f32(ggml_tensor* tensor, std::shared_ptr rng) { - uint32_t n = (uint32_t)ggml_nelements(tensor); - std::vector random_numbers = rng->randn(n); - for (uint32_t i = 0; i < n; i++) { - ggml_ext_im_set_f32_1d(tensor, i, random_numbers[i]); - } -} - -__STATIC_INLINE__ void ggml_ext_tensor_set_f32(ggml_tensor* tensor, float value, int64_t i0, int64_t i1 = 0, int64_t i2 = 0, int64_t i3 = 0) { - GGML_ASSERT(tensor->nb[0] == sizeof(float)); - *(float*)((char*)(tensor->data) + i3 * tensor->nb[3] + i2 * tensor->nb[2] + i1 * tensor->nb[1] + i0 * tensor->nb[0]) = value; -} - -__STATIC_INLINE__ float ggml_ext_tensor_get_f32(const ggml_tensor* tensor, int64_t i0, int64_t i1 = 0, int64_t i2 = 0, int64_t i3 = 0) { - if (tensor->buffer != nullptr) { - float value; - ggml_backend_tensor_get(tensor, &value, i3 * tensor->nb[3] + i2 * tensor->nb[2] + i1 * tensor->nb[1] + i0 * tensor->nb[0], sizeof(float)); - return value; - } - GGML_ASSERT(tensor->nb[0] == sizeof(float)); - return *(float*)((char*)(tensor->data) + i3 * tensor->nb[3] + i2 * tensor->nb[2] + i1 * tensor->nb[1] + i0 * tensor->nb[0]); -} - -__STATIC_INLINE__ int ggml_ext_tensor_get_i32(const ggml_tensor* tensor, int64_t i0, int64_t i1 = 0, int64_t i2 = 0, int64_t i3 = 0) { - if (tensor->buffer != nullptr) { - int value; - ggml_backend_tensor_get(tensor, &value, i3 * tensor->nb[3] + i2 * tensor->nb[2] + i1 * tensor->nb[1] + i0 * tensor->nb[0], sizeof(int)); - return value; - } - GGML_ASSERT(tensor->nb[0] == sizeof(int)); - return *(int*)((char*)(tensor->data) + i3 * tensor->nb[3] + i2 * tensor->nb[2] + i1 * tensor->nb[1] + i0 * tensor->nb[0]); -} - -__STATIC_INLINE__ ggml_fp16_t ggml_ext_tensor_get_f16(const ggml_tensor* tensor, int64_t i0, int64_t i1 = 0, int64_t i2 = 0, int64_t i3 = 0) { - GGML_ASSERT(tensor->nb[0] == sizeof(ggml_fp16_t)); - return *(ggml_fp16_t*)((char*)(tensor->data) + i3 * tensor->nb[3] + i2 * tensor->nb[2] + i1 * tensor->nb[1] + i0 * tensor->nb[0]); -} - -__STATIC_INLINE__ float sd_image_get_f32(sd_image_t image, int64_t iw, int64_t ih, int64_t ic, bool scale = true) { - float value = *(image.data + ih * image.width * image.channel + iw * image.channel + ic); - if (scale) { - value /= 255.f; - } - return value; -} - -__STATIC_INLINE__ void print_ggml_tensor(ggml_tensor* tensor, bool shape_only = false, const char* mark = "") { - printf("%s (%s): shape(%zu, %zu, %zu, %zu)\n", mark, ggml_type_name(tensor->type), tensor->ne[0], tensor->ne[1], tensor->ne[2], tensor->ne[3]); - fflush(stdout); - if (shape_only) { - return; - } - int range = 3; - for (int i3 = 0; i3 < tensor->ne[3]; i3++) { - if (i3 >= range && i3 + range < tensor->ne[3]) { - continue; - } - for (int i2 = 0; i2 < tensor->ne[2]; i2++) { - if (i2 >= range && i2 + range < tensor->ne[2]) { - continue; - } - for (int i1 = 0; i1 < tensor->ne[1]; i1++) { - if (i1 >= range && i1 + range < tensor->ne[1]) { - continue; - } - for (int i0 = 0; i0 < tensor->ne[0]; i0++) { - if (i0 >= range && i0 + range < tensor->ne[0]) { - continue; - } - if (tensor->type == GGML_TYPE_F32) { - printf(" [%d, %d, %d, %d] = %f\n", i3, i2, i1, i0, ggml_ext_tensor_get_f32(tensor, i0, i1, i2, i3)); - } else if (tensor->type == GGML_TYPE_F16) { - printf(" [%d, %d, %d, %d] = %f\n", i3, i2, i1, i0, ggml_fp16_to_fp32(ggml_ext_tensor_get_f16(tensor, i0, i1, i2, i3))); - } else if (tensor->type == GGML_TYPE_I32) { - printf(" [%d, %d, %d, %d] = %i3\n", i3, i2, i1, i0, ggml_ext_tensor_get_i32(tensor, i0, i1, i2, i3)); - } - fflush(stdout); - } - } - } - } -} - -template -__STATIC_INLINE__ void print_sd_tensor(const sd::Tensor& tensor, bool shape_only = false, const char* mark = "") { - printf("%s: shape(", mark); - for (size_t i = 0; i < static_cast(tensor.dim()); ++i) { - printf("%s%lld", i == 0 ? "" : ", ", static_cast(tensor.shape()[i])); - } - printf(")\n"); - fflush(stdout); - if (shape_only) { - return; - } - if (tensor.empty()) { - return; - } - int range = 3; - std::vector shape = tensor.shape(); - while (shape.size() < 4) { - shape.push_back(1); - } - for (int64_t i3 = 0; i3 < shape[3]; i3++) { - if (i3 >= range && i3 + range < shape[3]) { - continue; - } - for (int64_t i2 = 0; i2 < shape[2]; i2++) { - if (i2 >= range && i2 + range < shape[2]) { - continue; - } - for (int64_t i1 = 0; i1 < shape[1]; i1++) { - if (i1 >= range && i1 + range < shape[1]) { - continue; - } - for (int64_t i0 = 0; i0 < shape[0]; i0++) { - if (i0 >= range && i0 + range < shape[0]) { - continue; - } - size_t offset = static_cast(i0 + shape[0] * (i1 + shape[1] * (i2 + shape[2] * i3))); - printf(" [%lld, %lld, %lld, %lld] = ", static_cast(i3), static_cast(i2), static_cast(i1), static_cast(i0)); - if constexpr (std::is_same_v) { - printf("%f\n", tensor[static_cast(offset)]); - } else if constexpr (std::is_same_v) { - printf("%f\n", ggml_fp16_to_fp32(tensor[static_cast(offset)])); - } else if constexpr (std::is_same_v) { - printf("%d\n", tensor[static_cast(offset)]); - } else if constexpr (std::is_same_v) { - printf("%lld\n", static_cast(tensor[static_cast(offset)])); - } - fflush(stdout); - } - } - } - } -} - -__STATIC_INLINE__ void ggml_ext_tensor_iter( - ggml_tensor* tensor, - const std::function& fn) { - int64_t n0 = tensor->ne[0]; - int64_t n1 = tensor->ne[1]; - int64_t n2 = tensor->ne[2]; - int64_t n3 = tensor->ne[3]; - - for (int64_t i3 = 0; i3 < n3; i3++) { - for (int64_t i2 = 0; i2 < n2; i2++) { - for (int64_t i1 = 0; i1 < n1; i1++) { - for (int64_t i0 = 0; i0 < n0; i0++) { - fn(tensor, i0, i1, i2, i3); - } - } - } - } -} - -__STATIC_INLINE__ void ggml_ext_tensor_iter( - ggml_tensor* tensor, - const std::function& fn) { - int64_t n0 = tensor->ne[0]; - int64_t n1 = tensor->ne[1]; - int64_t n2 = tensor->ne[2]; - int64_t n3 = tensor->ne[3]; - - for (int64_t i = 0; i < ggml_nelements(tensor); i++) { - fn(tensor, i); - } -} - -__STATIC_INLINE__ void ggml_ext_tensor_diff( - ggml_tensor* a, - ggml_tensor* b, - float gap = 0.1f) { - GGML_ASSERT(ggml_nelements(a) == ggml_nelements(b)); - ggml_ext_tensor_iter(a, [&](ggml_tensor* a, int64_t i0, int64_t i1, int64_t i2, int64_t i3) { - float a_value = ggml_ext_tensor_get_f32(a, i0, i1, i2, i3); - float b_value = ggml_ext_tensor_get_f32(b, i0, i1, i2, i3); - if (abs(a_value - b_value) > gap) { - LOG_WARN("[%ld, %ld, %ld, %ld] %f %f", i3, i2, i1, i0, a_value, b_value); - } - }); -} - -__STATIC_INLINE__ ggml_tensor* load_tensor_from_file(ggml_context* ctx, const std::string& file_path) { - std::ifstream file(sd_get_u8path(file_path), std::ios::binary); - if (!file.is_open()) { - LOG_ERROR("failed to open '%s'", file_path.c_str()); - return nullptr; - } - int32_t n_dims; - int32_t length; - int32_t ttype; - - file.read(reinterpret_cast(&n_dims), sizeof(n_dims)); - file.read(reinterpret_cast(&length), sizeof(length)); - file.read(reinterpret_cast(&ttype), sizeof(ttype)); - - LOG_DEBUG("load_tensor_from_file %d %d %d", n_dims, length, ttype); - - if (file.eof()) { - LOG_ERROR("incomplete file '%s'", file_path.c_str()); - return nullptr; - } - - int32_t nelements = 1; - int32_t ne[4] = {1, 1, 1, 1}; - for (int i = 0; i < n_dims; ++i) { - file.read(reinterpret_cast(&ne[i]), sizeof(ne[i])); - nelements *= ne[i]; - } - std::string name(length, 0); - file.read(&name[0], length); - ggml_tensor* tensor = ggml_new_tensor_4d(ctx, (ggml_type)ttype, ne[0], ne[1], ne[2], ne[3]); - const size_t bpe = ggml_type_size(ggml_type(ttype)); - file.read(reinterpret_cast(tensor->data), ggml_nbytes(tensor)); - return tensor; -} - -// __STATIC_INLINE__ void save_tensor_to_file(const std::string& file_name, ggml_tensor* tensor, const std::string & name) { -// std::string file_name_ = file_name + ".tensor"; -// std::string name_ = name; -// std::ofstream file("./" + file_name_, std::ios::binary); -// file.write(reinterpret_cast(&tensor->n_dims), sizeof(tensor->n_dims)); -// int len = (int)name_.size(); -// file.write(reinterpret_cast(&len), sizeof(len)); -// int ttype = (int)tensor->type; -// file.write(reinterpret_cast(&ttype), sizeof(ttype)); -// for (int i = 0; i < tensor->n_dims; ++i) { -// int ne_ = (int) tensor->ne[i]; -// file.write(reinterpret_cast(&ne_), sizeof(ne_)); -// } -// file.write(&name_[0], len); -// char* data = nullptr; -// file.write((char*)tensor->data, ggml_nbytes(tensor)); -// file.close(); -// } - -__STATIC_INLINE__ float sigmoid(float x) { - return 1 / (1.0f + expf(-x)); -} - -// SPECIAL OPERATIONS WITH TENSORS - -__STATIC_INLINE__ uint8_t* ggml_tensor_to_sd_image(ggml_tensor* input, uint8_t* image_data = nullptr) { - int64_t width = input->ne[0]; - int64_t height = input->ne[1]; - int64_t channels = input->ne[2]; - GGML_ASSERT(input->type == GGML_TYPE_F32); - if (image_data == nullptr) { - image_data = (uint8_t*)malloc(width * height * channels); - } - for (int iy = 0; iy < height; iy++) { - for (int ix = 0; ix < width; ix++) { - for (int k = 0; k < channels; k++) { - float value = ggml_ext_tensor_get_f32(input, ix, iy, k); - *(image_data + iy * width * channels + ix * channels + k) = (uint8_t)(value * 255.0f); - } - } - } - return image_data; -} - -__STATIC_INLINE__ uint8_t* ggml_tensor_to_sd_image(ggml_tensor* input, int idx, bool video = false) { - int64_t width = input->ne[0]; - int64_t height = input->ne[1]; - int64_t channels; - if (video) { - channels = input->ne[3]; - } else { - channels = input->ne[2]; - } - GGML_ASSERT(channels == 3 && input->type == GGML_TYPE_F32); - uint8_t* image_data = (uint8_t*)malloc(width * height * channels); - for (int ih = 0; ih < height; ih++) { - for (int iw = 0; iw < width; iw++) { - for (int ic = 0; ic < channels; ic++) { - float value; - if (video) { - value = ggml_ext_tensor_get_f32(input, iw, ih, idx, ic); - } else { - value = ggml_ext_tensor_get_f32(input, iw, ih, ic, idx); - } - *(image_data + ih * width * channels + iw * channels + ic) = (uint8_t)(value * 255.0f); - } - } - } - return image_data; -} - -__STATIC_INLINE__ void sd_image_to_ggml_tensor(sd_image_t image, - ggml_tensor* tensor, - bool scale = true) { - GGML_ASSERT(image.width == tensor->ne[0]); - GGML_ASSERT(image.height == tensor->ne[1]); - GGML_ASSERT(image.channel == tensor->ne[2]); - GGML_ASSERT(1 == tensor->ne[3]); - GGML_ASSERT(tensor->type == GGML_TYPE_F32); - ggml_ext_tensor_iter(tensor, [&](ggml_tensor* tensor, int64_t i0, int64_t i1, int64_t i2, int64_t i3) { - float value = sd_image_get_f32(image, i0, i1, i2, scale); - ggml_ext_tensor_set_f32(tensor, value, i0, i1, i2, i3); - }); -} - -__STATIC_INLINE__ void ggml_ext_tensor_apply_mask(ggml_tensor* image_data, - ggml_tensor* mask, - ggml_tensor* output, - float masked_value = 0.5f) { - int64_t width = output->ne[0]; - int64_t height = output->ne[1]; - int64_t channels = output->ne[2]; - float rescale_mx = 1.f * mask->ne[0] / output->ne[0]; - float rescale_my = 1.f * mask->ne[1] / output->ne[1]; - GGML_ASSERT(output->type == GGML_TYPE_F32); - for (int ix = 0; ix < width; ix++) { - for (int iy = 0; iy < height; iy++) { - int mx = (int)(ix * rescale_mx); - int my = (int)(iy * rescale_my); - float m = ggml_ext_tensor_get_f32(mask, mx, my); - m = round(m); // inpaint models need binary masks - ggml_ext_tensor_set_f32(mask, m, mx, my); - for (int k = 0; k < channels; k++) { - float value = ggml_ext_tensor_get_f32(image_data, ix, iy, k); - value = (1 - m) * (value - masked_value) + masked_value; - ggml_ext_tensor_set_f32(output, value, ix, iy, k); - } - } - } -} - -__STATIC_INLINE__ float ggml_ext_tensor_mean(ggml_tensor* src) { - float mean = 0.0f; - int64_t nelements = ggml_nelements(src); - float* data = (float*)src->data; - for (int i = 0; i < nelements; i++) { - mean += data[i] / nelements * 1.0f; - } - return mean; -} - -// a = a+b -__STATIC_INLINE__ void ggml_ext_tensor_add_inplace(ggml_tensor* a, ggml_tensor* b) { - GGML_ASSERT(ggml_nelements(a) == ggml_nelements(b)); - int64_t nelements = ggml_nelements(a); - float* vec_a = (float*)a->data; - float* vec_b = (float*)b->data; - for (int i = 0; i < nelements; i++) { - vec_a[i] = vec_a[i] + vec_b[i]; - } -} - -__STATIC_INLINE__ void ggml_ext_tensor_scale_inplace(ggml_tensor* src, float scale) { - int64_t nelements = ggml_nelements(src); - float* data = (float*)src->data; - for (int i = 0; i < nelements; i++) { - data[i] = data[i] * scale; - } -} - -__STATIC_INLINE__ void ggml_ext_tensor_clamp_inplace(ggml_tensor* src, float min, float max) { - int64_t nelements = ggml_nelements(src); - float* data = (float*)src->data; - for (int i = 0; i < nelements; i++) { - float val = data[i]; - data[i] = val < min ? min : (val > max ? max : val); - } -} - -__STATIC_INLINE__ ggml_tensor* ggml_ext_tensor_concat(ggml_context* ctx, - ggml_tensor* a, - ggml_tensor* b, - int dim) { - int64_t ne[GGML_MAX_DIMS]; - for (int d = 0; d < GGML_MAX_DIMS; ++d) { - if (d == dim) { - ne[d] = a->ne[d] + b->ne[d]; - continue; - } - GGML_ASSERT(a->ne[d] == b->ne[d]); - ne[d] = a->ne[d]; - } - ggml_tensor* result = ggml_new_tensor(ctx, a->type, GGML_MAX_DIMS, ne); - int64_t o[4] = {0, 0, 0, 0}; - o[dim] = a->ne[dim]; - - float v; - for (int i3 = 0; i3 < result->ne[3]; i3++) { - for (int i2 = 0; i2 < result->ne[2]; i2++) { - for (int i1 = 0; i1 < result->ne[1]; i1++) { - for (int i0 = 0; i0 < result->ne[0]; i0++) { - if (i0 < a->ne[0] && i1 < a->ne[1] && i2 < a->ne[2] && i3 < a->ne[3]) { - v = ggml_ext_tensor_get_f32(a, i0, i1, i2, i3); - } else { - v = ggml_ext_tensor_get_f32(b, i0 - o[0], i1 - o[1], i2 - o[2], i3 - o[3]); - } - - ggml_ext_tensor_set_f32(result, v, i0, i1, i2, i3); - } - } - } - } - return result; -} - -// convert values from [0, 1] to [-1, 1] -__STATIC_INLINE__ void scale_to_minus1_1(ggml_tensor* src) { - int64_t nelements = ggml_nelements(src); - float* data = (float*)src->data; - for (int i = 0; i < nelements; i++) { - float val = data[i]; - data[i] = val * 2.0f - 1.0f; - } -} - -// convert values from [-1, 1] to [0, 1] -__STATIC_INLINE__ void scale_to_0_1(ggml_tensor* src) { - int64_t nelements = ggml_nelements(src); - float* data = (float*)src->data; - for (int i = 0; i < nelements; i++) { - float val = data[i]; - data[i] = (val + 1.0f) * 0.5f; - } -} - -__STATIC_INLINE__ ggml_tensor* ggml_ext_cont(ggml_context* ctx, - ggml_tensor* x) { - if (ggml_is_contiguous(x)) { - return x; - } - return ggml_cont(ctx, x); -} - -// torch like permute -__STATIC_INLINE__ ggml_tensor* ggml_ext_torch_permute(ggml_context* ctx, - ggml_tensor* x, - int axis0, - int axis1, - int axis2, - int axis3) { - int torch_axes[4] = {axis0, axis1, axis2, axis3}; - - int ggml_axes[4] = {0}; - for (int i = 0; i < 4; ++i) { - int found = 0; - for (int j = 0; j < 4; ++j) { - if (torch_axes[j] == i) { - ggml_axes[i] = j; - found = 1; - break; - } - } - GGML_ASSERT(found && "Invalid permute input: must be a permutation of 0-3"); - } - - return ggml_permute(ctx, x, ggml_axes[0], ggml_axes[1], ggml_axes[2], ggml_axes[3]); -} - -__STATIC_INLINE__ ggml_tensor* ggml_ext_slice(ggml_context* ctx, - ggml_tensor* x, - int dim, - int64_t start, - int64_t end, - bool cont = true) { - GGML_ASSERT(dim >= 0 && dim < 4); - if (x->ne[dim] == 1) { - return x; - } - while (start < 0) { - start = x->ne[dim] + start; - } - while (end < 0) { - end = x->ne[dim] + end; - } - GGML_ASSERT(end > start); - GGML_ASSERT(start >= 0 && start < x->ne[dim]); - GGML_ASSERT(end > start && end <= x->ne[dim]); - - int64_t slice_size = end - start; - int64_t slice_ne[4] = {x->ne[0], x->ne[1], x->ne[2], x->ne[3]}; - slice_ne[dim] = slice_size; - - x = ggml_view_4d(ctx, x, - slice_ne[0], slice_ne[1], slice_ne[2], slice_ne[3], - x->nb[1], x->nb[2], x->nb[3], start * x->nb[dim]); - - if (cont) { - x = ggml_cont(ctx, x); - } - - return x; -} - -// example: [N, 3*C, H, W] => ([N, C, H, W], [N, C, H, W], [N, C, H, W]) -__STATIC_INLINE__ std::vector ggml_ext_chunk(ggml_context* ctx, - ggml_tensor* x, - int num, - int64_t dim, - bool cont = true) { - GGML_ASSERT(dim >= 0 && dim < 4); - GGML_ASSERT(x->ne[dim] % num == 0); - - std::vector chunks; - int64_t chunk_size = x->ne[dim] / num; - int64_t stride = chunk_size * x->nb[dim]; - int64_t chunk_ne[4] = {x->ne[0], x->ne[1], x->ne[2], x->ne[3]}; - chunk_ne[dim] = chunk_size; - for (int i = 0; i < num; i++) { - auto chunk = ggml_view_4d( - ctx, x, - chunk_ne[0], chunk_ne[1], chunk_ne[2], chunk_ne[3], - x->nb[1], x->nb[2], x->nb[3], stride * i); - if (cont) { - chunk = ggml_cont(ctx, chunk); - } - chunks.push_back(chunk); - } - - return chunks; -} - -__STATIC_INLINE__ ggml_tensor* ggml_ext_silu_act(ggml_context* ctx, ggml_tensor* x, bool gate_first = true) { - // x: [ne3, ne2, ne1, ne0] - // return: [ne3, ne2, ne1, ne0/2] - - auto x_vec = ggml_ext_chunk(ctx, x, 2, 0, false); - ggml_tensor* gate; - if (gate_first) { - gate = x_vec[0]; - x = x_vec[1]; - } else { - x = x_vec[0]; - gate = x_vec[1]; - } - gate = ggml_cont(ctx, gate); - gate = ggml_silu_inplace(ctx, gate); - - x = ggml_mul(ctx, x, gate); // [ne3, ne2, ne1, ne0/2] - - return x; -} - -typedef std::function on_tile_process; - -__STATIC_INLINE__ void sd_tiling_calc_tiles(int& num_tiles_dim, - float& tile_overlap_factor_dim, - int small_dim, - int tile_size, - const float tile_overlap_factor, - bool circular) { - int tile_overlap = static_cast(tile_size * tile_overlap_factor); - int non_tile_overlap = tile_size - tile_overlap; - - if (circular) { - // circular means the last and first tile are overlapping (wraping around) - num_tiles_dim = small_dim / non_tile_overlap; - - if (num_tiles_dim < 1) { - num_tiles_dim = 1; - } - - tile_overlap_factor_dim = (tile_size - small_dim / num_tiles_dim) / (float)tile_size; - - // if single tile and tile_overlap_factor is not 0, add one to ensure we have at least two overlapping tiles - if (num_tiles_dim == 1 && tile_overlap_factor_dim > 0) { - num_tiles_dim++; - tile_overlap_factor_dim = 0.5; - } - - return; - } - // else, non-circular means the last and first tile are not overlapping - - num_tiles_dim = (small_dim - tile_overlap) / non_tile_overlap; - int overshoot_dim = ((num_tiles_dim + 1) * non_tile_overlap + tile_overlap) % small_dim; - - if ((overshoot_dim != non_tile_overlap) && (overshoot_dim <= num_tiles_dim * (tile_size / 2 - tile_overlap))) { - // if tiles don't fit perfectly using the desired overlap - // and there is enough room to squeeze an extra tile without overlap becoming >0.5 - num_tiles_dim++; - } - - tile_overlap_factor_dim = (float)(tile_size * num_tiles_dim - small_dim) / (float)(tile_size * (num_tiles_dim - 1)); - if (num_tiles_dim <= 2) { - if (small_dim <= tile_size) { - num_tiles_dim = 1; - tile_overlap_factor_dim = 0; - } else { - num_tiles_dim = 2; - tile_overlap_factor_dim = (2 * tile_size - small_dim) / (float)tile_size; - } - } -} - -// Tiling - -__STATIC_INLINE__ int64_t sd_tensor_plane_size(const sd::Tensor& tensor) { - GGML_ASSERT(tensor.dim() >= 2); - return tensor.shape()[0] * tensor.shape()[1]; -} - -__STATIC_INLINE__ sd::Tensor sd_tensor_split_2d(const sd::Tensor& input, int width, int height, int x, int y) { - GGML_ASSERT(input.dim() >= 4); - std::vector output_shape = input.shape(); - output_shape[0] = width; - output_shape[1] = height; - sd::Tensor output(std::move(output_shape)); - int64_t input_width = input.shape()[0]; - int64_t input_height = input.shape()[1]; - int64_t input_plane = sd_tensor_plane_size(input); - int64_t output_plane = sd_tensor_plane_size(output); - int64_t plane_count = input.numel() / input_plane; - for (int iy = 0; iy < height; iy++) { - for (int ix = 0; ix < width; ix++) { - int64_t src_xy = (ix + x) % input_width + input_width * ((iy + y) % input_height); - int64_t dst_xy = ix + width * iy; - for (int64_t plane = 0; plane < plane_count; ++plane) { - output[plane * output_plane + dst_xy] = input[plane * input_plane + src_xy]; - } - } - } - return output; -} - -__STATIC_INLINE__ void sd_tensor_merge_2d(const sd::Tensor& input, - sd::Tensor* output, - int x, - int y, - int overlap_x, - int overlap_y, - bool circular_x, - bool circular_y, - int x_skip = 0, - int y_skip = 0) { - GGML_ASSERT(output != nullptr); - int64_t width = input.shape()[0]; - int64_t height = input.shape()[1]; - int64_t img_width = output->shape()[0]; - int64_t img_height = output->shape()[1]; - int64_t input_plane = sd_tensor_plane_size(input); - int64_t output_plane = sd_tensor_plane_size(*output); - int64_t plane_count = input.numel() / input_plane; - GGML_ASSERT(output->numel() / output_plane == plane_count); - - // unclamped -> expects x in the range [0-1] - auto smootherstep_f32 = [](const float x) -> float { - GGML_ASSERT(x >= 0.f && x <= 1.f); - return x * x * x * (x * (6.0f * x - 15.0f) + 10.0f); - }; - - for (int iy = y_skip; iy < height; iy++) { - for (int ix = x_skip; ix < width; ix++) { - int64_t src_xy = ix + width * iy; - int64_t ox = (x + ix) % img_width; - int64_t oy = (y + iy) % img_height; - int64_t dst_xy = ox + img_width * oy; - for (int64_t plane = 0; plane < plane_count; ++plane) { - float new_value = input[plane * input_plane + src_xy]; - if (overlap_x > 0 || overlap_y > 0) { - float old_value = (*output)[plane * output_plane + dst_xy]; - const float x_f_0 = (circular_x || (overlap_x > 0 && x > 0)) ? (ix - x_skip) / float(overlap_x) : 1.f; - const float x_f_1 = (circular_x || (overlap_x > 0 && x < (img_width - width))) ? (width - ix) / float(overlap_x) : 1.f; - const float y_f_0 = (circular_y || (overlap_y > 0 && y > 0)) ? (iy - y_skip) / float(overlap_y) : 1.f; - const float y_f_1 = (circular_y || (overlap_y > 0 && y < (img_height - height))) ? (height - iy) / float(overlap_y) : 1.f; - const float x_f = std::min(std::min(x_f_0, x_f_1), 1.f); - const float y_f = std::min(std::min(y_f_0, y_f_1), 1.f); - (*output)[plane * output_plane + dst_xy] = - old_value + new_value * smootherstep_f32(y_f) * smootherstep_f32(x_f); - } else { - (*output)[plane * output_plane + dst_xy] = new_value; - } - } - } - } -} - -template -__STATIC_INLINE__ sd::Tensor process_tiles_2d(const sd::Tensor& input, - int output_width, - int output_height, - int scale, - int p_tile_size_x, - int p_tile_size_y, - float tile_overlap_factor, - bool circular_x, - bool circular_y, - Fn&& on_processing, - bool silent = false) { - sd::Tensor output; - int input_width = static_cast(input.shape()[0]); - int input_height = static_cast(input.shape()[1]); - - GGML_ASSERT(((input_width / output_width) == (input_height / output_height)) && - ((output_width / input_width) == (output_height / input_height))); - GGML_ASSERT(((input_width / output_width) == scale) || - ((output_width / input_width) == scale)); - - int small_width = output_width; - int small_height = output_height; - bool decode = output_width > input_width; - if (decode) { - small_width = input_width; - small_height = input_height; - } - - int num_tiles_x; - float tile_overlap_factor_x; - sd_tiling_calc_tiles(num_tiles_x, tile_overlap_factor_x, small_width, p_tile_size_x, tile_overlap_factor, circular_x); - - int num_tiles_y; - float tile_overlap_factor_y; - sd_tiling_calc_tiles(num_tiles_y, tile_overlap_factor_y, small_height, p_tile_size_y, tile_overlap_factor, circular_y); - - int tile_overlap_x = static_cast(p_tile_size_x * tile_overlap_factor_x); - int non_tile_overlap_x = p_tile_size_x - tile_overlap_x; - int tile_overlap_y = static_cast(p_tile_size_y * tile_overlap_factor_y); - int non_tile_overlap_y = p_tile_size_y - tile_overlap_y; - int tile_size_x = p_tile_size_x < small_width ? p_tile_size_x : small_width; - int tile_size_y = p_tile_size_y < small_height ? p_tile_size_y : small_height; - int input_tile_size_x = tile_size_x; - int input_tile_size_y = tile_size_y; - int output_tile_size_x = tile_size_x; - int output_tile_size_y = tile_size_y; - if (decode) { - output_tile_size_x *= scale; - output_tile_size_y *= scale; - } else { - input_tile_size_x *= scale; - input_tile_size_y *= scale; - } - - int num_tiles = num_tiles_x * num_tiles_y; - int tile_count = 1; - bool last_y = false; - bool last_x = false; - float last_time = 0.0f; - if (!silent) { - LOG_DEBUG("num tiles : %d, %d ", num_tiles_x, num_tiles_y); - LOG_DEBUG("optimal overlap : %f, %f (targeting %f)", tile_overlap_factor_x, tile_overlap_factor_y, tile_overlap_factor); - LOG_DEBUG("processing %i tiles", num_tiles); - pretty_progress(0, num_tiles, 0.0f); - } - for (int y = 0; y < small_height && !last_y; y += non_tile_overlap_y) { - int dy = 0; - if (!circular_y && y + tile_size_y >= small_height) { - int original_y = y; - y = small_height - tile_size_y; - dy = original_y - y; - if (decode) { - dy *= scale; - } - last_y = true; - } - for (int x = 0; x < small_width && !last_x; x += non_tile_overlap_x) { - int dx = 0; - if (!circular_x && x + tile_size_x >= small_width) { - int original_x = x; - x = small_width - tile_size_x; - dx = original_x - x; - if (decode) { - dx *= scale; - } - last_x = true; - } - - int x_in = decode ? x : scale * x; - int y_in = decode ? y : scale * y; - int x_out = decode ? x * scale : x; - int y_out = decode ? y * scale : y; - - int overlap_x_out = decode ? tile_overlap_x * scale : tile_overlap_x; - int overlap_y_out = decode ? tile_overlap_y * scale : tile_overlap_y; - - int64_t t1 = ggml_time_ms(); - auto input_tile = sd_tensor_split_2d(input, input_tile_size_x, input_tile_size_y, x_in, y_in); - auto output_tile = on_processing(input_tile); - if (output_tile.empty()) { - return {}; - } - GGML_ASSERT(output_tile.shape()[0] == output_tile_size_x && output_tile.shape()[1] == output_tile_size_y); - if (output.empty()) { - std::vector output_shape = output_tile.shape(); - output_shape[0] = output_width; - output_shape[1] = output_height; - output = sd::Tensor::zeros(std::move(output_shape)); - } - sd_tensor_merge_2d(output_tile, &output, x_out, y_out, overlap_x_out, overlap_y_out, circular_x, circular_y, dx, dy); - - if (!silent) { - int64_t t2 = ggml_time_ms(); - last_time = (t2 - t1) / 1000.0f; - pretty_progress(tile_count, num_tiles, last_time); - } - tile_count++; - } - last_x = false; - } - if (!silent && tile_count < num_tiles) { - pretty_progress(num_tiles, num_tiles, last_time); - } - if (output.empty()) { - return {}; - } - return output; -} - -__STATIC_INLINE__ ggml_tensor* ggml_ext_group_norm_32(ggml_context* ctx, - ggml_tensor* a) { - const float eps = 1e-6f; // default eps parameter - return ggml_group_norm(ctx, a, 32, eps); -} - -__STATIC_INLINE__ bool ggml_ext_is_padded_1d(const ggml_tensor* x) { - return x->nb[0] == ggml_type_size(x->type) && - x->nb[2] == x->nb[1] * x->ne[1] && - x->nb[3] == x->nb[2] * x->ne[2]; -} - -__STATIC_INLINE__ ggml_tensor* ggml_ext_scale(ggml_context* ctx, - ggml_tensor* x, - float factor, - bool inplace = false) { - if (!ggml_ext_is_padded_1d(x)) { - x = ggml_cont(ctx, x); - } - if (inplace) { - x = ggml_scale_inplace(ctx, x, factor); - } else { - x = ggml_scale(ctx, x, factor); - } - return x; -} - -__STATIC_INLINE__ ggml_tensor* ggml_ext_gelu(ggml_context* ctx, - ggml_tensor* x, - bool inplace = false) { - if (!ggml_is_contiguous(x)) { - x = ggml_cont(ctx, x); - } - if (inplace) { - x = ggml_gelu_inplace(ctx, x); - } else { - x = ggml_gelu(ctx, x); - } - return x; -} - -__STATIC_INLINE__ ggml_tensor* ggml_ext_gelu_quick(ggml_context* ctx, - ggml_tensor* x, - bool inplace = false) { - if (!ggml_is_contiguous(x)) { - x = ggml_cont(ctx, x); - } - if (inplace) { - x = ggml_gelu_quick_inplace(ctx, x); - } else { - x = ggml_gelu_quick(ctx, x); - } - return x; -} - -__STATIC_INLINE__ ggml_tensor* ggml_ext_linear(ggml_context* ctx, - ggml_tensor* x, - ggml_tensor* w, - ggml_tensor* b, - bool force_prec_f32 = false, - float scale = 1.f) { - if (scale != 1.f) { - x = ggml_ext_scale(ctx, x, scale); - } - if (x->ne[2] * x->ne[3] > 1024) { - // workaround: avoid ggml cuda error - int64_t ne2 = x->ne[2]; - int64_t ne3 = x->ne[3]; - x = ggml_reshape_2d(ctx, x, x->ne[0], x->ne[1] * x->ne[2] * x->ne[3]); - x = ggml_mul_mat(ctx, w, x); - if (force_prec_f32) { - ggml_mul_mat_set_prec(x, GGML_PREC_F32); - } - x = ggml_reshape_4d(ctx, x, x->ne[0], x->ne[1] / ne2 / ne3, ne2, ne3); - } else { - x = ggml_mul_mat(ctx, w, x); - if (force_prec_f32) { - ggml_mul_mat_set_prec(x, GGML_PREC_F32); - } - } - if (scale != 1.f) { - x = ggml_ext_scale(ctx, x, 1.f / scale); - } - if (b != nullptr) { - x = ggml_add_inplace(ctx, x, b); - } - return x; -} - -#if KCPP_MAINLINE_INT8_CONVROT -__STATIC_INLINE__ ggml_tensor* ggml_ext_linear_i8_tensorwise(ggml_context* ctx, - ggml_tensor* x, - ggml_tensor* w, - ggml_tensor* weight_scale, - ggml_tensor* b, - int convrot_group_size, - float scale = 1.f) { - GGML_ASSERT(x->type == GGML_TYPE_F32 || (x->type == GGML_TYPE_I8 && scale == 1.f)); - if (scale != 1.f) { - x = ggml_ext_scale(ctx, x, scale); - } - - ggml_tensor* fused_bias = scale == 1.f ? b : nullptr; - if (x->ne[2] * x->ne[3] > 1024) { - int64_t ne2 = x->ne[2]; - int64_t ne3 = x->ne[3]; - x = ggml_reshape_2d(ctx, x, x->ne[0], x->ne[1] * x->ne[2] * x->ne[3]); - x = ggml_mul_mat_i8_tensorwise(ctx, w, x, weight_scale, fused_bias, convrot_group_size); - x = ggml_reshape_4d(ctx, x, x->ne[0], x->ne[1] / ne2 / ne3, ne2, ne3); - } else { - x = ggml_mul_mat_i8_tensorwise(ctx, w, x, weight_scale, fused_bias, convrot_group_size); - } - - if (scale != 1.f) { - x = ggml_ext_scale(ctx, x, 1.f / scale); - if (b != nullptr) { - x = ggml_add_inplace(ctx, x, b); - } - } - return x; -} -#endif - -__STATIC_INLINE__ ggml_tensor* ggml_ext_pad_ext(ggml_context* ctx, - ggml_backend_t backend, - ggml_tensor* x, - int lp0, - int rp0, - int lp1, - int rp1, - int lp2, - int rp2, - int lp3, - int rp3, - bool circular_x = false, - bool circular_y = false) { - if (circular_x && circular_y) { - return ggml_pad_ext_circular(ctx, x, lp0, rp0, lp1, rp1, lp2, rp2, lp3, rp3); - } - - if (circular_x && (lp0 != 0 || rp0 != 0)) { - x = ggml_pad_ext_circular(ctx, x, lp0, rp0, 0, 0, 0, 0, 0, 0); - lp0 = rp0 = 0; - } - if (circular_y && (lp1 != 0 || rp1 != 0)) { - x = ggml_pad_ext_circular(ctx, x, 0, 0, lp1, rp1, 0, 0, 0, 0); - lp1 = rp1 = 0; - } - - if (lp0 != 0 || rp0 != 0 || lp1 != 0 || rp1 != 0 || lp2 != 0 || rp2 != 0 || lp3 != 0 || rp3 != 0) { - ggml_tensor* padded = ggml_pad_ext(ctx, x, lp0, rp0, lp1, rp1, lp2, rp2, lp3, rp3); - if (backend == nullptr || ggml_backend_supports_op(backend, padded)) { - x = padded; - } else { - // Some backends (e.g. Metal) only implement right-padding for - // GGML_OP_PAD (see #850): pad right by lp+rp instead, then roll - // the padding around to the left. shift < ne always holds because - // ne grew by lp+rp. - x = ggml_pad_ext(ctx, x, 0, lp0 + rp0, 0, lp1 + rp1, 0, lp2 + rp2, 0, lp3 + rp3); - x = ggml_roll(ctx, x, lp0, lp1, lp2, lp3); - } - } - return x; -} - -__STATIC_INLINE__ ggml_tensor* ggml_ext_pad(ggml_context* ctx, - ggml_tensor* x, - int p0, - int p1, - int p2 = 0, - int p3 = 0, - bool circular_x = false, - bool circular_y = false) { - return ggml_ext_pad_ext(ctx, nullptr, x, 0, p0, 0, p1, 0, p2, 0, p3, circular_x, circular_y); -} - -// w: [OC，IC, KH, KW] -// x: [N, IC, IH, IW] -// b: [OC,] -// result: [N, OC, OH, OW] -__STATIC_INLINE__ ggml_tensor* ggml_ext_conv_2d(ggml_context* ctx, - ggml_tensor* x, - ggml_tensor* w, - ggml_tensor* b, - int s0 = 1, - int s1 = 1, - int p0 = 0, - int p1 = 0, - int d0 = 1, - int d1 = 1, - bool direct = false, - bool circular_x = false, - bool circular_y = false, - float scale = 1.f) { - if (scale != 1.f) { - x = ggml_ext_scale(ctx, x, scale); - } - if (w->ne[2] != x->ne[2] && ggml_n_dims(w) == 2) { - w = ggml_reshape_4d(ctx, w, 1, 1, w->ne[0], w->ne[1]); - } - - if ((p0 != 0 || p1 != 0) && (circular_x || circular_y)) { - x = ggml_ext_pad_ext(ctx, nullptr, x, p0, p0, p1, p1, 0, 0, 0, 0, circular_x, circular_y); - p0 = 0; - p1 = 0; - } - - if (direct) { - x = ggml_conv_2d_direct(ctx, w, x, s0, s1, p0, p1, d0, d1); - } else { - x = ggml_conv_2d(ctx, w, x, s0, s1, p0, p1, d0, d1); - } - if (scale != 1.f) { - x = ggml_ext_scale(ctx, x, 1.f / scale); - } - if (b != nullptr) { - b = ggml_reshape_4d(ctx, b, 1, 1, b->ne[0], 1); - x = ggml_add_inplace(ctx, x, b); - } - return x; -} - -// w: [OC，IC, KD, 1 * 1] -// x: [N, IC, IH, IW] -// b: [OC,] -// result: [N*OC, OD, OH, OW] -__STATIC_INLINE__ ggml_tensor* ggml_ext_conv_3d(ggml_context* ctx, - ggml_backend_t backend, - ggml_tensor* x, - ggml_tensor* w, - ggml_tensor* b, - int64_t IC, - int s0 = 1, - int s1 = 1, - int s2 = 1, - int p0 = 0, - int p1 = 0, - int p2 = 0, - int d0 = 1, - int d1 = 1, - int d2 = 1, - bool force_prec_f32 = false) { - if (force_prec_f32) { - ggml_tensor* im2col = ggml_im2col_3d(ctx, w, x, IC, s0, s1, s2, p0, p1, p2, d0, d1, d2, w->type); - - int64_t OC = w->ne[3] / IC; - int64_t N = x->ne[3] / IC; - x = ggml_mul_mat(ctx, - ggml_reshape_2d(ctx, im2col, im2col->ne[0], im2col->ne[3] * im2col->ne[2] * im2col->ne[1]), - ggml_reshape_2d(ctx, w, w->ne[0] * w->ne[1] * w->ne[2] * IC, OC)); - ggml_mul_mat_set_prec(x, GGML_PREC_F32); - - int64_t OD = im2col->ne[3] / N; - x = ggml_reshape_4d(ctx, x, im2col->ne[1] * im2col->ne[2], OD, N, OC); - x = ggml_cont(ctx, ggml_permute(ctx, x, 0, 1, 3, 2)); - x = ggml_reshape_4d(ctx, x, im2col->ne[1], im2col->ne[2], OD, OC * N); - } else { - // ggml_conv_3d decomposes into GGML_OP_IM2COL_3D, which some backends - // (e.g. Metal, see #850) do not implement. Fall back to - // GGML_OP_CONV_3D on those backends. - bool im2col_3d_supported = true; - if (backend != nullptr) { - ggml_tensor* im2col = ggml_im2col_3d(ctx, w, x, IC, s0, s1, s2, p0, p1, p2, d0, d1, d2, w->type); - im2col_3d_supported = ggml_backend_supports_op(backend, im2col); - } - if (im2col_3d_supported) { - x = ggml_conv_3d(ctx, w, x, IC, s0, s1, s2, p0, p1, p2, d0, d1, d2); - } else { - int64_t OC = w->ne[3] / IC; - int64_t N = x->ne[3] / IC; - x = ggml_conv_3d_direct(ctx, w, x, s0, s1, s2, p0, p1, p2, d0, d1, d2, (int)IC, (int)N, (int)OC); - } - } - - if (b != nullptr) { - b = ggml_reshape_4d(ctx, b, 1, 1, 1, b->ne[0]); // [OC, 1, 1, 1] - x = ggml_add_inplace(ctx, x, b); - } - return x; -} - -// w: [OC，IC, KD, 1 * 1] -// x: [N, IC, ID, IH*IW] -// b: [OC,] -// result: [N, OC, OD, OH*OW] -__STATIC_INLINE__ ggml_tensor* ggml_ext_conv_3d_nx1x1(ggml_context* ctx, - ggml_tensor* x, - ggml_tensor* w, - ggml_tensor* b, - int s2 = 1, - int p2 = 1, - int d2 = 1) { - x = ggml_conv_2d(ctx, w, x, 1, s2, 0, p2, 1, d2); // [N, OC, T, OH * OW] - if (b != nullptr) { - b = ggml_reshape_4d(ctx, b, 1, 1, b->ne[0], 1); - x = ggml_add(ctx, x, b); - } - return x; // [N, OC, T, OH * OW] -} - -// qkv: [N, L, 3*C] -// return: ([N, L, C], [N, L, C], [N, L, C]) -__STATIC_INLINE__ std::vector split_qkv(ggml_context* ctx, - ggml_tensor* qkv) { - qkv = ggml_reshape_4d(ctx, qkv, qkv->ne[0] / 3, 3, qkv->ne[1], qkv->ne[2]); // [N, L, 3, C] - qkv = ggml_cont(ctx, ggml_permute(ctx, qkv, 0, 3, 1, 2)); // [3, N, L, C] - - int64_t offset = qkv->nb[2] * qkv->ne[2]; - auto q = ggml_view_3d(ctx, qkv, qkv->ne[0], qkv->ne[1], qkv->ne[2], qkv->nb[1], qkv->nb[2], offset * 0); // [N, L, C] - auto k = ggml_view_3d(ctx, qkv, qkv->ne[0], qkv->ne[1], qkv->ne[2], qkv->nb[1], qkv->nb[2], offset * 1); // [N, L, C] - auto v = ggml_view_3d(ctx, qkv, qkv->ne[0], qkv->ne[1], qkv->ne[2], qkv->nb[1], qkv->nb[2], offset * 2); // [N, L, C] - return {q, k, v}; -} - -// qkv: [N, 3*C, H, W] -// return: ([N, C, H, W], [N, C, H, W], [N, C, H, W]) -__STATIC_INLINE__ std::vector split_image_qkv(ggml_context* ctx, - ggml_tensor* qkv) { - int64_t W = qkv->ne[0]; - int64_t H = qkv->ne[1]; - int64_t C = qkv->ne[2] / 3; - int64_t N = qkv->ne[3]; - int64_t nb1 = qkv->nb[1]; - int64_t nb2 = qkv->nb[2]; - qkv = ggml_reshape_4d(ctx, qkv, W * H, C, 3, N); // [N, 3, C, H*W] - qkv = ggml_cont(ctx, ggml_ext_torch_permute(ctx, qkv, 0, 1, 3, 2)); // [3, N, C, H*W] - - int64_t offset = qkv->nb[2] * qkv->ne[2]; - auto q = ggml_view_4d(ctx, qkv, W, H, C, N, nb1, nb2, qkv->nb[3], offset * 0); // [N, C, H, W] - auto k = ggml_view_4d(ctx, qkv, W, H, C, N, nb1, nb2, qkv->nb[3], offset * 1); // [N, C, H, W] - auto v = ggml_view_4d(ctx, qkv, W, H, C, N, nb1, nb2, qkv->nb[3], offset * 2); // [N, C, H, W] - return {q, k, v}; -} - -__STATIC_INLINE__ ggml_tensor* ggml_ext_full(ggml_context* ctx, - float value, - int64_t ne0, - int64_t ne1, - int64_t ne2, - int64_t ne3) { - auto one = ggml_get_tensor(ctx, "ggml_runner_build_in_tensor:one"); - auto t = ggml_ext_scale(ctx, one, value); // [1,] - t = ggml_repeat_4d(ctx, t, ne0, ne1, ne2, ne3); // [ne0, ne1, ne2, ne3] - return t; -} - -__STATIC_INLINE__ ggml_tensor* ggml_ext_zeros(ggml_context* ctx, - int64_t ne0, - int64_t ne1, - int64_t ne2, - int64_t ne3) { - return ggml_ext_full(ctx, 0.f, ne0, ne1, ne2, ne3); -} - -__STATIC_INLINE__ ggml_tensor* ggml_ext_zeros_like(ggml_context* ctx, - ggml_tensor* x) { - return ggml_ext_zeros(ctx, x->ne[0], x->ne[1], x->ne[2], x->ne[3]); -} - -__STATIC_INLINE__ ggml_tensor* ggml_ext_ones(ggml_context* ctx, - int64_t ne0, - int64_t ne1, - int64_t ne2, - int64_t ne3) { - return ggml_ext_full(ctx, 1.f, ne0, ne1, ne2, ne3); -} - -__STATIC_INLINE__ ggml_tensor* ggml_ext_ones_like(ggml_context* ctx, - ggml_tensor* x) { - return ggml_ext_ones(ctx, x->ne[0], x->ne[1], x->ne[2], x->ne[3]); -} - -__STATIC_INLINE__ ggml_tensor* ggml_ext_cast_f32(ggml_context* ctx, ggml_backend_t backend, ggml_tensor* a) { - if (sd_backend_is(backend, "Vulkan")) { - auto zero_index = ggml_get_tensor(ctx, "ggml_runner_build_in_tensor:zero_int"); - auto out = ggml_reshape_1d(ctx, a, ggml_nelements(a)); - out = ggml_get_rows(ctx, out, zero_index); - out = ggml_reshape(ctx, out, a); - // auto out = ggml_cast(ctx, a, GGML_TYPE_F32); - return out; - } else { - auto out = ggml_reshape_2d(ctx, a, 1, ggml_nelements(a)); - ggml_tensor* one = ggml_ext_ones(ctx, 1, 1, 1, 1); // [1,] - if (ggml_is_transposed(out)) { - out = ggml_mul_mat(ctx, one, out); - } else { - out = ggml_mul_mat(ctx, out, one); - } - out = ggml_reshape(ctx, out, a); - return out; - } -} - -// q: [N, L_q, C(n_head*d_head)] or [N*n_head, L_q, d_head] -// k: [N, L_k, n_kv_head*d_head] or [N*n_kv_head, L_k, d_head] -// v: [N, L_k, n_kv_head*d_head] or [N, L_k, n_kv_head, d_head] -// mask: [N, L_q, L_k] -// return: [N, L_q, C] -__STATIC_INLINE__ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx, - ggml_backend_t backend, - ggml_tensor* q, - ggml_tensor* k, - ggml_tensor* v, - int64_t n_head, - ggml_tensor* mask = nullptr, - bool skip_reshape = false, - bool flash_attn = false, - float kv_scale = 1.0f) { // avoid overflow - int64_t L_q; - int64_t L_k; - int64_t C; - int64_t N; - int64_t d_head; - int64_t n_kv_head; - if (!skip_reshape) { - L_q = q->ne[1]; - L_k = k->ne[1]; - C = q->ne[0]; - N = q->ne[2]; - d_head = C / n_head; - n_kv_head = k->ne[0] / d_head; - - q = ggml_reshape_4d(ctx, q, d_head, n_head, L_q, N); // [N, L_q, n_head, d_head] - q = ggml_ext_cont(ctx, ggml_permute(ctx, q, 0, 2, 1, 3)); // [N, n_head, L_q, d_head] - q = ggml_reshape_3d(ctx, q, d_head, L_q, n_head * N); // [N * n_head, L_q, d_head] - - k = ggml_reshape_4d(ctx, k, d_head, n_kv_head, L_k, N); // [N, L_k, n_kv_head, d_head] - k = ggml_ext_cont(ctx, ggml_permute(ctx, k, 0, 2, 1, 3)); // [N, n_kv_head, L_k, d_head] - k = ggml_reshape_3d(ctx, k, d_head, L_k, n_kv_head * N); // [N * n_kv_head, L_k, d_head] - - v = ggml_reshape_4d(ctx, v, d_head, n_kv_head, L_k, N); // [N, L_k, n_kv_head, d_head] - } else { - L_q = q->ne[1]; - L_k = k->ne[1]; - d_head = v->ne[0]; - N = v->ne[3]; - n_kv_head = k->ne[2] / N; - C = d_head * n_head; - } - - float scale = (1.0f / sqrt((float)d_head)); - - ggml_tensor* kqv = nullptr; - - auto build_kqv = [&](ggml_tensor* q_in, ggml_tensor* k_in, ggml_tensor* v_in, ggml_tensor* mask_in) -> ggml_tensor* { - if (kv_scale != 1.0f) { - k_in = ggml_ext_scale(ctx, k_in, kv_scale); - } - k_in = ggml_cast(ctx, k_in, GGML_TYPE_F16); - - v_in = ggml_ext_cont(ctx, ggml_permute(ctx, v_in, 0, 2, 1, 3)); - v_in = ggml_reshape_3d(ctx, v_in, d_head, L_k, n_kv_head * N); - if (kv_scale != 1.0f) { - v_in = ggml_ext_scale(ctx, v_in, kv_scale); - } - v_in = ggml_cast(ctx, v_in, GGML_TYPE_F16); - - if (mask_in != nullptr) { - // ggml_flash_attn_ext expects the mask as a contiguous F16 tensor shaped - // [n_kv, n_q, (heads), (batch)] (ne0 = key length, ne1 = query length) and, - // unlike the manual-attention path, does not broadcast the query dimension. - // Some callers (e.g. Chroma/T5) pass a per-key padding mask broadcast over - // queries ([n_kv, 1, ...]); materialize the query dimension to L_q so the - // kernel indexes it correctly. (A bare ggml_transpose here produced a - // [1, n_kv, ...] mask that the kernel silently misreads, yielding NaN/blank - // output for masked flash attention.) - if (mask_in->ne[1] != L_q) { - mask_in = ggml_repeat(ctx, mask_in, - ggml_new_tensor_4d(ctx, mask_in->type, mask_in->ne[0], L_q, mask_in->ne[2], mask_in->ne[3])); - } - mask_in = ggml_cast(ctx, mask_in, GGML_TYPE_F16); - } - - auto out = ggml_flash_attn_ext(ctx, q_in, k_in, v_in, mask_in, scale / kv_scale, 0, 0); - if (!ggml_backend_supports_op(backend, out)) { - return nullptr; - } - ggml_flash_attn_ext_set_prec(out, GGML_PREC_F32); - if (kv_scale != 1.0f) { - out = ggml_ext_scale(ctx, out, 1.0f / kv_scale); - } - return out; - }; - - if (flash_attn) { - // LOG_DEBUG("attention_ext L_q:%d L_k:%d n_head:%d C:%d d_head:%d N:%d", L_q, L_k, n_head, C, d_head, N); - bool can_use_flash_attn = true; - if (mask != nullptr) { - // TODO: figure out if we can bend t5 to work too - can_use_flash_attn = can_use_flash_attn && mask->ne[3] == 1; - } - - if (can_use_flash_attn) { - kqv = build_kqv(q, k, v, mask); - if (kqv != nullptr) { - kqv = ggml_view_4d(ctx, - kqv, - d_head, - n_head, - L_q, - N, - kqv->nb[1], - kqv->nb[2], - kqv->nb[1] * n_head, - 0); - } - } - } - - if (kqv == nullptr) { - // if (flash_attn) { - // LOG_DEBUG("fallback to default attention, L_q:%d L_k:%d n_head:%d C:%d d_head:%d N:%d", L_q, L_k, n_head, C, d_head, N); - // } - v = ggml_ext_cont(ctx, ggml_permute(ctx, v, 1, 2, 0, 3)); // [N, n_kv_head, d_head, L_k] - v = ggml_reshape_3d(ctx, v, L_k, d_head, n_kv_head * N); // [N * n_kv_head, d_head, L_k] - - auto kq = ggml_mul_mat(ctx, k, q); // [N * n_head, L_q, L_k] - ggml_mul_mat_set_prec(kq, GGML_PREC_F32); - kq = ggml_scale_inplace(ctx, kq, scale); - if (mask) { - kq = ggml_add_inplace(ctx, kq, mask); - } - kq = ggml_soft_max_inplace(ctx, kq); - - kqv = ggml_mul_mat(ctx, v, kq); // [N * n_head, L_q, d_head] - - kqv = ggml_reshape_4d(ctx, kqv, d_head, L_q, n_head, N); // [N, n_head, L_q, d_head] - kqv = ggml_permute(ctx, kqv, 0, 2, 1, 3); // [N, L_q, n_head, d_head] - } - - kqv = ggml_ext_cont(ctx, kqv); - kqv = ggml_reshape_3d(ctx, kqv, d_head * n_head, L_q, N); // [N, L_q, C] - - return kqv; -} - -__STATIC_INLINE__ ggml_tensor* ggml_ext_layer_norm(ggml_context* ctx, - ggml_tensor* x, - ggml_tensor* w, - ggml_tensor* b, - float eps = EPS) { - x = ggml_norm(ctx, x, eps); - if (w != nullptr) { - x = ggml_mul_inplace(ctx, x, w); - if (b != nullptr) { - x = ggml_add_inplace(ctx, x, b); - } - } - return x; -} - -__STATIC_INLINE__ ggml_tensor* ggml_ext_group_norm(ggml_context* ctx, - ggml_tensor* x, - ggml_tensor* w, - ggml_tensor* b, - int num_groups = 32) { - if (ggml_n_dims(x) >= 3 && w != nullptr && b != nullptr) { - w = ggml_reshape_4d(ctx, w, 1, 1, w->ne[0], 1); - b = ggml_reshape_4d(ctx, b, 1, 1, b->ne[0], 1); - } - - const float eps = 1e-6f; // default eps parameter - x = ggml_group_norm(ctx, x, num_groups, eps); - if (w != nullptr && b != nullptr) { - x = ggml_mul_inplace(ctx, x, w); - // b = ggml_repeat(ctx, b, x); - x = ggml_add_inplace(ctx, x, b); - } - return x; -} - -__STATIC_INLINE__ void ggml_ext_backend_tensor_get_and_sync(ggml_backend_t backend, const ggml_tensor* tensor, void* data, size_t offset, size_t size) { - if ((sd_backend_is(backend, "ROCm") || sd_backend_is(backend, "CUDA") || sd_backend_is(backend, "SYCL")) && - !sd_backend_is_cpu(backend)) { - ggml_backend_tensor_get_async(backend, tensor, data, offset, size); - ggml_backend_synchronize(backend); - return; - } - - ggml_backend_tensor_get(tensor, data, offset, size); -} - -__STATIC_INLINE__ float ggml_ext_backend_tensor_get_f32(ggml_tensor* tensor) { - GGML_ASSERT(tensor->type == GGML_TYPE_F32 || tensor->type == GGML_TYPE_F16 || tensor->type == GGML_TYPE_I32 || tensor->type == GGML_TYPE_BF16); - float value; - if (tensor->type == GGML_TYPE_F32) { - ggml_backend_tensor_get(tensor, &value, 0, sizeof(value)); - } else if (tensor->type == GGML_TYPE_BF16) { - ggml_bf16_t bf16_value; - ggml_backend_tensor_get(tensor, &bf16_value, 0, sizeof(bf16_value)); - value = ggml_bf16_to_fp32(bf16_value); - } else if (tensor->type == GGML_TYPE_F16) { - ggml_fp16_t f16_value; - ggml_backend_tensor_get(tensor, &f16_value, 0, sizeof(f16_value)); - value = ggml_fp16_to_fp32(f16_value); - } else { // GGML_TYPE_I32 - int int32_value; - ggml_backend_tensor_get(tensor, &int32_value, 0, sizeof(int32_value)); - value = (float)int32_value; - } - return value; -} - -__STATIC_INLINE__ ggml_tensor* vector_to_ggml_tensor(ggml_context* ctx, - const std::vector& vec) { - ggml_tensor* t = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, vec.size()); - memcpy(t->data, (const void*)vec.data(), ggml_nbytes(t)); - return t; -} - -__STATIC_INLINE__ ggml_tensor* vector_to_ggml_tensor_i32(ggml_context* ctx, - const std::vector& vec) { - ggml_tensor* t = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, vec.size()); - memcpy(t->data, (const void*)vec.data(), ggml_nbytes(t)); - return t; -} - -__STATIC_INLINE__ std::vector arange(float start, float end, float step = 1.f) { - std::vector result; - - for (float value = start; value < end; value += step) { - result.push_back(value); - } - - return result; -} - -// Ref: https://github.com/CompVis/stable-diffusion/blob/main/ldm/modules/diffusionmodules/util.py#L151 -__STATIC_INLINE__ std::vector timestep_embedding(std::vector timesteps, - int dim, - int max_period = 10000, - bool flip_sin_to_cos = true, - float scale = 1.f) { - // timesteps: [N,] - // embedding: [N, dim] - size_t N = timesteps.size(); - std::vector embedding(N * dim, 0.f); - int half = dim / 2; - std::vector freqs(half); - for (int i = 0; i < half; ++i) { - freqs[i] = (float)std::exp(-std::log(max_period) * i / half); - } - for (int i = 0; i < N; ++i) { - for (int j = 0; j < half; ++j) { - float arg = timesteps[i] * freqs[j] * scale; - if (flip_sin_to_cos) { - embedding[i * dim + j] = std::cos(arg); - embedding[i * dim + j + half] = std::sin(arg); - } else { - embedding[i * dim + j] = std::sin(arg); - embedding[i * dim + j + half] = std::cos(arg); - } - } - } - return embedding; -} - -__STATIC_INLINE__ void set_timestep_embedding(std::vector timesteps, - ggml_tensor* embedding, - int dim, - int max_period = 10000) { - std::vector embedding_vec = timestep_embedding(timesteps, dim, max_period); - memcpy(((char*)embedding->data), ((char*)embedding_vec.data()), ggml_nbytes(embedding)); -} - -__STATIC_INLINE__ void set_timestep_embedding(std::vector timesteps, - sd::Tensor* embedding, - int dim, - int max_period = 10000) { - GGML_ASSERT(embedding != nullptr); - std::vector embedding_vec = timestep_embedding(timesteps, dim, max_period); - if (embedding->numel() != static_cast(embedding_vec.size())) { - embedding->resize({dim, static_cast(timesteps.size())}); - } - std::copy(embedding_vec.begin(), embedding_vec.end(), embedding->values().begin()); -} - -__STATIC_INLINE__ ggml_tensor* new_timestep_embedding(ggml_context* ctx, - std::vector timesteps, - int dim, - int max_period = 10000) { - // timesteps: [N,] - // embedding: [N, dim] - std::vector embedding_vec = timestep_embedding(timesteps, dim, max_period); - ggml_tensor* embedding = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, dim, timesteps.size()); - if (embedding->data != nullptr) { - memcpy(((char*)embedding->data), ((char*)embedding_vec.data()), ggml_nbytes(embedding)); - } else { - ggml_backend_tensor_set(embedding, embedding_vec.data(), 0, ggml_nbytes(embedding)); - } - return embedding; -} - -__STATIC_INLINE__ ggml_tensor* ggml_ext_timestep_embedding( - ggml_context* ctx, - ggml_tensor* timesteps, - int dim, - int max_period = 10000, - float time_factor = 1.0f) { - timesteps = ggml_ext_scale(ctx, timesteps, time_factor); - return ggml_timestep_embedding(ctx, timesteps, dim, max_period); -} - -__STATIC_INLINE__ size_t ggml_tensor_num(ggml_context* ctx) { - size_t num = 0; - for (ggml_tensor* t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) { - num++; - } - return num; -} - -__STATIC_INLINE__ ggml_tensor* ggml_ext_vec_concat(ggml_context* ctx, - std::vector& tensors, - int dim) { - while (tensors.size() > 1) { - std::vector next_level; - for (size_t i = 0; i < tensors.size(); i += 2) { - if (i + 1 < tensors.size()) { - next_level.push_back(ggml_concat(ctx, tensors[i], tensors[i + 1], dim)); - } else { - next_level.push_back(tensors[i]); - } - } - tensors = std::move(next_level); - } - return tensors[0]; -} - -/* SDXL with LoRA requires more space */ -#define MAX_PARAMS_TENSOR_NUM 32768 -#define MAX_GRAPH_SIZE 327680 - -struct WeightAdapter { - struct ForwardParams { - enum class op_type_t { - OP_LINEAR, - OP_CONV2D, - } op_type; - struct { - bool force_prec_f32 = false; - float scale = 1.f; - } linear; - struct conv2d_params_t { - int s0 = 1; - int s1 = 1; - int p0 = 0; - int p1 = 0; - int d0 = 1; - int d1 = 1; - bool direct = false; - bool circular_x = false; - bool circular_y = false; - float scale = 1.f; - } conv2d; - }; - virtual ggml_tensor* patch_weight(ggml_context* ctx, ggml_backend_t backend, ggml_tensor* weight, const std::string& weight_name) = 0; - virtual ggml_tensor* forward_with_lora(ggml_context* ctx, - ggml_backend_t backend, - ggml_tensor* x, - ggml_tensor* w, - ggml_tensor* b, - const std::string& prefix, - ForwardParams forward_params) = 0; - virtual ggml_tensor* add_lora_to_output(ggml_context* ctx, - ggml_backend_t backend, - ggml_tensor* x, - ggml_tensor* w, - ggml_tensor* output, - const std::string& prefix, - ForwardParams forward_params) = 0; - virtual size_t get_extra_graph_size() = 0; -}; - -struct GGMLRunnerContext { - ggml_backend_t backend = nullptr; - ggml_context* ggml_ctx = nullptr; - bool flash_attn_enabled = false; - bool conv2d_direct_enabled = false; - bool circular_x_enabled = false; - bool circular_y_enabled = false; - ggml_tensor* ip_context = nullptr; - float ip_scale = 1.0f; - std::shared_ptr weight_adapter = nullptr; - std::vector>* debug_tensors = nullptr; - std::function get_cache_tensor; - std::function cache_tensor; - std::function set_backend_tensor_data; - std::map, ggml_tensor*> int8_convrot_cache; - - void capture_tensor(const std::string& name, ggml_tensor* tensor) { - if (debug_tensors == nullptr || tensor == nullptr) { - return; - } - ggml_tensor* snapshot = ggml_cont(ggml_ctx, tensor); - ggml_tensor* dst = ggml_dup_tensor(ggml_ctx, snapshot); - snapshot = ggml_cpy(ggml_ctx, snapshot, dst); - ggml_set_output(snapshot); - debug_tensors->push_back({snapshot, name}); - } - - ggml_tensor* load_cache_tensor(const std::string& name) const { - if (!get_cache_tensor) { - return nullptr; - } - return get_cache_tensor(name); - } - - void persist_cache_tensor(const std::string& name, ggml_tensor* tensor) const { - if (!cache_tensor || tensor == nullptr) { - return; - } - cache_tensor(name, tensor); - } - - void bind_backend_tensor_data(ggml_tensor* tensor, const void* data) const { - if (!set_backend_tensor_data || tensor == nullptr || data == nullptr) { - return; - } - set_backend_tensor_data(tensor, data); - } -}; - -struct GGMLRunner { -protected: - typedef std::function get_graph_cb_t; - using GraphCutSegment = sd::ggml_graph_cut::Segment; - using GraphCutPlan = sd::ggml_graph_cut::Plan; - - ggml_backend_t runtime_backend = nullptr; - - ggml_context* params_ctx = nullptr; - - ggml_context* cache_ctx = nullptr; - ggml_backend_buffer_t cache_buffer = nullptr; - - ggml_context* compute_ctx = nullptr; - ggml_gallocr* compute_allocr = nullptr; - - size_t max_graph_vram_bytes = 0; - bool stream_layers_enabled = false; - size_t observed_max_effective_budget_ = 0; - bool graph_cut_layer_split_enabled = false; - std::vector graph_cut_layer_split_backend_vram_limits_; - - std::vector extra_runtime_backends; // borrowed (SDBackendManager-owned) - ggml_backend_sched_t sched = nullptr; // owned - size_t sched_graph_capacity = 0; - ggml_backend_t cpu_fallback_backend = nullptr; // owned, sched requires a trailing CPU backend - bool multi_device_eval_callback_warned = false; - - std::shared_ptr weight_adapter = nullptr; - std::weak_ptr weight_manager; - std::unordered_set kept_compute_param_tensor_set; - std::vector runner_param_tensors; - std::unordered_set runner_param_tensor_set; - bool params_tensor_set_dirty_ = true; - - std::vector one_vec = {1.f}; - ggml_tensor* one_tensor = nullptr; - - std::vector zero_int_vec = {0}; - ggml_tensor* zero_int_tensor = nullptr; - - std::map backend_tensor_data_map; - std::map cache_tensor_map; // name -> tensor - std::vector> debug_tensors; - const std::string final_result_name = "ggml_runner_final_result_tensor"; - - bool flash_attn_enabled = false; - bool conv2d_direct_enabled = false; - bool circular_x_enabled = false; - bool circular_y_enabled = false; - - sd::ggml_graph_cut::PlanCache graph_cut_plan_cache_; - std::unordered_set params_tensor_set_; - std::unordered_map graph_cut_layer_split_assignments_; - std::unordered_map graph_cut_layer_split_node_assignments_; - bool graph_cut_layer_split_primary_notice_logged_ = false; - - template - static sd::Tensor take_or_empty(std::optional> tensor) { - if (!tensor.has_value()) { - return {}; - } - return std::move(*tensor); - } - - template - static sd::Tensor restore_trailing_singleton_dims(std::optional> tensor, - size_t expected_dim) { - return restore_trailing_singleton_dims(take_or_empty(std::move(tensor)), expected_dim); - } - - template - static sd::Tensor restore_trailing_singleton_dims(sd::Tensor tensor, - size_t expected_dim) { - if (tensor.empty()) { - return tensor; - } - while (static_cast(tensor.dim()) < expected_dim) { - tensor.unsqueeze_(tensor.dim()); - } - return tensor; - } - - void alloc_params_ctx() { - ggml_init_params params; - params.mem_size = static_cast(MAX_PARAMS_TENSOR_NUM * ggml_tensor_overhead()); - params.mem_buffer = nullptr; - params.no_alloc = true; - - params_ctx = ggml_init(params); - GGML_ASSERT(params_ctx != nullptr); - params_tensor_set_.clear(); - params_tensor_set_dirty_ = true; - } - - void free_params_ctx() { - if (params_ctx != nullptr) { - ggml_free(params_ctx); - params_ctx = nullptr; - } - params_tensor_set_.clear(); - params_tensor_set_dirty_ = true; - } - - void alloc_cache_ctx() { - ggml_init_params params; - params.mem_size = static_cast(MAX_PARAMS_TENSOR_NUM * ggml_tensor_overhead()); - params.mem_buffer = nullptr; - params.no_alloc = true; - - cache_ctx = ggml_init(params); - GGML_ASSERT(cache_ctx != nullptr); - } - - void free_cache_ctx() { - if (cache_ctx != nullptr) { - ggml_free(cache_ctx); - cache_ctx = nullptr; - } - } - - void alloc_compute_ctx() { - ggml_init_params params; - params.mem_size = static_cast(ggml_tensor_overhead() * MAX_GRAPH_SIZE + ggml_graph_overhead()); - params.mem_buffer = nullptr; - params.no_alloc = true; - - compute_ctx = ggml_init(params); - GGML_ASSERT(compute_ctx != nullptr); - } - - void free_compute_ctx() { - debug_tensors.clear(); - if (compute_ctx != nullptr) { - ggml_free(compute_ctx); - compute_ctx = nullptr; - } - backend_tensor_data_map.clear(); - } - - void rebuild_params_tensor_set() { - if (!params_tensor_set_dirty_) { - return; - } - params_tensor_set_.clear(); - if (params_ctx == nullptr) { - return; - } - for (ggml_tensor* t = ggml_get_first_tensor(params_ctx); t != nullptr; t = ggml_get_next_tensor(params_ctx, t)) { - params_tensor_set_.insert(t); - } - params_tensor_set_dirty_ = false; - } - - ggml_tensor* canonical_param_tensor(ggml_tensor* tensor) { - if (tensor == nullptr) { - return nullptr; - } - if (params_tensor_set_.find(tensor) != params_tensor_set_.end()) { - return tensor; - } - if (tensor->view_src != nullptr && - params_tensor_set_.find(tensor->view_src) != params_tensor_set_.end()) { - return tensor->view_src; - } - return nullptr; - } - - std::vector collect_used_param_tensors(ggml_cgraph* gf) { - std::vector used_params; - rebuild_params_tensor_set(); - if (gf == nullptr || params_tensor_set_.empty()) { - return used_params; - } - - std::unordered_set seen_params; - const int n_leafs = sd::ggml_graph_cut::leaf_count(gf); - seen_params.reserve(static_cast(n_leafs)); - for (int i = 0; i < n_leafs; ++i) { - ggml_tensor* leaf = sd::ggml_graph_cut::leaf_tensor(gf, i); - ggml_tensor* param_leaf = canonical_param_tensor(leaf); - if (param_leaf != nullptr && - seen_params.insert(param_leaf).second) { - used_params.push_back(param_leaf); - } - } - return used_params; - } - - bool prepare_execute_graph_weights(ggml_cgraph* gf, - std::vector& graph_param_tensors, - std::vector& params_to_prepare, - bool keep_compute_params) { - graph_param_tensors = collect_used_param_tensors(gf); - params_to_prepare.clear(); - params_to_prepare.reserve(graph_param_tensors.size()); - for (ggml_tensor* param : graph_param_tensors) { - if (param == nullptr) { - continue; - } - if (keep_compute_params && - kept_compute_param_tensor_set.find(param) != kept_compute_param_tensor_set.end()) { - continue; - } - params_to_prepare.push_back(param); - } - auto manager = weight_manager.lock(); - if (manager == nullptr) { - if (!params_to_prepare.empty()) { - LOG_ERROR("%s weight manager is not set for graph params", get_desc().c_str()); - return false; - } - return true; - } - - if (!manager->prepare_params(params_to_prepare)) { - LOG_ERROR("%s prepare graph weights failed", get_desc().c_str()); - return false; - } - for (ggml_tensor* param : params_to_prepare) { - if (param == nullptr) { - continue; - } - if (runner_param_tensor_set.insert(param).second) { - runner_param_tensors.push_back(param); - } - } - return true; - } - - void free_compute_backend_param_tensors(const std::vector& tensors) { - if (tensors.empty()) { - return; - } - auto manager = weight_manager.lock(); - if (manager != nullptr) { - manager->release_compute_backend_params(tensors); - } - } - - void free_params_backend_param_tensors(const std::vector& tensors) { - if (tensors.empty()) { - return; - } - auto manager = weight_manager.lock(); - if (manager != nullptr) { - manager->release_params_backend_params(tensors); - } - } - - void prepare_build_in_tensor_before() { - one_tensor = ggml_new_tensor_1d(compute_ctx, GGML_TYPE_F32, 1); - ggml_set_name(one_tensor, "ggml_runner_build_in_tensor:one"); - set_backend_tensor_data(one_tensor, one_vec.data()); - - zero_int_tensor = ggml_new_tensor_1d(compute_ctx, GGML_TYPE_I32, 1); - ggml_set_name(zero_int_tensor, "ggml_runner_build_in_tensor:zero_int"); - set_backend_tensor_data(zero_int_tensor, zero_int_vec.data()); - } - - void prepare_build_in_tensor_after(ggml_cgraph* gf) { - ggml_build_forward_expand(gf, one_tensor); - ggml_build_forward_expand(gf, zero_int_tensor); - } - - ggml_cgraph* new_graph_custom(size_t graph_size) { - if (weight_adapter) { - graph_size += weight_adapter->get_extra_graph_size(); - } - return ggml_new_graph_custom(compute_ctx, graph_size, false); - } - - ggml_cgraph* get_compute_graph(get_graph_cb_t get_graph) { - prepare_build_in_tensor_before(); - ggml_cgraph* gf = get_graph(); - if (ggml_graph_n_nodes(gf) > 0) { - auto result = ggml_graph_node(gf, -1); - ggml_set_name(result, final_result_name.c_str()); - } - for (const auto& entry : debug_tensors) { - if (entry.first != nullptr) { - ggml_build_forward_expand(gf, entry.first); - } - } - for (const auto& entry : cache_tensor_map) { - if (entry.second != nullptr) { - ggml_build_forward_expand(gf, entry.second); - } - } - prepare_build_in_tensor_after(gf); - return gf; - } - - bool prepare_compute_graph(get_graph_cb_t get_graph, - ggml_cgraph** gf_out) { - GGML_ASSERT(gf_out != nullptr); - - reset_compute_ctx(); - ggml_cgraph* gf = get_compute_graph(get_graph); - if (gf == nullptr) { - free_compute_ctx(); - return false; - } - - *gf_out = gf; - return true; - } - - // Pass explicit buffer types: synthesized defaults can make CUDA devices - // report supporting each other's buffers and skip a required copy. - bool ensure_sched(ggml_cgraph* gf) { - const size_t required_graph_size = gf != nullptr - ? std::max(1, - (size_t)ggml_graph_n_nodes(gf) + - sd::ggml_graph_cut::leaf_count(gf)) - : 1; - if (sched != nullptr && sched_graph_capacity >= required_graph_size) { - return true; - } - if (sched != nullptr) { - ggml_backend_sched_free(sched); - sched = nullptr; - sched_graph_capacity = 0; - } - std::vector backends; - backends.reserve(extra_runtime_backends.size() + 2); - backends.push_back(runtime_backend); - for (ggml_backend_t backend : extra_runtime_backends) { - backends.push_back(backend); - } - if (cpu_fallback_backend == nullptr && !sd_backend_is_cpu(runtime_backend)) { - cpu_fallback_backend = sd_backend_cpu_init(); - } - if (cpu_fallback_backend != nullptr) { - backends.push_back(cpu_fallback_backend); - } - - std::vector bufts; - bufts.reserve(backends.size()); - ggml_backend_dev_t main_dev = ggml_backend_get_device(runtime_backend); - for (ggml_backend_t backend : backends) { - ggml_backend_buffer_type_t buft = nullptr; - if (backend == cpu_fallback_backend && main_dev != nullptr) { - buft = ggml_backend_dev_host_buffer_type(main_dev); - } - if (buft == nullptr) { - buft = ggml_backend_get_default_buffer_type(backend); - } - bufts.push_back(buft); - } - - sched = ggml_backend_sched_new(backends.data(), - bufts.data(), - (int)backends.size(), - required_graph_size, - /*parallel=*/false, - /*op_offload=*/false); - if (sched == nullptr) { - LOG_ERROR("%s: failed to create backend sched", get_desc().c_str()); - return false; - } - sched_graph_capacity = required_graph_size; - return true; - } - - ggml_backend_t backend_for_weight(const ggml_tensor* tensor) const { - if (tensor == nullptr || tensor->buffer == nullptr) { - return nullptr; - } - if (ggml_backend_buffer_get_usage(tensor->buffer) != GGML_BACKEND_BUFFER_USAGE_WEIGHTS || - ggml_backend_buffer_is_host(tensor->buffer)) { - return nullptr; - } - ggml_backend_dev_t dev = ggml_backend_buft_get_device(ggml_backend_buffer_get_type(tensor->buffer)); - if (dev == nullptr) { - return nullptr; - } - if (ggml_backend_get_device(runtime_backend) == dev) { - return runtime_backend; - } - for (ggml_backend_t backend : extra_runtime_backends) { - if (ggml_backend_get_device(backend) == dev) { - return backend; - } - } - return nullptr; - } - - // Weightless ops have no scheduler anchor, so pin them to the most recent - // weight device. Views must stay unpinned or cross-device copies can be - // skipped for their consumers. - void pin_multi_device_nodes(ggml_cgraph* gf) { - if (sched == nullptr || gf == nullptr) { - return; - } - ggml_backend_t current = runtime_backend; - const int n_nodes = ggml_graph_n_nodes(gf); - for (int i = 0; i < n_nodes; i++) { - ggml_tensor* node = ggml_graph_node(gf, i); - auto node_assignment = graph_cut_layer_split_node_assignments_.find(node); - if (node_assignment != graph_cut_layer_split_node_assignments_.end()) { - current = node_assignment->second; - } - for (int s = 0; s < GGML_MAX_SRC; s++) { - ggml_backend_t weight_backend = backend_for_weight(node->src[s]); - if (weight_backend != nullptr) { - if (node_assignment == graph_cut_layer_split_node_assignments_.end()) { - current = weight_backend; - } - } - } - if (node->op == GGML_OP_NONE || node->op == GGML_OP_VIEW || node->op == GGML_OP_RESHAPE || - node->op == GGML_OP_PERMUTE || node->op == GGML_OP_TRANSPOSE) { - continue; - } - if (ggml_backend_supports_op(current, node)) { - ggml_backend_sched_set_tensor_backend(sched, node, current); - } - } - } - - bool is_multi_device() const { - return !extra_runtime_backends.empty(); - } - - bool graph_requires_backend_fallback(ggml_cgraph* gf) const { - if (gf == nullptr || sd_backend_is_cpu(runtime_backend)) { - return false; - } - const int n_nodes = ggml_graph_n_nodes(gf); - for (int i = 0; i < n_nodes; ++i) { - ggml_tensor* node = ggml_graph_node(gf, i); - if (node != nullptr && !ggml_backend_supports_op(runtime_backend, node)) { - return true; - } - } - return false; - } - - bool alloc_compute_buffer(ggml_cgraph* gf) { - if (sched != nullptr || is_multi_device() || graph_requires_backend_fallback(gf)) { - // The sched replaces the gallocr. Do NOT ggml_backend_sched_reserve - // the graph here: reserve runs split_graph, which rewires the - // graph's src pointers to sched-internal copy tensors, and the - // later ggml_backend_sched_alloc_graph would split the already - // rewired graph, silently corrupting every cross-backend input. A - // graph must be split at most once; the alloc in execute_graph - // performs the real allocation. - if (compute_allocr != nullptr) { - ggml_gallocr_free(compute_allocr); - compute_allocr = nullptr; - } - return ensure_sched(gf); - } - if (compute_allocr != nullptr) { - return true; - } - compute_allocr = ggml_gallocr_new(ggml_backend_get_default_buffer_type(runtime_backend)); - - if (!ggml_gallocr_reserve(compute_allocr, gf)) { - // failed to allocate the compute buffer - LOG_ERROR("%s: failed to allocate the compute buffer\n", get_desc().c_str()); - free_compute_buffer(); - return false; - } - - // compute the required memory - size_t compute_buffer_size = ggml_gallocr_get_buffer_size(compute_allocr, 0); - LOG_DEBUG("%s compute buffer size: %.2f MB(%s)", - get_desc().c_str(), - compute_buffer_size / 1024.0 / 1024.0, - sd_backend_is_cpu(runtime_backend) ? "RAM" : "VRAM"); - return true; - } - - void free_cache_buffer() { - if (cache_buffer != nullptr) { - ggml_backend_buffer_free(cache_buffer); - cache_buffer = nullptr; - } - } - - bool copy_cache_tensors_to_cache_buffer(const std::unordered_set* cache_keep_names = nullptr) { - if (cache_tensor_map.empty() && cache_keep_names == nullptr) { - return true; - } - - ggml_context* old_cache_ctx = cache_ctx; - ggml_backend_buffer_t old_cache_buffer = cache_buffer; - cache_ctx = nullptr; - cache_buffer = nullptr; - std::map merged_cache_sources; - if (old_cache_ctx != nullptr) { - for (ggml_tensor* tensor = ggml_get_first_tensor(old_cache_ctx); tensor != nullptr; tensor = ggml_get_next_tensor(old_cache_ctx, tensor)) { - if (cache_keep_names != nullptr && cache_keep_names->find(tensor->name) == cache_keep_names->end()) { - continue; - } - merged_cache_sources[tensor->name] = tensor; - } - } - for (const auto& kv : cache_tensor_map) { - if (cache_keep_names != nullptr && cache_keep_names->find(kv.first) == cache_keep_names->end()) { - continue; - } - merged_cache_sources[kv.first] = kv.second; - } - cache_tensor_map.clear(); - if (merged_cache_sources.empty()) { - if (old_cache_buffer != nullptr) { - ggml_backend_buffer_free(old_cache_buffer); - } - if (old_cache_ctx != nullptr) { - ggml_free(old_cache_ctx); - } - return true; - } - - alloc_cache_ctx(); - std::vector> source_to_cache_tensors; - source_to_cache_tensors.reserve(merged_cache_sources.size()); - for (const auto& kv : merged_cache_sources) { - ggml_tensor* source_tensor = sd::ggml_graph_cut::cache_source_tensor(kv.second); - auto cache_tensor = ggml_dup_tensor(cache_ctx, source_tensor); - ggml_set_name(cache_tensor, kv.first.c_str()); - source_to_cache_tensors.push_back({source_tensor, cache_tensor}); - } - size_t num_tensors = ggml_tensor_num(cache_ctx); - cache_buffer = ggml_backend_alloc_ctx_tensors(cache_ctx, runtime_backend); - GGML_ASSERT(cache_buffer != nullptr); - for (const auto& kv : source_to_cache_tensors) { - ggml_tensor* src = kv.first; - ggml_tensor* dst = kv.second; - ggml_backend_buffer_t src_buf = sd::ggml_graph_cut::tensor_buffer(src); - ggml_backend_buffer_t dst_buf = sd::ggml_graph_cut::tensor_buffer(dst); - if (src_buf == nullptr || dst_buf == nullptr) { - LOG_ERROR("%s cache copy tensor buffer missing: name=%s op=%s src0=%p src0_name=%s src0_buffer=%p src_buffer=%p src_view_src=%p src_view_src_buffer=%p dst_buffer=%p", - get_desc().c_str(), - src && src->name[0] != '\0' ? src->name : "", - src ? ggml_op_name(src->op) : "", - src ? src->src[0] : nullptr, - (src && src->src[0] && src->src[0]->name[0] != '\0') ? src->src[0]->name : "", - (src && src->src[0]) ? sd::ggml_graph_cut::tensor_buffer(src->src[0]) : nullptr, - src ? src->buffer : nullptr, - src ? src->view_src : nullptr, - (src && src->view_src) ? src->view_src->buffer : nullptr, - dst ? dst->buffer : nullptr); - return false; - } - const bool use_staging_copy = src->view_src != nullptr || !ggml_is_contiguous(src) || src->buffer == nullptr; - if (use_staging_copy) { - std::vector host_data(ggml_nbytes(src)); - ggml_backend_tensor_get(src, host_data.data(), 0, host_data.size()); - ggml_backend_tensor_set(dst, host_data.data(), 0, host_data.size()); - } else { - ggml_backend_tensor_copy(src, dst); - } - } - ggml_backend_synchronize(runtime_backend); - size_t cache_buffer_size = ggml_backend_buffer_get_size(cache_buffer); - LOG_DEBUG("%s cache backend buffer size = % 6.2f MB(%s) (%i tensors)", - get_desc().c_str(), - cache_buffer_size / (1024.f * 1024.f), - sd_backend_is_cpu(runtime_backend) ? "RAM" : "VRAM", - num_tensors); - if (old_cache_buffer != nullptr) { - ggml_backend_buffer_free(old_cache_buffer); - } - if (old_cache_ctx != nullptr) { - ggml_free(old_cache_ctx); - } - return true; - } - - template - std::optional> read_graph_tensor(ggml_tensor* tensor, const char* label) { - if (tensor == nullptr) { - LOG_ERROR("%s %s tensor is null", get_desc().c_str(), label); - return std::nullopt; - } - if (tensor->type != sd::GGMLTypeTraits::type) { - LOG_ERROR("%s %s tensor type mismatch: got %s", - get_desc().c_str(), - label, - ggml_type_name(tensor->type)); - return std::nullopt; - } - ggml_backend_buffer_t buf = sd::ggml_graph_cut::tensor_buffer(tensor); - if (buf == nullptr) { - LOG_ERROR("%s %s tensor buffer missing: name=%s op=%s buffer=%p view_src=%p view_src_buffer=%p data=%p", - get_desc().c_str(), - label, - tensor->name[0] != '\0' ? tensor->name : "", - ggml_op_name(tensor->op), - tensor->buffer, - tensor->view_src, - tensor->view_src ? tensor->view_src->buffer : nullptr, - tensor->data); - return std::nullopt; - } - - sd::Tensor result(sd::shape_from_ggml(tensor)); - if (tensor->view_src != nullptr || !ggml_is_contiguous(tensor) || tensor->buffer == nullptr) { - ggml_backend_tensor_get(tensor, result.data(), 0, ggml_nbytes(tensor)); - } else { - ggml_backend_tensor_get(tensor, result.data(), 0, ggml_nbytes(tensor)); - } - return result; - } - - void copy_data_to_backend_tensor(ggml_cgraph* gf, bool clear_after_copy = true) { - GGML_ASSERT(gf != nullptr); - std::unordered_set graph_tensor_set; - const int n_leafs = sd::ggml_graph_cut::leaf_count(gf); - const int n_nodes = ggml_graph_n_nodes(gf); - graph_tensor_set.reserve(static_cast(n_leafs + n_nodes)); - for (int i = 0; i < n_leafs; ++i) { - graph_tensor_set.insert(sd::ggml_graph_cut::leaf_tensor(gf, i)); - } - for (int i = 0; i < n_nodes; ++i) { - graph_tensor_set.insert(ggml_graph_node(gf, i)); - } - - for (auto& kv : backend_tensor_data_map) { - auto tensor = kv.first; - auto data = kv.second; - if (tensor == nullptr || data == nullptr) { - continue; - } - const char* name = ggml_get_name(tensor); - if (graph_tensor_set.find(tensor) == graph_tensor_set.end()) { - continue; - } - if (tensor->buffer == nullptr) { - LOG_WARN("%s skip backend tensor copy: tensor buffer not set, name='%s', ne=[%lld,%lld,%lld,%lld], type=%s", - get_desc().c_str(), - name != nullptr ? name : "", - (long long)tensor->ne[0], - (long long)tensor->ne[1], - (long long)tensor->ne[2], - (long long)tensor->ne[3], - ggml_type_name(tensor->type)); - continue; - } - - ggml_backend_buffer_t buf = tensor->view_src ? tensor->view_src->buffer : tensor->buffer; - if (buf == nullptr) { - LOG_WARN("%s graph exec skip tensor copy: name=%s op=%s reason=buffer_not_set data=%p view_src=%p view_src_buffer=%p", - get_desc().c_str(), - tensor && tensor->name[0] != '\0' ? tensor->name : "", - tensor ? ggml_op_name(tensor->op) : "", - data, - tensor ? tensor->view_src : nullptr, - (tensor && tensor->view_src) ? tensor->view_src->buffer : nullptr); - continue; - } - - ggml_backend_tensor_set(tensor, data, 0, ggml_nbytes(tensor)); - } - - if (clear_after_copy) { - backend_tensor_data_map.clear(); - } - } - - bool should_use_graph_cut_segmented_compute(const GraphCutPlan& plan) { - return plan.has_cuts && - plan.valid && - max_graph_vram_bytes > 0 && - plan.segments.size() > 1 && - !sd_backend_is_cpu(runtime_backend) && - !is_multi_device(); - } - - bool can_attempt_graph_cut_segmented_compute() const { - return max_graph_vram_bytes > 0 && - !sd_backend_is_cpu(runtime_backend) && - !is_multi_device(); - } - - bool resolve_graph_cut_plan(ggml_cgraph* gf, - GraphCutPlan* plan_out, - size_t* effective_budget_out = nullptr) { - GGML_ASSERT(plan_out != nullptr); - GGML_ASSERT(gf != nullptr); - - size_t effective_budget = max_graph_vram_bytes; - size_t free_clamp = SIZE_MAX; - if (stream_layers_enabled && max_graph_vram_bytes > 0 && runtime_backend != nullptr) { - ggml_backend_dev_t dev = ggml_backend_get_device(runtime_backend); - if (dev != nullptr && ggml_backend_dev_type(dev) != GGML_BACKEND_DEVICE_TYPE_CPU) { - size_t free_vram = 0, total_vram = 0; - ggml_backend_dev_memory(dev, &free_vram, &total_vram); - constexpr size_t safety_margin = 512ull * 1024 * 1024; - free_clamp = (free_vram > safety_margin) ? (free_vram - safety_margin) : 0; - if (free_clamp < effective_budget) { - LOG_DEBUG("%s clamping streaming budget: actual free VRAM %.2f MB < user cap %.2f MB", - get_desc().c_str(), - free_clamp / (1024.0 * 1024.0), - effective_budget / (1024.0 * 1024.0)); - effective_budget = free_clamp; - } - } - } - - bool budget_increased = false; - if (stream_layers_enabled) { - if (effective_budget > observed_max_effective_budget_) { - observed_max_effective_budget_ = effective_budget; - budget_increased = true; - } else { - // Keep the plan cache stable, but never plan above what is free now: - // another model or process can take VRAM after the first measurement. - effective_budget = std::min(observed_max_effective_budget_, free_clamp); - } - } - - if (effective_budget_out != nullptr) { - *effective_budget_out = effective_budget; - } - - // When streaming and the model dwarfs the budget, cap the planner at - // a quarter so it builds smaller merged segments and chunk-K can fit - // alongside. Without streaming the cap only adds dispatch overhead. - size_t planner_budget = effective_budget; - if (stream_layers_enabled) { - size_t total_params_bytes = 0; - for (const ggml_tensor* t : params_tensor_set_) { - if (t != nullptr) { - total_params_bytes += ggml_nbytes(t); - } - } - if (total_params_bytes * 4 > effective_budget * 3) { - planner_budget = effective_budget / 4; - } - } - - *plan_out = sd::ggml_graph_cut::resolve_plan(runtime_backend, - gf, - &graph_cut_plan_cache_, - planner_budget, - params_tensor_set_, - get_desc().c_str()); - if (stream_layers_enabled) { - sd::ggml_graph_cut::annotate_residency(*plan_out, effective_budget); - } - if (stream_layers_enabled) { - if (budget_increased) { - LOG_INFO("%s streaming budget = %.2f MB", - get_desc().c_str(), - effective_budget / (1024.0 * 1024.0)); - } else { - LOG_DEBUG("%s streaming budget = %.2f MB", - get_desc().c_str(), - effective_budget / (1024.0 * 1024.0)); - } - } - return true; - } - - bool resolve_graph_cut_layer_split_plan(ggml_cgraph* gf, - GraphCutPlan* plan_out) { - GGML_ASSERT(plan_out != nullptr); - GGML_ASSERT(gf != nullptr); - *plan_out = sd::ggml_graph_cut::resolve_plan(runtime_backend, - gf, - &graph_cut_plan_cache_, - 0, - params_tensor_set_, - get_desc().c_str()); - return true; - } - - bool assign_graph_cut_layer_split_backends(ggml_cgraph* gf) { - graph_cut_layer_split_node_assignments_.clear(); - if (!graph_cut_layer_split_enabled) { - return true; - } - if (!is_multi_device()) { - LOG_ERROR("%s graph-cut layer split requires multiple runtime backends", get_desc().c_str()); - return false; - } - - GraphCutPlan plan; - if (!resolve_graph_cut_layer_split_plan(gf, &plan)) { - return false; - } - if (!plan.valid || !plan.has_cuts || plan.segments.size() <= 1) { - auto manager = weight_manager.lock(); - if (manager == nullptr) { - LOG_ERROR("%s weight manager is not set for graph-cut layer split", get_desc().c_str()); - return false; - } - std::vector graph_params = collect_used_param_tensors(gf); - if (!graph_params.empty() && - !manager->assign_compute_backend(graph_params, runtime_backend)) { - LOG_ERROR("%s graph-cut layer split failed to assign unmarked graph params to %s", - get_desc().c_str(), - sd::layer_split_backend_device_display_name(runtime_backend).c_str()); - return false; - } - for (ggml_tensor* param : graph_params) { - if (param != nullptr) { - graph_cut_layer_split_assignments_[param] = runtime_backend; - } - } - const int n_nodes = ggml_graph_n_nodes(gf); - for (int i = 0; i < n_nodes; i++) { - ggml_tensor* node = ggml_graph_node(gf, i); - if (node != nullptr) { - graph_cut_layer_split_node_assignments_[node] = runtime_backend; - } - } - if (!graph_cut_layer_split_primary_notice_logged_) { - LOG_WARN("%s graph-cut layer split: graph has no mark_graph_cut segments; using primary backend %s for %zu graph params", - get_desc().c_str(), - sd::layer_split_backend_device_display_name(runtime_backend).c_str(), - graph_params.size()); - graph_cut_layer_split_primary_notice_logged_ = true; - } else { - LOG_DEBUG("%s graph-cut layer split: graph has no mark_graph_cut segments; using primary backend %s for %zu graph params", - get_desc().c_str(), - sd::layer_split_backend_device_display_name(runtime_backend).c_str(), - graph_params.size()); - } - return true; - } - - std::vector split_backends; - split_backends.reserve(extra_runtime_backends.size() + 1); - split_backends.push_back(runtime_backend); - for (ggml_backend_t backend : extra_runtime_backends) { - if (backend != nullptr) { - split_backends.push_back(backend); - } - } - - auto manager = weight_manager.lock(); - if (manager == nullptr) { - LOG_ERROR("%s weight manager is not set for graph-cut layer split", get_desc().c_str()); - return false; - } - - sd::GraphCutLayerSplitAssignment assignment; - auto canonicalize_param = [this](ggml_tensor* tensor) { - return canonical_param_tensor(tensor); - }; - if (!sd::partition_graph_cut_layer_split(get_desc().c_str(), - gf, - plan, - split_backends, - graph_cut_layer_split_backend_vram_limits_, - max_graph_vram_bytes, - graph_cut_layer_split_assignments_, - canonicalize_param, - &assignment)) { - return false; - } - - for (size_t i = 0; i < split_backends.size(); i++) { - if (assignment.tensors_by_backend[i].empty()) { - continue; - } - if (!manager->assign_compute_backend(assignment.tensors_by_backend[i], split_backends[i])) { - LOG_ERROR("%s graph-cut layer split failed to assign params to %s", - get_desc().c_str(), - sd::layer_split_backend_device_display_name(split_backends[i]).c_str()); - return false; - } - } - - graph_cut_layer_split_node_assignments_ = std::move(assignment.node_assignments); - sd::log_graph_cut_layer_split_assignment(get_desc().c_str(), split_backends, assignment); - - return true; - } - - struct PersistentExternalBinding { - ggml_backend_buffer_t buffer = nullptr; - void* data = nullptr; - void* extra = nullptr; - }; - - void snapshot_persistent_externals(const sd::ggml_graph_cut::Plan& plan, - ggml_cgraph* gf, - std::unordered_map& out) { - GGML_ASSERT(gf != nullptr); - out.clear(); - for (const auto& segment : plan.segments) { - for (const auto& input : segment.input_refs) { - if (input.type != GraphCutSegment::INPUT_EXTERNAL) { - continue; - } - ggml_tensor* tensor = sd::ggml_graph_cut::input_tensor(gf, input); - if (tensor == nullptr || tensor->buffer == nullptr) { - continue; - } - PersistentExternalBinding binding; - binding.buffer = tensor->buffer; - binding.data = tensor->data; - binding.extra = tensor->extra; - out[tensor] = binding; - } - } - } - - void reset_segment_runtime_tensors(const GraphCutSegment& segment, - ggml_cgraph* gf, - const std::unordered_map* persistent_externals = nullptr) { - GGML_ASSERT(gf != nullptr); - - for (const auto& input : segment.input_refs) { - ggml_tensor* input_tensor = sd::ggml_graph_cut::input_tensor(gf, input); - if (input_tensor == nullptr) { - continue; - } - switch (input.type) { - case GraphCutSegment::INPUT_PREVIOUS_CUT: - input_tensor->buffer = nullptr; - input_tensor->data = nullptr; - input_tensor->extra = nullptr; - break; - case GraphCutSegment::INPUT_EXTERNAL: { - if (persistent_externals != nullptr) { - auto it = persistent_externals->find(input_tensor); - if (it != persistent_externals->end()) { - input_tensor->buffer = it->second.buffer; - input_tensor->data = it->second.data; - input_tensor->extra = it->second.extra; - break; - } - } - input_tensor->buffer = nullptr; - input_tensor->data = nullptr; - input_tensor->extra = nullptr; - break; - } - case GraphCutSegment::INPUT_PARAM: - break; - } - } - - for (int node_idx : segment.internal_node_indices) { - ggml_tensor* node = ggml_graph_node(gf, node_idx); - if (node == nullptr) { - continue; - } - node->buffer = nullptr; - node->data = nullptr; - node->extra = nullptr; - } - } - - bool bind_segment_cached_inputs(ggml_cgraph* gf, const GraphCutSegment& segment) { - GGML_ASSERT(gf != nullptr); - for (const auto& input : segment.input_refs) { - ggml_tensor* input_tensor = sd::ggml_graph_cut::input_tensor(gf, input); - if (input_tensor == nullptr) { - continue; - } - switch (input.type) { - case GraphCutSegment::INPUT_PREVIOUS_CUT: { - ggml_tensor* cache_tensor = get_cache_tensor_by_name(input.display_name); - if (cache_tensor == nullptr) { - LOG_ERROR("%s missing graph cut cache tensor: %s", - get_desc().c_str(), - input.display_name.c_str()); - return false; - } - if (input_tensor->view_src != nullptr) { - input_tensor->view_src = cache_tensor; - input_tensor->buffer = nullptr; - input_tensor->data = cache_tensor->data == nullptr - ? nullptr - : static_cast(static_cast(cache_tensor->data) + input_tensor->view_offs); - input_tensor->extra = cache_tensor->extra; - } else { - input_tensor->buffer = cache_tensor->buffer; - input_tensor->data = cache_tensor->data; - input_tensor->extra = cache_tensor->extra; - } - for (int src_idx = 0; src_idx < GGML_MAX_SRC; ++src_idx) { - input_tensor->src[src_idx] = nullptr; - } - input_tensor->op = GGML_OP_NONE; - break; - } - case GraphCutSegment::INPUT_EXTERNAL: - case GraphCutSegment::INPUT_PARAM: - break; - } - } - return true; - } - - template - std::optional> execute_graph(ggml_cgraph* gf, - int n_threads, - bool free_compute_buffer, - bool free_compute_params, - bool preserve_backend_tensor_data_map, - bool no_return = false, - const std::unordered_set* cache_keep_names = nullptr) { - std::vector graph_param_tensors; - std::vector params_to_prepare; - if (!prepare_execute_graph_weights(gf, graph_param_tensors, params_to_prepare, !free_compute_params)) { - return std::nullopt; - } - struct GraphWeightDoneGuard { - GraphWeightDoneGuard(GGMLRunner* runner, const std::vector* tensors) - : runner(runner), - tensors(tensors) {} - - GGMLRunner* runner = nullptr; - const std::vector* tensors = nullptr; - bool enabled = true; - - ~GraphWeightDoneGuard() { - if (enabled && runner != nullptr && tensors != nullptr) { - runner->free_compute_backend_param_tensors(*tensors); - } - } - - void dismiss() { enabled = false; } - - GraphWeightDoneGuard(const GraphWeightDoneGuard&) = delete; - GraphWeightDoneGuard& operator=(const GraphWeightDoneGuard&) = delete; - }; - GraphWeightDoneGuard graph_weight_done_guard(this, ¶ms_to_prepare); - - if (!alloc_compute_buffer(gf)) { - LOG_ERROR("%s alloc compute buffer failed", get_desc().c_str()); - return std::nullopt; - } - struct ComputeBufferGuard { - ComputeBufferGuard(GGMLRunner* runner, bool enabled) - : runner(runner), - enabled(enabled) {} - - GGMLRunner* runner = nullptr; - bool enabled = false; - - ~ComputeBufferGuard() { - if (enabled && runner != nullptr) { - runner->free_compute_buffer(); - } - } - - ComputeBufferGuard(const ComputeBufferGuard&) = delete; - ComputeBufferGuard& operator=(const ComputeBufferGuard&) = delete; - }; - ComputeBufferGuard compute_buffer_guard(this, free_compute_buffer); - - if (sched != nullptr) { - ggml_backend_sched_reset(sched); - pin_multi_device_nodes(gf); // reset clears the pins; re-apply before alloc - if (!ggml_backend_sched_alloc_graph(sched, gf)) { - LOG_ERROR("%s sched alloc compute graph failed", get_desc().c_str()); - return std::nullopt; - } - } else if (!ggml_gallocr_alloc_graph(compute_allocr, gf)) { - LOG_ERROR("%s alloc compute graph failed", get_desc().c_str()); - return std::nullopt; - } - - copy_data_to_backend_tensor(gf, !preserve_backend_tensor_data_map); - if (sd_backend_is_cpu(runtime_backend)) { - sd_backend_cpu_set_n_threads(runtime_backend, n_threads); - } - if (cpu_fallback_backend != nullptr) { - sd_backend_cpu_set_n_threads(cpu_fallback_backend, n_threads); - } - - ggml_status status; - if (sched != nullptr) { - if (sd_get_backend_eval_callback() != nullptr && !multi_device_eval_callback_warned) { - LOG_WARN("%s: eval callback is not supported with the backend scheduler; ignoring", - get_desc().c_str()); - multi_device_eval_callback_warned = true; - } - status = ggml_backend_sched_graph_compute(sched, gf); - if (status == GGML_STATUS_SUCCESS) { - ggml_backend_sched_synchronize(sched); - } - } else { - status = sd_backend_graph_compute_with_eval_callback(runtime_backend, - gf, - sd_get_backend_eval_callback(), - sd_get_backend_eval_callback_data()); - } - if (status != GGML_STATUS_SUCCESS) { - LOG_ERROR("%s compute failed: %s", get_desc().c_str(), ggml_status_to_string(status)); - return std::nullopt; - } - - if (!debug_tensors.empty()) { - std::unordered_set debug_graph_tensor_set; - const int n_debug_leafs = sd::ggml_graph_cut::leaf_count(gf); - const int n_debug_nodes = ggml_graph_n_nodes(gf); - debug_graph_tensor_set.reserve(static_cast(n_debug_leafs + n_debug_nodes)); - for (int i = 0; i < n_debug_leafs; ++i) { - debug_graph_tensor_set.insert(sd::ggml_graph_cut::leaf_tensor(gf, i)); - } - for (int i = 0; i < n_debug_nodes; ++i) { - debug_graph_tensor_set.insert(ggml_graph_node(gf, i)); - } - - for (const auto& entry : debug_tensors) { - auto tensor = entry.first; - if (tensor == nullptr) { - continue; - } - if (debug_graph_tensor_set.find(tensor) == debug_graph_tensor_set.end()) { - continue; - } - ggml_backend_buffer_t tensor_buf = tensor->view_src ? tensor->view_src->buffer : tensor->buffer; - if (tensor_buf == nullptr) { - LOG_WARN("%s skip debug tensor '%s': tensor buffer not set", - get_desc().c_str(), - entry.second.c_str()); - continue; - } - if (tensor->type != GGML_TYPE_F32) { - LOG_WARN("%s skip debug tensor '%s': only GGML_TYPE_F32 is supported, got %s", - get_desc().c_str(), - entry.second.c_str(), - ggml_type_name(tensor->type)); - continue; - } - auto debug_tensor = sd::make_sd_tensor_from_ggml(tensor); - print_sd_tensor(debug_tensor, false, entry.second.c_str()); - } - } - - if (!copy_cache_tensors_to_cache_buffer(cache_keep_names)) { - return std::nullopt; - } - auto result = ggml_get_tensor(compute_ctx, final_result_name.c_str()); - std::optional> output; - if (!no_return) { - output = read_graph_tensor(result, "output"); - if (!output.has_value()) { - return std::nullopt; - } - } else { - output = sd::Tensor(); - } - - if (!free_compute_params) { - for (ggml_tensor* param : params_to_prepare) { - if (param == nullptr) { - continue; - } - kept_compute_param_tensor_set.insert(param); - } - graph_weight_done_guard.dismiss(); - } - return output; - } - - template - std::optional> compute_graph_cut_segments(ggml_cgraph* gf, - const GraphCutPlan& plan, - int n_threads, - bool log_residency, - bool no_return = false) { - GGML_ASSERT(gf != nullptr); - - free_compute_buffer(); - free_cache_ctx_and_buffer(); - - std::unordered_map persistent_externals; - snapshot_persistent_externals(plan, gf, persistent_externals); - - std::optional> output = sd::Tensor(); - for (size_t seg_idx = 0; seg_idx < plan.segments.size(); ++seg_idx) { - const auto& segment = plan.segments[seg_idx]; - const bool is_last = seg_idx + 1 == plan.segments.size(); - auto future_cut_names = sd::ggml_graph_cut::collect_future_input_names(gf, plan, seg_idx); - if (log_residency) { - LOG_DEBUG("%s graph cut executing segment %zu/%zu: %s (residency=%s)", - get_desc().c_str(), - seg_idx + 1, - plan.segments.size(), - segment.group_name.c_str(), - segment.residency == sd::ggml_graph_cut::SegmentResidency::RESIDENT ? "RESIDENT" : "STREAMED"); - } else { - LOG_DEBUG("%s graph cut executing segment %zu/%zu: %s", - get_desc().c_str(), - seg_idx + 1, - plan.segments.size(), - segment.group_name.c_str()); - } - - reset_segment_runtime_tensors(segment, gf, &persistent_externals); - if (!bind_segment_cached_inputs(gf, segment)) { - free_cache_ctx_and_buffer(); - free_compute_buffer(); - free_compute_ctx(); - return std::nullopt; - } - - if (!is_last) { - for (size_t output_idx = 0; output_idx < segment.output_node_indices.size(); ++output_idx) { - ggml_tensor* output_tensor = sd::ggml_graph_cut::output_tensor(gf, segment, output_idx); - if (output_tensor != nullptr && - sd::ggml_graph_cut::is_graph_cut_tensor(output_tensor) && - future_cut_names.find(output_tensor->name) != future_cut_names.end()) { - cache(output_tensor->name, output_tensor); - } - } - } - - ggml_context* segment_graph_ctx = nullptr; - ggml_cgraph* segment_graph = sd::ggml_graph_cut::build_segment_graph(gf, segment, &segment_graph_ctx); - const bool keep_segment_params = segment.residency == sd::ggml_graph_cut::SegmentResidency::RESIDENT; - auto segment_output = execute_graph(segment_graph, - n_threads, - true, - !keep_segment_params, - true, - !is_last || no_return, - &future_cut_names); - ggml_free(segment_graph_ctx); - if (!segment_output.has_value()) { - free_cache_ctx_and_buffer(); - free_compute_buffer(); - free_compute_ctx(); - return std::nullopt; - } - output = std::move(segment_output); - } - - backend_tensor_data_map.clear(); - free_cache_ctx_and_buffer(); - free_compute_ctx(); - return output; - } - -public: - void runner_done() { - free_compute_buffer(); - std::vector tensors_to_release = std::move(this->runner_param_tensors); - this->runner_param_tensors.clear(); - runner_param_tensor_set.clear(); - kept_compute_param_tensor_set.clear(); - free_compute_backend_param_tensors(tensors_to_release); - free_params_backend_param_tensors(tensors_to_release); - } - -public: - virtual std::string get_desc() = 0; - - GGMLRunner(ggml_backend_t backend, - std::shared_ptr manager = nullptr) - : runtime_backend(backend), - weight_manager(manager) { - GGML_ASSERT(runtime_backend != nullptr); - alloc_params_ctx(); - } - - virtual ~GGMLRunner() { - free_compute_buffer(); - free_params_ctx(); - free_compute_ctx(); - free_cache_ctx_and_buffer(); - if (cpu_fallback_backend != nullptr) { - ggml_backend_free(cpu_fallback_backend); - cpu_fallback_backend = nullptr; - } - } - - virtual GGMLRunnerContext get_context() { - GGMLRunnerContext runner_ctx; - runner_ctx.ggml_ctx = compute_ctx; - runner_ctx.backend = runtime_backend; - runner_ctx.flash_attn_enabled = flash_attn_enabled; - runner_ctx.conv2d_direct_enabled = conv2d_direct_enabled; - runner_ctx.circular_x_enabled = circular_x_enabled; - runner_ctx.circular_y_enabled = circular_y_enabled; - runner_ctx.weight_adapter = weight_adapter; - runner_ctx.debug_tensors = &debug_tensors; - runner_ctx.get_cache_tensor = [this](const std::string& name) { - return this->get_cache_tensor_by_name(name); - }; - runner_ctx.cache_tensor = [this](const std::string& name, ggml_tensor* tensor) { - this->cache(name, tensor); - }; - runner_ctx.set_backend_tensor_data = [this](ggml_tensor* tensor, const void* data) { - this->set_backend_tensor_data(tensor, data); - }; - return runner_ctx; - } - - void reset_compute_ctx() { - free_compute_ctx(); - alloc_compute_ctx(); - } - -public: - void free_cache_ctx_and_buffer() { - free_cache_buffer(); - free_cache_ctx(); - } - - void free_compute_buffer() { - if (compute_allocr != nullptr) { - ggml_gallocr_free(compute_allocr); - compute_allocr = nullptr; - } - if (sched != nullptr) { - ggml_backend_sched_free(sched); - sched = nullptr; - sched_graph_capacity = 0; - } - } - - // do copy after alloc graph - void set_backend_tensor_data(ggml_tensor* tensor, const void* data) { - // The scheduler only allocates standalone data tensors when they are - // marked as graph inputs. The flag is harmless for single-backend graphs. - ggml_set_input(tensor); - backend_tensor_data_map[tensor] = data; - } - - template - ggml_tensor* make_input(const sd::Tensor& tensor) { - ggml_tensor* input = sd::make_ggml_tensor(compute_ctx, tensor, false); - set_backend_tensor_data(input, tensor.data()); - return input; - } - - template - ggml_tensor* make_optional_input(const sd::Tensor& tensor) { - if (tensor.empty()) { - return nullptr; - } - return make_input(tensor); - } - - template - ggml_tensor* make_optional_input(const sd::Tensor* tensor) { - if (tensor == nullptr) { - return nullptr; - } - return make_input(*tensor); - } - - ggml_tensor* to_backend(ggml_tensor* tensor) { - GGML_ASSERT(compute_ctx != nullptr); - if (tensor == nullptr) { - return nullptr; - } - // it's performing a compute, check if backend isn't cpu - if (!sd_backend_is_cpu(runtime_backend) && (tensor->buffer == nullptr || ggml_backend_buffer_is_host(tensor->buffer))) { - // pass input tensors to gpu memory - auto backend_tensor = ggml_dup_tensor(compute_ctx, tensor); - - set_backend_tensor_data(backend_tensor, tensor->data); - return backend_tensor; - } else { - return tensor; - } - } - - void cache(const std::string name, ggml_tensor* tensor) { - if (tensor != nullptr && tensor->view_src != nullptr) { - tensor = ggml_cont(compute_ctx, tensor); - } - cache_tensor_map[name] = tensor; - } - - ggml_tensor* get_cache_tensor_by_name(const std::string& name) { - if (cache_ctx == nullptr) { - return nullptr; - } - return ggml_get_tensor(cache_ctx, name.c_str()); - } - - template - std::optional> compute(get_graph_cb_t get_graph, - int n_threads, - bool auto_free = true, - bool free_compute_buffer = true, - bool free_compute_params = true, - bool no_return = false) { - struct RunnerDoneGuard { - RunnerDoneGuard(GGMLRunner* runner, bool enabled) - : runner(runner), - enabled(enabled) {} - - ~RunnerDoneGuard() { - if (enabled && runner != nullptr) { - runner->runner_done(); - } - } - - RunnerDoneGuard(const RunnerDoneGuard&) = delete; - RunnerDoneGuard& operator=(const RunnerDoneGuard&) = delete; - - GGMLRunner* runner = nullptr; - bool enabled = false; - }; - RunnerDoneGuard runner_done_guard(this, auto_free); - - ggml_cgraph* gf = nullptr; - if (!prepare_compute_graph(get_graph, &gf)) { - return std::nullopt; - } - GGML_ASSERT(gf != nullptr); - rebuild_params_tensor_set(); - - if (!assign_graph_cut_layer_split_backends(gf)) { - free_compute_ctx(); - return std::nullopt; - } - - if (can_attempt_graph_cut_segmented_compute()) { - GraphCutPlan plan; - if (!resolve_graph_cut_plan(gf, &plan)) { - free_compute_ctx(); - return std::nullopt; - } - if (should_use_graph_cut_segmented_compute(plan)) { - return compute_graph_cut_segments(gf, - plan, - n_threads, - stream_layers_enabled, - no_return); - } - } - return execute_graph(gf, - n_threads, - free_compute_buffer, - free_compute_params, - false, - no_return, - nullptr); - } - - void set_flash_attention_enabled(bool enabled) { - flash_attn_enabled = enabled; - } - - void set_conv2d_direct_enabled(bool enabled) { - conv2d_direct_enabled = enabled; - } - - void set_circular_axes(bool circular_x, bool circular_y) { - circular_x_enabled = circular_x; - circular_y_enabled = circular_y; - } - - void set_weight_adapter(const std::shared_ptr& adapter) { - weight_adapter = adapter; - } - - void set_max_graph_vram_bytes(size_t max_vram_bytes) { - max_graph_vram_bytes = max_vram_bytes; - } - - void set_stream_layers_enabled(bool enabled) { - if (enabled && is_multi_device()) { - LOG_WARN("%s: --stream-layers is not supported with multiple runtime backends; ignoring", - get_desc().c_str()); - return; - } - stream_layers_enabled = enabled; - } - - void set_graph_cut_layer_split_enabled(bool enabled) { - graph_cut_layer_split_enabled = enabled; - if (!enabled) { - graph_cut_layer_split_assignments_.clear(); - graph_cut_layer_split_node_assignments_.clear(); - graph_cut_layer_split_primary_notice_logged_ = false; - } - } - - void set_graph_cut_layer_split_backend_vram_limits(const std::vector& limits) { - graph_cut_layer_split_backend_vram_limits_ = limits; - graph_cut_layer_split_assignments_.clear(); - graph_cut_layer_split_node_assignments_.clear(); - graph_cut_layer_split_primary_notice_logged_ = false; - } - - void set_runtime_backends(const std::vector& backends) { - extra_runtime_backends.clear(); - for (ggml_backend_t backend : backends) { - if (backend == nullptr || backend == runtime_backend) { - continue; - } - if (std::find(extra_runtime_backends.begin(), extra_runtime_backends.end(), backend) == - extra_runtime_backends.end()) { - extra_runtime_backends.push_back(backend); - } - } - graph_cut_layer_split_assignments_.clear(); - graph_cut_layer_split_node_assignments_.clear(); - graph_cut_layer_split_primary_notice_logged_ = false; - if (is_multi_device() && stream_layers_enabled) { - LOG_WARN("%s: --stream-layers is not supported with multiple runtime backends; ignoring", - get_desc().c_str()); - stream_layers_enabled = false; - } - } -}; - -class GGMLBlock { -protected: - typedef std::unordered_map ParameterMap; - typedef std::unordered_map> GGMLBlockMap; - GGMLBlockMap blocks; - ParameterMap params; - - ggml_type get_type(const std::string& name, const String2TensorStorage& tensor_storage_map, ggml_type default_type) { - ggml_type wtype = default_type; - auto iter = tensor_storage_map.find(name); - if (iter != tensor_storage_map.end()) { - const TensorStorage& tensor_storage = iter->second; - if (tensor_storage.expected_type != GGML_TYPE_COUNT) { - wtype = tensor_storage.expected_type; - } else { - wtype = tensor_storage.type; - } - } - return wtype; - } - - void init_blocks(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") { - for (auto& pair : blocks) { - auto& block = pair.second; - block->init(ctx, tensor_storage_map, prefix + pair.first); - } - } - - virtual void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") {} - - virtual enum ggml_op param_usage_op(const std::string& name) const { - (void)name; - return GGML_OP_NONE; - } - -public: - void init(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, std::string prefix = "") { - if (prefix.size() > 0) { - prefix = prefix + "."; - } - init_params(ctx, tensor_storage_map, prefix); - init_blocks(ctx, tensor_storage_map, prefix); - } - - size_t get_params_num() { - size_t num_tensors = params.size(); - for (auto& pair : blocks) { - auto& block = pair.second; - - num_tensors += block->get_params_num(); - } - return num_tensors; - }; - - size_t get_params_mem_size() { - size_t mem_size = 0; - for (auto& pair : blocks) { - auto& block = pair.second; - - mem_size += block->get_params_mem_size(); - } - - for (auto& pair : params) { - mem_size += ggml_nbytes(pair.second); - } - - return mem_size; - } - - void get_param_tensors(std::map& tensors, std::string prefix = "") { - if (prefix.size() > 0) { - prefix = prefix + "."; - } - for (auto& pair : blocks) { - auto& block = pair.second; - block->get_param_tensors(tensors, prefix + pair.first); - } - - for (auto& pair : params) { - ggml_tensor* param = pair.second; - tensors[prefix + pair.first] = pair.second; - ggml_set_name(param, (prefix + pair.first).c_str()); - } - } - - void get_param_tensor_ops(std::map& tensor_ops) { - for (auto& pair : blocks) { - pair.second->get_param_tensor_ops(tensor_ops); - } - for (auto& pair : params) { - enum ggml_op op = param_usage_op(pair.first); - if (op != GGML_OP_NONE) { - tensor_ops[pair.second] = op; - } - } - } - - virtual std::string get_desc() { - return "GGMLBlock"; - } - - void get_all_blocks(std::vector& result) { - result.push_back(this); - for (auto& block_iter : blocks) { - if (block_iter.second) { - block_iter.second->get_all_blocks(result); - } - } - } -}; - -class UnaryBlock : public GGMLBlock { -public: - virtual ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) = 0; -}; - -class Identity : public UnaryBlock { -public: - ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { - return x; - } -}; - -class Linear : public UnaryBlock { -protected: - int64_t in_features; - int64_t out_features; - bool bias; - bool force_f32; - bool force_prec_f32; - bool has_weight_scale = false; - bool int8_convrot = false; - int int8_convrot_group_size = 0; - float scale; - std::string prefix; - - void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override { - this->prefix = prefix; - has_weight_scale = false; - int8_convrot = false; - int8_convrot_group_size = 0; - enum ggml_type wtype = get_type(prefix + "weight", tensor_storage_map, GGML_TYPE_F32); - if (in_features % ggml_blck_size(wtype) != 0 || force_f32) { - wtype = GGML_TYPE_F32; - } - params["weight"] = ggml_new_tensor_2d(ctx, wtype, in_features, out_features); - if (bias) { - enum ggml_type wtype = GGML_TYPE_F32; - params["bias"] = ggml_new_tensor_1d(ctx, wtype, out_features); - } - auto weight_storage = tensor_storage_map.find(prefix + "weight"); - const bool is_int8_tensorwise = weight_storage != tensor_storage_map.end() && weight_storage->second.is_int8_tensorwise; - auto weight_scale_storage = tensor_storage_map.find(prefix + "weight_scale"); - if (weight_scale_storage != tensor_storage_map.end()) { - const int64_t scale_nelements = weight_scale_storage->second.nelements(); - GGML_ASSERT(scale_nelements == 1 || scale_nelements == out_features); - params["weight_scale"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, scale_nelements); - has_weight_scale = true; - } - if (is_int8_tensorwise) { - GGML_ASSERT(wtype == GGML_TYPE_I8); - GGML_ASSERT(has_weight_scale); - int8_convrot = weight_storage->second.int8_convrot; - int8_convrot_group_size = weight_storage->second.int8_convrot_group_size; - } - } - -public: - Linear(int64_t in_features, - int64_t out_features, - bool bias = true, - bool force_f32 = false, - bool force_prec_f32 = false, - float scale = 1.f) - : in_features(in_features), - out_features(out_features), - bias(bias), - force_f32(force_f32), - force_prec_f32(force_prec_f32), - scale(scale) {} - - void set_scale(float scale_) { - scale = scale_; - } - - void set_force_prec_f32(bool force_prec_f32_) { - force_prec_f32 = force_prec_f32_; - } - - ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { - ggml_tensor* w = params["weight"]; - ggml_tensor* weight_scale = has_weight_scale ? params["weight_scale"] : nullptr; - #if KCPP_MAINLINE_FP8_SCALED - if (w->type == GGML_TYPE_F8_E4M3 || w->type == GGML_TYPE_F8_E5M2) { - bool supports_fp8_matmul = false; - if (ctx->backend != nullptr) { - ggml_tensor* fp8_matmul = ggml_mul_mat(ctx->ggml_ctx, w, x); - if (force_prec_f32) { - ggml_mul_mat_set_prec(fp8_matmul, GGML_PREC_F32); - } - supports_fp8_matmul = ggml_backend_supports_op(ctx->backend, fp8_matmul); - } - if (!supports_fp8_matmul) { - w = ggml_cast(ctx->ggml_ctx, w, GGML_TYPE_BF16); - } - } - #endif - ggml_tensor* b = nullptr; - if (bias) { - b = params["bias"]; - } - ggml_tensor* linear_bias = has_weight_scale ? nullptr : b; - ggml_tensor* out = nullptr; - #if KCPP_MAINLINE_INT8_CONVROT - if (w->type == GGML_TYPE_I8) { - if (x->type != GGML_TYPE_F32) { - x = ggml_ext_cast_f32(ctx->ggml_ctx, ctx->backend, x); - } - if (!ggml_is_contiguous(x)) { - x = ggml_cont(ctx->ggml_ctx, x); - } - ggml_tensor* lora_input = x; - if (ctx->weight_adapter && b != nullptr) { - b = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, b, prefix + "bias"); - } - if (int8_convrot && scale == 1.f) { - const auto cache_key = std::make_pair(x, int8_convrot_group_size); - auto cached = ctx->int8_convrot_cache.find(cache_key); - if (cached == ctx->int8_convrot_cache.end()) { - x = ggml_quantize_i8_convrot(ctx->ggml_ctx, x, int8_convrot_group_size); - ctx->int8_convrot_cache.emplace(cache_key, x); - } else { - x = cached->second; - } - } - out = ggml_ext_linear_i8_tensorwise(ctx->ggml_ctx, - x, - w, - weight_scale, - b, - int8_convrot ? int8_convrot_group_size : 0, - scale); - if (ctx->weight_adapter) { - WeightAdapter::ForwardParams forward_params; - forward_params.op_type = WeightAdapter::ForwardParams::op_type_t::OP_LINEAR; - forward_params.linear.force_prec_f32 = force_prec_f32; - forward_params.linear.scale = scale; - out = ctx->weight_adapter->add_lora_to_output(ctx->ggml_ctx, - ctx->backend, - lora_input, - w, - out, - prefix, - forward_params); - } - return out; - } - #endif // kcpp - if (has_weight_scale) { - out = ggml_ext_linear(ctx->ggml_ctx, x, w, nullptr, force_prec_f32, scale); - out = ggml_mul(ctx->ggml_ctx, out, weight_scale); - if (ctx->weight_adapter) { - WeightAdapter::ForwardParams forward_params; - forward_params.op_type = WeightAdapter::ForwardParams::op_type_t::OP_LINEAR; - forward_params.linear.force_prec_f32 = force_prec_f32; - forward_params.linear.scale = scale; - out = ctx->weight_adapter->add_lora_to_output(ctx->ggml_ctx, - ctx->backend, - x, - w, - out, - prefix, - forward_params); - if (b != nullptr) { - b = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, b, prefix + "bias"); - } - } - if (b != nullptr) { - out = ggml_add_inplace(ctx->ggml_ctx, out, b); - } - return out; - } - if (ctx->weight_adapter) { - WeightAdapter::ForwardParams forward_params; - forward_params.op_type = WeightAdapter::ForwardParams::op_type_t::OP_LINEAR; - forward_params.linear.force_prec_f32 = force_prec_f32; - forward_params.linear.scale = scale; - out = ctx->weight_adapter->forward_with_lora(ctx->ggml_ctx, ctx->backend, x, w, linear_bias, prefix, forward_params); - } else { - out = ggml_ext_linear(ctx->ggml_ctx, x, w, linear_bias, force_prec_f32, scale); - } - return out; - } -}; - -__STATIC_INLINE__ bool support_get_rows(ggml_type wtype) { - std::set allow_types = {GGML_TYPE_F16, GGML_TYPE_Q8_0, GGML_TYPE_Q5_1, GGML_TYPE_Q5_0, GGML_TYPE_Q4_1, GGML_TYPE_Q4_0}; - if (allow_types.find(wtype) != allow_types.end()) { - return true; - } - return false; -} - -class Embedding : public UnaryBlock { -protected: - int64_t embedding_dim; - int64_t num_embeddings; - void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map, const std::string prefix = "") override { - enum ggml_type wtype = get_type(prefix + "weight", tensor_storage_map, GGML_TYPE_F32); - if (!support_get_rows(wtype)) { - wtype = GGML_TYPE_F32; - } - params["weight"] = ggml_new_tensor_2d(ctx, wtype, embedding_dim, num_embeddings); - } - - enum ggml_op param_usage_op(const std::string& name) const override { - return name == "weight" ? GGML_OP_GET_ROWS : GGML_OP_NONE; - } - -public: - Embedding(int64_t num_embeddings, int64_t embedding_dim) - : embedding_dim(embedding_dim), - num_embeddings(num_embeddings) { - } - - ggml_tensor* forward(GGMLRunnerContext* ctx, - ggml_tensor* input_ids) override { - // input_ids: [N, n_token] - auto weight = params["weight"]; - - // There are issues with ggml batch inference, so we are expanding it here first. - // TODO: fix ggml batch inference - int64_t n = input_ids->ne[1]; - input_ids = ggml_reshape_1d(ctx->ggml_ctx, input_ids, input_ids->ne[0] * input_ids->ne[1]); - - input_ids = ggml_reshape_3d(ctx->ggml_ctx, input_ids, input_ids->ne[0], 1, input_ids->ne[1]); - auto embedding = ggml_get_rows(ctx->ggml_ctx, weight, input_ids); - embedding = ggml_reshape_3d(ctx->ggml_ctx, embedding, embedding->ne[0], embedding->ne[1] / n, n); - - // [N, n_token, embedding_dim] - return embedding; - } -}; - -class Conv2d : public UnaryBlock { -protected: - int64_t in_channels; - int64_t out_channels; - std::pair kernel_size; - std::pair stride; - std::pair padding; - std::pair dilation; - bool bias; - float scale = 1.f; - std::string prefix; - - void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map, const std::string prefix = "") override { - this->prefix = prefix; - enum ggml_type wtype = GGML_TYPE_F16; - params["weight"] = ggml_new_tensor_4d(ctx, wtype, kernel_size.second, kernel_size.first, in_channels, out_channels); - if (bias) { - enum ggml_type wtype = GGML_TYPE_F32; - params["bias"] = ggml_new_tensor_1d(ctx, wtype, out_channels); - } - } - -public: - Conv2d(int64_t in_channels, - int64_t out_channels, - std::pair kernel_size, - std::pair stride = {1, 1}, - std::pair padding = {0, 0}, - std::pair dilation = {1, 1}, - bool bias = true) - : in_channels(in_channels), - out_channels(out_channels), - kernel_size(kernel_size), - stride(stride), - padding(padding), - dilation(dilation), - bias(bias) {} - - void set_scale(float scale_value) { - scale = scale_value; - } - - std::string get_desc() override { - return "Conv2d"; - } - - ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { - ggml_tensor* w = params["weight"]; - ggml_tensor* b = nullptr; - if (bias) { - b = params["bias"]; - } - if (ctx->weight_adapter) { - WeightAdapter::ForwardParams forward_params; - forward_params.op_type = WeightAdapter::ForwardParams::op_type_t::OP_CONV2D; - forward_params.conv2d.s0 = stride.second; - forward_params.conv2d.s1 = stride.first; - forward_params.conv2d.p0 = padding.second; - forward_params.conv2d.p1 = padding.first; - forward_params.conv2d.d0 = dilation.second; - forward_params.conv2d.d1 = dilation.first; - forward_params.conv2d.direct = ctx->conv2d_direct_enabled; - forward_params.conv2d.circular_x = ctx->circular_x_enabled; - forward_params.conv2d.circular_y = ctx->circular_y_enabled; - forward_params.conv2d.scale = scale; - return ctx->weight_adapter->forward_with_lora(ctx->ggml_ctx, ctx->backend, x, w, b, prefix, forward_params); - } - return ggml_ext_conv_2d(ctx->ggml_ctx, - x, - w, - b, - stride.second, - stride.first, - padding.second, - padding.first, - dilation.second, - dilation.first, - ctx->conv2d_direct_enabled, - ctx->circular_x_enabled, - ctx->circular_y_enabled, - scale); - } -}; - -class Conv2d_grouped : public UnaryBlock { -protected: - int64_t in_channels; - int64_t out_channels; - int groups; - std::pair kernel_size; - std::pair stride; - std::pair padding; - std::pair dilation; - bool bias; - float scale = 1.f; - std::string prefix; - - void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map, const std::string prefix = "") override { - this->prefix = prefix; - enum ggml_type wtype = GGML_TYPE_F16; - params["weight"] = ggml_new_tensor_4d(ctx, wtype, kernel_size.second, kernel_size.first, in_channels / groups, out_channels); - if (bias) { - enum ggml_type wtype = GGML_TYPE_F32; - params["bias"] = ggml_new_tensor_1d(ctx, wtype, out_channels); - } - } - -public: - Conv2d_grouped(int64_t in_channels, - int64_t out_channels, - int groups, - std::pair kernel_size, - std::pair stride = {1, 1}, - std::pair padding = {0, 0}, - std::pair dilation = {1, 1}, - bool bias = true) - : in_channels(in_channels), - out_channels(out_channels), - groups(groups), - kernel_size(kernel_size), - stride(stride), - padding(padding), - dilation(dilation), - bias(bias) {} - - void set_scale(float scale_value) { - scale = scale_value; - } - - std::string get_desc() override { - return "Conv2d_grouped"; - } - - ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { - ggml_tensor* w = params["weight"]; - ggml_tensor* b = nullptr; - if (bias) { - b = params["bias"]; - } - - if (groups == 1) { - if (ctx->weight_adapter) { - WeightAdapter::ForwardParams forward_params; - forward_params.op_type = WeightAdapter::ForwardParams::op_type_t::OP_CONV2D; - forward_params.conv2d.s0 = stride.second; - forward_params.conv2d.s1 = stride.first; - forward_params.conv2d.p0 = padding.second; - forward_params.conv2d.p1 = padding.first; - forward_params.conv2d.d0 = dilation.second; - forward_params.conv2d.d1 = dilation.first; - forward_params.conv2d.direct = ctx->conv2d_direct_enabled; - forward_params.conv2d.circular_x = ctx->circular_x_enabled; - forward_params.conv2d.circular_y = ctx->circular_y_enabled; - forward_params.conv2d.scale = scale; - return ctx->weight_adapter->forward_with_lora(ctx->ggml_ctx, ctx->backend, x, w, b, prefix, forward_params); - } - return ggml_ext_conv_2d(ctx->ggml_ctx, x, w, b, - stride.second, stride.first, - padding.second, padding.first, - dilation.second, dilation.first, - ctx->conv2d_direct_enabled, - ctx->circular_x_enabled, - ctx->circular_y_enabled, - scale); - } - - if (groups == in_channels && groups == out_channels) { - ggml_tensor* res; - if (ctx->conv2d_direct_enabled) { - res = ggml_conv_2d_dw_direct(ctx->ggml_ctx, w, x, - stride.second, stride.first, - padding.second, padding.first, - dilation.second, dilation.first); - } else { - res = ggml_conv_2d_dw(ctx->ggml_ctx, w, x, - stride.second, stride.first, - padding.second, padding.first, - dilation.second, dilation.first); - } - if (b) { - b = ggml_reshape_4d(ctx->ggml_ctx, b, 1, 1, b->ne[0], 1); - res = ggml_add_inplace(ctx->ggml_ctx, res, b); - } - return res; - } - - int64_t ic_g = in_channels / groups; - int64_t oc_g = out_channels / groups; - - std::vector out_slices(groups); - - for (int i = 0; i < groups; ++i) { - size_t x_offset = i * ic_g * x->nb[2]; - ggml_tensor* x_i = ggml_view_4d(ctx->ggml_ctx, x, - x->ne[0], x->ne[1], ic_g, x->ne[3], - x->nb[1], x->nb[2], x->nb[3], - x_offset); - - size_t w_offset = i * oc_g * w->nb[3]; - ggml_tensor* w_i = ggml_view_4d(ctx->ggml_ctx, w, - w->ne[0], w->ne[1], w->ne[2], oc_g, - w->nb[1], w->nb[2], w->nb[3], - w_offset); - - ggml_tensor* b_i = nullptr; - if (b) { - size_t b_offset = i * oc_g * b->nb[0]; - b_i = ggml_view_1d(ctx->ggml_ctx, b, oc_g, b_offset); - } - - if (ctx->weight_adapter) { - WeightAdapter::ForwardParams forward_params; - forward_params.op_type = WeightAdapter::ForwardParams::op_type_t::OP_CONV2D; - forward_params.conv2d.s0 = stride.second; - forward_params.conv2d.s1 = stride.first; - forward_params.conv2d.p0 = padding.second; - forward_params.conv2d.p1 = padding.first; - forward_params.conv2d.d0 = dilation.second; - forward_params.conv2d.d1 = dilation.first; - forward_params.conv2d.direct = ctx->conv2d_direct_enabled; - forward_params.conv2d.circular_x = ctx->circular_x_enabled; - forward_params.conv2d.circular_y = ctx->circular_y_enabled; - forward_params.conv2d.scale = scale; - out_slices[i] = ctx->weight_adapter->forward_with_lora(ctx->ggml_ctx, ctx->backend, x_i, w_i, b_i, prefix, forward_params); - } else { - out_slices[i] = ggml_ext_conv_2d(ctx->ggml_ctx, x_i, w_i, b_i, - stride.second, stride.first, - padding.second, padding.first, - dilation.second, dilation.first, - ctx->conv2d_direct_enabled, - ctx->circular_x_enabled, - ctx->circular_y_enabled, - scale); - } - } - - ggml_tensor* out = ggml_ext_vec_concat(ctx->ggml_ctx, out_slices, 2); - - return out; - } -}; - -class Conv3d : public UnaryBlock { -protected: - int64_t in_channels; - int64_t out_channels; - std::tuple kernel_size; - std::tuple stride; - std::tuple padding; - std::tuple dilation; - bool bias; - bool force_prec_f32; - std::string prefix; - - void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map, const std::string prefix = "") override { - this->prefix = prefix; - enum ggml_type wtype = GGML_TYPE_F16; - params["weight"] = ggml_new_tensor_4d(ctx, - wtype, - std::get<2>(kernel_size), - std::get<1>(kernel_size), - std::get<0>(kernel_size), - in_channels * out_channels); - if (bias) { - params["bias"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, out_channels); - } - } - -public: - Conv3d(int64_t in_channels, - int64_t out_channels, - std::tuple kernel_size, - std::tuple stride = {1, 1, 1}, - std::tuple padding = {0, 0, 0}, - std::tuple dilation = {1, 1, 1}, - bool bias = true, - bool force_prec_f32 = false) - : in_channels(in_channels), - out_channels(out_channels), - kernel_size(kernel_size), - stride(stride), - padding(padding), - dilation(dilation), - bias(bias), - force_prec_f32(force_prec_f32) {} - - ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { - ggml_tensor* w = params["weight"]; - ggml_tensor* b = nullptr; - if (ctx->weight_adapter) { - w = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, w, prefix + "weight"); - if (w->type != GGML_TYPE_F16) { - w = ggml_cast(ctx->ggml_ctx, w, GGML_TYPE_F16); - } - } - if (bias) { - b = params["bias"]; - if (ctx->weight_adapter) { - b = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, b, prefix + "bias"); - } - } - return ggml_ext_conv_3d(ctx->ggml_ctx, ctx->backend, x, w, b, in_channels, - std::get<2>(stride), std::get<1>(stride), std::get<0>(stride), - std::get<2>(padding), std::get<1>(padding), std::get<0>(padding), - std::get<2>(dilation), std::get<1>(dilation), std::get<0>(dilation), - force_prec_f32); - } -}; - -class LayerNorm : public UnaryBlock { -protected: - int64_t normalized_shape; - float eps; - bool elementwise_affine; - bool bias; - std::string prefix; - - void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override { - this->prefix = prefix; - if (elementwise_affine) { - enum ggml_type wtype = GGML_TYPE_F32; - params["weight"] = ggml_new_tensor_1d(ctx, wtype, normalized_shape); - if (bias) { - enum ggml_type wtype = GGML_TYPE_F32; - params["bias"] = ggml_new_tensor_1d(ctx, wtype, normalized_shape); - } - } - } - -public: - LayerNorm(int64_t normalized_shape, - float eps = 1e-05f, - bool elementwise_affine = true, - bool bias = true) - : normalized_shape(normalized_shape), - eps(eps), - elementwise_affine(elementwise_affine), - bias(bias) {} - - ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { - ggml_tensor* w = nullptr; - ggml_tensor* b = nullptr; - - if (elementwise_affine) { - w = params["weight"]; - if (ctx->weight_adapter) { - w = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, w, prefix + "weight"); - } - if (bias) { - b = params["bias"]; - if (ctx->weight_adapter) { - b = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, b, prefix + "bias"); - } - } - } - return ggml_ext_layer_norm(ctx->ggml_ctx, x, w, b, eps); - } -}; - -class GroupNorm : public GGMLBlock { -protected: - int num_groups; - int64_t num_channels; - float eps; - bool affine; - std::string prefix; - - void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override { - this->prefix = prefix; - if (affine) { - enum ggml_type wtype = GGML_TYPE_F32; - enum ggml_type bias_wtype = GGML_TYPE_F32; - params["weight"] = ggml_new_tensor_1d(ctx, wtype, num_channels); - params["bias"] = ggml_new_tensor_1d(ctx, bias_wtype, num_channels); - } - } - -public: - GroupNorm(int num_groups, - int64_t num_channels, - float eps = 1e-05f, - bool affine = true) - : num_groups(num_groups), - num_channels(num_channels), - eps(eps), - affine(affine) {} - - ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { - ggml_tensor* w = nullptr; - ggml_tensor* b = nullptr; - if (affine) { - w = params["weight"]; - b = params["bias"]; - if (ctx->weight_adapter) { - w = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, w, prefix + "weight"); - b = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, b, prefix + "bias"); - } - } - return ggml_ext_group_norm(ctx->ggml_ctx, x, w, b, num_groups); - } -}; - -class GroupNorm32 : public GroupNorm { -public: - GroupNorm32(int64_t num_channels) - : GroupNorm(32, num_channels, 1e-06f) {} -}; - -class RMSNorm : public UnaryBlock { -protected: - int64_t hidden_size; - float eps; - std::string prefix; - - void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, std::string prefix = "") override { - this->prefix = prefix; - enum ggml_type wtype = GGML_TYPE_F32; - params["weight"] = ggml_new_tensor_1d(ctx, wtype, hidden_size); - } - -public: - RMSNorm(int64_t hidden_size, - float eps = 1e-06f) - : hidden_size(hidden_size), - eps(eps) {} - - ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { - ggml_tensor* w = params["weight"]; - if (ctx->weight_adapter) { - w = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, w, prefix + "weight"); - } - x = ggml_rms_norm(ctx->ggml_ctx, x, eps); - x = ggml_mul_inplace(ctx->ggml_ctx, x, w); - return x; - } -}; - -class MultiheadAttention : public GGMLBlock { -protected: - int64_t embed_dim; - int64_t n_head; - bool proj_in; - std::string q_proj_name; - std::string k_proj_name; - std::string v_proj_name; - std::string in_proj_name; - std::string out_proj_name; - -public: - MultiheadAttention(int64_t embed_dim, - int64_t n_head, - bool qkv_proj_bias = true, - bool out_proj_bias = true, - bool proj_in = false, - std::string q_proj_name = "q_proj", - std::string k_proj_name = "k_proj", - std::string v_proj_name = "v_proj", - std::string in_proj_name = "in_proj", - std::string out_proj_name = "out_proj") - : embed_dim(embed_dim), - n_head(n_head), - proj_in(proj_in), - q_proj_name(q_proj_name), - k_proj_name(k_proj_name), - v_proj_name(v_proj_name), - in_proj_name(in_proj_name), - out_proj_name(out_proj_name) { - if (proj_in) { - blocks[in_proj_name] = std::shared_ptr(new Linear(embed_dim, embed_dim * 3, qkv_proj_bias)); - } else { - blocks[q_proj_name] = std::shared_ptr(new Linear(embed_dim, embed_dim, qkv_proj_bias)); - blocks[k_proj_name] = std::shared_ptr(new Linear(embed_dim, embed_dim, qkv_proj_bias)); - blocks[v_proj_name] = std::shared_ptr(new Linear(embed_dim, embed_dim, qkv_proj_bias)); - } - blocks[out_proj_name] = std::shared_ptr(new Linear(embed_dim, embed_dim, out_proj_bias)); - } - - // x: [N, n_token, embed_dim] - ggml_tensor* forward(GGMLRunnerContext* ctx, - ggml_tensor* x, - ggml_tensor* mask = nullptr) { - auto out_proj = std::dynamic_pointer_cast(blocks[out_proj_name]); - - ggml_tensor* q; - ggml_tensor* k; - ggml_tensor* v; - if (proj_in) { - auto in_proj = std::dynamic_pointer_cast(blocks[in_proj_name]); - auto qkv = in_proj->forward(ctx, x); - auto qkv_vec = split_qkv(ctx->ggml_ctx, qkv); - q = qkv_vec[0]; - k = qkv_vec[1]; - v = qkv_vec[2]; - } else { - auto q_proj = std::dynamic_pointer_cast(blocks[q_proj_name]); - auto k_proj = std::dynamic_pointer_cast(blocks[k_proj_name]); - auto v_proj = std::dynamic_pointer_cast(blocks[v_proj_name]); - - q = q_proj->forward(ctx, x); - k = k_proj->forward(ctx, x); - v = v_proj->forward(ctx, x); - } - - x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, n_head, mask, false); // [N, n_token, embed_dim] - - x = out_proj->forward(ctx, x); // [N, n_token, embed_dim] - return x; - } -}; - -__STATIC_INLINE__ ggml_tensor* ggml_ext_lokr_forward( - ggml_context* ctx, - ggml_backend_t backend, - ggml_tensor* h, // Input: [q, batch] or [W, H, q, batch] - ggml_tensor* w1, // Outer C (Full rank) - ggml_tensor* w1a, // Outer A (Low rank part 1) - ggml_tensor* w1b, // Outer B (Low rank part 2) - ggml_tensor* w2, // Inner BA (Full rank) - ggml_tensor* w2a, // Inner A (Low rank part 1) - ggml_tensor* w2b, // Inner B (Low rank part 2) - bool is_conv, - WeightAdapter::ForwardParams::conv2d_params_t conv_params, - float scale) { - GGML_ASSERT((w1 != nullptr || (w1a != nullptr && w1b != nullptr))); - GGML_ASSERT((w2 != nullptr || (w2a != nullptr && w2b != nullptr))); - - int uq = (w1 != nullptr) ? (int)w1->ne[0] : (int)w1a->ne[0]; - int up = (w1 != nullptr) ? (int)w1->ne[1] : (int)w1b->ne[1]; - - int q_actual = is_conv ? (int)h->ne[2] : (int)h->ne[0]; - int vq = q_actual / uq; - - int vp = (w2 != nullptr) ? (is_conv ? (int)w2->ne[3] : (int)w2->ne[1]) - : (int)w2a->ne[1]; - GGML_ASSERT(q_actual == (uq * vq) && "Input dimension mismatch for LoKR split"); - - ggml_tensor* hb; - - if (!is_conv) { - int batch = (int)h->ne[1]; - int merge_batch_uq = batch; - int merge_batch_vp = batch; - - if (sd_backend_is(backend, "Vulkan")) { - if (batch > 1) { - // no access to backend here, worst case is slightly worse perfs for other backends when built alongside Vulkan backend - int max_batch = 65535; - int max_batch_uq = max_batch / uq; - merge_batch_uq = 1; - for (int i = max_batch_uq; i > 0; i--) { - if (batch % i == 0) { - merge_batch_uq = i; - break; - } - } - - int max_batch_vp = max_batch / vp; - merge_batch_vp = 1; - for (int i = max_batch_vp; i > 0; i--) { - if (batch % i == 0) { - merge_batch_vp = i; - break; - } - } - } - } - - ggml_tensor* h_split = ggml_reshape_3d(ctx, h, vq, uq * merge_batch_uq, batch / merge_batch_uq); - if (w2 != nullptr) { - hb = ggml_mul_mat(ctx, w2, h_split); - } else { - hb = ggml_mul_mat(ctx, w2b, ggml_mul_mat(ctx, w2a, h_split)); - } - - if (batch > 1) { - hb = ggml_reshape_3d(ctx, hb, vp, uq, batch); - } - ggml_tensor* hb_t = ggml_cont(ctx, ggml_transpose(ctx, hb)); - hb_t = ggml_reshape_3d(ctx, hb_t, uq, vp * merge_batch_vp, batch / merge_batch_vp); - - ggml_tensor* hc_t; - if (w1 != nullptr) { - hc_t = ggml_mul_mat(ctx, w1, hb_t); - } else { - hc_t = ggml_mul_mat(ctx, w1b, ggml_mul_mat(ctx, w1a, hb_t)); - } - - if (batch > 1) { - hc_t = ggml_reshape_3d(ctx, hc_t, up, vp, batch); - } - - ggml_tensor* hc = ggml_transpose(ctx, hc_t); - ggml_tensor* out = ggml_reshape_2d(ctx, ggml_cont(ctx, hc), up * vp, batch); - return ggml_ext_scale(ctx, out, scale); - } else { - int batch = (int)h->ne[3]; - // 1. Reshape input: [W, H, vq*uq, batch] -> [W, H, vq, uq * batch] - ggml_tensor* h_split = ggml_reshape_4d(ctx, h, h->ne[0], h->ne[1], vq, uq * batch); - - if (w2 != nullptr) { - hb = ggml_ext_conv_2d(ctx, h_split, w2, nullptr, - conv_params.s0, - conv_params.s1, - conv_params.p0, - conv_params.p1, - conv_params.d0, - conv_params.d1, - conv_params.direct, - conv_params.circular_x, - conv_params.circular_y, - conv_params.scale); - } else { - // swap a and b order for conv lora - ggml_tensor* a = w2b; - ggml_tensor* b = w2a; - - // unpack conv2d weights if needed - if (ggml_n_dims(a) < 4) { - int k = (int)sqrt(a->ne[0] / h_split->ne[2]); - GGML_ASSERT(k * k * h_split->ne[2] == a->ne[0]); - a = ggml_reshape_4d(ctx, a, k, k, a->ne[0] / (k * k), a->ne[1]); - } else if (a->ne[2] != h_split->ne[2]) { - int k = (int)sqrt(a->ne[2] / h_split->ne[2]); - GGML_ASSERT(k * k * h_split->ne[2] == a->ne[2]); - a = ggml_reshape_4d(ctx, a, a->ne[0] * k, a->ne[1] * k, a->ne[2] / (k * k), a->ne[3]); - } - ggml_tensor* ha = ggml_ext_conv_2d(ctx, h_split, a, nullptr, - conv_params.s0, - conv_params.s1, - conv_params.p0, - conv_params.p1, - conv_params.d0, - conv_params.d1, - conv_params.direct, - conv_params.circular_x, - conv_params.circular_y, - conv_params.scale); - - // not supporting lora_mid here - hb = ggml_ext_conv_2d(ctx, - ha, - b, - nullptr, - 1, - 1, - 0, - 0, - 1, - 1, - conv_params.direct, - conv_params.circular_x, - conv_params.circular_y, - conv_params.scale); - } - - // Current hb shape: [W_out, H_out, vp, uq * batch] - int w_out = (int)hb->ne[0]; - int h_out = (int)hb->ne[1]; - - // ggml_tensor* hb_cat = ggml_reshape_4d(ctx, hb, w_out , h_out , vp * uq, batch); - // [W_out, H_out, vp * uq, batch] - // Now left to compute (W1 kr Id) * hb_cat == (W1 kr W2) cv h - - // merge the uq groups of size vp*w_out*h_out - ggml_tensor* hb_merged = ggml_reshape_2d(ctx, hb, w_out * h_out * vp, uq * batch); - ggml_tensor* hc_t; - ggml_tensor* hb_merged_t = ggml_cont(ctx, ggml_transpose(ctx, hb_merged)); - if (w1 != nullptr) { - // Would be great to be able to transpose w1 instead to avoid transposing both hb and hc - hc_t = ggml_mul_mat(ctx, w1, hb_merged_t); - } else { - hc_t = ggml_mul_mat(ctx, w1b, ggml_mul_mat(ctx, w1a, hb_merged_t)); - } - ggml_tensor* hc = ggml_transpose(ctx, hc_t); - // ungroup - ggml_tensor* out = ggml_reshape_4d(ctx, ggml_cont(ctx, hc), w_out, h_out, up * vp, batch); - return ggml_ext_scale(ctx, out, scale); - } -} - -#endif // __SD_CORE_GGML_EXTEND_HPP__ diff --git a/otherarch/sdcpp/src/core/ggml_extend_backend.cpp b/otherarch/sdcpp/src/core/ggml_extend_backend.cpp index a83166438..9c4cfa42a 100644 --- a/otherarch/sdcpp/src/core/ggml_extend_backend.cpp +++ b/otherarch/sdcpp/src/core/ggml_extend_backend.cpp @@ -392,7 +392,7 @@ static bool backend_name_exists(const std::string& name) { static ggml_backend_t init_named_backend(const std::string& name) { ggml_backend_load_all_once(); - LOG_DEBUG("Initializing backend: %s", name.c_str()); + LOG_VERBOSE("Initializing backend: %s", name.c_str()); if (trim_copy(name).empty()) { return ggml_backend_init_best(); } @@ -542,10 +542,10 @@ static ggml_backend_t sd_get_default_backend() { if (dev_count == 0) { LOG_ERROR("No devices found!"); } else { - LOG_DEBUG("Found %zu backend devices:", dev_count); + LOG_VERBOSE("Found %zu backend devices:", dev_count); for (size_t i = 0; i < dev_count; ++i) { auto dev = ggml_backend_dev_get(i); - LOG_DEBUG("#%zu: %s", i, ggml_backend_dev_name(dev)); + LOG_VERBOSE("#%zu: %s", i, ggml_backend_dev_name(dev)); } } }); @@ -587,7 +587,7 @@ static ggml_backend_t sd_get_default_backend() { } if (sd_backend_is_cpu(backend)) { - LOG_DEBUG("Using CPU backend"); + LOG_VERBOSE("Using CPU backend"); } return backend; @@ -965,3 +965,35 @@ const char* sd_backend_module_name(SDBackendModule module) { } return "unknown"; } + +void ggml_ext_backend_tensor_get_and_sync(ggml_backend_t backend, const ggml_tensor* tensor, void* data, size_t offset, size_t size) { + if ((sd_backend_is(backend, "ROCm") || sd_backend_is(backend, "CUDA") || sd_backend_is(backend, "SYCL")) && + !sd_backend_is_cpu(backend)) { + ggml_backend_tensor_get_async(backend, tensor, data, offset, size); + ggml_backend_synchronize(backend); + return; + } + + ggml_backend_tensor_get(tensor, data, offset, size); +} + +float ggml_ext_backend_tensor_get_f32(ggml_tensor* tensor) { + GGML_ASSERT(tensor->type == GGML_TYPE_F32 || tensor->type == GGML_TYPE_F16 || tensor->type == GGML_TYPE_I32 || tensor->type == GGML_TYPE_BF16); + float value; + if (tensor->type == GGML_TYPE_F32) { + ggml_backend_tensor_get(tensor, &value, 0, sizeof(value)); + } else if (tensor->type == GGML_TYPE_BF16) { + ggml_bf16_t bf16_value; + ggml_backend_tensor_get(tensor, &bf16_value, 0, sizeof(bf16_value)); + value = ggml_bf16_to_fp32(bf16_value); + } else if (tensor->type == GGML_TYPE_F16) { + ggml_fp16_t f16_value; + ggml_backend_tensor_get(tensor, &f16_value, 0, sizeof(f16_value)); + value = ggml_fp16_to_fp32(f16_value); + } else { // GGML_TYPE_I32 + int int32_value; + ggml_backend_tensor_get(tensor, &int32_value, 0, sizeof(int32_value)); + value = (float)int32_value; + } + return value; +} diff --git a/otherarch/sdcpp/src/core/ggml_extend_backend.h b/otherarch/sdcpp/src/core/ggml_extend_backend.h index d5498e8f4..01652cbfb 100644 --- a/otherarch/sdcpp/src/core/ggml_extend_backend.h +++ b/otherarch/sdcpp/src/core/ggml_extend_backend.h @@ -96,4 +96,6 @@ std::string sd_backend_resolve_name(const std::string& name); const char* sd_backend_module_name(SDBackendModule module); void ggml_ext_im_set_f32_1d(const struct ggml_tensor* tensor, int i, float value); bool add_rpc_devices(const std::string& servers); +void ggml_ext_backend_tensor_get_and_sync(ggml_backend_t backend, const ggml_tensor* tensor, void* data, size_t offset, size_t size); +float ggml_ext_backend_tensor_get_f32(ggml_tensor* tensor); #endif // __SD_CORE_GGML_EXTEND_BACKEND_H__ diff --git a/otherarch/sdcpp/src/core/ggml_graph_cut.cpp b/otherarch/sdcpp/src/core/ggml_graph_cut.cpp index 542b7fe16..797100587 100644 --- a/otherarch/sdcpp/src/core/ggml_graph_cut.cpp +++ b/otherarch/sdcpp/src/core/ggml_graph_cut.cpp @@ -2,6 +2,7 @@ #include #include +#include #include #include #include @@ -67,25 +68,6 @@ namespace sd::ggml_graph_cut { return -1; } - static Plan::InputShape input_shape(const ggml_tensor* tensor) { - Plan::InputShape shape; - if (tensor == nullptr) { - return shape; - } - shape.type = tensor->type; - for (int i = 0; i < GGML_MAX_DIMS; ++i) { - shape.ne[static_cast(i)] = tensor->ne[i]; - } - return shape; - } - - static size_t graph_cut_segment_vram_bytes(const Segment& segment) { - return segment.compute_buffer_size + - segment.input_param_bytes + - segment.input_previous_cut_bytes + - segment.output_bytes; - } - static std::string lower_ascii_copy(std::string value) { std::transform(value.begin(), value.end(), value.begin(), [](unsigned char c) { return static_cast(std::tolower(c)); @@ -291,55 +273,6 @@ 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) { - GGML_ASSERT(start_segment_index < plan.segments.size()); - GGML_ASSERT(end_segment_index < plan.segments.size()); - GGML_ASSERT(start_segment_index <= end_segment_index); - - Segment seed; - const auto& start_segment = plan.segments[start_segment_index]; - const auto& target_segment = plan.segments[end_segment_index]; - std::unordered_set 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 ((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); - } - } - } - if (start_segment_index == end_segment_index) { - seed.group_name = target_segment.group_name; - } else { - seed.group_name = sd_format("%s..%s", - start_segment.group_name.c_str(), - target_segment.group_name.c_str()); - } - return seed; - } - static void build_segment(ggml_cgraph* gf, Plan& plan, Segment& segment, @@ -416,31 +349,7 @@ namespace sd::ggml_graph_cut { } return a.display_name < b.display_name; }); - segment.input_refs = input_refs; - for (const auto& input : input_refs) { - ggml_tensor* current_input = input_tensor(gf, input); - size_t tensor_bytes = current_input == nullptr - ? 0 - : (input.type == Segment::INPUT_PREVIOUS_CUT - ? cache_tensor_bytes(current_input) - : ggml_nbytes(current_input)); - switch (input.type) { - case Segment::INPUT_PREVIOUS_CUT: - segment.input_previous_cut_bytes += tensor_bytes; - break; - case Segment::INPUT_PARAM: - segment.input_param_bytes += tensor_bytes; - break; - case Segment::INPUT_EXTERNAL: - default: - segment.input_external_bytes += tensor_bytes; - break; - } - } - for (int output_node_index : segment.output_node_indices) { - ggml_tensor* output = ggml_graph_node(gf, output_node_index); - segment.output_bytes += cache_tensor_bytes(output); - } + segment.input_refs = input_refs; segment.compute_buffer_size = measure_segment_compute_buffer(backend, gf, segment, log_desc); for (int output_node_index : segment.output_node_indices) { @@ -449,6 +358,70 @@ namespace sd::ggml_graph_cut { plan.segments.push_back(std::move(segment)); } + static bool validate_plan(ggml_cgraph* gf, + const Plan& plan, + std::string* validation_error) { + auto fail = [&](const std::string& reason) { + if (validation_error != nullptr) { + *validation_error = reason; + } + return false; + }; + if (!plan.has_cuts) { + return true; + } + if (plan.segments.size() <= 1) { + return fail("fewer than two segments"); + } + const int n_nodes = ggml_graph_n_nodes(gf); + std::unordered_set completed_outputs; + for (size_t segment_index = 0; segment_index < plan.segments.size(); ++segment_index) { + const Segment& segment = plan.segments[segment_index]; + const std::string segment_label = "segment " + std::to_string(segment_index) + + " ('" + segment.group_name + "')"; + if (segment.internal_node_indices.empty() || segment.output_node_indices.empty()) { + return fail(segment_label + " has no internal nodes or outputs"); + } + for (const Segment::InputRef& input : segment.input_refs) { + if (input.type == Segment::INPUT_PREVIOUS_CUT) { + if (input.node_index < 0 || input.node_index >= n_nodes || + completed_outputs.find(input.node_index) == completed_outputs.end()) { + return fail(segment_label + " references an unavailable cut node " + + std::to_string(input.node_index)); + } + } else if (input.leaf_index < 0 || input.leaf_index >= gf->n_leafs) { + return fail(segment_label + " references an invalid leaf " + + std::to_string(input.leaf_index)); + } + } + std::unordered_set segment_nodes; + segment_nodes.reserve(segment.internal_node_indices.size()); + for (int node_index : segment.internal_node_indices) { + if (node_index < 0 || node_index >= n_nodes) { + return fail(segment_label + " contains an invalid node " + + std::to_string(node_index)); + } + if (!segment_nodes.insert(node_index).second) { + return fail(segment_label + " contains duplicate node " + + std::to_string(node_index)); + } + } + for (int output_index : segment.output_node_indices) { + if (output_index < 0 || output_index >= n_nodes || + segment_nodes.find(output_index) == segment_nodes.end()) { + return fail(segment_label + " has an output outside its node set: " + + std::to_string(output_index)); + } + if (completed_outputs.find(output_index) != completed_outputs.end()) { + return fail(segment_label + " repeats output node " + + std::to_string(output_index)); + } + completed_outputs.insert(output_index); + } + } + return true; + } + bool is_graph_cut_tensor(const ggml_tensor* tensor) { if (tensor == nullptr || tensor->name[0] == '\0') { return false; @@ -509,26 +482,87 @@ namespace sd::ggml_graph_cut { return ggml_nbytes(cache_src); } - bool plan_matches_graph(ggml_cgraph* gf, const Plan& plan) { - GGML_ASSERT(gf != nullptr); - if (ggml_graph_n_nodes(gf) != plan.n_nodes || gf->n_leafs != plan.n_leafs) { - return false; - } - for (const auto& input_shape_ref : plan.input_shapes) { - if (input_shape_ref.leaf_index < 0 || input_shape_ref.leaf_index >= gf->n_leafs) { - return false; + std::vector graph_layout(ggml_cgraph* graph, bool include_bindings) { + std::vector tensors; + std::unordered_map indices; + auto add = [&](const ggml_tensor* tensor) { + if (tensor != nullptr && indices.emplace(tensor, tensors.size() + 1).second) { + tensors.push_back(tensor); } - ggml_tensor* leaf = gf->leafs[input_shape_ref.leaf_index]; - if (leaf == nullptr || input_shape_ref.type != leaf->type) { - return false; + }; + for (int i = 0; i < graph->n_leafs; ++i) { + add(graph->leafs[i]); + } + for (int i = 0; i < graph->n_nodes; ++i) { + add(graph->nodes[i]); + } + for (size_t i = 0; i < tensors.size(); ++i) { + add(tensors[i]->view_src); + for (auto source : tensors[i]->src) { + add(source); + } + } + std::vector signature; + signature.reserve(tensors.size() * 24); + signature.push_back(graph->n_nodes); + signature.push_back(graph->n_leafs); + for (int i = 0; i < graph->n_leafs; ++i) { + signature.push_back(indices.at(graph->leafs[i])); + } + for (int i = 0; i < graph->n_nodes; ++i) { + signature.push_back(indices.at(graph->nodes[i])); + } + for (auto tensor : tensors) { + signature.push_back(tensor->op); + signature.push_back(tensor->type); + signature.push_back(tensor->flags); + signature.push_back(tensor->view_offs); + if (include_bindings) { + signature.push_back(tensor->data != nullptr); + auto buffer = tensor_buffer(tensor); + signature.push_back(reinterpret_cast(buffer == nullptr ? nullptr : ggml_backend_buffer_get_type(buffer))); } for (int d = 0; d < GGML_MAX_DIMS; ++d) { - if (input_shape_ref.ne[static_cast(d)] != leaf->ne[d]) { - return false; - } + signature.push_back(tensor->ne[d]); + signature.push_back(tensor->nb[d]); + } + signature.push_back(tensor->view_src == nullptr ? 0 : indices.at(tensor->view_src)); + for (auto source : tensor->src) { + signature.push_back(source == nullptr ? 0 : indices.at(source)); + } + for (int value : tensor->op_params) { + signature.push_back(static_cast(value)); } } - return true; + return signature; + } + + static bool plan_matches_graph(ggml_cgraph* gf, + const Plan& plan, + const std::vector& layout) { + GGML_ASSERT(gf != nullptr); + if (plan.leaf_names.size() != static_cast(gf->n_leafs) || + plan.layout != layout) { + return false; + } + for (int i = 0; i < gf->n_leafs; ++i) { + if (plan.leaf_names[i] != gf->leafs[i]->name) { + return false; + } + } + std::vector> cut_markers; + for (int i = 0; i < ggml_graph_n_nodes(gf); ++i) { + auto node = ggml_graph_node(gf, i); + if (is_graph_cut_tensor(node)) { + cut_markers.emplace_back(i, node->name); + } + } + return cut_markers == plan.cut_markers; + } + + bool plan_matches_graph(ggml_cgraph* gf, const Plan& plan) { + GGML_ASSERT(gf != nullptr); + return plan_matches_graph(gf, plan, graph_layout(gf, false)); } ggml_tensor* output_tensor(ggml_cgraph* gf, const Segment& segment, size_t output_index) { @@ -578,26 +612,6 @@ namespace sd::ggml_graph_cut { return tensors; } - std::unordered_set collect_future_input_names(ggml_cgraph* gf, - const Plan& plan, - size_t current_segment_index) { - GGML_ASSERT(gf != nullptr); - std::unordered_set future_input_names; - for (size_t seg_idx = current_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) { - continue; - } - ggml_tensor* current_input = input_tensor(gf, input_ref); - if (current_input != nullptr && current_input->name[0] != '\0') { - future_input_names.insert(current_input->name); - } - } - } - return future_input_names; - } - ggml_cgraph* build_segment_graph(ggml_cgraph* gf, const Segment& segment, ggml_context** graph_ctx_out) { @@ -662,6 +676,10 @@ namespace sd::ggml_graph_cut { continue; } ggml_set_output(output); + if (output->view_src != nullptr) { + // A consumed output view does not keep its storage alive in gallocr. + ggml_set_output(output->view_src); + } } for (int node_idx : segment.internal_node_indices) { ggml_graph_add_node(segment_graph, ggml_graph_node(gf, node_idx)); @@ -716,6 +734,10 @@ namespace sd::ggml_graph_cut { if (output != nullptr && saved_output_flags.find(output) == saved_output_flags.end()) { saved_output_flags[output] = output->flags; } + if (output != nullptr && output->view_src != nullptr && + saved_output_flags.find(output->view_src) == saved_output_flags.end()) { + saved_output_flags[output->view_src] = output->view_src->flags; + } } ggml_context* graph_ctx = nullptr; @@ -744,6 +766,46 @@ namespace sd::ggml_graph_cut { return buffer_size; } + static size_t measure_graph_compute_buffer( + ggml_backend_t backend, + ggml_cgraph* gf, + const std::unordered_set& params_tensor_set) { + struct TensorRuntimeBinding { + ggml_backend_buffer_t buffer = nullptr; + void* data = nullptr; + void* extra = nullptr; + }; + std::unordered_map saved_bindings; + auto mark_external = [&](ggml_tensor* tensor) { + if (tensor == nullptr || saved_bindings.find(tensor) != saved_bindings.end()) { + return; + } + saved_bindings[tensor] = {tensor->buffer, tensor->data, tensor->extra}; + tensor->data = reinterpret_cast(static_cast(1)); + }; + for (int i = 0; i < leaf_count(gf); ++i) { + ggml_tensor* leaf = leaf_tensor(gf, i); + if (!is_params_tensor(params_tensor_set, leaf)) { + continue; + } + mark_external(leaf); + mark_external(leaf->view_src); + } + + ggml_gallocr_t allocr = ggml_gallocr_new( + ggml_backend_get_default_buffer_type(backend)); + size_t sizes[1] = {0}; + ggml_gallocr_reserve_n_size(allocr, gf, nullptr, nullptr, sizes); + ggml_gallocr_free(allocr); + + 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 sizes[0]; + } + Plan build_plan(ggml_backend_t backend, ggml_cgraph* gf, const std::unordered_set& params_tensor_set, @@ -756,24 +818,22 @@ namespace sd::ggml_graph_cut { if (n_nodes <= 0) { return plan; } - plan.n_nodes = n_nodes; - plan.n_leafs = gf->n_leafs; + plan.layout = graph_layout(gf, false); for (int i = 0; i < gf->n_leafs; ++i) { - ggml_tensor* leaf = gf->leafs[i]; - if (is_params_tensor(params_tensor_set, leaf)) { - continue; - } - auto shape = input_shape(leaf); - shape.leaf_index = i; - plan.input_shapes.push_back(shape); + plan.leaf_names.emplace_back(gf->leafs[i]->name); } + plan.compute_buffer_size = + measure_graph_compute_buffer(backend, gf, params_tensor_set); std::unordered_map producer_index; producer_index.reserve(static_cast(n_nodes)); for (int i = 0; i < n_nodes; ++i) { - producer_index[ggml_graph_node(gf, i)] = i; + ggml_tensor* node = ggml_graph_node(gf, i); + producer_index[node] = i; + if (is_graph_cut_tensor(node)) { + plan.cut_markers.push_back({i, node->name}); + } } - std::vector grouped_segments; std::unordered_map group_to_segment; for (int i = 0; i < n_nodes; ++i) { @@ -824,11 +884,24 @@ namespace sd::ggml_graph_cut { if (final_output_index < 0) { final_output_index = n_nodes - 1; } - ggml_tensor* final_output = final_output_index >= 0 ? ggml_graph_node(gf, final_output_index) : nullptr; - if (final_output != nullptr && available_cut_output_node_indices.find(final_output_index) == available_cut_output_node_indices.end()) { - Segment final_segment; - final_segment.group_name = "ggml_runner.final"; + Segment final_segment; + final_segment.group_name = "ggml_runner.final"; + if (final_output_index >= 0 && + available_cut_output_node_indices.find(final_output_index) == + available_cut_output_node_indices.end()) { final_segment.output_node_indices.push_back(final_output_index); + } + for (int i = 0; i < n_nodes; ++i) { + ggml_tensor* node = ggml_graph_node(gf, i); + if (i == final_output_index || node == nullptr || + (node->flags & GGML_TENSOR_FLAG_OUTPUT) == 0 || + available_cut_output_node_indices.find(i) != + available_cut_output_node_indices.end()) { + continue; + } + final_segment.output_node_indices.push_back(i); + } + if (!final_segment.output_node_indices.empty()) { build_segment(gf, plan, final_segment, @@ -839,210 +912,59 @@ namespace sd::ggml_graph_cut { log_desc); } - return plan; - } - - Plan apply_max_vram_budget(ggml_cgraph* gf, - const Plan& base_plan, - size_t max_graph_vram_bytes, - ggml_backend_t backend, - const std::unordered_set& params_tensor_set, - const char* log_desc) { - GGML_ASSERT(backend != nullptr); - GGML_ASSERT(gf != nullptr); - int64_t t_budget_begin = ggml_time_ms(); - if (max_graph_vram_bytes == 0 || !base_plan.has_cuts || base_plan.segments.size() <= 1) { - return base_plan; - } - - const int n_nodes = ggml_graph_n_nodes(gf); - std::unordered_map producer_index; - producer_index.reserve(static_cast(n_nodes)); - for (int i = 0; i < n_nodes; ++i) { - producer_index[ggml_graph_node(gf, i)] = i; - } - - Plan merged_plan; - merged_plan.available = true; - merged_plan.has_cuts = base_plan.has_cuts; - merged_plan.valid = base_plan.valid; - merged_plan.n_nodes = base_plan.n_nodes; - merged_plan.n_leafs = base_plan.n_leafs; - - std::unordered_set available_cut_output_node_indices; - available_cut_output_node_indices.reserve(static_cast(n_nodes)); - - size_t start_segment_index = 0; - while (start_segment_index < base_plan.segments.size()) { - Plan single_plan; - auto single_available_cut_output_node_indices = available_cut_output_node_indices; - auto single_seed = make_segment_seed(base_plan, - start_segment_index, - start_segment_index); - build_segment(gf, - single_plan, - single_seed, - producer_index, - single_available_cut_output_node_indices, - backend, - params_tensor_set, - log_desc); - GGML_ASSERT(!single_plan.segments.empty()); - - size_t best_end_segment_index = start_segment_index; - bool can_merge_next_segment = graph_cut_segment_vram_bytes(single_plan.segments.back()) <= max_graph_vram_bytes; - - while (can_merge_next_segment && best_end_segment_index + 1 < base_plan.segments.size()) { - const size_t next_end_segment_index = best_end_segment_index + 1; - Plan candidate_plan; - auto candidate_available_cut_output_node_indices = available_cut_output_node_indices; - auto candidate_seed = make_segment_seed(base_plan, - start_segment_index, - next_end_segment_index); - build_segment(gf, - candidate_plan, - candidate_seed, - producer_index, - candidate_available_cut_output_node_indices, - backend, - params_tensor_set, - log_desc); - GGML_ASSERT(!candidate_plan.segments.empty()); - - const auto& candidate_segment = candidate_plan.segments.back(); - const size_t candidate_bytes = graph_cut_segment_vram_bytes(candidate_segment); - if (candidate_bytes > max_graph_vram_bytes) { - break; + std::unordered_set future_cut_names; + for (auto segment = plan.segments.rbegin(); segment != plan.segments.rend(); ++segment) { + segment->future_cut_names = future_cut_names; + segment->live_cut_names = future_cut_names; + for (const auto& input : segment->input_refs) { + if (input.type != Segment::INPUT_PREVIOUS_CUT) { + continue; } - - best_end_segment_index = next_end_segment_index; + segment->live_cut_names.insert(input.display_name); + future_cut_names.insert(input.display_name); } - - auto best_seed = make_segment_seed(base_plan, - start_segment_index, - best_end_segment_index); - build_segment(gf, - merged_plan, - best_seed, - producer_index, - available_cut_output_node_indices, - backend, - params_tensor_set, - log_desc); - start_segment_index = best_end_segment_index + 1; } - if (log_desc != nullptr && merged_plan.segments.size() != base_plan.segments.size()) { - LOG_INFO("%s graph cut max_vram=%.2f MB merged %zu segments -> %zu segments", + std::string plan_validation_error; + plan.valid = validate_plan(gf, plan, &plan_validation_error); + if (!plan.valid && log_desc != nullptr) { + LOG_WARN("%s graph cut plan validation failed (%s); using monolithic execution", log_desc, - max_graph_vram_bytes / 1024.0 / 1024.0, - base_plan.segments.size(), - merged_plan.segments.size()); + plan_validation_error.c_str()); } - if (log_desc != nullptr) { - LOG_DEBUG("%s graph cut max_vram budget merge took %lld ms", - log_desc, - ggml_time_ms() - t_budget_begin); - } - - return merged_plan; + return plan; } Plan resolve_plan(ggml_backend_t backend, ggml_cgraph* gf, PlanCache* cache, - size_t max_graph_vram_bytes, const std::unordered_set& params_tensor_set, const char* log_desc) { GGML_ASSERT(backend != nullptr); GGML_ASSERT(gf != nullptr); GGML_ASSERT(cache != nullptr); - int64_t t_prepare_begin = ggml_time_ms(); - Plan base_plan; + const auto layout = graph_layout(gf, false); + auto& plans = cache->graph_cut_plans; + for (auto it = plans.begin(); it != plans.end(); ++it) { + if (it->available && plan_matches_graph(gf, *it, layout)) { + plans.splice(plans.begin(), plans, it); + return plans.front(); + } + } + int64_t t_plan_begin = ggml_time_ms(); - if (cache->graph_cut_plan.available && plan_matches_graph(gf, cache->graph_cut_plan)) { - base_plan = cache->graph_cut_plan; - } else { - base_plan = build_plan(backend, gf, params_tensor_set, log_desc); - cache->graph_cut_plan = base_plan; - cache->graph_cut_plan.available = true; - cache->budgeted_graph_cut_plan.available = false; - if (log_desc != nullptr) { - LOG_INFO("%s build cached graph cut plan done (taking %lld ms)", log_desc, ggml_time_ms() - t_plan_begin); - } + plans.push_front(build_plan(backend, gf, params_tensor_set, log_desc)); + if (plans.size() > PlanCache::MAX_PLANS) { + plans.pop_back(); } - - Plan resolved_plan = base_plan; - if (max_graph_vram_bytes > 0 && base_plan.has_cuts) { - if (cache->budgeted_graph_cut_plan.available && - cache->budgeted_graph_cut_plan_max_vram_bytes == max_graph_vram_bytes && - plan_matches_graph(gf, cache->budgeted_graph_cut_plan)) { - resolved_plan = cache->budgeted_graph_cut_plan; - } else { - resolved_plan = apply_max_vram_budget(gf, - base_plan, - max_graph_vram_bytes, - backend, - params_tensor_set, - log_desc); - cache->budgeted_graph_cut_plan = resolved_plan; - cache->budgeted_graph_cut_plan.available = true; - cache->budgeted_graph_cut_plan_max_vram_bytes = max_graph_vram_bytes; - } - } - return resolved_plan; - } - - void annotate_residency(Plan& plan, size_t max_graph_vram_bytes) { - // Cached plans may be reused with a smaller live budget. - for (auto& seg : plan.segments) { - seg.residency = SegmentResidency::STREAMED; - } - if (max_graph_vram_bytes == 0 || plan.segments.size() < 2) { - return; - } - - bool any_param_bearing = false; - for (const auto& seg : plan.segments) { - if (seg.input_param_bytes > 0) { - any_param_bearing = true; - break; - } - } - if (!any_param_bearing) { - return; - } - - // Leave room for the largest active streamed segment. - size_t worst_streamed_footprint = 0; - for (const auto& seg : plan.segments) { - const size_t seg_footprint = seg.input_param_bytes + - seg.compute_buffer_size + - seg.output_bytes + - seg.input_previous_cut_bytes + - seg.input_external_bytes; - if (seg_footprint > worst_streamed_footprint) { - worst_streamed_footprint = seg_footprint; - } - } - constexpr size_t safety = 512ull * 1024 * 1024; - const size_t reserved = safety + worst_streamed_footprint; - - if (max_graph_vram_bytes <= reserved) { - return; - } - const size_t available = max_graph_vram_bytes - reserved; - - size_t cumulative = 0; - for (auto& seg : plan.segments) { - if (cumulative + seg.input_param_bytes > available) { - break; - } - seg.residency = SegmentResidency::RESIDENT; - cumulative += seg.input_param_bytes; + if (log_desc != nullptr) { + LOG_INFO("%s build cached graph cut plan done (taking %lld ms)", + log_desc, + ggml_time_ms() - t_plan_begin); } + return plans.front(); } } // namespace sd::ggml_graph_cut diff --git a/otherarch/sdcpp/src/core/ggml_graph_cut.h b/otherarch/sdcpp/src/core/ggml_graph_cut.h index d955049d8..c93f368ce 100644 --- a/otherarch/sdcpp/src/core/ggml_graph_cut.h +++ b/otherarch/sdcpp/src/core/ggml_graph_cut.h @@ -3,22 +3,17 @@ #include #include +#include #include #include #include +#include #include #include "ggml-backend.h" #include "ggml.h" namespace sd::ggml_graph_cut { - - // Streaming residency for a segment's params. - enum class SegmentResidency : uint8_t { - STREAMED = 0, - RESIDENT = 1, - }; - struct Segment { enum InputType { INPUT_EXTERNAL = 0, @@ -33,38 +28,29 @@ namespace sd::ggml_graph_cut { int node_index = -1; }; - size_t compute_buffer_size = 0; - size_t output_bytes = 0; - size_t input_external_bytes = 0; - size_t input_previous_cut_bytes = 0; - size_t input_param_bytes = 0; + size_t compute_buffer_size = 0; std::string group_name; std::vector internal_node_indices; std::vector output_node_indices; std::vector input_refs; - SegmentResidency residency = SegmentResidency::STREAMED; + std::unordered_set future_cut_names; + std::unordered_set live_cut_names; }; struct Plan { - struct InputShape { - int leaf_index = -1; - ggml_type type = GGML_TYPE_COUNT; - std::array ne = {0, 0, 0, 0}; - }; - - bool available = false; - bool has_cuts = false; - bool valid = true; - int n_nodes = 0; - int n_leafs = 0; - std::vector input_shapes; + bool available = false; + bool has_cuts = false; + bool valid = true; + size_t compute_buffer_size = 0; + std::vector layout; + std::vector leaf_names; + std::vector> cut_markers; std::vector segments; }; struct PlanCache { - Plan graph_cut_plan; - Plan budgeted_graph_cut_plan; - size_t budgeted_graph_cut_plan_max_vram_bytes = 0; + static constexpr size_t MAX_PLANS = 4; + std::list graph_cut_plans; }; static constexpr const char* GGML_RUNNER_CUT_PREFIX = "ggml_runner_cut:"; @@ -89,13 +75,12 @@ namespace sd::ggml_graph_cut { ggml_backend_buffer_t tensor_buffer(const ggml_tensor* tensor); ggml_tensor* cache_source_tensor(ggml_tensor* tensor); size_t cache_tensor_bytes(const ggml_tensor* tensor); + // Plans ignore runtime bindings; allocator reservations must include them. + std::vector graph_layout(ggml_cgraph* graph, bool include_bindings); bool plan_matches_graph(ggml_cgraph* gf, const Plan& plan); 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 param_tensors(ggml_cgraph* gf, const Segment& segment); - std::unordered_set collect_future_input_names(ggml_cgraph* gf, - const Plan& plan, - size_t current_segment_index); ggml_cgraph* build_segment_graph(ggml_cgraph* gf, const Segment& segment, ggml_context** graph_ctx_out); @@ -109,21 +94,12 @@ namespace sd::ggml_graph_cut { ggml_cgraph* gf, const std::unordered_set& params_tensor_set, const char* log_desc); - Plan apply_max_vram_budget(ggml_cgraph* gf, - const Plan& base_plan, - size_t max_graph_vram_bytes, - ggml_backend_t backend, - const std::unordered_set& params_tensor_set, - const char* log_desc); Plan resolve_plan(ggml_backend_t backend, ggml_cgraph* gf, PlanCache* cache, - size_t max_graph_vram_bytes, const std::unordered_set& params_tensor_set, const char* log_desc); - // Mark leading segments resident when they fit after streamed-segment headroom. - void annotate_residency(Plan& plan, size_t max_graph_vram_bytes); } // namespace sd::ggml_graph_cut #endif // __SD_CORE_GGML_GRAPH_CUT_H__ diff --git a/otherarch/sdcpp/src/core/ggml_runner.cpp b/otherarch/sdcpp/src/core/ggml_runner.cpp new file mode 100644 index 000000000..93faa0e43 --- /dev/null +++ b/otherarch/sdcpp/src/core/ggml_runner.cpp @@ -0,0 +1,952 @@ +#include +#include +#include + +#include "core/ggml_extend_backend.h" +#include "core/ggml_runner.h" +#include "core/ggml_tensor_utils.h" +#include "core/layer_split_partition.h" +#include "core/segment_graph_bindings.h" +#include "core/segment_weight_pipeline.h" + +using namespace sd; + +void GGMLRunner::alloc_params_ctx() { + ggml_init_params params; + params.mem_size = static_cast(MAX_PARAMS_TENSOR_NUM * ggml_tensor_overhead()); + params.mem_buffer = nullptr; + params.no_alloc = true; + + params_ctx = ggml_init(params); + GGML_ASSERT(params_ctx != nullptr); + params_tensor_set_.clear(); + params_tensor_set_dirty_ = true; +} + +void GGMLRunner::free_params_ctx() { + if (params_ctx != nullptr) { + ggml_free(params_ctx); + params_ctx = nullptr; + } + params_tensor_set_.clear(); + params_tensor_set_dirty_ = true; +} + +void GGMLRunner::alloc_compute_ctx() { + ggml_init_params params; + params.mem_size = static_cast(ggml_tensor_overhead() * MAX_GRAPH_SIZE + ggml_graph_overhead()); + params.mem_buffer = nullptr; + params.no_alloc = true; + + compute_ctx = ggml_init(params); + GGML_ASSERT(compute_ctx != nullptr); +} + +void GGMLRunner::free_compute_ctx() { + debug_tensors.clear(); + if (compute_ctx != nullptr) { + ggml_free(compute_ctx); + compute_ctx = nullptr; + } + backend_tensor_data_map.clear(); +} + +void GGMLRunner::rebuild_params_tensor_set() { + if (!params_tensor_set_dirty_) { + return; + } + params_tensor_set_.clear(); + if (params_ctx == nullptr) { + return; + } + for (ggml_tensor* t = ggml_get_first_tensor(params_ctx); t != nullptr; t = ggml_get_next_tensor(params_ctx, t)) { + params_tensor_set_.insert(t); + } + params_tensor_set_dirty_ = false; +} + +ggml_tensor* GGMLRunner::canonical_param_tensor(ggml_tensor* tensor) { + if (tensor == nullptr) { + return nullptr; + } + if (params_tensor_set_.find(tensor) != params_tensor_set_.end()) { + return tensor; + } + if (tensor->view_src != nullptr && + params_tensor_set_.find(tensor->view_src) != params_tensor_set_.end()) { + return tensor->view_src; + } + return nullptr; +} + +std::vector GGMLRunner::collect_used_param_tensors(ggml_cgraph* gf) { + std::vector used_params; + rebuild_params_tensor_set(); + if (gf == nullptr || params_tensor_set_.empty()) { + return used_params; + } + + std::unordered_set seen_params; + const int n_leafs = sd::ggml_graph_cut::leaf_count(gf); + seen_params.reserve(static_cast(n_leafs)); + for (int i = 0; i < n_leafs; ++i) { + ggml_tensor* leaf = sd::ggml_graph_cut::leaf_tensor(gf, i); + ggml_tensor* param_leaf = canonical_param_tensor(leaf); + if (param_leaf != nullptr && + seen_params.insert(param_leaf).second) { + used_params.push_back(param_leaf); + } + } + return used_params; +} + +void GGMLRunner::evict_compute_backend_param_tensors(const std::vector& tensors) { + if (tensors.empty()) { + return; + } + auto manager = residency_manager.lock(); + if (manager != nullptr) { + manager->evict_compute_backend_params(tensors); + } +} + +void GGMLRunner::prepare_build_in_tensor_before() { + one_tensor = ggml_new_tensor_1d(compute_ctx, GGML_TYPE_F32, 1); + ggml_set_name(one_tensor, "ggml_runner_build_in_tensor:one"); + set_backend_tensor_data(one_tensor, one_vec.data()); + + zero_int_tensor = ggml_new_tensor_1d(compute_ctx, GGML_TYPE_I32, 1); + ggml_set_name(zero_int_tensor, "ggml_runner_build_in_tensor:zero_int"); + set_backend_tensor_data(zero_int_tensor, zero_int_vec.data()); +} + +void GGMLRunner::prepare_build_in_tensor_after(ggml_cgraph* gf) { + ggml_build_forward_expand(gf, one_tensor); + ggml_build_forward_expand(gf, zero_int_tensor); +} + +ggml_cgraph* GGMLRunner::new_graph_custom(size_t graph_size) { + if (weight_adapter) { + graph_size += weight_adapter->get_extra_graph_size(); + } + return ggml_new_graph_custom(compute_ctx, graph_size, false); +} + +ggml_cgraph* GGMLRunner::get_compute_graph(get_graph_cb_t get_graph) { + prepare_build_in_tensor_before(); + ggml_cgraph* gf = get_graph(); + if (gf == nullptr) { + return nullptr; + } + if (ggml_graph_n_nodes(gf) > 0) { + auto result = ggml_graph_node(gf, -1); + ggml_set_name(result, final_result_name.c_str()); + } + for (const auto& entry : debug_tensors) { + if (entry.first != nullptr) { + ggml_build_forward_expand(gf, entry.first); + } + } + for (const auto& entry : cache_.outputs()) { + if (entry.second != nullptr) { + ggml_build_forward_expand(gf, entry.second); + } + } + prepare_build_in_tensor_after(gf); + return gf; +} + +bool GGMLRunner::prepare_compute_graph(get_graph_cb_t get_graph, + ggml_cgraph** gf_out) { + GGML_ASSERT(gf_out != nullptr); + + reset_compute_ctx(); + ggml_cgraph* gf = get_compute_graph(get_graph); + if (gf == nullptr) { + free_compute_ctx(); + return false; + } + + *gf_out = gf; + return true; +} + +ggml_backend_t GGMLRunner::backend_for_weight(const ggml_tensor* tensor) const { + if (tensor == nullptr || tensor->buffer == nullptr) { + return nullptr; + } + if (ggml_backend_buffer_get_usage(tensor->buffer) != GGML_BACKEND_BUFFER_USAGE_WEIGHTS || + ggml_backend_buffer_is_host(tensor->buffer)) { + return nullptr; + } + ggml_backend_dev_t dev = ggml_backend_buft_get_device(ggml_backend_buffer_get_type(tensor->buffer)); + if (dev == nullptr) { + return nullptr; + } + if (ggml_backend_get_device(runtime_backend) == dev) { + return runtime_backend; + } + for (ggml_backend_t backend : extra_runtime_backends) { + if (ggml_backend_get_device(backend) == dev) { + return backend; + } + } + return nullptr; +} + +void GGMLRunner::pin_multi_device_nodes(ggml_backend_sched_t sched, ggml_cgraph* gf, ggml_cgraph* original_graph) { + if (sched == nullptr || gf == nullptr) { + return; + } + ggml_backend_t current = runtime_backend; + const int n_nodes = ggml_graph_n_nodes(gf); + for (int i = 0; i < n_nodes; i++) { + ggml_tensor* node = ggml_graph_node(gf, i); + auto node_assignment = graph_cut_layer_split_node_assignments_.find(original_graph == nullptr ? node : ggml_graph_node(original_graph, i)); + if (node_assignment != graph_cut_layer_split_node_assignments_.end()) { + current = node_assignment->second; + } + for (int s = 0; s < GGML_MAX_SRC; s++) { + ggml_backend_t weight_backend = backend_for_weight(node->src[s]); + if (weight_backend != nullptr) { + if (node_assignment == graph_cut_layer_split_node_assignments_.end()) { + current = weight_backend; + } + } + } + if (node->op == GGML_OP_NONE || node->op == GGML_OP_VIEW || node->op == GGML_OP_RESHAPE || + node->op == GGML_OP_PERMUTE || node->op == GGML_OP_TRANSPOSE) { + continue; + } + if (ggml_backend_supports_op(current, node)) { + ggml_backend_sched_set_tensor_backend(sched, node, current); + } + } +} + +size_t GGMLRunner::retained_runtime_buffer_bytes(ggml_backend_t backend) const { + backend = backend == nullptr ? runtime_backend : backend; + size_t bytes = workspace_.bytes(backend); + if (backend == runtime_backend) { + const size_t cache_bytes = cache_.resident_bytes(ggml_backend_get_device(backend)); + bytes = cache_bytes > SIZE_MAX - bytes ? SIZE_MAX : bytes + cache_bytes; + const size_t cut_bytes = cut_cache_.resident_bytes(ggml_backend_get_device(backend)); + bytes = cut_bytes > SIZE_MAX - bytes ? SIZE_MAX : bytes + cut_bytes; + } + return bytes; +} + +void GGMLRunner::sync_runtime_residency() { + if (auto manager = residency_manager.lock()) { + manager->update_runtime_residency(reinterpret_cast(this), + runtime_backend, retained_runtime_buffer_bytes()); + for (auto backend : extra_runtime_backends) { + manager->update_runtime_residency(reinterpret_cast(this), + backend, retained_runtime_buffer_bytes(backend)); + } + } +} + +std::optional> GGMLRunner::read_graph_tensor(ggml_tensor* tensor, const char* label) { + if (tensor == nullptr) { + LOG_ERROR("%s %s tensor is null", get_desc().c_str(), label); + return std::nullopt; + } + if (tensor->type != GGML_TYPE_F32) { + LOG_ERROR("%s %s tensor type mismatch: got %s", + get_desc().c_str(), + label, + ggml_type_name(tensor->type)); + return std::nullopt; + } + ggml_backend_buffer_t buf = sd::ggml_graph_cut::tensor_buffer(tensor); + if (buf == nullptr) { + LOG_ERROR("%s %s tensor buffer missing: name=%s op=%s buffer=%p view_src=%p view_src_buffer=%p data=%p", + get_desc().c_str(), + label, + tensor->name[0] != '\0' ? tensor->name : "", + ggml_op_name(tensor->op), + tensor->buffer, + tensor->view_src, + tensor->view_src ? tensor->view_src->buffer : nullptr, + tensor->data); + return std::nullopt; + } + + return sd::make_sd_tensor_from_ggml(tensor); +} + +void GGMLRunner::copy_data_to_backend_tensor(ggml_cgraph* gf, bool clear_after_copy) { + GGML_ASSERT(gf != nullptr); + std::unordered_set graph_tensor_set; + const int n_leafs = sd::ggml_graph_cut::leaf_count(gf); + const int n_nodes = ggml_graph_n_nodes(gf); + graph_tensor_set.reserve(static_cast(n_leafs + n_nodes)); + for (int i = 0; i < n_leafs; ++i) { + graph_tensor_set.insert(sd::ggml_graph_cut::leaf_tensor(gf, i)); + } + for (int i = 0; i < n_nodes; ++i) { + graph_tensor_set.insert(ggml_graph_node(gf, i)); + } + + for (auto& kv : backend_tensor_data_map) { + auto tensor = kv.first; + auto data = kv.second; + if (tensor == nullptr || data == nullptr) { + continue; + } + const char* name = ggml_get_name(tensor); + if (graph_tensor_set.find(tensor) == graph_tensor_set.end()) { + continue; + } + if (tensor->buffer == nullptr) { + LOG_WARN("%s skip backend tensor copy: tensor buffer not set, name='%s', ne=[%lld,%lld,%lld,%lld], type=%s", + get_desc().c_str(), + name != nullptr ? name : "", + (long long)tensor->ne[0], + (long long)tensor->ne[1], + (long long)tensor->ne[2], + (long long)tensor->ne[3], + ggml_type_name(tensor->type)); + continue; + } + + ggml_backend_buffer_t buf = tensor->view_src ? tensor->view_src->buffer : tensor->buffer; + if (buf == nullptr) { + LOG_WARN("%s graph exec skip tensor copy: name=%s op=%s reason=buffer_not_set data=%p view_src=%p view_src_buffer=%p", + get_desc().c_str(), + tensor && tensor->name[0] != '\0' ? tensor->name : "", + tensor ? ggml_op_name(tensor->op) : "", + data, + tensor ? tensor->view_src : nullptr, + (tensor && tensor->view_src) ? tensor->view_src->buffer : nullptr); + continue; + } + + ggml_backend_tensor_set(tensor, data, 0, ggml_nbytes(tensor)); + } + + if (clear_after_copy) { + backend_tensor_data_map.clear(); + } +} + +bool GGMLRunner::resolve_graph_cut_plan(ggml_cgraph* gf, + GraphCutPlan* plan_out) { + GGML_ASSERT(plan_out != nullptr); + GGML_ASSERT(gf != nullptr); + *plan_out = sd::ggml_graph_cut::resolve_plan(runtime_backend, + gf, + &graph_cut_plan_cache_, + params_tensor_set_, + get_desc().c_str()); + return true; +} + +bool GGMLRunner::resolve_graph_cut_layer_split_plan(ggml_cgraph* gf, + GraphCutPlan* plan_out) { + return resolve_graph_cut_plan(gf, plan_out); +} + +bool GGMLRunner::assign_graph_cut_layer_split_backends(ggml_cgraph* gf) { + graph_cut_layer_split_node_assignments_.clear(); + if (!graph_cut_layer_split_enabled) { + return true; + } + if (!is_multi_device()) { + LOG_ERROR("%s graph-cut layer split requires multiple runtime backends", get_desc().c_str()); + return false; + } + + GraphCutPlan plan; + if (!resolve_graph_cut_layer_split_plan(gf, &plan)) { + return false; + } + if (!plan.valid || !plan.has_cuts || plan.segments.size() <= 1) { + auto manager = residency_manager.lock(); + if (manager == nullptr) { + LOG_ERROR("%s weight manager is not set for graph-cut layer split", get_desc().c_str()); + return false; + } + std::vector graph_params = collect_used_param_tensors(gf); + if (!graph_params.empty() && + !manager->assign_compute_backend(graph_params, runtime_backend)) { + LOG_ERROR("%s graph-cut layer split failed to assign unmarked graph params to %s", + get_desc().c_str(), + sd::layer_split_backend_device_display_name(runtime_backend).c_str()); + return false; + } + for (ggml_tensor* param : graph_params) { + if (param != nullptr) { + graph_cut_layer_split_assignments_[param] = runtime_backend; + } + } + const int n_nodes = ggml_graph_n_nodes(gf); + for (int i = 0; i < n_nodes; i++) { + ggml_tensor* node = ggml_graph_node(gf, i); + if (node != nullptr) { + graph_cut_layer_split_node_assignments_[node] = runtime_backend; + } + } + if (!graph_cut_layer_split_primary_notice_logged_) { + LOG_WARN("%s graph-cut layer split: graph has no mark_graph_cut segments; using primary backend %s for %zu graph params", + get_desc().c_str(), + sd::layer_split_backend_device_display_name(runtime_backend).c_str(), + graph_params.size()); + graph_cut_layer_split_primary_notice_logged_ = true; + } else { + LOG_VERBOSE("%s graph-cut layer split: graph has no mark_graph_cut segments; using primary backend %s for %zu graph params", + get_desc().c_str(), + sd::layer_split_backend_device_display_name(runtime_backend).c_str(), + graph_params.size()); + } + return true; + } + + std::vector split_backends; + split_backends.reserve(extra_runtime_backends.size() + 1); + split_backends.push_back(runtime_backend); + for (ggml_backend_t backend : extra_runtime_backends) { + if (backend != nullptr) { + split_backends.push_back(backend); + } + } + + auto manager = residency_manager.lock(); + if (manager == nullptr) { + LOG_ERROR("%s weight manager is not set for graph-cut layer split", get_desc().c_str()); + return false; + } + + sd::GraphCutLayerSplitAssignment assignment; + auto canonicalize_param = [this](ggml_tensor* tensor) { + return canonical_param_tensor(tensor); + }; + if (!sd::partition_graph_cut_layer_split(get_desc().c_str(), + gf, + plan, + split_backends, + graph_cut_layer_split_backend_vram_limits_, + max_graph_vram_bytes, + graph_cut_layer_split_assignments_, + canonicalize_param, + &assignment)) { + return false; + } + + for (size_t i = 0; i < split_backends.size(); i++) { + if (assignment.tensors_by_backend[i].empty()) { + continue; + } + if (!manager->assign_compute_backend(assignment.tensors_by_backend[i], split_backends[i])) { + LOG_ERROR("%s graph-cut layer split failed to assign params to %s", + get_desc().c_str(), + sd::layer_split_backend_device_display_name(split_backends[i]).c_str()); + return false; + } + } + + graph_cut_layer_split_node_assignments_ = std::move(assignment.node_assignments); + sd::log_graph_cut_layer_split_assignment(get_desc().c_str(), split_backends, assignment); + + return true; +} + +bool GGMLRunner::runner_start() { + if (runner_started_) { + return true; + } + cache_.clear(); + workspace_.set_extra_backends(extra_runtime_backends); + if (auto manager = residency_manager.lock()) { + manager->set_workspace_reclaimer(reinterpret_cast(this), [this]() { + if (!workspace_.release()) { + return false; + } + sync_runtime_residency(); + return true; + }); + } + runner_started_ = true; + return true; +} + +void GGMLRunner::runner_end() { + GGML_ASSERT(!graph_active_); + if (!runner_started_) { + return; + } + workspace_.release(); + cache_.clear(); + logged_compute_bytes_.clear(); + logged_segment_count_ = 0; + if (auto manager = residency_manager.lock()) { + manager->clear_prefetched_params(reinterpret_cast(this)); + std::vector tensors; + for (auto tensor = ggml_get_first_tensor(params_ctx); tensor != nullptr; + tensor = ggml_get_next_tensor(params_ctx, tensor)) { + tensors.push_back(tensor); + } + manager->evict_compute_backend_params(tensors); + manager->remove_runtime_owner(reinterpret_cast(this)); + } + runner_started_ = false; +} + +GGMLRunner::GGMLRunner(ggml_backend_t backend, + std::shared_ptr manager) + : runtime_backend(backend), + cache_(backend), + cut_cache_(backend), + workspace_(backend), + residency_manager(manager) { + GGML_ASSERT(runtime_backend != nullptr); + alloc_params_ctx(); +} + +GGMLRunner::~GGMLRunner() { + runner_end(); + free_compute_ctx(); + free_params_ctx(); +} + +GGMLRunnerContext GGMLRunner::get_context() { + GGMLRunnerContext runner_ctx; + runner_ctx.ggml_ctx = compute_ctx; + runner_ctx.backend = runtime_backend; + runner_ctx.flash_attn_enabled = flash_attn_enabled; + runner_ctx.conv2d_direct_enabled = conv2d_direct_enabled; + runner_ctx.circular_x_enabled = circular_x_enabled; + runner_ctx.circular_y_enabled = circular_y_enabled; + runner_ctx.weight_adapter = weight_adapter; + runner_ctx.debug_tensors = &debug_tensors; + runner_ctx.get_cache_tensor = [this](const std::string& name) { + return this->get_cache_tensor_by_name(name); + }; + runner_ctx.cache_tensor = [this](const std::string& name, ggml_tensor* tensor) { + this->cache(name, tensor); + }; + runner_ctx.set_backend_tensor_data = [this](ggml_tensor* tensor, const void* data) { + this->set_backend_tensor_data(tensor, data); + }; + return runner_ctx; +} + +void GGMLRunner::reset_compute_ctx() { + free_compute_ctx(); + alloc_compute_ctx(); +} + +void GGMLRunner::free_cache_ctx_and_buffer() { + cache_.clear(); + sync_runtime_residency(); +} + +void GGMLRunner::set_backend_tensor_data(ggml_tensor* tensor, const void* data) { + // The scheduler only allocates standalone data tensors when they are + // marked as graph inputs. The flag is harmless for single-backend graphs. + ggml_set_input(tensor); + backend_tensor_data_map[tensor] = data; +} + +ggml_tensor* GGMLRunner::to_backend(ggml_tensor* tensor) { + GGML_ASSERT(compute_ctx != nullptr); + if (tensor == nullptr) { + return nullptr; + } + // it's performing a compute, check if backend isn't cpu + if (!sd_backend_is_cpu(runtime_backend) && (tensor->buffer == nullptr || ggml_backend_buffer_is_host(tensor->buffer))) { + // pass input tensors to gpu memory + auto backend_tensor = ggml_dup_tensor(compute_ctx, tensor); + + set_backend_tensor_data(backend_tensor, tensor->data); + return backend_tensor; + } else { + return tensor; + } +} + +void GGMLRunner::cache(const std::string name, ggml_tensor* tensor) { + if (tensor != nullptr && tensor->view_src != nullptr) { + tensor = ggml_cont(compute_ctx, tensor); + } + if (tensor != nullptr) { + ggml_set_output(tensor); + } + cache_.stage(name, tensor); +} + +std::optional> GGMLRunner::compute(get_graph_cb_t get_graph, + int n_threads, + bool auto_runner_end, + bool no_return, + const std::function& read_outputs) { + if (graph_active_) { + LOG_ERROR("%s does not support reentrant graph execution", get_desc().c_str()); + return std::nullopt; + } + if (!runner_start()) { + runner_end(); + return std::nullopt; + } + struct RunnerEndGuard { + GGMLRunner& runner; + bool enabled; + ~RunnerEndGuard() { + if (enabled) { + runner.runner_end(); + } + } + } runner_guard{*this, auto_runner_end}; + graph_active_ = true; + bool success = false; + struct GraphEndGuard { + GGMLRunner& runner; + const bool& success; + ~GraphEndGuard() { + runner.workspace_.segment_end(); + runner.cache_.graph_end(false); + runner.cut_cache_.clear(); + runner.free_compute_ctx(); + runner.graph_active_ = false; + if (!success) { + runner.workspace_.release(); + } + runner.sync_runtime_residency(); + } + } graph_guard{*this, success}; + + ggml_cgraph* graph = nullptr; + if (!prepare_compute_graph(get_graph, &graph)) { + return std::nullopt; + } + rebuild_params_tensor_set(); + auto output = execute_graph(graph, n_threads, no_return, read_outputs); + success = output.has_value(); + if (success) { + cache_.graph_end(true); + } + return output; +} + +void GGMLRunner::set_graph_cut_layer_split_enabled(bool enabled) { + graph_cut_layer_split_enabled = enabled; + if (!enabled) { + graph_cut_layer_split_assignments_.clear(); + graph_cut_layer_split_node_assignments_.clear(); + graph_cut_layer_split_primary_notice_logged_ = false; + } +} + +void GGMLRunner::set_graph_cut_layer_split_backend_vram_limits(const std::vector& limits) { + graph_cut_layer_split_backend_vram_limits_ = limits; + graph_cut_layer_split_assignments_.clear(); + graph_cut_layer_split_node_assignments_.clear(); + graph_cut_layer_split_primary_notice_logged_ = false; +} + +void GGMLRunner::set_runtime_backends(const std::vector& backends) { + extra_runtime_backends.clear(); + for (ggml_backend_t backend : backends) { + if (backend == nullptr || backend == runtime_backend) { + continue; + } + if (std::find(extra_runtime_backends.begin(), extra_runtime_backends.end(), backend) == + extra_runtime_backends.end()) { + extra_runtime_backends.push_back(backend); + } + } + workspace_.set_extra_backends(extra_runtime_backends); + graph_cut_layer_split_assignments_.clear(); + graph_cut_layer_split_node_assignments_.clear(); + graph_cut_layer_split_primary_notice_logged_ = false; +} + +static size_t add_bytes(size_t a, size_t b) { + return b > SIZE_MAX - a ? SIZE_MAX : a + b; +} + +ComputeWorkspace::Measurement GGMLRunner::measure(ggml_cgraph* graph, size_t direct_bytes) { + auto external_backend = [&](const ggml_tensor* tensor) -> ggml_backend_t { + if (!params_tensor_set_.count(tensor)) { + return nullptr; + } + auto placement = graph_cut_layer_split_assignments_.find(tensor); + return placement == graph_cut_layer_split_assignments_.end() ? runtime_backend : placement->second; + }; + auto assign_nodes = [&](ggml_backend_sched_t scheduler, ggml_cgraph* copy) { + pin_multi_device_nodes(scheduler, copy, graph); + }; + return workspace_.measure(graph, direct_bytes, external_backend, assign_nodes); +} + +std::vector GGMLRunner::memory_requests( + const std::vector& sizes, + size_t pending_cache_bytes) const { + std::vector requests; + for (const auto& size : sizes) { + const size_t retained = retained_runtime_buffer_bytes(size.backend); + const size_t reusable = workspace_.bytes(size.backend); + const size_t cache_bytes = size.backend == runtime_backend ? pending_cache_bytes : 0; + const size_t pending = add_bytes(size.bytes > reusable ? size.bytes - reusable : 0, cache_bytes); + size_t limit = max_graph_vram_bytes; + if (is_multi_device()) { + size_t index = 0; + if (size.backend != runtime_backend) { + auto position = std::find(extra_runtime_backends.begin(), extra_runtime_backends.end(), size.backend); + index = static_cast(position - extra_runtime_backends.begin()) + 1; + } + if (index < graph_cut_layer_split_backend_vram_limits_.size()) { + limit = graph_cut_layer_split_backend_vram_limits_[index]; + } + } + requests.push_back({size.backend, reinterpret_cast(this), pending, + retained, limit}); + } + return requests; +} + +bool GGMLRunner::fits(const std::vector& requests, + const std::vector& params) const { + auto manager = residency_manager.lock(); + if (manager == nullptr) { + return params.empty(); + } + for (const auto& request : requests) { + if (!manager->fits_compute_backend_capacity(request, params)) { + return false; + } + } + return true; +} + +bool GGMLRunner::execute_segment(ggml_cgraph* graph, int n_threads) { + if (sd_backend_is_cpu(runtime_backend)) { + sd_backend_cpu_set_n_threads(runtime_backend, n_threads); + } + if (workspace_.cpu_backend() != nullptr) { + sd_backend_cpu_set_n_threads(workspace_.cpu_backend(), n_threads); + } + auto scheduler = workspace_.scheduler(); + ggml_status status; + if (scheduler != nullptr) { + if (sd_get_backend_eval_callback() != nullptr && !multi_device_eval_callback_warned) { + LOG_WARN("%s: eval callback is not supported with the backend scheduler; ignoring", get_desc().c_str()); + multi_device_eval_callback_warned = true; + } + status = ggml_backend_sched_graph_compute(scheduler, graph); + } else { + status = sd_backend_graph_compute_with_eval_callback(runtime_backend, graph, + sd_get_backend_eval_callback(), + sd_get_backend_eval_callback_data()); + } + workspace_.synchronize(); + if (status != GGML_STATUS_SUCCESS) { + LOG_ERROR("%s compute failed: %s", get_desc().c_str(), ggml_status_to_string(status)); + return false; + } + const std::string description = get_desc(); + if (!debug_tensors.empty()) { + std::unordered_set graph_tensors; + const int leaf_count = ggml_graph_cut::leaf_count(graph); + const int node_count = ggml_graph_n_nodes(graph); + graph_tensors.reserve(static_cast(leaf_count + node_count)); + for (int index = 0; index < leaf_count; ++index) { + graph_tensors.insert(ggml_graph_cut::leaf_tensor(graph, index)); + } + for (int index = 0; index < node_count; ++index) { + graph_tensors.insert(ggml_graph_node(graph, index)); + } + + for (const auto& entry : debug_tensors) { + ggml_tensor* tensor = entry.first; + if (tensor == nullptr || graph_tensors.find(tensor) == graph_tensors.end()) { + continue; + } + ggml_backend_buffer_t buffer = + tensor->view_src != nullptr ? tensor->view_src->buffer : tensor->buffer; + if (buffer == nullptr) { + LOG_WARN("%s skip debug tensor '%s': tensor buffer not set", + description.c_str(), + entry.second.c_str()); + continue; + } + if (tensor->type != GGML_TYPE_F32) { + LOG_WARN("%s skip debug tensor '%s': only GGML_TYPE_F32 is supported, got %s", + description.c_str(), + entry.second.c_str(), + ggml_type_name(tensor->type)); + continue; + } + auto debug_tensor = make_sd_tensor_from_ggml(tensor); + print_sd_tensor(debug_tensor, false, entry.second.c_str()); + } + } + + return true; +} + +std::optional> GGMLRunner::execute_graph(ggml_cgraph* graph, int n_threads, bool no_return, const std::function& read_outputs) { + if (!assign_graph_cut_layer_split_backends(graph)) { + return std::nullopt; + } + const auto params = collect_used_param_tensors(graph); + ggml_graph_cut::Plan plan; + if (!resolve_graph_cut_plan(graph, &plan)) { + return std::nullopt; + } + const auto full_measurement = measure(graph, plan.compute_buffer_size); + if (full_measurement.buffers.empty()) { + return std::nullopt; + } + auto manager = residency_manager.lock(); + const bool segmented = !is_multi_device() && !sd_backend_is_cpu(runtime_backend) && + manager != nullptr && manager->segmented_compute_enabled() && + plan.valid && plan.has_cuts && plan.segments.size() > 1 && + !fits(memory_requests(full_measurement.buffers, cache_.pending_bytes(graph)), params); + if (!segmented) { + ggml_graph_cut::Segment segment; + segment.group_name = "graph"; + segment.compute_buffer_size = plan.compute_buffer_size; + for (int i = 0; i < ggml_graph_n_nodes(graph); ++i) { + segment.internal_node_indices.push_back(i); + } + for (int i = 0; i < ggml_graph_cut::leaf_count(graph); ++i) { + auto tensor = ggml_graph_cut::leaf_tensor(graph, i); + ggml_graph_cut::Segment::InputRef input; + input.leaf_index = i; + input.type = canonical_param_tensor(tensor) != nullptr + ? ggml_graph_cut::Segment::INPUT_PARAM + : ggml_graph_cut::Segment::INPUT_EXTERNAL; + segment.input_refs.push_back(input); + } + plan.segments = {std::move(segment)}; + } + const bool segments_changed = plan.segments.size() != logged_segment_count_; + if (segments_changed && (segmented || logged_segment_count_ > 1)) { + LOG_VERBOSE("%s using %zu segment%s", get_desc().c_str(), + plan.segments.size(), plan.segments.size() == 1 ? "" : "s"); + } + SegmentGraphBindings bindings(cut_cache_, plan, graph); + SegmentWeightPipeline weights(manager, runtime_backend, reinterpret_cast(this), + graph, plan, params_tensor_set_, + segmented && manager != nullptr && manager->prefetch_enabled()); + + std::map peak_compute_bytes; + auto track_compute_buffer = [&](ggml_backend_t backend) { + if (backend != nullptr) { + auto& peak = peak_compute_bytes[backend]; + peak = std::max(peak, workspace_.bytes(backend)); + } + }; + std::optional> output = Tensor(); + for (size_t index = 0; index < plan.segments.size(); ++index) { + const auto& segment = plan.segments[index]; + const bool last = index + 1 == plan.segments.size(); + auto fail_segment = [&](const char* phase) { + LOG_ERROR("%s segment %zu/%zu (%s) failed during %s", get_desc().c_str(), + index + 1, plan.segments.size(), segment.group_name.c_str(), phase); + return std::nullopt; + }; + cut_cache_.prune(segment.live_cut_names); + bindings.reset(segment); + if (!bindings.bind_cached_inputs(segment, get_desc().c_str())) { + return fail_segment("input binding"); + } + ggml_context* segment_context = nullptr; + auto segment_graph = segmented + ? ggml_graph_cut::build_segment_graph(graph, segment, &segment_context) + : graph; + struct SegmentCleanup { + GGMLRunner& runner; + SegmentWeightPipeline& weights; + SegmentGraphBindings& bindings; + ggml_context* context; + ~SegmentCleanup() { + runner.workspace_.segment_end(); + bindings.restore(); + weights.segment_end(); + ggml_free(context); + runner.sync_runtime_residency(); + } + } segment_cleanup{*this, weights, bindings, segment_context}; + + auto measurement = segmented ? measure(segment_graph, segment.compute_buffer_size) : full_measurement; + if (!workspace_.prepare(measurement)) { + return fail_segment("workspace preparation"); + } + const size_t cut_bytes = last ? 0 : cut_cache_.estimate_output_bytes(graph, segment); + const size_t new_cache_bytes = add_bytes(cut_bytes, cache_.pending_bytes(segment_graph)); + auto ensure_capacity = [&]() { + sync_runtime_residency(); + auto requests = memory_requests(measurement.buffers, new_cache_bytes); + if (!fits(requests, weights.params(index)) && workspace_.release_excess(measurement)) { + sync_runtime_residency(); + requests = memory_requests(measurement.buffers, new_cache_bytes); + } + return weights.ensure_segment_capacity(index, requests); + }; + if (!weights.segment_start(index, ensure_capacity)) { + return fail_segment("weight preparation"); + } + // Preparing weights can execute LoRA graphs and reclaim an idle workspace. + if (!workspace_.measurement_matches(segment_graph, measurement)) { + measurement = measure(segment_graph, segment.compute_buffer_size); + } + if (!workspace_.prepare(measurement) || !ensure_capacity()) { + return fail_segment("workspace capacity check"); + } + if (!workspace_.allocate(segment_graph, [&](ggml_backend_sched_t scheduler, ggml_cgraph* current) { + pin_multi_device_nodes(scheduler, current); + })) { + return fail_segment("workspace allocation"); + } + for (const auto& size : measurement.buffers) { + track_compute_buffer(size.backend); + } + if (workspace_.scheduler() != nullptr) { + track_compute_buffer(workspace_.cpu_backend()); + } + if (!ensure_capacity()) { + return fail_segment("allocated capacity check"); + } + copy_data_to_backend_tensor(segment_graph, false); + auto prefetch_requests = memory_requests(measurement.buffers, new_cache_bytes); + if (!prefetch_requests.empty()) { + weights.enqueue_next(index, prefetch_requests.front()); + } + LOG_DEBUG("%s executing segment %zu/%zu: %s", get_desc().c_str(), + index + 1, plan.segments.size(), segment.group_name.c_str()); + if (!execute_segment(segment_graph, n_threads) || + !cache_.capture(segment_graph) || + !cut_cache_.capture(graph, segment, get_desc().c_str())) { + return fail_segment("execution or output caching"); + } + sync_runtime_residency(); + if (last) { + if (read_outputs && !read_outputs()) { + return fail_segment("output finalization"); + } + if (!no_return) { + auto result = ggml_get_tensor(compute_ctx, final_result_name.c_str()); + output = read_graph_tensor(result, "output"); + if (!output.has_value()) { + return fail_segment("output readback"); + } + } + } + // Final outputs and their callbacks may still be views of consumed cuts. + cut_cache_.prune(segment.future_cut_names); + } + if (segments_changed || peak_compute_bytes != logged_compute_bytes_) { + for (const auto& entry : peak_compute_bytes) { + LOG_VERBOSE("%s compute buffer size: %.2f MB(%s) on %s (peak across %zu segment%s)", + get_desc().c_str(), entry.second / (1024.0 * 1024.0), + sd_backend_is_cpu(entry.first) ? "RAM" : "VRAM", ggml_backend_name(entry.first), + plan.segments.size(), plan.segments.size() == 1 ? "" : "s"); + } + logged_compute_bytes_ = std::move(peak_compute_bytes); + logged_segment_count_ = plan.segments.size(); + } + return output; +} diff --git a/otherarch/sdcpp/src/core/ggml_runner.h b/otherarch/sdcpp/src/core/ggml_runner.h new file mode 100644 index 000000000..7960f5aaa --- /dev/null +++ b/otherarch/sdcpp/src/core/ggml_runner.h @@ -0,0 +1,350 @@ +#ifndef __SD_CORE_GGML_RUNNER_H__ +#define __SD_CORE_GGML_RUNNER_H__ + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "core/compute_workspace.h" +#include "core/ggml_graph_cut.h" +#include "core/runner_cache.h" +#include "core/tensor_ggml.hpp" +#include "core/util.h" +#include "device_residency_manager.h" + +/* SDXL with LoRA requires more space */ +#define MAX_PARAMS_TENSOR_NUM 32768 +#define MAX_GRAPH_SIZE 327680 + +struct WeightAdapter { + struct ForwardParams { + enum class op_type_t { + OP_LINEAR, + OP_CONV2D, + } op_type; + struct { + bool force_prec_f32 = false; + float scale = 1.f; + } linear; + struct conv2d_params_t { + int s0 = 1; + int s1 = 1; + int p0 = 0; + int p1 = 0; + int d0 = 1; + int d1 = 1; + bool direct = false; + bool circular_x = false; + bool circular_y = false; + float scale = 1.f; + } conv2d; + }; + virtual ggml_tensor* patch_weight(ggml_context* ctx, ggml_backend_t backend, ggml_tensor* weight, const std::string& weight_name) = 0; + virtual ggml_tensor* forward_with_lora(ggml_context* ctx, + ggml_backend_t backend, + ggml_tensor* x, + ggml_tensor* w, + ggml_tensor* b, + const std::string& prefix, + ForwardParams forward_params) = 0; + virtual ggml_tensor* add_lora_to_output(ggml_context* ctx, + ggml_backend_t backend, + ggml_tensor* x, + ggml_tensor* w, + ggml_tensor* output, + const std::string& prefix, + ForwardParams forward_params) = 0; + virtual size_t get_extra_graph_size() = 0; +}; + +struct GGMLRunnerContext { + ggml_backend_t backend = nullptr; + ggml_context* ggml_ctx = nullptr; + bool flash_attn_enabled = false; + bool conv2d_direct_enabled = false; + bool circular_x_enabled = false; + bool circular_y_enabled = false; + ggml_tensor* ip_context = nullptr; + float ip_scale = 1.0f; + std::shared_ptr weight_adapter = nullptr; + std::vector>* debug_tensors = nullptr; + std::function get_cache_tensor; + std::function cache_tensor; + std::function set_backend_tensor_data; + std::map, ggml_tensor*> int8_convrot_cache; + + void capture_tensor(const std::string& name, ggml_tensor* tensor) { + if (debug_tensors == nullptr || tensor == nullptr) { + return; + } + ggml_tensor* snapshot = ggml_cont(ggml_ctx, tensor); + ggml_tensor* dst = ggml_dup_tensor(ggml_ctx, snapshot); + snapshot = ggml_cpy(ggml_ctx, snapshot, dst); + ggml_set_output(snapshot); + debug_tensors->push_back({snapshot, name}); + } + + ggml_tensor* load_cache_tensor(const std::string& name) const { + if (!get_cache_tensor) { + return nullptr; + } + return get_cache_tensor(name); + } + + void persist_cache_tensor(const std::string& name, ggml_tensor* tensor) const { + if (!cache_tensor || tensor == nullptr) { + return; + } + cache_tensor(name, tensor); + } + + void bind_backend_tensor_data(ggml_tensor* tensor, const void* data) const { + if (!set_backend_tensor_data || tensor == nullptr || data == nullptr) { + return; + } + set_backend_tensor_data(tensor, data); + } +}; + +struct GGMLRunner { +private: + std::map logged_compute_bytes_; + size_t logged_segment_count_ = 0; + + sd::ComputeWorkspace::Measurement measure(ggml_cgraph* graph, size_t direct_bytes); + std::vector memory_requests(const std::vector& sizes, + size_t pending_cache_bytes) const; + bool fits(const std::vector& requests, + const std::vector& params) const; + bool execute_segment(ggml_cgraph* graph, int n_threads); + std::optional> execute_graph(ggml_cgraph* graph, int n_threads, bool no_return, const std::function& read_outputs); + +protected: + typedef std::function get_graph_cb_t; + using GraphCutPlan = sd::ggml_graph_cut::Plan; + + ggml_backend_t runtime_backend = nullptr; + + ggml_context* params_ctx = nullptr; + + sd::RunnerCache cache_; + sd::GraphCutTensorCache cut_cache_; + sd::ComputeWorkspace workspace_; + ggml_context* compute_ctx = nullptr; + bool runner_started_ = false; + bool graph_active_ = false; + + size_t max_graph_vram_bytes = 0; + bool graph_cut_layer_split_enabled = false; + std::vector graph_cut_layer_split_backend_vram_limits_; + + std::vector extra_runtime_backends; // borrowed (SDBackendManager-owned) + bool multi_device_eval_callback_warned = false; + + std::shared_ptr weight_adapter = nullptr; + std::weak_ptr residency_manager; + bool params_tensor_set_dirty_ = true; + + std::vector one_vec = {1.f}; + ggml_tensor* one_tensor = nullptr; + + std::vector zero_int_vec = {0}; + ggml_tensor* zero_int_tensor = nullptr; + + std::map backend_tensor_data_map; + std::vector> debug_tensors; + const std::string final_result_name = "ggml_runner_final_result_tensor"; + + bool flash_attn_enabled = false; + bool conv2d_direct_enabled = false; + bool circular_x_enabled = false; + bool circular_y_enabled = false; + + sd::ggml_graph_cut::PlanCache graph_cut_plan_cache_; + std::unordered_set params_tensor_set_; + std::unordered_map graph_cut_layer_split_assignments_; + std::unordered_map graph_cut_layer_split_node_assignments_; + bool graph_cut_layer_split_primary_notice_logged_ = false; + + template + static sd::Tensor take_or_empty(std::optional> tensor) { + if (!tensor.has_value()) { + return {}; + } + return std::move(*tensor); + } + + template + static sd::Tensor restore_trailing_singleton_dims(std::optional> tensor, + size_t expected_dim) { + return restore_trailing_singleton_dims(take_or_empty(std::move(tensor)), expected_dim); + } + + template + static sd::Tensor restore_trailing_singleton_dims(sd::Tensor tensor, + size_t expected_dim) { + if (tensor.empty()) { + return tensor; + } + while (static_cast(tensor.dim()) < expected_dim) { + tensor.unsqueeze_(tensor.dim()); + } + return tensor; + } + + void alloc_params_ctx(); + + void free_params_ctx(); + + void alloc_compute_ctx(); + + void free_compute_ctx(); + + void rebuild_params_tensor_set(); + + ggml_tensor* canonical_param_tensor(ggml_tensor* tensor); + + std::vector collect_used_param_tensors(ggml_cgraph* gf); + + void evict_compute_backend_param_tensors(const std::vector& tensors); + + void prepare_build_in_tensor_before(); + + void prepare_build_in_tensor_after(ggml_cgraph* gf); + + ggml_cgraph* new_graph_custom(size_t graph_size); + + ggml_cgraph* get_compute_graph(get_graph_cb_t get_graph); + + bool prepare_compute_graph(get_graph_cb_t get_graph, + ggml_cgraph** gf_out); + + ggml_backend_t backend_for_weight(const ggml_tensor* tensor) const; + + // Weightless ops have no scheduler anchor, so pin them to the most recent + // weight device. Views must stay unpinned or cross-device copies can be + // skipped for their consumers. + void pin_multi_device_nodes(ggml_backend_sched_t sched, ggml_cgraph* gf, ggml_cgraph* original_graph = nullptr); + + bool is_multi_device() const { + return !extra_runtime_backends.empty(); + } + + size_t reusable_compute_buffer_bytes() const { + return workspace_.bytes(runtime_backend); + } + + size_t retained_runtime_buffer_bytes(ggml_backend_t backend = nullptr) const; + + void sync_runtime_residency(); + + std::optional> read_graph_tensor(ggml_tensor* tensor, const char* label); + + void copy_data_to_backend_tensor(ggml_cgraph* gf, bool clear_after_copy = true); + + bool resolve_graph_cut_plan(ggml_cgraph* gf, + GraphCutPlan* plan_out); + + bool resolve_graph_cut_layer_split_plan(ggml_cgraph* gf, + GraphCutPlan* plan_out); + + bool assign_graph_cut_layer_split_backends(ggml_cgraph* gf); + +public: + bool runner_start(); + + bool runner_started() const { return runner_started_; } + + void runner_end(); + +public: + virtual std::string get_desc() = 0; + + GGMLRunner(ggml_backend_t backend, + std::shared_ptr manager = nullptr); + + virtual ~GGMLRunner(); + + virtual GGMLRunnerContext get_context(); + + void reset_compute_ctx(); + +public: + void free_cache_ctx_and_buffer(); + + // do copy after alloc graph + void set_backend_tensor_data(ggml_tensor* tensor, const void* data); + + template + ggml_tensor* make_input(const sd::Tensor& tensor) { + ggml_tensor* input = sd::make_ggml_tensor(compute_ctx, tensor, false); + set_backend_tensor_data(input, tensor.data()); + return input; + } + + template + ggml_tensor* make_optional_input(const sd::Tensor& tensor) { + if (tensor.empty()) { + return nullptr; + } + return make_input(tensor); + } + + template + ggml_tensor* make_optional_input(const sd::Tensor* tensor) { + if (tensor == nullptr) { + return nullptr; + } + return make_input(*tensor); + } + + ggml_tensor* to_backend(ggml_tensor* tensor); + + void cache(const std::string name, ggml_tensor* tensor); + + ggml_tensor* get_cache_tensor_by_name(const std::string& name) { + return cache_.get(name); + } + + std::optional> compute(get_graph_cb_t get_graph, + int n_threads, + bool auto_runner_end = true, + bool no_return = false, + const std::function& read_outputs = {}); + + void set_flash_attention_enabled(bool enabled) { + flash_attn_enabled = enabled; + } + + void set_conv2d_direct_enabled(bool enabled) { + conv2d_direct_enabled = enabled; + } + + void set_circular_axes(bool circular_x, bool circular_y) { + circular_x_enabled = circular_x; + circular_y_enabled = circular_y; + } + + void set_weight_adapter(const std::shared_ptr& adapter) { + weight_adapter = adapter; + } + + void set_max_graph_vram_bytes(size_t max_vram_bytes) { + max_graph_vram_bytes = max_vram_bytes; + } + + void set_graph_cut_layer_split_enabled(bool enabled); + + void set_graph_cut_layer_split_backend_vram_limits(const std::vector& limits); + + void set_runtime_backends(const std::vector& backends); +}; + +#endif // __SD_CORE_GGML_RUNNER_H__ diff --git a/otherarch/sdcpp/src/core/ggml_tensor_utils.cpp b/otherarch/sdcpp/src/core/ggml_tensor_utils.cpp new file mode 100644 index 000000000..1028972fc --- /dev/null +++ b/otherarch/sdcpp/src/core/ggml_tensor_utils.cpp @@ -0,0 +1,428 @@ +#include "core/ggml_tensor_utils.h" + +#include +#include +#include +#include + +#include "core/ggml_extend_backend.h" +#include "core/rng.hpp" + +void ggml_ext_im_set_randn_f32(ggml_tensor* tensor, std::shared_ptr rng) { + uint32_t n = (uint32_t)ggml_nelements(tensor); + std::vector random_numbers = rng->randn(n); + for (uint32_t i = 0; i < n; i++) { + ggml_ext_im_set_f32_1d(tensor, i, random_numbers[i]); + } +} + +void print_ggml_tensor(ggml_tensor* tensor, bool shape_only, const char* mark) { + printf("%s (%s): shape(%zu, %zu, %zu, %zu)\n", mark, ggml_type_name(tensor->type), tensor->ne[0], tensor->ne[1], tensor->ne[2], tensor->ne[3]); + fflush(stdout); + if (shape_only) { + return; + } + int range = 3; + for (int i3 = 0; i3 < tensor->ne[3]; i3++) { + if (i3 >= range && i3 + range < tensor->ne[3]) { + continue; + } + for (int i2 = 0; i2 < tensor->ne[2]; i2++) { + if (i2 >= range && i2 + range < tensor->ne[2]) { + continue; + } + for (int i1 = 0; i1 < tensor->ne[1]; i1++) { + if (i1 >= range && i1 + range < tensor->ne[1]) { + continue; + } + for (int i0 = 0; i0 < tensor->ne[0]; i0++) { + if (i0 >= range && i0 + range < tensor->ne[0]) { + continue; + } + if (tensor->type == GGML_TYPE_F32) { + printf(" [%d, %d, %d, %d] = %f\n", i3, i2, i1, i0, ggml_ext_tensor_get_f32(tensor, i0, i1, i2, i3)); + } else if (tensor->type == GGML_TYPE_F16) { + printf(" [%d, %d, %d, %d] = %f\n", i3, i2, i1, i0, ggml_fp16_to_fp32(ggml_ext_tensor_get_f16(tensor, i0, i1, i2, i3))); + } else if (tensor->type == GGML_TYPE_I32) { + printf(" [%d, %d, %d, %d] = %i3\n", i3, i2, i1, i0, ggml_ext_tensor_get_i32(tensor, i0, i1, i2, i3)); + } + fflush(stdout); + } + } + } + } +} + +void ggml_ext_tensor_iter( + ggml_tensor* tensor, + const std::function& fn) { + int64_t n0 = tensor->ne[0]; + int64_t n1 = tensor->ne[1]; + int64_t n2 = tensor->ne[2]; + int64_t n3 = tensor->ne[3]; + + for (int64_t i3 = 0; i3 < n3; i3++) { + for (int64_t i2 = 0; i2 < n2; i2++) { + for (int64_t i1 = 0; i1 < n1; i1++) { + for (int64_t i0 = 0; i0 < n0; i0++) { + fn(tensor, i0, i1, i2, i3); + } + } + } + } +} + +void ggml_ext_tensor_iter( + ggml_tensor* tensor, + const std::function& fn) { + int64_t n0 = tensor->ne[0]; + int64_t n1 = tensor->ne[1]; + int64_t n2 = tensor->ne[2]; + int64_t n3 = tensor->ne[3]; + + for (int64_t i = 0; i < ggml_nelements(tensor); i++) { + fn(tensor, i); + } +} + +void ggml_ext_tensor_diff( + ggml_tensor* a, + ggml_tensor* b, + float gap) { + GGML_ASSERT(ggml_nelements(a) == ggml_nelements(b)); + ggml_ext_tensor_iter(a, [&](ggml_tensor* a, int64_t i0, int64_t i1, int64_t i2, int64_t i3) { + float a_value = ggml_ext_tensor_get_f32(a, i0, i1, i2, i3); + float b_value = ggml_ext_tensor_get_f32(b, i0, i1, i2, i3); + if (abs(a_value - b_value) > gap) { + LOG_WARN("[%ld, %ld, %ld, %ld] %f %f", i3, i2, i1, i0, a_value, b_value); + } + }); +} + +ggml_tensor* load_tensor_from_file(ggml_context* ctx, const std::string& file_path) { + std::ifstream file(sd_get_u8path(file_path), std::ios::binary); + if (!file.is_open()) { + LOG_ERROR("failed to open '%s'", file_path.c_str()); + return nullptr; + } + int32_t n_dims; + int32_t length; + int32_t ttype; + + file.read(reinterpret_cast(&n_dims), sizeof(n_dims)); + file.read(reinterpret_cast(&length), sizeof(length)); + file.read(reinterpret_cast(&ttype), sizeof(ttype)); + + LOG_VERBOSE("load_tensor_from_file %d %d %d", n_dims, length, ttype); + + if (file.eof()) { + LOG_ERROR("incomplete file '%s'", file_path.c_str()); + return nullptr; + } + + int32_t nelements = 1; + int32_t ne[4] = {1, 1, 1, 1}; + for (int i = 0; i < n_dims; ++i) { + file.read(reinterpret_cast(&ne[i]), sizeof(ne[i])); + nelements *= ne[i]; + } + std::string name(length, 0); + file.read(&name[0], length); + ggml_tensor* tensor = ggml_new_tensor_4d(ctx, (ggml_type)ttype, ne[0], ne[1], ne[2], ne[3]); + const size_t bpe = ggml_type_size(ggml_type(ttype)); + file.read(reinterpret_cast(tensor->data), ggml_nbytes(tensor)); + return tensor; +} + +// __STATIC_INLINE__ void save_tensor_to_file(const std::string& file_name, ggml_tensor* tensor, const std::string & name) { +// std::string file_name_ = file_name + ".tensor"; +// std::string name_ = name; +// std::ofstream file("./" + file_name_, std::ios::binary); +// file.write(reinterpret_cast(&tensor->n_dims), sizeof(tensor->n_dims)); +// int len = (int)name_.size(); +// file.write(reinterpret_cast(&len), sizeof(len)); +// int ttype = (int)tensor->type; +// file.write(reinterpret_cast(&ttype), sizeof(ttype)); +// for (int i = 0; i < tensor->n_dims; ++i) { +// int ne_ = (int) tensor->ne[i]; +// file.write(reinterpret_cast(&ne_), sizeof(ne_)); +// } +// file.write(&name_[0], len); +// char* data = nullptr; +// file.write((char*)tensor->data, ggml_nbytes(tensor)); +// file.close(); +// } + +uint8_t* ggml_tensor_to_sd_image(ggml_tensor* input, uint8_t* image_data) { + int64_t width = input->ne[0]; + int64_t height = input->ne[1]; + int64_t channels = input->ne[2]; + GGML_ASSERT(input->type == GGML_TYPE_F32); + if (image_data == nullptr) { + image_data = (uint8_t*)malloc(width * height * channels); + } + for (int iy = 0; iy < height; iy++) { + for (int ix = 0; ix < width; ix++) { + for (int k = 0; k < channels; k++) { + float value = ggml_ext_tensor_get_f32(input, ix, iy, k); + *(image_data + iy * width * channels + ix * channels + k) = (uint8_t)(value * 255.0f); + } + } + } + return image_data; +} + +uint8_t* ggml_tensor_to_sd_image(ggml_tensor* input, int idx, bool video) { + int64_t width = input->ne[0]; + int64_t height = input->ne[1]; + int64_t channels; + if (video) { + channels = input->ne[3]; + } else { + channels = input->ne[2]; + } + GGML_ASSERT(channels == 3 && input->type == GGML_TYPE_F32); + uint8_t* image_data = (uint8_t*)malloc(width * height * channels); + for (int ih = 0; ih < height; ih++) { + for (int iw = 0; iw < width; iw++) { + for (int ic = 0; ic < channels; ic++) { + float value; + if (video) { + value = ggml_ext_tensor_get_f32(input, iw, ih, idx, ic); + } else { + value = ggml_ext_tensor_get_f32(input, iw, ih, ic, idx); + } + *(image_data + ih * width * channels + iw * channels + ic) = (uint8_t)(value * 255.0f); + } + } + } + return image_data; +} + +void sd_image_to_ggml_tensor(sd_image_t image, + ggml_tensor* tensor, + bool scale) { + GGML_ASSERT(image.width == tensor->ne[0]); + GGML_ASSERT(image.height == tensor->ne[1]); + GGML_ASSERT(image.channel == tensor->ne[2]); + GGML_ASSERT(1 == tensor->ne[3]); + GGML_ASSERT(tensor->type == GGML_TYPE_F32); + ggml_ext_tensor_iter(tensor, [&](ggml_tensor* tensor, int64_t i0, int64_t i1, int64_t i2, int64_t i3) { + float value = sd_image_get_f32(image, i0, i1, i2, scale); + ggml_ext_tensor_set_f32(tensor, value, i0, i1, i2, i3); + }); +} + +void ggml_ext_tensor_apply_mask(ggml_tensor* image_data, + ggml_tensor* mask, + ggml_tensor* output, + float masked_value) { + int64_t width = output->ne[0]; + int64_t height = output->ne[1]; + int64_t channels = output->ne[2]; + float rescale_mx = 1.f * mask->ne[0] / output->ne[0]; + float rescale_my = 1.f * mask->ne[1] / output->ne[1]; + GGML_ASSERT(output->type == GGML_TYPE_F32); + for (int ix = 0; ix < width; ix++) { + for (int iy = 0; iy < height; iy++) { + int mx = (int)(ix * rescale_mx); + int my = (int)(iy * rescale_my); + float m = ggml_ext_tensor_get_f32(mask, mx, my); + m = round(m); // inpaint models need binary masks + ggml_ext_tensor_set_f32(mask, m, mx, my); + for (int k = 0; k < channels; k++) { + float value = ggml_ext_tensor_get_f32(image_data, ix, iy, k); + value = (1 - m) * (value - masked_value) + masked_value; + ggml_ext_tensor_set_f32(output, value, ix, iy, k); + } + } + } +} + +float ggml_ext_tensor_mean(ggml_tensor* src) { + float mean = 0.0f; + int64_t nelements = ggml_nelements(src); + float* data = (float*)src->data; + for (int i = 0; i < nelements; i++) { + mean += data[i] / nelements * 1.0f; + } + return mean; +} + +void ggml_ext_tensor_add_inplace(ggml_tensor* a, ggml_tensor* b) { + GGML_ASSERT(ggml_nelements(a) == ggml_nelements(b)); + int64_t nelements = ggml_nelements(a); + float* vec_a = (float*)a->data; + float* vec_b = (float*)b->data; + for (int i = 0; i < nelements; i++) { + vec_a[i] = vec_a[i] + vec_b[i]; + } +} + +void ggml_ext_tensor_scale_inplace(ggml_tensor* src, float scale) { + int64_t nelements = ggml_nelements(src); + float* data = (float*)src->data; + for (int i = 0; i < nelements; i++) { + data[i] = data[i] * scale; + } +} + +void ggml_ext_tensor_clamp_inplace(ggml_tensor* src, float min, float max) { + int64_t nelements = ggml_nelements(src); + float* data = (float*)src->data; + for (int i = 0; i < nelements; i++) { + float val = data[i]; + data[i] = val < min ? min : (val > max ? max : val); + } +} + +ggml_tensor* ggml_ext_tensor_concat(ggml_context* ctx, + ggml_tensor* a, + ggml_tensor* b, + int dim) { + int64_t ne[GGML_MAX_DIMS]; + for (int d = 0; d < GGML_MAX_DIMS; ++d) { + if (d == dim) { + ne[d] = a->ne[d] + b->ne[d]; + continue; + } + GGML_ASSERT(a->ne[d] == b->ne[d]); + ne[d] = a->ne[d]; + } + ggml_tensor* result = ggml_new_tensor(ctx, a->type, GGML_MAX_DIMS, ne); + int64_t o[4] = {0, 0, 0, 0}; + o[dim] = a->ne[dim]; + + float v; + for (int i3 = 0; i3 < result->ne[3]; i3++) { + for (int i2 = 0; i2 < result->ne[2]; i2++) { + for (int i1 = 0; i1 < result->ne[1]; i1++) { + for (int i0 = 0; i0 < result->ne[0]; i0++) { + if (i0 < a->ne[0] && i1 < a->ne[1] && i2 < a->ne[2] && i3 < a->ne[3]) { + v = ggml_ext_tensor_get_f32(a, i0, i1, i2, i3); + } else { + v = ggml_ext_tensor_get_f32(b, i0 - o[0], i1 - o[1], i2 - o[2], i3 - o[3]); + } + + ggml_ext_tensor_set_f32(result, v, i0, i1, i2, i3); + } + } + } + } + return result; +} + +void scale_to_minus1_1(ggml_tensor* src) { + int64_t nelements = ggml_nelements(src); + float* data = (float*)src->data; + for (int i = 0; i < nelements; i++) { + float val = data[i]; + data[i] = val * 2.0f - 1.0f; + } +} + +void scale_to_0_1(ggml_tensor* src) { + int64_t nelements = ggml_nelements(src); + float* data = (float*)src->data; + for (int i = 0; i < nelements; i++) { + float val = data[i]; + data[i] = (val + 1.0f) * 0.5f; + } +} + +ggml_tensor* vector_to_ggml_tensor(ggml_context* ctx, + const std::vector& vec) { + ggml_tensor* t = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, vec.size()); + memcpy(t->data, (const void*)vec.data(), ggml_nbytes(t)); + return t; +} + +ggml_tensor* vector_to_ggml_tensor_i32(ggml_context* ctx, + const std::vector& vec) { + ggml_tensor* t = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, vec.size()); + memcpy(t->data, (const void*)vec.data(), ggml_nbytes(t)); + return t; +} + +std::vector arange(float start, float end, float step) { + std::vector result; + + for (float value = start; value < end; value += step) { + result.push_back(value); + } + + return result; +} + +std::vector timestep_embedding(std::vector timesteps, + int dim, + int max_period, + bool flip_sin_to_cos, + float scale) { + // timesteps: [N,] + // embedding: [N, dim] + size_t N = timesteps.size(); + std::vector embedding(N * dim, 0.f); + int half = dim / 2; + std::vector freqs(half); + for (int i = 0; i < half; ++i) { + freqs[i] = (float)std::exp(-std::log(max_period) * i / half); + } + for (int i = 0; i < N; ++i) { + for (int j = 0; j < half; ++j) { + float arg = timesteps[i] * freqs[j] * scale; + if (flip_sin_to_cos) { + embedding[i * dim + j] = std::cos(arg); + embedding[i * dim + j + half] = std::sin(arg); + } else { + embedding[i * dim + j] = std::sin(arg); + embedding[i * dim + j + half] = std::cos(arg); + } + } + } + return embedding; +} + +void set_timestep_embedding(std::vector timesteps, + ggml_tensor* embedding, + int dim, + int max_period) { + std::vector embedding_vec = timestep_embedding(timesteps, dim, max_period); + memcpy(((char*)embedding->data), ((char*)embedding_vec.data()), ggml_nbytes(embedding)); +} + +void set_timestep_embedding(std::vector timesteps, + sd::Tensor* embedding, + int dim, + int max_period) { + GGML_ASSERT(embedding != nullptr); + std::vector embedding_vec = timestep_embedding(timesteps, dim, max_period); + if (embedding->numel() != static_cast(embedding_vec.size())) { + embedding->resize({dim, static_cast(timesteps.size())}); + } + std::copy(embedding_vec.begin(), embedding_vec.end(), embedding->values().begin()); +} + +ggml_tensor* new_timestep_embedding(ggml_context* ctx, + std::vector timesteps, + int dim, + int max_period) { + // timesteps: [N,] + // embedding: [N, dim] + std::vector embedding_vec = timestep_embedding(timesteps, dim, max_period); + ggml_tensor* embedding = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, dim, timesteps.size()); + if (embedding->data != nullptr) { + memcpy(((char*)embedding->data), ((char*)embedding_vec.data()), ggml_nbytes(embedding)); + } else { + ggml_backend_tensor_set(embedding, embedding_vec.data(), 0, ggml_nbytes(embedding)); + } + return embedding; +} + +size_t ggml_tensor_num(ggml_context* ctx) { + size_t num = 0; + for (ggml_tensor* t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) { + num++; + } + return num; +} diff --git a/otherarch/sdcpp/src/core/ggml_tensor_utils.h b/otherarch/sdcpp/src/core/ggml_tensor_utils.h new file mode 100644 index 000000000..65cc5f2f4 --- /dev/null +++ b/otherarch/sdcpp/src/core/ggml_tensor_utils.h @@ -0,0 +1,210 @@ +#ifndef __SD_CORE_GGML_TENSOR_UTILS_H__ +#define __SD_CORE_GGML_TENSOR_UTILS_H__ + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "core/tensor.hpp" +#include "core/util.h" +#include "ggml-backend.h" +#include "ggml.h" +#include "stable-diffusion.h" + +class RNG; + +__STATIC_INLINE__ int align_up_offset(int n, int multiple) { + return (multiple - n % multiple) % multiple; +} + +__STATIC_INLINE__ int align_up(int n, int multiple) { + return n + align_up_offset(n, multiple); +} + +void ggml_ext_im_set_randn_f32(ggml_tensor* tensor, std::shared_ptr rng); + +__STATIC_INLINE__ void ggml_ext_tensor_set_f32(ggml_tensor* tensor, float value, int64_t i0, int64_t i1 = 0, int64_t i2 = 0, int64_t i3 = 0) { + GGML_ASSERT(tensor->nb[0] == sizeof(float)); + *(float*)((char*)(tensor->data) + i3 * tensor->nb[3] + i2 * tensor->nb[2] + i1 * tensor->nb[1] + i0 * tensor->nb[0]) = value; +} + +__STATIC_INLINE__ float ggml_ext_tensor_get_f32(const ggml_tensor* tensor, int64_t i0, int64_t i1 = 0, int64_t i2 = 0, int64_t i3 = 0) { + if (tensor->buffer != nullptr) { + float value; + ggml_backend_tensor_get(tensor, &value, i3 * tensor->nb[3] + i2 * tensor->nb[2] + i1 * tensor->nb[1] + i0 * tensor->nb[0], sizeof(float)); + return value; + } + GGML_ASSERT(tensor->nb[0] == sizeof(float)); + return *(float*)((char*)(tensor->data) + i3 * tensor->nb[3] + i2 * tensor->nb[2] + i1 * tensor->nb[1] + i0 * tensor->nb[0]); +} + +__STATIC_INLINE__ int ggml_ext_tensor_get_i32(const ggml_tensor* tensor, int64_t i0, int64_t i1 = 0, int64_t i2 = 0, int64_t i3 = 0) { + if (tensor->buffer != nullptr) { + int value; + ggml_backend_tensor_get(tensor, &value, i3 * tensor->nb[3] + i2 * tensor->nb[2] + i1 * tensor->nb[1] + i0 * tensor->nb[0], sizeof(int)); + return value; + } + GGML_ASSERT(tensor->nb[0] == sizeof(int)); + return *(int*)((char*)(tensor->data) + i3 * tensor->nb[3] + i2 * tensor->nb[2] + i1 * tensor->nb[1] + i0 * tensor->nb[0]); +} + +__STATIC_INLINE__ ggml_fp16_t ggml_ext_tensor_get_f16(const ggml_tensor* tensor, int64_t i0, int64_t i1 = 0, int64_t i2 = 0, int64_t i3 = 0) { + GGML_ASSERT(tensor->nb[0] == sizeof(ggml_fp16_t)); + return *(ggml_fp16_t*)((char*)(tensor->data) + i3 * tensor->nb[3] + i2 * tensor->nb[2] + i1 * tensor->nb[1] + i0 * tensor->nb[0]); +} + +__STATIC_INLINE__ float sd_image_get_f32(sd_image_t image, int64_t iw, int64_t ih, int64_t ic, bool scale = true) { + float value = *(image.data + ih * image.width * image.channel + iw * image.channel + ic); + if (scale) { + value /= 255.f; + } + return value; +} + +void print_ggml_tensor(ggml_tensor* tensor, bool shape_only = false, const char* mark = ""); + +template +__STATIC_INLINE__ void print_sd_tensor(const sd::Tensor& tensor, bool shape_only = false, const char* mark = "") { + printf("%s: shape(", mark); + for (size_t i = 0; i < static_cast(tensor.dim()); ++i) { + printf("%s%lld", i == 0 ? "" : ", ", static_cast(tensor.shape()[i])); + } + printf(")\n"); + fflush(stdout); + if (shape_only) { + return; + } + if (tensor.empty()) { + return; + } + int range = 3; + std::vector shape = tensor.shape(); + while (shape.size() < 4) { + shape.push_back(1); + } + for (int64_t i3 = 0; i3 < shape[3]; i3++) { + if (i3 >= range && i3 + range < shape[3]) { + continue; + } + for (int64_t i2 = 0; i2 < shape[2]; i2++) { + if (i2 >= range && i2 + range < shape[2]) { + continue; + } + for (int64_t i1 = 0; i1 < shape[1]; i1++) { + if (i1 >= range && i1 + range < shape[1]) { + continue; + } + for (int64_t i0 = 0; i0 < shape[0]; i0++) { + if (i0 >= range && i0 + range < shape[0]) { + continue; + } + size_t offset = static_cast(i0 + shape[0] * (i1 + shape[1] * (i2 + shape[2] * i3))); + printf(" [%lld, %lld, %lld, %lld] = ", static_cast(i3), static_cast(i2), static_cast(i1), static_cast(i0)); + if constexpr (std::is_same_v) { + printf("%f\n", tensor[static_cast(offset)]); + } else if constexpr (std::is_same_v) { + printf("%f\n", ggml_fp16_to_fp32(tensor[static_cast(offset)])); + } else if constexpr (std::is_same_v) { + printf("%d\n", tensor[static_cast(offset)]); + } else if constexpr (std::is_same_v) { + printf("%lld\n", static_cast(tensor[static_cast(offset)])); + } + fflush(stdout); + } + } + } + } +} + +void ggml_ext_tensor_iter( + ggml_tensor* tensor, + const std::function& fn); + +void ggml_ext_tensor_iter( + ggml_tensor* tensor, + const std::function& fn); + +void ggml_ext_tensor_diff( + ggml_tensor* a, + ggml_tensor* b, + float gap = 0.1f); + +ggml_tensor* load_tensor_from_file(ggml_context* ctx, const std::string& file_path); + +__STATIC_INLINE__ float sigmoid(float x) { + return 1 / (1.0f + expf(-x)); +} + +// SPECIAL OPERATIONS WITH TENSORS + +uint8_t* ggml_tensor_to_sd_image(ggml_tensor* input, uint8_t* image_data = nullptr); + +uint8_t* ggml_tensor_to_sd_image(ggml_tensor* input, int idx, bool video = false); + +void sd_image_to_ggml_tensor(sd_image_t image, + ggml_tensor* tensor, + bool scale = true); + +void ggml_ext_tensor_apply_mask(ggml_tensor* image_data, + ggml_tensor* mask, + ggml_tensor* output, + float masked_value = 0.5f); + +float ggml_ext_tensor_mean(ggml_tensor* src); + +// a = a+b +void ggml_ext_tensor_add_inplace(ggml_tensor* a, ggml_tensor* b); + +void ggml_ext_tensor_scale_inplace(ggml_tensor* src, float scale); + +void ggml_ext_tensor_clamp_inplace(ggml_tensor* src, float min, float max); + +ggml_tensor* ggml_ext_tensor_concat(ggml_context* ctx, + ggml_tensor* a, + ggml_tensor* b, + int dim); + +// convert values from [0, 1] to [-1, 1] +void scale_to_minus1_1(ggml_tensor* src); + +// convert values from [-1, 1] to [0, 1] +void scale_to_0_1(ggml_tensor* src); + +ggml_tensor* vector_to_ggml_tensor(ggml_context* ctx, + const std::vector& vec); + +ggml_tensor* vector_to_ggml_tensor_i32(ggml_context* ctx, + const std::vector& vec); + +std::vector arange(float start, float end, float step = 1.f); + +// Ref: https://github.com/CompVis/stable-diffusion/blob/main/ldm/modules/diffusionmodules/util.py#L151 +std::vector timestep_embedding(std::vector timesteps, + int dim, + int max_period = 10000, + bool flip_sin_to_cos = true, + float scale = 1.f); + +void set_timestep_embedding(std::vector timesteps, + ggml_tensor* embedding, + int dim, + int max_period = 10000); + +void set_timestep_embedding(std::vector timesteps, + sd::Tensor* embedding, + int dim, + int max_period = 10000); + +ggml_tensor* new_timestep_embedding(ggml_context* ctx, + std::vector timesteps, + int dim, + int max_period = 10000); + +size_t ggml_tensor_num(ggml_context* ctx); + +#endif // __SD_CORE_GGML_TENSOR_UTILS_H__ diff --git a/otherarch/sdcpp/src/core/layer_split_partition.cpp b/otherarch/sdcpp/src/core/layer_split_partition.cpp index 654b3056b..e19f47b26 100644 --- a/otherarch/sdcpp/src/core/layer_split_partition.cpp +++ b/otherarch/sdcpp/src/core/layer_split_partition.cpp @@ -145,19 +145,24 @@ namespace sd { std::vector backend_capacities = graph_cut_layer_split_backend_capacities(split_backends, backend_vram_limits, primary_backend_vram_limit); + // Existing placements may already occupy the reported free VRAM. Reuse + // them; execution checks missing weights and reclaims memory as needed. + const bool reuse_assignments = std::all_of(seen_params.begin(), seen_params.end(), [&](ggml_tensor* param) { + return param_assignments.count(param) != 0; + }); std::vector backend_by_segment(plan.segments.size(), split_backends[0]); size_t current_backend = 0; int64_t current_used = 0; for (size_t seg_idx = 0; seg_idx < plan.segments.size(); seg_idx++) { int64_t bytes = segment_param_bytes[seg_idx]; - while (current_backend + 1 < split_backends.size() && + while (!reuse_assignments && current_backend + 1 < split_backends.size() && bytes > 0 && current_used + bytes > backend_capacities[current_backend]) { current_backend++; current_used = 0; } - if (bytes > 0 && current_used + bytes > backend_capacities[current_backend]) { + if (!reuse_assignments && bytes > 0 && current_used + bytes > backend_capacities[current_backend]) { LOG_ERROR("%s graph-cut layer split: segment %zu needs %.1f MB on %s, but only %.1f MB is available under current VRAM limits", desc, seg_idx, @@ -167,7 +172,6 @@ namespace sd { return false; } current_used += bytes; - backend_by_segment[seg_idx] = split_backends[current_backend]; for (ggml_tensor* param : segment_params[seg_idx]) { ggml_backend_t target_backend = split_backends[current_backend]; @@ -186,12 +190,16 @@ namespace sd { ggml_get_name(param)); return false; } - size_t backend_idx = (size_t)std::distance(split_backends.begin(), backend_it); + size_t backend_idx = (size_t)std::distance(split_backends.begin(), backend_it); + if (reuse_assignments) { + current_backend = backend_idx; + } assignment.first_segment_by_backend[backend_idx] = std::min(assignment.first_segment_by_backend[backend_idx], seg_idx); assignment.last_segment_by_backend[backend_idx] = std::max(assignment.last_segment_by_backend[backend_idx], seg_idx + 1); assignment.tensors_by_backend[backend_idx].push_back(param); assignment.bytes_by_backend[backend_idx] += (int64_t)ggml_nbytes(param); } + backend_by_segment[seg_idx] = split_backends[current_backend]; } const int n_nodes = ggml_graph_n_nodes(gf); @@ -243,13 +251,13 @@ namespace sd { assignment.tensors_by_backend[i].size(), assignment.bytes_by_backend[i] / (1024.0 * 1024.0)); } else { - LOG_DEBUG("%s graph-cut layer split: %s <- segments [%zu, %zu), %zu tensors, %.1f MB", - desc, - layer_split_backend_device_display_name(split_backends[i]).c_str(), - first_segment, - last_segment, - assignment.tensors_by_backend[i].size(), - assignment.bytes_by_backend[i] / (1024.0 * 1024.0)); + LOG_VERBOSE("%s graph-cut layer split: %s <- segments [%zu, %zu), %zu tensors, %.1f MB", + desc, + layer_split_backend_device_display_name(split_backends[i]).c_str(), + first_segment, + last_segment, + assignment.tensors_by_backend[i].size(), + assignment.bytes_by_backend[i] / (1024.0 * 1024.0)); } } } diff --git a/otherarch/sdcpp/src/core/runner_cache.cpp b/otherarch/sdcpp/src/core/runner_cache.cpp new file mode 100644 index 000000000..095621068 --- /dev/null +++ b/otherarch/sdcpp/src/core/runner_cache.cpp @@ -0,0 +1,211 @@ +#include "core/runner_cache.h" + +#include +#include +#include + +#include "core/ggml_graph_cut.h" +#include "core/util.h" + +namespace sd { + static std::unordered_set cache_graph_tensors(ggml_cgraph* graph) { + std::unordered_set tensors; + for (int i = 0; i < ggml_graph_n_nodes(graph); ++i) { + tensors.insert(ggml_graph_node(graph, i)); + } + for (int i = 0; i < ggml_graph_cut::leaf_count(graph); ++i) { + tensors.insert(ggml_graph_cut::leaf_tensor(graph, i)); + } + return tensors; + } + + CachedTensor::~CachedTensor() { + ggml_backend_buffer_free(buffer); + ggml_free(context); + } + + std::unique_ptr CachedTensor::copy(ggml_backend_t backend, + const std::string& name, + ggml_tensor* source) { + if (ggml_graph_cut::tensor_buffer(source) == nullptr) { + return nullptr; + } + auto entry = std::make_unique(); + entry->context = ggml_init({2 * ggml_tensor_overhead(), nullptr, true}); + if (entry->context == nullptr) { + return nullptr; + } + entry->tensor = ggml_dup_tensor(entry->context, source); + // Cut views are rebound with their original strides and offsets. + std::copy(std::begin(source->nb), std::end(source->nb), std::begin(entry->tensor->nb)); + ggml_set_name(entry->tensor, name.c_str()); + entry->buffer = ggml_backend_alloc_ctx_tensors(entry->context, backend); + if (entry->buffer == nullptr) { + return nullptr; + } + if (source->view_src != nullptr || !ggml_is_contiguous(source) || source->buffer == nullptr) { + std::vector data(ggml_nbytes(source)); + ggml_backend_tensor_get(source, data.data(), 0, data.size()); + ggml_backend_tensor_set(entry->tensor, data.data(), 0, data.size()); + } else { + ggml_backend_tensor_copy(source, entry->tensor); + } + return entry; + } + + static ggml_tensor* cached_tensor(const CachedTensors& tensors, const std::string& name) { + auto entry = tensors.find(name); + return entry == tensors.end() ? nullptr : entry->second->tensor; + } + + static size_t resident_bytes(const CachedTensors& tensors, ggml_backend_dev_t device) { + size_t bytes = 0; + for (const auto& entry : tensors) { + auto buffer = entry.second->buffer; + if (!ggml_backend_buffer_is_host(buffer) && + ggml_backend_buft_get_device(ggml_backend_buffer_get_type(buffer)) == device) { + const size_t size = ggml_backend_buffer_get_size(buffer); + bytes = size > SIZE_MAX - bytes ? SIZE_MAX : bytes + size; + } + } + return bytes; + } + + ggml_tensor* RunnerCache::get(const std::string& name) const { + return cached_tensor(committed_, name); + } + + void RunnerCache::stage(const std::string& name, ggml_tensor* tensor) { + if (tensor != nullptr) { + ggml_set_output(tensor); + outputs_[name] = tensor; + } + } + + size_t RunnerCache::pending_bytes(ggml_cgraph* graph) const { + if (outputs_.empty()) { + return 0; + } + auto tensors = cache_graph_tensors(graph); + auto buft = ggml_backend_get_default_buffer_type(backend_); + size_t bytes = 0; + for (const auto& output : outputs_) { + if (pending_.count(output.first) || !tensors.count(output.second)) { + continue; + } + const size_t size = GGML_PAD(ggml_backend_buft_get_alloc_size(buft, output.second), + ggml_backend_buft_get_alignment(buft)); + bytes = size > SIZE_MAX - bytes ? SIZE_MAX : bytes + size; + } + return bytes; + } + + size_t RunnerCache::resident_bytes(ggml_backend_dev_t device) const { + const size_t committed = sd::resident_bytes(committed_, device); + const size_t pending = sd::resident_bytes(pending_, device); + return pending > SIZE_MAX - committed ? SIZE_MAX : committed + pending; + } + + bool RunnerCache::capture(ggml_cgraph* graph) { + if (outputs_.empty()) { + return true; + } + const auto tensors = cache_graph_tensors(graph); + for (const auto& output : outputs_) { + if (pending_.count(output.first) || !tensors.count(output.second)) { + continue; + } + GGML_ASSERT(ggml_is_contiguous(output.second)); + auto entry = CachedTensor::copy(backend_, output.first, output.second); + if (entry == nullptr) { + return false; + } + pending_[output.first] = std::move(entry); + } + ggml_backend_synchronize(backend_); + return true; + } + + void RunnerCache::graph_end(bool success) { + // Graph inputs can still reference the previous generation until graph end. + if (success) { + for (auto& entry : pending_) { + committed_[entry.first] = std::move(entry.second); + } + } + pending_.clear(); + outputs_.clear(); + } + + void RunnerCache::clear() { + graph_end(false); + committed_.clear(); + } + + ggml_tensor* GraphCutTensorCache::get(const std::string& name) const { + return cached_tensor(tensors_, name); + } + + size_t GraphCutTensorCache::resident_bytes(ggml_backend_dev_t device) const { + return sd::resident_bytes(tensors_, device); + } + + size_t GraphCutTensorCache::estimate_output_bytes( + ggml_cgraph* graph, + const ggml_graph_cut::Segment& segment) const { + ggml_backend_buffer_type_t buffer_type = + ggml_backend_get_default_buffer_type(backend_); + if (buffer_type == nullptr) { + return SIZE_MAX; + } + const size_t alignment = ggml_backend_buft_get_alignment(buffer_type); + size_t total_size = 0; + for (size_t output_idx = 0; output_idx < segment.output_node_indices.size(); ++output_idx) { + ggml_tensor* output = ggml_graph_cut::output_tensor(graph, segment, output_idx); + if (output == nullptr || !ggml_graph_cut::is_graph_cut_tensor(output) || + !segment.future_cut_names.count(output->name)) { + continue; + } + ggml_tensor* source = ggml_graph_cut::cache_source_tensor(output); + const size_t tensor_size = GGML_PAD( + ggml_backend_buft_get_alloc_size(buffer_type, source), alignment); + total_size = tensor_size > SIZE_MAX - total_size ? SIZE_MAX : total_size + tensor_size; + } + return total_size; + } + + void GraphCutTensorCache::prune(const std::unordered_set& keep_names) { + for (auto it = tensors_.begin(); it != tensors_.end();) { + it = keep_names.count(it->first) ? std::next(it) : tensors_.erase(it); + } + } + + bool GraphCutTensorCache::capture(ggml_cgraph* graph, + const ggml_graph_cut::Segment& segment, + const char* log_desc) { + size_t copied_bytes = 0; + size_t copied_count = 0; + for (int index : segment.output_node_indices) { + auto output = ggml_graph_node(graph, index); + if (!ggml_graph_cut::is_graph_cut_tensor(output) || + !segment.future_cut_names.count(output->name)) { + continue; + } + auto entry = CachedTensor::copy(backend_, output->name, ggml_graph_cut::cache_source_tensor(output)); + if (entry == nullptr) { + LOG_ERROR("%s failed to capture graph cut tensor: %s", log_desc, output->name); + return false; + } + const size_t size = ggml_backend_buffer_get_size(entry->buffer); + copied_bytes = size > SIZE_MAX - copied_bytes ? SIZE_MAX : copied_bytes + size; + ++copied_count; + tensors_[output->name] = std::move(entry); + } + ggml_backend_synchronize(backend_); + if (copied_count > 0) { + LOG_DEBUG("%s graph cut cache added %6.2f MB (%zu tensors)", + log_desc, copied_bytes / (1024.f * 1024.f), copied_count); + } + return true; + } +} diff --git a/otherarch/sdcpp/src/core/runner_cache.h b/otherarch/sdcpp/src/core/runner_cache.h new file mode 100644 index 000000000..bd331a4bb --- /dev/null +++ b/otherarch/sdcpp/src/core/runner_cache.h @@ -0,0 +1,66 @@ +#ifndef __SD_CORE_RUNNER_CACHE_H__ +#define __SD_CORE_RUNNER_CACHE_H__ + +#include +#include +#include +#include + +#include "ggml-backend.h" + +namespace sd::ggml_graph_cut { + struct Segment; +} + +namespace sd { + struct CachedTensor { + ggml_context* context = nullptr; + ggml_backend_buffer_t buffer = nullptr; + ggml_tensor* tensor = nullptr; + ~CachedTensor(); + static std::unique_ptr copy(ggml_backend_t backend, + const std::string& name, + ggml_tensor* source); + }; + using CachedTensors = std::map>; + + class RunnerCache { + ggml_backend_t backend_; + CachedTensors committed_; + CachedTensors pending_; + std::map outputs_; + + public: + explicit RunnerCache(ggml_backend_t backend) + : backend_(backend) {} + RunnerCache(const RunnerCache&) = delete; + RunnerCache& operator=(const RunnerCache&) = delete; + + ggml_tensor* get(const std::string& name) const; + void stage(const std::string& name, ggml_tensor* tensor); + const std::map& outputs() const { return outputs_; } + size_t pending_bytes(ggml_cgraph* graph) const; + size_t resident_bytes(ggml_backend_dev_t device) const; + bool capture(ggml_cgraph* graph); + void graph_end(bool success); + void clear(); + }; + + class GraphCutTensorCache { + ggml_backend_t backend_; + CachedTensors tensors_; + + public: + explicit GraphCutTensorCache(ggml_backend_t backend) + : backend_(backend) {} + ggml_tensor* get(const std::string& name) const; + size_t resident_bytes(ggml_backend_dev_t device) const; + size_t estimate_output_bytes(ggml_cgraph* graph, + const ggml_graph_cut::Segment& segment) const; + bool capture(ggml_cgraph* graph, const ggml_graph_cut::Segment& segment, const char* log_desc); + void prune(const std::unordered_set& keep_names); + void clear() { tensors_.clear(); } + }; +} + +#endif // __SD_CORE_RUNNER_CACHE_H__ diff --git a/otherarch/sdcpp/src/core/segment_graph_bindings.cpp b/otherarch/sdcpp/src/core/segment_graph_bindings.cpp new file mode 100644 index 000000000..3c40a3b03 --- /dev/null +++ b/otherarch/sdcpp/src/core/segment_graph_bindings.cpp @@ -0,0 +1,137 @@ +#include "core/segment_graph_bindings.h" + +#include +#include + +#include "core/ggml_graph_cut.h" +#include "core/runner_cache.h" +#include "core/util.h" +#include "ggml.h" + +namespace sd { + SegmentGraphBindings::SegmentGraphBindings(GraphCutTensorCache& tensor_cache, + const ggml_graph_cut::Plan& plan, + ggml_cgraph* graph) + : tensor_cache_(tensor_cache), + graph_(graph) { + GGML_ASSERT(graph_ != nullptr); + for (int i = 0; i < ggml_graph_n_nodes(graph_); ++i) { + ggml_tensor* tensor = ggml_graph_node(graph_, i); + Topology topology{tensor->op, {}, tensor->view_src, tensor->flags}; + std::copy(std::begin(tensor->src), std::end(tensor->src), topology.sources.begin()); + topology_[tensor] = topology; + } + for (const auto& segment : plan.segments) { + for (const auto& input : segment.input_refs) { + if (input.type != ggml_graph_cut::Segment::INPUT_EXTERNAL) { + continue; + } + ggml_tensor* tensor = ggml_graph_cut::input_tensor(graph_, input); + if (tensor == nullptr || tensor->buffer == nullptr) { + continue; + } + external_bindings_[tensor] = {tensor->buffer, tensor->data, tensor->extra}; + } + } + } + + void SegmentGraphBindings::reset(const ggml_graph_cut::Segment& segment) { + restore(); + for (const auto& input : segment.input_refs) { + ggml_tensor* tensor = ggml_graph_cut::input_tensor(graph_, input); + if (tensor == nullptr) { + continue; + } + switch (input.type) { + case ggml_graph_cut::Segment::INPUT_PREVIOUS_CUT: + tensor->buffer = nullptr; + tensor->data = nullptr; + tensor->extra = nullptr; + break; + case ggml_graph_cut::Segment::INPUT_EXTERNAL: { + auto binding = external_bindings_.find(tensor); + if (binding != external_bindings_.end()) { + tensor->buffer = binding->second.buffer; + tensor->data = binding->second.data; + tensor->extra = binding->second.extra; + } else { + tensor->buffer = nullptr; + tensor->data = nullptr; + tensor->extra = nullptr; + } + break; + } + case ggml_graph_cut::Segment::INPUT_PARAM: + break; + } + } + + for (int node_index : segment.internal_node_indices) { + ggml_tensor* node = ggml_graph_node(graph_, node_index); + if (node == nullptr) { + continue; + } + node->buffer = nullptr; + node->data = nullptr; + node->extra = nullptr; + } + } + + void SegmentGraphBindings::restore() { + for (const auto& entry : topology_) { + entry.first->op = entry.second.op; + entry.first->view_src = entry.second.view_source; + entry.first->flags = entry.second.flags; + std::copy(entry.second.sources.begin(), entry.second.sources.end(), std::begin(entry.first->src)); + } + } + + bool SegmentGraphBindings::bind_cached_inputs( + const ggml_graph_cut::Segment& segment, + const char* log_desc) { + std::unordered_map cached_view_sources; + for (const auto& input : segment.input_refs) { + if (input.type != ggml_graph_cut::Segment::INPUT_PREVIOUS_CUT) { + continue; + } + ggml_tensor* input_tensor = ggml_graph_cut::input_tensor(graph_, input); + if (input_tensor == nullptr) { + continue; + } + ggml_tensor* cached_tensor = tensor_cache_.get(input.display_name); + if (cached_tensor == nullptr) { + LOG_ERROR("%s missing graph cut cache tensor: %s", + log_desc, + input.display_name.c_str()); + return false; + } + if (input_tensor->view_src != nullptr) { + cached_view_sources[topology_.at(input_tensor).view_source] = cached_tensor; + input_tensor->view_src = cached_tensor; + input_tensor->buffer = nullptr; + input_tensor->data = cached_tensor->data == nullptr + ? nullptr + : static_cast(static_cast(cached_tensor->data) + + input_tensor->view_offs); + input_tensor->extra = cached_tensor->extra; + } else { + input_tensor->buffer = cached_tensor->buffer; + input_tensor->data = cached_tensor->data; + input_tensor->extra = cached_tensor->extra; + } + for (int source_index = 0; source_index < GGML_MAX_SRC; ++source_index) { + input_tensor->src[source_index] = nullptr; + } + input_tensor->op = GGML_OP_NONE; + } + // ggml flattens view chains, so descendants also need the cached root. + for (int node_index : segment.internal_node_indices) { + ggml_tensor* node = ggml_graph_node(graph_, node_index); + auto cached_source = cached_view_sources.find(topology_.at(node).view_source); + if (cached_source != cached_view_sources.end()) { + node->view_src = cached_source->second; + } + } + return true; + } +} diff --git a/otherarch/sdcpp/src/core/segment_graph_bindings.h b/otherarch/sdcpp/src/core/segment_graph_bindings.h new file mode 100644 index 000000000..08aba156b --- /dev/null +++ b/otherarch/sdcpp/src/core/segment_graph_bindings.h @@ -0,0 +1,52 @@ +#ifndef __SD_CORE_SEGMENT_GRAPH_BINDINGS_H__ +#define __SD_CORE_SEGMENT_GRAPH_BINDINGS_H__ + +#include +#include +#include "ggml.h" + +struct ggml_backend_buffer; +struct ggml_cgraph; +struct ggml_tensor; + +namespace sd { + class GraphCutTensorCache; + + namespace ggml_graph_cut { + struct Plan; + struct Segment; + } + + class SegmentGraphBindings { + public: + SegmentGraphBindings(GraphCutTensorCache& tensor_cache, + const ggml_graph_cut::Plan& plan, + ggml_cgraph* graph); + + void reset(const ggml_graph_cut::Segment& segment); + void restore(); + ~SegmentGraphBindings() { restore(); } + bool bind_cached_inputs(const ggml_graph_cut::Segment& segment, + const char* log_desc); + + private: + struct ExternalBinding { + ggml_backend_buffer* buffer = nullptr; + void* data = nullptr; + void* extra = nullptr; + }; + + GraphCutTensorCache& tensor_cache_; + ggml_cgraph* graph_ = nullptr; + std::unordered_map external_bindings_; + struct Topology { + ggml_op op; + std::array sources; + ggml_tensor* view_source; + int flags; + }; + std::unordered_map topology_; + }; +} + +#endif // __SD_CORE_SEGMENT_GRAPH_BINDINGS_H__ diff --git a/otherarch/sdcpp/src/core/segment_weight_pipeline.cpp b/otherarch/sdcpp/src/core/segment_weight_pipeline.cpp new file mode 100644 index 000000000..f769794df --- /dev/null +++ b/otherarch/sdcpp/src/core/segment_weight_pipeline.cpp @@ -0,0 +1,205 @@ +#include "core/segment_weight_pipeline.h" + +#include + +#include "core/ggml_graph_cut.h" +#include "device_residency_manager.h" + +namespace sd { + static ggml_tensor* canonical_param( + ggml_tensor* tensor, + const std::unordered_set& params) { + for (ggml_tensor* current = tensor; current != nullptr; current = current->view_src) { + if (params.find(current) != params.end()) { + return current; + } + } + return nullptr; + } + + SegmentWeightPipeline::SegmentWeightPipeline( + const std::shared_ptr& residency_manager, + ggml_backend_t compute_backend, + uintptr_t owner_id, + ggml_cgraph* graph, + const ggml_graph_cut::Plan& plan, + const std::unordered_set& params, + bool enabled) + : residency_manager_(residency_manager), + compute_backend_(compute_backend), + owner_id_(owner_id), + enabled_(enabled && residency_manager != nullptr) { + segment_params_.resize(plan.segments.size()); + for (size_t segment_index = 0; segment_index < plan.segments.size(); ++segment_index) { + std::unordered_set seen; + for (ggml_tensor* tensor : + ggml_graph_cut::param_tensors(graph, plan.segments[segment_index])) { + ggml_tensor* param = canonical_param(tensor, params); + if (param != nullptr && seen.insert(param).second) { + segment_params_[segment_index].push_back(param); + } + } + } + } + + SegmentWeightPipeline::~SegmentWeightPipeline() { + segment_end(); + clear(); + } + + size_t SegmentWeightPipeline::next_parameter_segment(size_t segment_index) const { + for (size_t next = segment_index + 1; next < segment_params_.size(); ++next) { + if (!segment_params_[next].empty()) { + return next; + } + } + return SIZE_MAX; + } + + std::vector> SegmentWeightPipeline::preferred_eviction_order() const { + return {segment_params_.rbegin(), segment_params_.rend()}; + } + + void SegmentWeightPipeline::disable() { + clear(); + enabled_ = false; + } + + void SegmentWeightPipeline::activate(size_t segment_index) { + if (!enabled_ || queued_segment_ == SIZE_MAX || queued_segment_ != segment_index) { + return; + } + + auto manager = residency_manager_.lock(); + if (manager == nullptr || + !manager->activate_prefetched_params(owner_id_, queued_params_)) { + disable(); + return; + } + queued_params_.clear(); + queued_segment_ = SIZE_MAX; + } + + bool SegmentWeightPipeline::ensure_segment_capacity( + size_t segment_index, + const std::vector& requests) { + if (segment_index >= segment_params_.size()) { + return false; + } + auto manager = residency_manager_.lock(); + if (manager == nullptr) { + return segment_params_[segment_index].empty(); + } + std::vector protected_params = segment_params_[segment_index]; + protected_params.insert(protected_params.end(), queued_params_.begin(), queued_params_.end()); + for (const auto& request : requests) { + if (!manager->ensure_compute_backend_capacity(request, segment_params_[segment_index], + preferred_eviction_order(), protected_params)) { + return false; + } + } + return true; + } + + bool SegmentWeightPipeline::segment_start(size_t segment_index, const std::function& ensure_capacity) { + GGML_ASSERT(pinned_params_.empty()); + activate(segment_index); + if (!ensure_capacity()) { + return false; + } + auto manager = residency_manager_.lock(); + if (manager == nullptr) { + return segment_params_[segment_index].empty(); + } + if (!manager->prepare_params(segment_params_[segment_index])) { + return false; + } + pinned_params_ = segment_params_[segment_index]; + return true; + } + + void SegmentWeightPipeline::segment_end() { + if (auto manager = residency_manager_.lock()) { + manager->release_compute_backend_params(pinned_params_); + } + pinned_params_.clear(); + } + + void SegmentWeightPipeline::enqueue_next( + size_t segment_index, + const DeviceMemoryRequest& request) { + if (!enabled_ || queued_segment_ != SIZE_MAX) { + return; + } + + const size_t next_segment = next_parameter_segment(segment_index); + if (next_segment == SIZE_MAX) { + return; + } + + std::unordered_set active_params( + segment_params_[segment_index].begin(), + segment_params_[segment_index].end()); + std::vector params; + params.reserve(segment_params_[next_segment].size()); + for (ggml_tensor* param : segment_params_[next_segment]) { + if (active_params.find(param) == active_params.end()) { + params.push_back(param); + } + } + if (params.empty()) { + return; + } + + auto manager = residency_manager_.lock(); + if (manager == nullptr) { + disable(); + return; + } + const WeightResidencyInfo residency = + manager->inspect_compute_backend_params(params); + if (residency.missing_bytes == 0) { + return; + } + if (!residency.async_prefetch_supported) { + disable(); + return; + } + + DeviceMemoryRequest backend_request = request; + backend_request.compute_backend = compute_backend_; + backend_request.owner_id = owner_id_; + std::vector protected_params = segment_params_[segment_index]; + protected_params.insert(protected_params.end(), params.begin(), params.end()); + if (!manager->ensure_compute_backend_capacity(backend_request, + params, + preferred_eviction_order(), + protected_params)) { + disable(); + return; + } + switch (manager->prefetch_params(owner_id_, params)) { + case WeightPrefetchResult::Scheduled: + queued_params_ = std::move(params); + queued_segment_ = next_segment; + return; + case WeightPrefetchResult::AlreadyResident: + return; + case WeightPrefetchResult::Unsupported: + disable(); + return; + case WeightPrefetchResult::Failed: + disable(); + return; + } + disable(); + } + + void SegmentWeightPipeline::clear() { + if (auto manager = residency_manager_.lock()) { + manager->clear_prefetched_params(owner_id_); + } + queued_params_.clear(); + queued_segment_ = SIZE_MAX; + } +} diff --git a/otherarch/sdcpp/src/core/segment_weight_pipeline.h b/otherarch/sdcpp/src/core/segment_weight_pipeline.h new file mode 100644 index 000000000..598bf9f40 --- /dev/null +++ b/otherarch/sdcpp/src/core/segment_weight_pipeline.h @@ -0,0 +1,62 @@ +#ifndef __SD_CORE_SEGMENT_WEIGHT_PIPELINE_H__ +#define __SD_CORE_SEGMENT_WEIGHT_PIPELINE_H__ + +#include +#include +#include +#include +#include +#include + +#include "ggml-backend.h" + +struct DeviceMemoryRequest; +struct DeviceResidencyManager; +struct ggml_cgraph; +struct ggml_tensor; + +namespace sd::ggml_graph_cut { + struct Plan; +} + +namespace sd { + class SegmentWeightPipeline { + private: + std::weak_ptr residency_manager_; + ggml_backend_t compute_backend_ = nullptr; + uintptr_t owner_id_ = 0; + std::vector> segment_params_; + std::vector queued_params_; + std::vector pinned_params_; + size_t queued_segment_ = SIZE_MAX; + bool enabled_ = true; + + size_t next_parameter_segment(size_t segment_index) const; + std::vector> preferred_eviction_order() const; + void disable(); + void activate(size_t segment_index); + void clear(); + + public: + SegmentWeightPipeline( + const std::shared_ptr& residency_manager, + ggml_backend_t compute_backend, + uintptr_t owner_id, + ggml_cgraph* graph, + const ggml_graph_cut::Plan& plan, + const std::unordered_set& params, + bool enabled = true); + ~SegmentWeightPipeline(); + + const std::vector& params(size_t index) const { return segment_params_[index]; } + bool ensure_segment_capacity(size_t segment_index, + const std::vector& requests); + bool segment_start(size_t segment_index, const std::function& ensure_capacity); + void segment_end(); + // Prefetch is best effort; segment_start falls back to synchronous loading. + void enqueue_next(size_t segment_index, + const DeviceMemoryRequest& request); + }; +} + +#endif // __SD_CORE_SEGMENT_WEIGHT_PIPELINE_H__ diff --git a/otherarch/sdcpp/src/core/tensor_ggml.hpp b/otherarch/sdcpp/src/core/tensor_ggml.hpp index 774574f79..66ea22206 100644 --- a/otherarch/sdcpp/src/core/tensor_ggml.hpp +++ b/otherarch/sdcpp/src/core/tensor_ggml.hpp @@ -9,6 +9,7 @@ #include #include "core/tensor.hpp" +#include "ggml-backend.h" #include "ggml.h" namespace sd { @@ -54,10 +55,28 @@ namespace sd { GGML_ABORT("ggml tensor type does not match sd::Tensor type"); } Tensor result(shape_from_ggml(tensor)); - if (tensor->buffer != nullptr) { - ggml_backend_tensor_get(tensor, result.data(), 0, ggml_nbytes(tensor)); + std::vector strided_data; + void* destination = result.data(); + if (!ggml_is_contiguous(tensor)) { + strided_data.resize(ggml_nbytes(tensor)); + destination = strided_data.data(); + } + auto buffer = tensor->view_src != nullptr ? tensor->view_src->buffer : tensor->buffer; + if (buffer != nullptr) { + ggml_backend_tensor_get(tensor, destination, 0, ggml_nbytes(tensor)); } else { - std::memcpy(result.data(), tensor->data, ggml_nbytes(tensor)); + std::memcpy(destination, tensor->data, ggml_nbytes(tensor)); + } + if (!strided_data.empty()) { + for (int64_t i = 0; i < result.numel(); ++i) { + int64_t index = i; + size_t offset = 0; + for (int d = 0; d < GGML_MAX_DIMS; ++d) { + offset += static_cast(index % tensor->ne[d]) * tensor->nb[d]; + index /= tensor->ne[d]; + } + std::memcpy(result.data() + i, strided_data.data() + offset, sizeof(T)); + } } return result; } diff --git a/otherarch/sdcpp/src/core/util.cpp b/otherarch/sdcpp/src/core/util.cpp index 4672b139f..2ba895d90 100644 --- a/otherarch/sdcpp/src/core/util.cpp +++ b/otherarch/sdcpp/src/core/util.cpp @@ -1,6 +1,7 @@ #include "core/util.h" #include #include +#include #include #include #include @@ -16,6 +17,7 @@ #include #include #include +#include "core/ggml_tensor_utils.h" #include "runtime/preprocessing.hpp" #include @@ -518,7 +520,7 @@ bool parse_strict_bool(const std::string& text, bool& value) { } // { kcpp -static int sdloglevel = 0; //-1 = hide all, 0 = normal, 1 = showall +static int sdloglevel = INT_MAX; // -1 = hide all, 0 = normal, 1 = showall, INT_MAX = sdcpp static bool sdquiet = false; // } kcpp @@ -618,18 +620,17 @@ void* sd_log_cb_data = nullptr; #define LOG_BUFFER_SIZE 4096 -void log_message(const char* format, ...) { - if (sdloglevel>0) { +void log_printf(sd_log_level_t level, const char* file, int line, const char* format, ...) { + if (sdloglevel > 0 && sdloglevel != INT_MAX) { printf("\n"); va_list args; va_start(args, format); vprintf(format, args); va_end(args); fflush(stdout); + return; } -} -void log_printf(sd_log_level_t level, const char* file, int line, const char* format, ...) { va_list args; va_start(args, format); @@ -651,6 +652,25 @@ void log_printf(sd_log_level_t level, const char* file, int line, const char* fo va_end(args); } +void sd_ggml_log_callback(ggml_log_level level, const char* text, void*) { + switch (level) { + case GGML_LOG_LEVEL_DEBUG: + LOG_VERBOSE(text); + break; + case GGML_LOG_LEVEL_INFO: + LOG_INFO(text); + break; + case GGML_LOG_LEVEL_WARN: + LOG_WARN(text); + break; + case GGML_LOG_LEVEL_ERROR: + LOG_ERROR(text); + break; + default: + LOG_VERBOSE(text); + } +} + void sd_set_log_callback(sd_log_cb_t cb, void* data) { sd_log_cb = cb; sd_log_cb_data = data; diff --git a/otherarch/sdcpp/src/core/util.h b/otherarch/sdcpp/src/core/util.h index 07968cbcd..bb4dc9d7a 100644 --- a/otherarch/sdcpp/src/core/util.h +++ b/otherarch/sdcpp/src/core/util.h @@ -11,6 +11,14 @@ #include "ggml-backend.h" #include "stable-diffusion.h" +#ifndef __STATIC_INLINE__ +#define __STATIC_INLINE__ static inline +#endif + +#ifndef SD_UNUSED +#define SD_UNUSED(x) (void)(x) +#endif + #define SAFE_STR(s) ((s) ? (s) : "") #define BOOL_STR(b) ((b) ? "true" : "false") @@ -36,7 +44,6 @@ std::string sd_get_u8path(const std::string& file_path); sd_image_t tensor_to_sd_image(const sd::Tensor& tensor, int frame_index = 0); - sd::Tensor sd_image_to_tensor(sd_image_t image, int target_width = -1, int target_height = -1, @@ -82,6 +89,7 @@ void pretty_progress(int step, int steps, float time); void pretty_bytes_progress(int step, int steps, uint64_t bytes_processed, float elapsed_seconds); void log_printf(sd_log_level_t level, const char* file, int line, const char* format, ...); +void sd_ggml_log_callback(ggml_log_level level, const char* text, void*); ggml_type sd_type_to_ggml_type(sd_type_t sdtype); @@ -107,13 +115,9 @@ void* sd_get_backend_eval_callback_data(); // test if the backend is a specific one, e.g. "CUDA", "ROCm", "Vulkan" etc. bool sd_backend_is(ggml_backend_t backend, const std::string& name); -void log_message(const char* format, ...); -#define LOG_DEBUG(...) log_message(__VA_ARGS__) -#define LOG_INFO(...) log_message(__VA_ARGS__) -#define LOG_WARN(...) log_message(__VA_ARGS__) -#define LOG_ERROR(...) log_message(__VA_ARGS__) -// #define LOG_DEBUG(format, ...) log_printf(SD_LOG_DEBUG, __FILE__, __LINE__, format, ##__VA_ARGS__) -// #define LOG_INFO(format, ...) log_printf(SD_LOG_INFO, __FILE__, __LINE__, format, ##__VA_ARGS__) -// #define LOG_WARN(format, ...) log_printf(SD_LOG_WARN, __FILE__, __LINE__, format, ##__VA_ARGS__) -// #define LOG_ERROR(format, ...) log_printf(SD_LOG_ERROR, __FILE__, __LINE__, format, ##__VA_ARGS__) +#define LOG_DEBUG(format, ...) log_printf(SD_LOG_DEBUG, __FILE__, __LINE__, format, ##__VA_ARGS__) +#define LOG_VERBOSE(format, ...) log_printf(SD_LOG_VERBOSE, __FILE__, __LINE__, format, ##__VA_ARGS__) +#define LOG_INFO(format, ...) log_printf(SD_LOG_INFO, __FILE__, __LINE__, format, ##__VA_ARGS__) +#define LOG_WARN(format, ...) log_printf(SD_LOG_WARN, __FILE__, __LINE__, format, ##__VA_ARGS__) +#define LOG_ERROR(format, ...) log_printf(SD_LOG_ERROR, __FILE__, __LINE__, format, ##__VA_ARGS__) #endif // __SD_CORE_UTIL_H__ diff --git a/otherarch/sdcpp/src/detailer.cpp b/otherarch/sdcpp/src/detailer.cpp index fa9884be3..2523dd513 100644 --- a/otherarch/sdcpp/src/detailer.cpp +++ b/otherarch/sdcpp/src/detailer.cpp @@ -715,7 +715,7 @@ std::vector ADetailerGGML::predict(sd_image_t image, LetterboxInput input = make_letterbox_input(image, params.input_size); int64_t start = ggml_time_ms(); sd::Tensor raw = detector->compute(n_threads, input.tensor); - detector->free_compute_buffer(); + detector->runner_end(); if (raw.empty()) { LOG_ERROR("YOLOv8 detector inference failed"); return {}; diff --git a/otherarch/sdcpp/src/device_residency_manager.h b/otherarch/sdcpp/src/device_residency_manager.h new file mode 100644 index 000000000..2fb2a342d --- /dev/null +++ b/otherarch/sdcpp/src/device_residency_manager.h @@ -0,0 +1,76 @@ +#ifndef __DEVICE_RESIDENCY_MANAGER_H__ +#define __DEVICE_RESIDENCY_MANAGER_H__ + +#include +#include +#include + +#include "ggml-backend.h" + +struct ggml_tensor; + +enum class WeightPrefetchResult { + Scheduled, + AlreadyResident, + Unsupported, + Failed, +}; + +struct WeightResidencyInfo { + bool async_prefetch_supported = false; + size_t missing_bytes = 0; +}; + +struct DeviceMemoryRequest { + ggml_backend_t compute_backend = nullptr; + uintptr_t owner_id = 0; + size_t pending_allocation_bytes = 0; + size_t runtime_resident_bytes = 0; + size_t max_backend_bytes = 0; + + // Runtime buffers only; the manager accounts for weights separately. + size_t runtime_peak_bytes() const { + return pending_allocation_bytes > SIZE_MAX - runtime_resident_bytes + ? SIZE_MAX + : runtime_resident_bytes + pending_allocation_bytes; + } +}; + +struct DeviceResidencyManager { + virtual ~DeviceResidencyManager() = default; + + virtual bool segmented_compute_enabled() const = 0; + virtual bool prefetch_enabled() const = 0; + virtual void set_workspace_reclaimer(uintptr_t owner_id, std::function reclaim) = 0; + virtual void remove_runtime_owner(uintptr_t owner_id) = 0; + // Capacity requests select their backend's weights; protection spans all backends. + virtual bool fits_compute_backend_capacity(const DeviceMemoryRequest& request, + const std::vector& required_params) const = 0; + virtual bool assign_compute_backend(const std::vector& tensors, + ggml_backend_t compute_backend) = 0; + virtual bool prepare_params(const std::vector& tensors) = 0; + virtual void release_compute_backend_params(const std::vector& tensors) = 0; + virtual void evict_compute_backend_params(const std::vector& tensors) = 0; + virtual WeightResidencyInfo inspect_compute_backend_params( + const std::vector& tensors) const = 0; + virtual void update_runtime_residency(uintptr_t owner_id, + ggml_backend_t compute_backend, + size_t resident_bytes) = 0; + virtual bool ensure_compute_backend_capacity( + const DeviceMemoryRequest& request, + const std::vector& required_params, + const std::vector>& preferred_eviction_order, + const std::vector& protected_params) = 0; + virtual WeightPrefetchResult prefetch_params( + uintptr_t owner_id, + const std::vector& tensors) = 0; + virtual bool activate_prefetched_params(uintptr_t owner_id, + const std::vector& tensors) = 0; + virtual void clear_prefetched_params(uintptr_t owner_id) = 0; +}; + +// Transitional alias for model constructors that have not yet adopted the +// residency-oriented name. It does not introduce a second implementation. +using RunnerWeightManager = DeviceResidencyManager; + +#endif // __DEVICE_RESIDENCY_MANAGER_H__ diff --git a/otherarch/sdcpp/src/extensions/generation_extension.h b/otherarch/sdcpp/src/extensions/generation_extension.h index 67085c157..b1ed9729b 100644 --- a/otherarch/sdcpp/src/extensions/generation_extension.h +++ b/otherarch/sdcpp/src/extensions/generation_extension.h @@ -49,7 +49,7 @@ struct GenerationExtension { virtual void get_param_tensors(std::map&) {} virtual void collect_loras(std::vector&) {} virtual void add_ignore_tensors(std::set&) const {} - virtual void runner_done() {} + virtual void runner_end() {} virtual void reset_runtime_condition() {} virtual bool prepare_condition(GenerationExtensionConditionContext&) { return false; diff --git a/otherarch/sdcpp/src/extensions/photomaker_extension.cpp b/otherarch/sdcpp/src/extensions/photomaker_extension.cpp index 48b023f99..b3bd6ec26 100644 --- a/otherarch/sdcpp/src/extensions/photomaker_extension.cpp +++ b/otherarch/sdcpp/src/extensions/photomaker_extension.cpp @@ -1,3 +1,4 @@ +#include #include "extensions/generation_extension.h" #include @@ -175,9 +176,9 @@ struct PhotoMakerExtension : public GenerationExtension { ignore_tensors.insert("pmid.unet."); } - void runner_done() override { + void runner_end() override { if (pmid_model != nullptr) { - pmid_model->runner_done(); + pmid_model->runner_end(); } } diff --git a/otherarch/sdcpp/src/model/adapter/ip_adapter.hpp b/otherarch/sdcpp/src/model/adapter/ip_adapter.hpp index cc4933e75..2cc660b59 100644 --- a/otherarch/sdcpp/src/model/adapter/ip_adapter.hpp +++ b/otherarch/sdcpp/src/model/adapter/ip_adapter.hpp @@ -1,8 +1,10 @@ #ifndef __SD_MODEL_ADAPTER_IP_ADAPTER_HPP__ #define __SD_MODEL_ADAPTER_IP_ADAPTER_HPP__ -#include "core/ggml_extend.hpp" +#include "core/ggml_extend.h" +#include "core/ggml_runner.h" #include "model/common/block.hpp" +#include "model/common/ggml_block.hpp" #include "model_loader.h" namespace IPAdapter { @@ -200,7 +202,7 @@ namespace IPAdapter { auto get_graph = [&]() -> ggml_cgraph* { return build_graph(image_embeds); }; - return take_or_empty(GGMLRunner::compute(get_graph, n_threads, true, true, true)); + return take_or_empty(GGMLRunner::compute(get_graph, n_threads, true)); } }; diff --git a/otherarch/sdcpp/src/model/adapter/lora.hpp b/otherarch/sdcpp/src/model/adapter/lora.hpp index f6b39d5a6..c72cd8bda 100644 --- a/otherarch/sdcpp/src/model/adapter/lora.hpp +++ b/otherarch/sdcpp/src/model/adapter/lora.hpp @@ -2,7 +2,13 @@ #define __SD_MODEL_ADAPTER_LORA_HPP__ #include -#include "core/ggml_extend.hpp" +#include "core/ggml_extend.h" +#include "core/ggml_extend_backend.h" +#include "core/ggml_runner.h" +#include "core/ggml_tensor_utils.h" +#include "core/util.h" +#include "model.h" +#include "model/adapter/lora_ops.h" #include "model_loader.h" #include "model_manager.h" @@ -120,18 +126,17 @@ struct LoraModel : public GGMLRunner { return false; } - LOG_DEBUG("finished loaded lora"); + LOG_VERBOSE("finished loaded lora"); return true; } void release_loaded_tensors() { - runner_done(); - free_compute_buffer(); + runner_end(); model_manager.reset(); free_params_ctx(); alloc_params_ctx(); - model_manager = std::make_shared(); - weight_manager = model_manager; + model_manager = std::make_shared(); + residency_manager = model_manager; lora_tensors.clear(); original_tensor_to_final_tensor.clear(); applied_lora_tensors.clear(); @@ -243,7 +248,7 @@ struct LoraModel : public GGMLRunner { if (iter != lora_tensors.end()) { float alpha = ggml_ext_backend_tensor_get_f32(iter->second); scale_value = alpha / rank; - // LOG_DEBUG("rank %s %ld %.2f %.2f", alpha_name.c_str(), rank, alpha, scale_value); + // LOG_VERBOSE("rank %s %ld %.2f %.2f", alpha_name.c_str(), rank, alpha, scale_value); applied_lora_tensors.insert(alpha_name); } } @@ -799,7 +804,7 @@ struct LoraModel : public GGMLRunner { float alpha = ggml_ext_backend_tensor_get_f32(iter->second); scale_value = alpha / rank; scale_tensor_name = alpha_name; - // LOG_DEBUG("rank %s %ld %.2f %.2f", alpha_name.c_str(), rank, alpha, scale_value); + // LOG_VERBOSE("rank %s %ld %.2f %.2f", alpha_name.c_str(), rank, alpha, scale_value); } } scale_value *= multiplier; @@ -952,16 +957,19 @@ struct LoraModel : public GGMLRunner { auto get_graph = [&]() -> ggml_cgraph* { return build_lora_graph(model_tensors, model_tensor_names, version); }; - GGMLRunner::compute(get_graph, n_threads, false, false, false, true); - stat(!warn_unused); - for (auto item : original_tensor_to_final_tensor) { - ggml_tensor* original_tensor = item.first; - ggml_tensor* final_tensor = item.second; - - ggml_backend_tensor_copy(final_tensor, original_tensor); + auto read_outputs = [&]() { + for (const auto& item : original_tensor_to_final_tensor) { + ggml_backend_tensor_copy(item.second, item.first); + } + return true; + }; + auto result = GGMLRunner::compute(get_graph, n_threads, false, true, read_outputs); + if (!result.has_value()) { + LOG_ERROR("LoRA graph execution failed"); } + stat(!warn_unused); original_tensor_to_final_tensor.clear(); - GGMLRunner::free_compute_buffer(); + runner_end(); } void apply(std::map model_tensors, SDVersion version, int n_threads, bool warn_unused = true) { diff --git a/otherarch/sdcpp/src/model/adapter/lora_ops.cpp b/otherarch/sdcpp/src/model/adapter/lora_ops.cpp new file mode 100644 index 000000000..7dacfd164 --- /dev/null +++ b/otherarch/sdcpp/src/model/adapter/lora_ops.cpp @@ -0,0 +1,207 @@ +#include "model/adapter/lora_ops.h" + +#include + +#include "core/ggml_extend.h" +#include "core/ggml_extend_backend.h" + +ggml_tensor* ggml_ext_merge_lora(ggml_context* ctx, + ggml_tensor* lora_down, + ggml_tensor* lora_up, + ggml_tensor* lora_mid) { + ggml_tensor* updown; + // flat lora tensors to multiply it + int64_t lora_up_rows = lora_up->ne[ggml_n_dims(lora_up) - 1]; + lora_up = ggml_reshape_2d(ctx, lora_up, ggml_nelements(lora_up) / lora_up_rows, lora_up_rows); + auto lora_down_n_dims = ggml_n_dims(lora_down); + // assume n_dims should always be a multiple of 2 (otherwise rank 1 doesn't work) + lora_down_n_dims = (lora_down_n_dims + lora_down_n_dims % 2); + int64_t lora_down_rows = lora_down->ne[lora_down_n_dims - 1]; + lora_down = ggml_reshape_2d(ctx, lora_down, ggml_nelements(lora_down) / lora_down_rows, lora_down_rows); + + // ggml_mul_mat requires tensor b transposed + lora_down = ggml_cont(ctx, ggml_transpose(ctx, lora_down)); + if (lora_mid == nullptr) { + updown = ggml_mul_mat(ctx, lora_up, lora_down); + updown = ggml_cont(ctx, ggml_transpose(ctx, updown)); + } else { + // undoing tucker decomposition for conv layers. + // lora_mid has shape (3, 3, Rank, Rank) + // lora_down has shape (Rank, In, 1, 1) + // lora_up has shape (Rank, Out, 1, 1) + // conv layer shape is (3, 3, Out, In) + updown = ggml_ext_mul_n_mode(ctx, ggml_ext_mul_n_mode(ctx, lora_mid, lora_down, 3), lora_up, 2); + updown = ggml_cont(ctx, updown); + } + return updown; +} + +ggml_tensor* ggml_ext_lokr_forward( + ggml_context* ctx, + ggml_backend_t backend, + ggml_tensor* h, // Input: [q, batch] or [W, H, q, batch] + ggml_tensor* w1, // Outer C (Full rank) + ggml_tensor* w1a, // Outer A (Low rank part 1) + ggml_tensor* w1b, // Outer B (Low rank part 2) + ggml_tensor* w2, // Inner BA (Full rank) + ggml_tensor* w2a, // Inner A (Low rank part 1) + ggml_tensor* w2b, // Inner B (Low rank part 2) + bool is_conv, + WeightAdapter::ForwardParams::conv2d_params_t conv_params, + float scale) { + GGML_ASSERT((w1 != nullptr || (w1a != nullptr && w1b != nullptr))); + GGML_ASSERT((w2 != nullptr || (w2a != nullptr && w2b != nullptr))); + + int uq = (w1 != nullptr) ? (int)w1->ne[0] : (int)w1a->ne[0]; + int up = (w1 != nullptr) ? (int)w1->ne[1] : (int)w1b->ne[1]; + + int q_actual = is_conv ? (int)h->ne[2] : (int)h->ne[0]; + int vq = q_actual / uq; + + int vp = (w2 != nullptr) ? (is_conv ? (int)w2->ne[3] : (int)w2->ne[1]) + : (int)w2a->ne[1]; + GGML_ASSERT(q_actual == (uq * vq) && "Input dimension mismatch for LoKR split"); + + ggml_tensor* hb; + + if (!is_conv) { + int batch = (int)h->ne[1]; + int merge_batch_uq = batch; + int merge_batch_vp = batch; + + if (sd_backend_is(backend, "Vulkan")) { + if (batch > 1) { + // no access to backend here, worst case is slightly worse perfs for other backends when built alongside Vulkan backend + int max_batch = 65535; + int max_batch_uq = max_batch / uq; + merge_batch_uq = 1; + for (int i = max_batch_uq; i > 0; i--) { + if (batch % i == 0) { + merge_batch_uq = i; + break; + } + } + + int max_batch_vp = max_batch / vp; + merge_batch_vp = 1; + for (int i = max_batch_vp; i > 0; i--) { + if (batch % i == 0) { + merge_batch_vp = i; + break; + } + } + } + } + + ggml_tensor* h_split = ggml_reshape_3d(ctx, h, vq, uq * merge_batch_uq, batch / merge_batch_uq); + if (w2 != nullptr) { + hb = ggml_mul_mat(ctx, w2, h_split); + } else { + hb = ggml_mul_mat(ctx, w2b, ggml_mul_mat(ctx, w2a, h_split)); + } + + if (batch > 1) { + hb = ggml_reshape_3d(ctx, hb, vp, uq, batch); + } + ggml_tensor* hb_t = ggml_cont(ctx, ggml_transpose(ctx, hb)); + hb_t = ggml_reshape_3d(ctx, hb_t, uq, vp * merge_batch_vp, batch / merge_batch_vp); + + ggml_tensor* hc_t; + if (w1 != nullptr) { + hc_t = ggml_mul_mat(ctx, w1, hb_t); + } else { + hc_t = ggml_mul_mat(ctx, w1b, ggml_mul_mat(ctx, w1a, hb_t)); + } + + if (batch > 1) { + hc_t = ggml_reshape_3d(ctx, hc_t, up, vp, batch); + } + + ggml_tensor* hc = ggml_transpose(ctx, hc_t); + ggml_tensor* out = ggml_reshape_2d(ctx, ggml_cont(ctx, hc), up * vp, batch); + return ggml_ext_scale(ctx, out, scale); + } else { + int batch = (int)h->ne[3]; + // 1. Reshape input: [W, H, vq*uq, batch] -> [W, H, vq, uq * batch] + ggml_tensor* h_split = ggml_reshape_4d(ctx, h, h->ne[0], h->ne[1], vq, uq * batch); + + if (w2 != nullptr) { + hb = ggml_ext_conv_2d(ctx, h_split, w2, nullptr, + conv_params.s0, + conv_params.s1, + conv_params.p0, + conv_params.p1, + conv_params.d0, + conv_params.d1, + conv_params.direct, + conv_params.circular_x, + conv_params.circular_y, + conv_params.scale); + } else { + // swap a and b order for conv lora + ggml_tensor* a = w2b; + ggml_tensor* b = w2a; + + // unpack conv2d weights if needed + if (ggml_n_dims(a) < 4) { + int k = (int)sqrt(a->ne[0] / h_split->ne[2]); + GGML_ASSERT(k * k * h_split->ne[2] == a->ne[0]); + a = ggml_reshape_4d(ctx, a, k, k, a->ne[0] / (k * k), a->ne[1]); + } else if (a->ne[2] != h_split->ne[2]) { + int k = (int)sqrt(a->ne[2] / h_split->ne[2]); + GGML_ASSERT(k * k * h_split->ne[2] == a->ne[2]); + a = ggml_reshape_4d(ctx, a, a->ne[0] * k, a->ne[1] * k, a->ne[2] / (k * k), a->ne[3]); + } + ggml_tensor* ha = ggml_ext_conv_2d(ctx, h_split, a, nullptr, + conv_params.s0, + conv_params.s1, + conv_params.p0, + conv_params.p1, + conv_params.d0, + conv_params.d1, + conv_params.direct, + conv_params.circular_x, + conv_params.circular_y, + conv_params.scale); + + // not supporting lora_mid here + hb = ggml_ext_conv_2d(ctx, + ha, + b, + nullptr, + 1, + 1, + 0, + 0, + 1, + 1, + conv_params.direct, + conv_params.circular_x, + conv_params.circular_y, + conv_params.scale); + } + + // Current hb shape: [W_out, H_out, vp, uq * batch] + int w_out = (int)hb->ne[0]; + int h_out = (int)hb->ne[1]; + + // ggml_tensor* hb_cat = ggml_reshape_4d(ctx, hb, w_out , h_out , vp * uq, batch); + // [W_out, H_out, vp * uq, batch] + // Now left to compute (W1 kr Id) * hb_cat == (W1 kr W2) cv h + + // merge the uq groups of size vp*w_out*h_out + ggml_tensor* hb_merged = ggml_reshape_2d(ctx, hb, w_out * h_out * vp, uq * batch); + ggml_tensor* hc_t; + ggml_tensor* hb_merged_t = ggml_cont(ctx, ggml_transpose(ctx, hb_merged)); + if (w1 != nullptr) { + // Would be great to be able to transpose w1 instead to avoid transposing both hb and hc + hc_t = ggml_mul_mat(ctx, w1, hb_merged_t); + } else { + hc_t = ggml_mul_mat(ctx, w1b, ggml_mul_mat(ctx, w1a, hb_merged_t)); + } + ggml_tensor* hc = ggml_transpose(ctx, hc_t); + // ungroup + ggml_tensor* out = ggml_reshape_4d(ctx, ggml_cont(ctx, hc), w_out, h_out, up * vp, batch); + return ggml_ext_scale(ctx, out, scale); + } +} diff --git a/otherarch/sdcpp/src/model/adapter/lora_ops.h b/otherarch/sdcpp/src/model/adapter/lora_ops.h new file mode 100644 index 000000000..b32debfa6 --- /dev/null +++ b/otherarch/sdcpp/src/model/adapter/lora_ops.h @@ -0,0 +1,25 @@ +#ifndef __SD_MODEL_ADAPTER_LORA_OPS_H__ +#define __SD_MODEL_ADAPTER_LORA_OPS_H__ + +#include "core/ggml_runner.h" + +ggml_tensor* ggml_ext_merge_lora(ggml_context* ctx, + ggml_tensor* lora_down, + ggml_tensor* lora_up, + ggml_tensor* lora_mid = nullptr); + +ggml_tensor* ggml_ext_lokr_forward( + ggml_context* ctx, + ggml_backend_t backend, + ggml_tensor* h, // Input: [q, batch] or [W, H, q, batch] + ggml_tensor* w1, // Outer C (Full rank) + ggml_tensor* w1a, // Outer A (Low rank part 1) + ggml_tensor* w1b, // Outer B (Low rank part 2) + ggml_tensor* w2, // Inner BA (Full rank) + ggml_tensor* w2a, // Inner A (Low rank part 1) + ggml_tensor* w2b, // Inner B (Low rank part 2) + bool is_conv, + WeightAdapter::ForwardParams::conv2d_params_t conv_params, + float scale); + +#endif // __SD_MODEL_ADAPTER_LORA_OPS_H__ diff --git a/otherarch/sdcpp/src/model/adapter/pmid.hpp b/otherarch/sdcpp/src/model/adapter/pmid.hpp index 8f7d4dbde..f16d7ba1f 100644 --- a/otherarch/sdcpp/src/model/adapter/pmid.hpp +++ b/otherarch/sdcpp/src/model/adapter/pmid.hpp @@ -1,7 +1,10 @@ #ifndef __SD_MODEL_ADAPTER_PMID_HPP__ #define __SD_MODEL_ADAPTER_PMID_HPP__ -#include "core/ggml_extend.hpp" +#include "core/ggml_extend.h" +#include "core/ggml_runner.h" +#include "core/util.h" +#include "model/common/ggml_block.hpp" #include "model/adapter/lora.hpp" #include "model/common/block.hpp" @@ -558,7 +561,7 @@ public: return build_graph(id_pixel_values, prompt_embeds, class_tokens_mask, id_embeds); }; - return take_or_empty(GGMLRunner::compute(get_graph, n_threads, true, true, true)); + return take_or_empty(GGMLRunner::compute(get_graph, n_threads, true)); } }; @@ -639,7 +642,7 @@ struct PhotoMakerIDEmbed : public GGMLRunner { return false; } - LOG_DEBUG("finished loading PhotoMaker ID Embeds "); + LOG_VERBOSE("finished loading PhotoMaker ID Embeds "); return true; } diff --git a/otherarch/sdcpp/src/model/adapter/pulid.hpp b/otherarch/sdcpp/src/model/adapter/pulid.hpp index 442c5b8b2..65b3064a2 100644 --- a/otherarch/sdcpp/src/model/adapter/pulid.hpp +++ b/otherarch/sdcpp/src/model/adapter/pulid.hpp @@ -1,8 +1,10 @@ #ifndef __PULID_HPP__ #define __PULID_HPP__ -#include "core/ggml_extend.hpp" +#include "core/ggml_extend.h" +#include "core/ggml_runner.h" #include "model/common/block.hpp" +#include "model/common/ggml_block.hpp" class PuLIDPerceiverAttentionCA : public GGMLBlock { public: diff --git a/otherarch/sdcpp/src/model/common/block.hpp b/otherarch/sdcpp/src/model/common/block.hpp index 6eb387d9f..09db9250a 100644 --- a/otherarch/sdcpp/src/model/common/block.hpp +++ b/otherarch/sdcpp/src/model/common/block.hpp @@ -1,9 +1,11 @@ #ifndef __SD_MODEL_COMMON_BLOCK_HPP__ #define __SD_MODEL_COMMON_BLOCK_HPP__ -#include "core/ggml_extend.hpp" +#include "core/ggml_extend.h" +#include "core/ggml_runner.h" #include "core/util.h" #include "ggml-backend.h" +#include "model/common/ggml_block.hpp" class DownSampleBlock : public GGMLBlock { protected: @@ -340,7 +342,7 @@ public: enable_ip(enable_ip) { int64_t inner_dim = d_head * n_head; if (context_dim == 320 && d_head == 320) { - // LOG_DEBUG("CrossAttention: temp set dim to 1024 for sdxs_09"); + // LOG_VERBOSE("CrossAttention: temp set dim to 1024 for sdxs_09"); xtra_dim = true; context_dim = 1024; } @@ -370,7 +372,7 @@ public: auto q = to_q->forward(ctx, x); // [N, n_token, inner_dim] if (xtra_dim) { - // LOG_DEBUG("CrossAttention: temp set dim to 1024 for sdxs_09"); + // LOG_VERBOSE("CrossAttention: temp set dim to 1024 for sdxs_09"); context->ne[0] = 1024; // patch dim } auto k = to_k->forward(ctx, context); // [N, n_context, inner_dim] diff --git a/otherarch/sdcpp/src/model/common/ggml_block.hpp b/otherarch/sdcpp/src/model/common/ggml_block.hpp new file mode 100644 index 000000000..6822922ec --- /dev/null +++ b/otherarch/sdcpp/src/model/common/ggml_block.hpp @@ -0,0 +1,883 @@ +#ifndef __SD_MODEL_COMMON_GGML_BLOCK_HPP__ +#define __SD_MODEL_COMMON_GGML_BLOCK_HPP__ + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "core/ggml_extend.h" +#include "core/ggml_runner.h" +#include "model.h" + +class GGMLBlock { +protected: + typedef std::unordered_map ParameterMap; + typedef std::unordered_map> GGMLBlockMap; + GGMLBlockMap blocks; + ParameterMap params; + + ggml_type get_type(const std::string& name, const String2TensorStorage& tensor_storage_map, ggml_type default_type) { + ggml_type wtype = default_type; + auto iter = tensor_storage_map.find(name); + if (iter != tensor_storage_map.end()) { + const TensorStorage& tensor_storage = iter->second; + if (tensor_storage.expected_type != GGML_TYPE_COUNT) { + wtype = tensor_storage.expected_type; + } else { + wtype = tensor_storage.type; + } + } + return wtype; + } + + void init_blocks(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") { + for (auto& pair : blocks) { + auto& block = pair.second; + block->init(ctx, tensor_storage_map, prefix + pair.first); + } + } + + virtual void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") {} + + virtual enum ggml_op param_usage_op(const std::string& name) const { + (void)name; + return GGML_OP_NONE; + } + +public: + void init(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, std::string prefix = "") { + if (prefix.size() > 0) { + prefix = prefix + "."; + } + init_params(ctx, tensor_storage_map, prefix); + init_blocks(ctx, tensor_storage_map, prefix); + } + + size_t get_params_num() { + size_t num_tensors = params.size(); + for (auto& pair : blocks) { + auto& block = pair.second; + + num_tensors += block->get_params_num(); + } + return num_tensors; + }; + + size_t get_params_mem_size() { + size_t mem_size = 0; + for (auto& pair : blocks) { + auto& block = pair.second; + + mem_size += block->get_params_mem_size(); + } + + for (auto& pair : params) { + mem_size += ggml_nbytes(pair.second); + } + + return mem_size; + } + + void get_param_tensors(std::map& tensors, std::string prefix = "") { + if (prefix.size() > 0) { + prefix = prefix + "."; + } + for (auto& pair : blocks) { + auto& block = pair.second; + block->get_param_tensors(tensors, prefix + pair.first); + } + + for (auto& pair : params) { + ggml_tensor* param = pair.second; + tensors[prefix + pair.first] = pair.second; + ggml_set_name(param, (prefix + pair.first).c_str()); + } + } + + void get_param_tensor_ops(std::map& tensor_ops) { + for (auto& pair : blocks) { + pair.second->get_param_tensor_ops(tensor_ops); + } + for (auto& pair : params) { + enum ggml_op op = param_usage_op(pair.first); + if (op != GGML_OP_NONE) { + tensor_ops[pair.second] = op; + } + } + } + + virtual std::string get_desc() { + return "GGMLBlock"; + } + + void get_all_blocks(std::vector& result) { + result.push_back(this); + for (auto& block_iter : blocks) { + if (block_iter.second) { + block_iter.second->get_all_blocks(result); + } + } + } +}; + +class UnaryBlock : public GGMLBlock { +public: + virtual ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) = 0; +}; + +class Identity : public UnaryBlock { +public: + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { + return x; + } +}; + +class Linear : public UnaryBlock { +protected: + int64_t in_features; + int64_t out_features; + bool bias; + bool force_f32; + bool force_prec_f32; + bool has_weight_scale = false; + bool int8_convrot = false; + int int8_convrot_group_size = 0; + float scale; + std::string prefix; + + void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override { + this->prefix = prefix; + has_weight_scale = false; + int8_convrot = false; + int8_convrot_group_size = 0; + enum ggml_type wtype = get_type(prefix + "weight", tensor_storage_map, GGML_TYPE_F32); + if (in_features % ggml_blck_size(wtype) != 0 || force_f32) { + wtype = GGML_TYPE_F32; + } + params["weight"] = ggml_new_tensor_2d(ctx, wtype, in_features, out_features); + if (bias) { + enum ggml_type wtype = GGML_TYPE_F32; + params["bias"] = ggml_new_tensor_1d(ctx, wtype, out_features); + } + auto weight_storage = tensor_storage_map.find(prefix + "weight"); + const bool is_int8_tensorwise = weight_storage != tensor_storage_map.end() && weight_storage->second.is_int8_tensorwise; + auto weight_scale_storage = tensor_storage_map.find(prefix + "weight_scale"); + if (weight_scale_storage != tensor_storage_map.end()) { + const int64_t scale_nelements = weight_scale_storage->second.nelements(); + GGML_ASSERT(scale_nelements == 1 || scale_nelements == out_features); + params["weight_scale"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, scale_nelements); + has_weight_scale = true; + } + if (is_int8_tensorwise) { + GGML_ASSERT(wtype == GGML_TYPE_I8); + GGML_ASSERT(has_weight_scale); + int8_convrot = weight_storage->second.int8_convrot; + int8_convrot_group_size = weight_storage->second.int8_convrot_group_size; + } + } + +public: + Linear(int64_t in_features, + int64_t out_features, + bool bias = true, + bool force_f32 = false, + bool force_prec_f32 = false, + float scale = 1.f) + : in_features(in_features), + out_features(out_features), + bias(bias), + force_f32(force_f32), + force_prec_f32(force_prec_f32), + scale(scale) {} + + void set_scale(float scale_) { + scale = scale_; + } + + void set_force_prec_f32(bool force_prec_f32_) { + force_prec_f32 = force_prec_f32_; + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { + ggml_tensor* w = params["weight"]; + ggml_tensor* weight_scale = has_weight_scale ? params["weight_scale"] : nullptr; + #if KCPP_MAINLINE_FP8_SCALED + if (w->type == GGML_TYPE_F8_E4M3 || w->type == GGML_TYPE_F8_E5M2) { + bool supports_fp8_matmul = false; + if (ctx->backend != nullptr) { + ggml_tensor* fp8_matmul = ggml_mul_mat(ctx->ggml_ctx, w, x); + if (force_prec_f32) { + ggml_mul_mat_set_prec(fp8_matmul, GGML_PREC_F32); + } + supports_fp8_matmul = ggml_backend_supports_op(ctx->backend, fp8_matmul); + } + if (!supports_fp8_matmul) { + w = ggml_cast(ctx->ggml_ctx, w, GGML_TYPE_BF16); + } + } + #endif //kcpp + ggml_tensor* b = nullptr; + if (bias) { + b = params["bias"]; + } + ggml_tensor* linear_bias = has_weight_scale ? nullptr : b; + ggml_tensor* out = nullptr; + #if KCPP_MAINLINE_INT8_CONVROT + if (w->type == GGML_TYPE_I8) { + if (x->type != GGML_TYPE_F32) { + x = ggml_ext_cast_f32(ctx->ggml_ctx, ctx->backend, x); + } + if (!ggml_is_contiguous(x)) { + x = ggml_cont(ctx->ggml_ctx, x); + } + ggml_tensor* lora_input = x; + if (ctx->weight_adapter && b != nullptr) { + b = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, b, prefix + "bias"); + } + if (int8_convrot && scale == 1.f) { + const auto cache_key = std::make_pair(x, int8_convrot_group_size); + auto cached = ctx->int8_convrot_cache.find(cache_key); + if (cached == ctx->int8_convrot_cache.end()) { + x = ggml_quantize_i8_convrot(ctx->ggml_ctx, x, int8_convrot_group_size); + ctx->int8_convrot_cache.emplace(cache_key, x); + } else { + x = cached->second; + } + } + out = ggml_ext_linear_i8_tensorwise(ctx->ggml_ctx, + x, + w, + weight_scale, + b, + int8_convrot ? int8_convrot_group_size : 0, + scale); + if (ctx->weight_adapter) { + WeightAdapter::ForwardParams forward_params; + forward_params.op_type = WeightAdapter::ForwardParams::op_type_t::OP_LINEAR; + forward_params.linear.force_prec_f32 = force_prec_f32; + forward_params.linear.scale = scale; + out = ctx->weight_adapter->add_lora_to_output(ctx->ggml_ctx, + ctx->backend, + lora_input, + w, + out, + prefix, + forward_params); + } + return out; + } + #endif //kcpp + if (has_weight_scale) { + out = ggml_ext_linear(ctx->ggml_ctx, x, w, nullptr, force_prec_f32, scale); + out = ggml_mul(ctx->ggml_ctx, out, weight_scale); + if (ctx->weight_adapter) { + WeightAdapter::ForwardParams forward_params; + forward_params.op_type = WeightAdapter::ForwardParams::op_type_t::OP_LINEAR; + forward_params.linear.force_prec_f32 = force_prec_f32; + forward_params.linear.scale = scale; + out = ctx->weight_adapter->add_lora_to_output(ctx->ggml_ctx, + ctx->backend, + x, + w, + out, + prefix, + forward_params); + if (b != nullptr) { + b = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, b, prefix + "bias"); + } + } + if (b != nullptr) { + out = ggml_add_inplace(ctx->ggml_ctx, out, b); + } + return out; + } + if (ctx->weight_adapter) { + WeightAdapter::ForwardParams forward_params; + forward_params.op_type = WeightAdapter::ForwardParams::op_type_t::OP_LINEAR; + forward_params.linear.force_prec_f32 = force_prec_f32; + forward_params.linear.scale = scale; + out = ctx->weight_adapter->forward_with_lora(ctx->ggml_ctx, ctx->backend, x, w, linear_bias, prefix, forward_params); + } else { + out = ggml_ext_linear(ctx->ggml_ctx, x, w, linear_bias, force_prec_f32, scale); + } + return out; + } +}; + +__STATIC_INLINE__ bool support_get_rows(ggml_type wtype) { + switch (wtype) { + case GGML_TYPE_F16: + case GGML_TYPE_Q8_0: + case GGML_TYPE_Q5_1: + case GGML_TYPE_Q5_0: + case GGML_TYPE_Q4_1: + case GGML_TYPE_Q4_0: + case GGML_TYPE_Q2_K: + case GGML_TYPE_Q3_K: + case GGML_TYPE_Q4_K: + case GGML_TYPE_Q5_K: + case GGML_TYPE_Q6_K: + return true; + default: + return false; + } +} + +class Embedding : public UnaryBlock { +protected: + int64_t embedding_dim; + int64_t num_embeddings; + void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map, const std::string prefix = "") override { + enum ggml_type wtype = get_type(prefix + "weight", tensor_storage_map, GGML_TYPE_F32); + if (!support_get_rows(wtype)) { + wtype = GGML_TYPE_F32; + } + params["weight"] = ggml_new_tensor_2d(ctx, wtype, embedding_dim, num_embeddings); + } + + enum ggml_op param_usage_op(const std::string& name) const override { + return name == "weight" ? GGML_OP_GET_ROWS : GGML_OP_NONE; + } + +public: + Embedding(int64_t num_embeddings, int64_t embedding_dim) + : embedding_dim(embedding_dim), + num_embeddings(num_embeddings) { + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, + ggml_tensor* input_ids) override { + // input_ids: [N, n_token] + auto weight = params["weight"]; + + // There are issues with ggml batch inference, so we are expanding it here first. + // TODO: fix ggml batch inference + int64_t n = input_ids->ne[1]; + input_ids = ggml_reshape_1d(ctx->ggml_ctx, input_ids, input_ids->ne[0] * input_ids->ne[1]); + + input_ids = ggml_reshape_3d(ctx->ggml_ctx, input_ids, input_ids->ne[0], 1, input_ids->ne[1]); + auto embedding = ggml_get_rows(ctx->ggml_ctx, weight, input_ids); + embedding = ggml_reshape_3d(ctx->ggml_ctx, embedding, embedding->ne[0], embedding->ne[1] / n, n); + + // [N, n_token, embedding_dim] + return embedding; + } +}; + +class Conv2d : public UnaryBlock { +protected: + int64_t in_channels; + int64_t out_channels; + std::pair kernel_size; + std::pair stride; + std::pair padding; + std::pair dilation; + bool bias; + float scale = 1.f; + std::string prefix; + + void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map, const std::string prefix = "") override { + this->prefix = prefix; + enum ggml_type wtype = GGML_TYPE_F16; + params["weight"] = ggml_new_tensor_4d(ctx, wtype, kernel_size.second, kernel_size.first, in_channels, out_channels); + if (bias) { + enum ggml_type wtype = GGML_TYPE_F32; + params["bias"] = ggml_new_tensor_1d(ctx, wtype, out_channels); + } + } + +public: + Conv2d(int64_t in_channels, + int64_t out_channels, + std::pair kernel_size, + std::pair stride = {1, 1}, + std::pair padding = {0, 0}, + std::pair dilation = {1, 1}, + bool bias = true) + : in_channels(in_channels), + out_channels(out_channels), + kernel_size(kernel_size), + stride(stride), + padding(padding), + dilation(dilation), + bias(bias) {} + + void set_scale(float scale_value) { + scale = scale_value; + } + + std::string get_desc() override { + return "Conv2d"; + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { + ggml_tensor* w = params["weight"]; + ggml_tensor* b = nullptr; + if (bias) { + b = params["bias"]; + } + if (ctx->weight_adapter) { + WeightAdapter::ForwardParams forward_params; + forward_params.op_type = WeightAdapter::ForwardParams::op_type_t::OP_CONV2D; + forward_params.conv2d.s0 = stride.second; + forward_params.conv2d.s1 = stride.first; + forward_params.conv2d.p0 = padding.second; + forward_params.conv2d.p1 = padding.first; + forward_params.conv2d.d0 = dilation.second; + forward_params.conv2d.d1 = dilation.first; + forward_params.conv2d.direct = ctx->conv2d_direct_enabled; + forward_params.conv2d.circular_x = ctx->circular_x_enabled; + forward_params.conv2d.circular_y = ctx->circular_y_enabled; + forward_params.conv2d.scale = scale; + return ctx->weight_adapter->forward_with_lora(ctx->ggml_ctx, ctx->backend, x, w, b, prefix, forward_params); + } + return ggml_ext_conv_2d(ctx->ggml_ctx, + x, + w, + b, + stride.second, + stride.first, + padding.second, + padding.first, + dilation.second, + dilation.first, + ctx->conv2d_direct_enabled, + ctx->circular_x_enabled, + ctx->circular_y_enabled, + scale); + } +}; + +class Conv2d_grouped : public UnaryBlock { +protected: + int64_t in_channels; + int64_t out_channels; + int groups; + std::pair kernel_size; + std::pair stride; + std::pair padding; + std::pair dilation; + bool bias; + float scale = 1.f; + std::string prefix; + + void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map, const std::string prefix = "") override { + this->prefix = prefix; + enum ggml_type wtype = GGML_TYPE_F16; + params["weight"] = ggml_new_tensor_4d(ctx, wtype, kernel_size.second, kernel_size.first, in_channels / groups, out_channels); + if (bias) { + enum ggml_type wtype = GGML_TYPE_F32; + params["bias"] = ggml_new_tensor_1d(ctx, wtype, out_channels); + } + } + +public: + Conv2d_grouped(int64_t in_channels, + int64_t out_channels, + int groups, + std::pair kernel_size, + std::pair stride = {1, 1}, + std::pair padding = {0, 0}, + std::pair dilation = {1, 1}, + bool bias = true) + : in_channels(in_channels), + out_channels(out_channels), + groups(groups), + kernel_size(kernel_size), + stride(stride), + padding(padding), + dilation(dilation), + bias(bias) {} + + void set_scale(float scale_value) { + scale = scale_value; + } + + std::string get_desc() override { + return "Conv2d_grouped"; + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { + ggml_tensor* w = params["weight"]; + ggml_tensor* b = nullptr; + if (bias) { + b = params["bias"]; + } + + if (groups == 1) { + if (ctx->weight_adapter) { + WeightAdapter::ForwardParams forward_params; + forward_params.op_type = WeightAdapter::ForwardParams::op_type_t::OP_CONV2D; + forward_params.conv2d.s0 = stride.second; + forward_params.conv2d.s1 = stride.first; + forward_params.conv2d.p0 = padding.second; + forward_params.conv2d.p1 = padding.first; + forward_params.conv2d.d0 = dilation.second; + forward_params.conv2d.d1 = dilation.first; + forward_params.conv2d.direct = ctx->conv2d_direct_enabled; + forward_params.conv2d.circular_x = ctx->circular_x_enabled; + forward_params.conv2d.circular_y = ctx->circular_y_enabled; + forward_params.conv2d.scale = scale; + return ctx->weight_adapter->forward_with_lora(ctx->ggml_ctx, ctx->backend, x, w, b, prefix, forward_params); + } + return ggml_ext_conv_2d(ctx->ggml_ctx, x, w, b, + stride.second, stride.first, + padding.second, padding.first, + dilation.second, dilation.first, + ctx->conv2d_direct_enabled, + ctx->circular_x_enabled, + ctx->circular_y_enabled, + scale); + } + + if (groups == in_channels && groups == out_channels) { + ggml_tensor* res; + if (ctx->conv2d_direct_enabled) { + res = ggml_conv_2d_dw_direct(ctx->ggml_ctx, w, x, + stride.second, stride.first, + padding.second, padding.first, + dilation.second, dilation.first); + } else { + res = ggml_conv_2d_dw(ctx->ggml_ctx, w, x, + stride.second, stride.first, + padding.second, padding.first, + dilation.second, dilation.first); + } + if (b) { + b = ggml_reshape_4d(ctx->ggml_ctx, b, 1, 1, b->ne[0], 1); + res = ggml_add_inplace(ctx->ggml_ctx, res, b); + } + return res; + } + + int64_t ic_g = in_channels / groups; + int64_t oc_g = out_channels / groups; + + std::vector out_slices(groups); + + for (int i = 0; i < groups; ++i) { + size_t x_offset = i * ic_g * x->nb[2]; + ggml_tensor* x_i = ggml_view_4d(ctx->ggml_ctx, x, + x->ne[0], x->ne[1], ic_g, x->ne[3], + x->nb[1], x->nb[2], x->nb[3], + x_offset); + + size_t w_offset = i * oc_g * w->nb[3]; + ggml_tensor* w_i = ggml_view_4d(ctx->ggml_ctx, w, + w->ne[0], w->ne[1], w->ne[2], oc_g, + w->nb[1], w->nb[2], w->nb[3], + w_offset); + + ggml_tensor* b_i = nullptr; + if (b) { + size_t b_offset = i * oc_g * b->nb[0]; + b_i = ggml_view_1d(ctx->ggml_ctx, b, oc_g, b_offset); + } + + if (ctx->weight_adapter) { + WeightAdapter::ForwardParams forward_params; + forward_params.op_type = WeightAdapter::ForwardParams::op_type_t::OP_CONV2D; + forward_params.conv2d.s0 = stride.second; + forward_params.conv2d.s1 = stride.first; + forward_params.conv2d.p0 = padding.second; + forward_params.conv2d.p1 = padding.first; + forward_params.conv2d.d0 = dilation.second; + forward_params.conv2d.d1 = dilation.first; + forward_params.conv2d.direct = ctx->conv2d_direct_enabled; + forward_params.conv2d.circular_x = ctx->circular_x_enabled; + forward_params.conv2d.circular_y = ctx->circular_y_enabled; + forward_params.conv2d.scale = scale; + out_slices[i] = ctx->weight_adapter->forward_with_lora(ctx->ggml_ctx, ctx->backend, x_i, w_i, b_i, prefix, forward_params); + } else { + out_slices[i] = ggml_ext_conv_2d(ctx->ggml_ctx, x_i, w_i, b_i, + stride.second, stride.first, + padding.second, padding.first, + dilation.second, dilation.first, + ctx->conv2d_direct_enabled, + ctx->circular_x_enabled, + ctx->circular_y_enabled, + scale); + } + } + + ggml_tensor* out = ggml_ext_vec_concat(ctx->ggml_ctx, out_slices, 2); + + return out; + } +}; + +class Conv3d : public UnaryBlock { +protected: + int64_t in_channels; + int64_t out_channels; + std::tuple kernel_size; + std::tuple stride; + std::tuple padding; + std::tuple dilation; + bool bias; + bool force_prec_f32; + std::string prefix; + + void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map, const std::string prefix = "") override { + this->prefix = prefix; + enum ggml_type wtype = GGML_TYPE_F16; + params["weight"] = ggml_new_tensor_4d(ctx, + wtype, + std::get<2>(kernel_size), + std::get<1>(kernel_size), + std::get<0>(kernel_size), + in_channels * out_channels); + if (bias) { + params["bias"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, out_channels); + } + } + +public: + Conv3d(int64_t in_channels, + int64_t out_channels, + std::tuple kernel_size, + std::tuple stride = {1, 1, 1}, + std::tuple padding = {0, 0, 0}, + std::tuple dilation = {1, 1, 1}, + bool bias = true, + bool force_prec_f32 = false) + : in_channels(in_channels), + out_channels(out_channels), + kernel_size(kernel_size), + stride(stride), + padding(padding), + dilation(dilation), + bias(bias), + force_prec_f32(force_prec_f32) {} + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { + ggml_tensor* w = params["weight"]; + ggml_tensor* b = nullptr; + if (ctx->weight_adapter) { + w = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, w, prefix + "weight"); + if (w->type != GGML_TYPE_F16) { + w = ggml_cast(ctx->ggml_ctx, w, GGML_TYPE_F16); + } + } + if (bias) { + b = params["bias"]; + if (ctx->weight_adapter) { + b = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, b, prefix + "bias"); + } + } + return ggml_ext_conv_3d(ctx->ggml_ctx, ctx->backend, x, w, b, in_channels, + std::get<2>(stride), std::get<1>(stride), std::get<0>(stride), + std::get<2>(padding), std::get<1>(padding), std::get<0>(padding), + std::get<2>(dilation), std::get<1>(dilation), std::get<0>(dilation), + force_prec_f32); + } +}; + +class LayerNorm : public UnaryBlock { +protected: + int64_t normalized_shape; + float eps; + bool elementwise_affine; + bool bias; + std::string prefix; + + void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override { + this->prefix = prefix; + if (elementwise_affine) { + enum ggml_type wtype = GGML_TYPE_F32; + params["weight"] = ggml_new_tensor_1d(ctx, wtype, normalized_shape); + if (bias) { + enum ggml_type wtype = GGML_TYPE_F32; + params["bias"] = ggml_new_tensor_1d(ctx, wtype, normalized_shape); + } + } + } + +public: + LayerNorm(int64_t normalized_shape, + float eps = 1e-05f, + bool elementwise_affine = true, + bool bias = true) + : normalized_shape(normalized_shape), + eps(eps), + elementwise_affine(elementwise_affine), + bias(bias) {} + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { + ggml_tensor* w = nullptr; + ggml_tensor* b = nullptr; + + if (elementwise_affine) { + w = params["weight"]; + if (ctx->weight_adapter) { + w = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, w, prefix + "weight"); + } + if (bias) { + b = params["bias"]; + if (ctx->weight_adapter) { + b = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, b, prefix + "bias"); + } + } + } + return ggml_ext_layer_norm(ctx->ggml_ctx, x, w, b, eps); + } +}; + +class GroupNorm : public GGMLBlock { +protected: + int num_groups; + int64_t num_channels; + float eps; + bool affine; + std::string prefix; + + void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override { + this->prefix = prefix; + if (affine) { + enum ggml_type wtype = GGML_TYPE_F32; + enum ggml_type bias_wtype = GGML_TYPE_F32; + params["weight"] = ggml_new_tensor_1d(ctx, wtype, num_channels); + params["bias"] = ggml_new_tensor_1d(ctx, bias_wtype, num_channels); + } + } + +public: + GroupNorm(int num_groups, + int64_t num_channels, + float eps = 1e-05f, + bool affine = true) + : num_groups(num_groups), + num_channels(num_channels), + eps(eps), + affine(affine) {} + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + ggml_tensor* w = nullptr; + ggml_tensor* b = nullptr; + if (affine) { + w = params["weight"]; + b = params["bias"]; + if (ctx->weight_adapter) { + w = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, w, prefix + "weight"); + b = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, b, prefix + "bias"); + } + } + return ggml_ext_group_norm(ctx->ggml_ctx, x, w, b, num_groups); + } +}; + +class GroupNorm32 : public GroupNorm { +public: + GroupNorm32(int64_t num_channels) + : GroupNorm(32, num_channels, 1e-06f) {} +}; + +class RMSNorm : public UnaryBlock { +protected: + int64_t hidden_size; + float eps; + std::string prefix; + + void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, std::string prefix = "") override { + this->prefix = prefix; + enum ggml_type wtype = GGML_TYPE_F32; + params["weight"] = ggml_new_tensor_1d(ctx, wtype, hidden_size); + } + +public: + RMSNorm(int64_t hidden_size, + float eps = 1e-06f) + : hidden_size(hidden_size), + eps(eps) {} + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { + ggml_tensor* w = params["weight"]; + if (ctx->weight_adapter) { + w = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, w, prefix + "weight"); + } + x = ggml_rms_norm(ctx->ggml_ctx, x, eps); + x = ggml_mul_inplace(ctx->ggml_ctx, x, w); + return x; + } +}; + +class MultiheadAttention : public GGMLBlock { +protected: + int64_t embed_dim; + int64_t n_head; + bool proj_in; + std::string q_proj_name; + std::string k_proj_name; + std::string v_proj_name; + std::string in_proj_name; + std::string out_proj_name; + +public: + MultiheadAttention(int64_t embed_dim, + int64_t n_head, + bool qkv_proj_bias = true, + bool out_proj_bias = true, + bool proj_in = false, + std::string q_proj_name = "q_proj", + std::string k_proj_name = "k_proj", + std::string v_proj_name = "v_proj", + std::string in_proj_name = "in_proj", + std::string out_proj_name = "out_proj") + : embed_dim(embed_dim), + n_head(n_head), + proj_in(proj_in), + q_proj_name(q_proj_name), + k_proj_name(k_proj_name), + v_proj_name(v_proj_name), + in_proj_name(in_proj_name), + out_proj_name(out_proj_name) { + if (proj_in) { + blocks[in_proj_name] = std::shared_ptr(new Linear(embed_dim, embed_dim * 3, qkv_proj_bias)); + } else { + blocks[q_proj_name] = std::shared_ptr(new Linear(embed_dim, embed_dim, qkv_proj_bias)); + blocks[k_proj_name] = std::shared_ptr(new Linear(embed_dim, embed_dim, qkv_proj_bias)); + blocks[v_proj_name] = std::shared_ptr(new Linear(embed_dim, embed_dim, qkv_proj_bias)); + } + blocks[out_proj_name] = std::shared_ptr(new Linear(embed_dim, embed_dim, out_proj_bias)); + } + + // x: [N, n_token, embed_dim] + ggml_tensor* forward(GGMLRunnerContext* ctx, + ggml_tensor* x, + ggml_tensor* mask = nullptr) { + auto out_proj = std::dynamic_pointer_cast(blocks[out_proj_name]); + + ggml_tensor* q; + ggml_tensor* k; + ggml_tensor* v; + if (proj_in) { + auto in_proj = std::dynamic_pointer_cast(blocks[in_proj_name]); + auto qkv = in_proj->forward(ctx, x); + auto qkv_vec = split_qkv(ctx->ggml_ctx, qkv); + q = qkv_vec[0]; + k = qkv_vec[1]; + v = qkv_vec[2]; + } else { + auto q_proj = std::dynamic_pointer_cast(blocks[q_proj_name]); + auto k_proj = std::dynamic_pointer_cast(blocks[k_proj_name]); + auto v_proj = std::dynamic_pointer_cast(blocks[v_proj_name]); + + q = q_proj->forward(ctx, x); + k = k_proj->forward(ctx, x); + v = v_proj->forward(ctx, x); + } + + x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, n_head, mask, false); // [N, n_token, embed_dim] + + x = out_proj->forward(ctx, x); // [N, n_token, embed_dim] + return x; + } +}; + +#endif // __SD_MODEL_COMMON_GGML_BLOCK_HPP__ diff --git a/otherarch/sdcpp/src/model/common/rope.hpp b/otherarch/sdcpp/src/model/common/rope.hpp index c36cc7e9c..4f2580394 100644 --- a/otherarch/sdcpp/src/model/common/rope.hpp +++ b/otherarch/sdcpp/src/model/common/rope.hpp @@ -2,9 +2,13 @@ #define __SD_MODEL_COMMON_ROPE_HPP__ #include +#include #include +#include #include -#include "core/ggml_extend.hpp" +#include "core/ggml_extend.h" +#include "core/ggml_runner.h" +#include "core/util.h" namespace Rope { enum class EmbedNDLayout { diff --git a/otherarch/sdcpp/src/model/detector/yolov8.h b/otherarch/sdcpp/src/model/detector/yolov8.h index a90fdf9a6..980d08c97 100644 --- a/otherarch/sdcpp/src/model/detector/yolov8.h +++ b/otherarch/sdcpp/src/model/detector/yolov8.h @@ -8,8 +8,9 @@ #include #include -#include "core/ggml_extend.hpp" +#include "core/ggml_runner.h" #include "core/util.h" +#include "model/common/ggml_block.hpp" struct YOLOv8Config { std::array out_channels{}; @@ -68,12 +69,12 @@ struct YOLOv8Config { } if (config.valid) { - LOG_DEBUG("yolov8: classes=%d, reg_max=%d, p3=%d, p4=%d, p5=%d", - config.num_classes, - config.reg_max, - config.out_channels[15], - config.out_channels[18], - config.out_channels[21]); + LOG_VERBOSE("yolov8: classes=%d, reg_max=%d, p3=%d, p4=%d, p5=%d", + config.num_classes, + config.reg_max, + config.out_channels[15], + config.out_channels[18], + config.out_channels[21]); } return config; } @@ -355,7 +356,7 @@ struct YOLOv8Runner : public GGMLRunner { sd::Tensor compute(int n_threads, const sd::Tensor& input) { auto get_graph = [&]() { return build_graph(input); }; - return take_or_empty(GGMLRunner::compute(get_graph, n_threads, false, false, false)); + return take_or_empty(GGMLRunner::compute(get_graph, n_threads, false)); } }; diff --git a/otherarch/sdcpp/src/model/diffusion/anima.hpp b/otherarch/sdcpp/src/model/diffusion/anima.hpp index 4fe7e4650..930efacea 100644 --- a/otherarch/sdcpp/src/model/diffusion/anima.hpp +++ b/otherarch/sdcpp/src/model/diffusion/anima.hpp @@ -2,6 +2,7 @@ #define __SD_MODEL_DIFFUSION_ANIMA_HPP__ #include +#include #include #include #include @@ -46,11 +47,11 @@ namespace Anima { } if (detected_layers > 0) { config.num_layers = detected_layers; - LOG_DEBUG("anima: num_layers = %" PRId64 ", hidden_size = %" PRId64 ", num_heads = %" PRId64 ", head_dim = %" PRId64, - config.num_layers, - config.hidden_size, - config.num_heads, - config.head_dim); + LOG_VERBOSE("anima: num_layers = %" PRId64 ", hidden_size = %" PRId64 ", num_heads = %" PRId64 ", head_dim = %" PRId64, + config.num_layers, + config.hidden_size, + config.num_heads, + config.head_dim); } return config; } @@ -717,7 +718,7 @@ namespace Anima { auto get_graph = [&]() -> ggml_cgraph* { return build_graph(x, timesteps, context, t5_ids, t5_weights, ref_latents); }; - return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), x.dim()); + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim()); } sd::Tensor compute(int n_threads, diff --git a/otherarch/sdcpp/src/model/diffusion/animatediff.hpp b/otherarch/sdcpp/src/model/diffusion/animatediff.hpp index d470544fe..6a94e21df 100644 --- a/otherarch/sdcpp/src/model/diffusion/animatediff.hpp +++ b/otherarch/sdcpp/src/model/diffusion/animatediff.hpp @@ -1,8 +1,10 @@ #ifndef __SD_MODEL_DIFFUSION_ANIMATEDIFF_HPP__ #define __SD_MODEL_DIFFUSION_ANIMATEDIFF_HPP__ -#include "core/ggml_extend.hpp" +#include "core/ggml_extend.h" +#include "core/ggml_runner.h" #include "model/common/block.hpp" +#include "model/common/ggml_block.hpp" // AnimateDiff (https://arxiv.org/abs/2307.04725) SD 1.5 motion modules. namespace AnimateDiff { diff --git a/otherarch/sdcpp/src/model/diffusion/boogu.hpp b/otherarch/sdcpp/src/model/diffusion/boogu.hpp index 7bdbcd814..8c91bb55d 100644 --- a/otherarch/sdcpp/src/model/diffusion/boogu.hpp +++ b/otherarch/sdcpp/src/model/diffusion/boogu.hpp @@ -2,11 +2,15 @@ #define __SD_MODEL_DIFFUSION_BOOGU_HPP__ #include +#include #include #include #include -#include "core/ggml_extend.hpp" +#include "core/ggml_extend.h" +#include "core/ggml_runner.h" +#include "core/util.h" +#include "model/common/ggml_block.hpp" #include "model/common/rope.hpp" #include "model/diffusion/dit.hpp" #include "model/diffusion/model.hpp" @@ -109,16 +113,16 @@ namespace Boogu { } config.timestep_embed_dim = std::min(config.hidden_size, 1024); - LOG_DEBUG("boogu_image: layers=%" PRId64 ", double_stream_layers=%" PRId64 ", refiner_layers=%" PRId64 ", hidden=%" PRId64 ", heads=%" PRId64 ", kv_heads=%" PRId64 ", head_dim=%" PRId64 ", in_channels=%" PRId64 ", out_channels=%" PRId64, - config.num_layers, - config.num_double_stream_layers, - config.num_refiner_layers, - config.hidden_size, - config.num_attention_heads, - config.num_kv_heads, - config.head_dim, - config.in_channels, - config.out_channels); + LOG_VERBOSE("boogu_image: layers=%" PRId64 ", double_stream_layers=%" PRId64 ", refiner_layers=%" PRId64 ", hidden=%" PRId64 ", heads=%" PRId64 ", kv_heads=%" PRId64 ", head_dim=%" PRId64 ", in_channels=%" PRId64 ", out_channels=%" PRId64, + config.num_layers, + config.num_double_stream_layers, + config.num_refiner_layers, + config.hidden_size, + config.num_attention_heads, + config.num_kv_heads, + config.head_dim, + config.in_channels, + config.out_channels); return config; } }; @@ -815,7 +819,7 @@ namespace Boogu { auto get_graph = [&]() -> ggml_cgraph* { return build_graph(x, timesteps, context, ref_latents); }; - return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), x.dim()); + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim()); } sd::Tensor compute(int n_threads, diff --git a/otherarch/sdcpp/src/model/diffusion/control.hpp b/otherarch/sdcpp/src/model/diffusion/control.hpp index bf3c7e435..66cce5759 100644 --- a/otherarch/sdcpp/src/model/diffusion/control.hpp +++ b/otherarch/sdcpp/src/model/diffusion/control.hpp @@ -423,19 +423,26 @@ struct ControlNet : public GGMLRunner { return build_graph(x, hint, timesteps, context, y); }; - auto compute_result = GGMLRunner::compute(get_graph, n_threads, false, false, false, true); + auto read_outputs = [&]() { + controls.clear(); + controls.reserve(control_outputs_ggml.size()); + for (ggml_tensor* control : control_outputs_ggml) { + auto control_host = restore_trailing_singleton_dims(sd::make_sd_tensor_from_ggml(control), 4); + if (control_host.empty()) { + return false; + } + controls.push_back(std::move(control_host)); + } + return true; + }; + auto compute_result = GGMLRunner::compute(get_graph, n_threads, false, true, read_outputs); + control_outputs_ggml.clear(); + guided_hint_output_ggml = nullptr; if (!compute_result.has_value()) { + controls.clear(); return std::nullopt; } - guided_hint_cached = get_cache_tensor_by_name(guided_hint_cache_name()) != nullptr; - controls.clear(); - controls.reserve(control_outputs_ggml.size()); - for (ggml_tensor* control : control_outputs_ggml) { - auto control_host = restore_trailing_singleton_dims(sd::make_sd_tensor_from_ggml(control), 4); - GGML_ASSERT(!control_host.empty()); - controls.push_back(std::move(control_host)); - } return controls; } @@ -444,10 +451,10 @@ struct ControlNet : public GGMLRunner { std::map tensors; control_net.get_param_tensors(tensors); - auto manager = std::dynamic_pointer_cast(weight_manager.lock()); + auto manager = std::dynamic_pointer_cast(residency_manager.lock()); if (manager == nullptr) { owned_model_manager = std::make_shared(); - weight_manager = owned_model_manager; + residency_manager = owned_model_manager; manager = owned_model_manager; } diff --git a/otherarch/sdcpp/src/model/diffusion/dit.hpp b/otherarch/sdcpp/src/model/diffusion/dit.hpp index e7d7b67f8..c709b8ab4 100644 --- a/otherarch/sdcpp/src/model/diffusion/dit.hpp +++ b/otherarch/sdcpp/src/model/diffusion/dit.hpp @@ -1,7 +1,8 @@ #ifndef __SD_MODEL_DIFFUSION_DIT_HPP__ #define __SD_MODEL_DIFFUSION_DIT_HPP__ -#include "core/ggml_extend.hpp" +#include "core/ggml_extend.h" +#include "core/ggml_runner.h" namespace DiT { inline ggml_tensor* patchify(ggml_context* ctx, diff --git a/otherarch/sdcpp/src/model/diffusion/ernie_image.hpp b/otherarch/sdcpp/src/model/diffusion/ernie_image.hpp index 12fcada59..c4a5df192 100644 --- a/otherarch/sdcpp/src/model/diffusion/ernie_image.hpp +++ b/otherarch/sdcpp/src/model/diffusion/ernie_image.hpp @@ -1,6 +1,7 @@ #ifndef __SD_MODEL_DIFFUSION_ERNIE_IMAGE_HPP__ #define __SD_MODEL_DIFFUSION_ERNIE_IMAGE_HPP__ +#include #include #include @@ -72,13 +73,13 @@ namespace ErnieImage { for (int axis_dim : config.axes_dim) { config.axes_dim_sum += axis_dim; } - LOG_DEBUG("ernie_image: num_layers = %" PRId64 ", hidden_size = %" PRId64 ", num_heads = %" PRId64 ", ffn_hidden_size = %" PRId64 ", in_channels = %" PRId64 ", out_channels = %" PRId64, - config.num_layers, - config.hidden_size, - config.num_heads, - config.ffn_hidden_size, - config.in_channels, - config.out_channels); + LOG_VERBOSE("ernie_image: num_layers = %" PRId64 ", hidden_size = %" PRId64 ", num_heads = %" PRId64 ", ffn_hidden_size = %" PRId64 ", in_channels = %" PRId64 ", out_channels = %" PRId64, + config.num_layers, + config.hidden_size, + config.num_heads, + config.ffn_hidden_size, + config.in_channels, + config.out_channels); return config; } }; @@ -440,7 +441,7 @@ namespace ErnieImage { auto get_graph = [&]() -> ggml_cgraph* { return build_graph(x, timesteps, context); }; - return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), x.dim()); + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim()); } sd::Tensor compute(int n_threads, diff --git a/otherarch/sdcpp/src/model/diffusion/flux.hpp b/otherarch/sdcpp/src/model/diffusion/flux.hpp index 375576c14..0ac54a9a1 100644 --- a/otherarch/sdcpp/src/model/diffusion/flux.hpp +++ b/otherarch/sdcpp/src/model/diffusion/flux.hpp @@ -1,8 +1,11 @@ #ifndef __SD_MODEL_DIFFUSION_FLUX_HPP__ #define __SD_MODEL_DIFFUSION_FLUX_HPP__ +#include #include #include +#include "core/ggml_extend_backend.h" +#include "core/ggml_tensor_utils.h" #include "core/util.h" #include "model/adapter/pulid.hpp" @@ -123,16 +126,16 @@ namespace Flux { config.guidance_embed = true; } if (name.find("__x0__") != std::string::npos) { - LOG_DEBUG("using x0 prediction"); + LOG_VERBOSE("using x0 prediction"); config.chroma_radiance_params.use_x0 = true; } if (name.find("__32x32__") != std::string::npos) { - LOG_DEBUG("using patch size 32"); + LOG_VERBOSE("using patch size 32"); config.patch_size = 32; } if (name.find("img_in_patch.weight") != std::string::npos) { actual_radiance_patch_size = tensor_storage.ne[0]; - LOG_DEBUG("actual radiance patch size: %" PRId64, actual_radiance_patch_size); + LOG_VERBOSE("actual radiance patch size: %" PRId64, actual_radiance_patch_size); } if (name.find("distilled_guidance_layer.in_proj.weight") != std::string::npos) { config.is_chroma = true; @@ -169,7 +172,7 @@ namespace Flux { } if (actual_radiance_patch_size > 0 && actual_radiance_patch_size != config.patch_size) { GGML_ASSERT(config.patch_size == 2 * actual_radiance_patch_size); - LOG_DEBUG("using fake x2 patch size"); + LOG_VERBOSE("using fake x2 patch size"); config.chroma_radiance_params.fake_patch_size_x2 = true; } if (head_dim > 0) { @@ -179,13 +182,13 @@ namespace Flux { for (int axis_dim : config.axes_dim) { config.axes_dim_sum += axis_dim; } - LOG_DEBUG("flux: depth = %d, depth_single_blocks = %d, guidance_embed = %s, context_in_dim = %" PRId64 ", hidden_size = %" PRId64 ", num_heads = %d", - config.depth, - config.depth_single_blocks, - config.guidance_embed ? "true" : "false", - config.context_in_dim, - config.hidden_size, - config.num_heads); + LOG_VERBOSE("flux: depth = %d, depth_single_blocks = %d, guidance_embed = %s, context_in_dim = %" PRId64 ", hidden_size = %" PRId64 ", num_heads = %d", + config.depth, + config.depth_single_blocks, + config.guidance_embed ? "true" : "false", + config.context_in_dim, + config.hidden_size, + config.num_heads); return config; } }; @@ -1560,7 +1563,7 @@ namespace Flux { config.axes_dim, sd_version_is_longcat(version)); int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); - // LOG_DEBUG("pos_len %d", pos_len); + // LOG_VERBOSE("pos_len %d", pos_len); auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len); // pe->data = pe_vec.data(); // print_ggml_tensor(pe); @@ -1626,7 +1629,7 @@ namespace Flux { return build_graph(x, timesteps, context, c_concat, y, guidance, ref_latents, ref_index_mode, skip_layers, pulid_id, pulid_id_weight); }; - auto result = restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), x.dim()); + auto result = restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim()); return result; } @@ -1702,7 +1705,7 @@ namespace Flux { GGML_ASSERT(!out_opt.empty()); out = std::move(out_opt); print_sd_tensor(out); - LOG_DEBUG("flux test done in %lldms", t1 - t0); + LOG_VERBOSE("flux test done in %lldms", t1 - t0); } } diff --git a/otherarch/sdcpp/src/model/diffusion/hidream_o1.hpp b/otherarch/sdcpp/src/model/diffusion/hidream_o1.hpp index 9d3df0391..677f81e7d 100644 --- a/otherarch/sdcpp/src/model/diffusion/hidream_o1.hpp +++ b/otherarch/sdcpp/src/model/diffusion/hidream_o1.hpp @@ -325,13 +325,11 @@ namespace HiDreamO1 { sd::Tensor compute(int n_threads, const sd::Tensor& image, - bool auto_free = true, - bool free_compute_buffer = true, - bool free_compute_params = true) { + bool auto_runner_end = true) { auto get_graph = [&]() { return build_graph(image); }; - auto output = GGMLRunner::compute(get_graph, n_threads, auto_free, free_compute_buffer, free_compute_params); + auto output = GGMLRunner::compute(get_graph, n_threads, auto_runner_end); return output.has_value() ? std::move(output.value()) : sd::Tensor(); } }; @@ -459,7 +457,7 @@ namespace HiDreamO1 { auto get_graph = [&]() { return build_graph(x, timestep, input_ids, input_pos, token_types, vinput_mask, image_embeds, ref_images); }; - return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), x.dim()); + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim()); } sd::Tensor compute(int n_threads, @@ -510,8 +508,8 @@ namespace HiDreamO1 { vision_runner->set_weight_adapter(adapter); } - void runner_done() override { - vision_runner->runner_done(); + void runner_end() override { + vision_runner->runner_end(); } SDCondition get_learned_condition(int n_threads, @@ -659,7 +657,7 @@ namespace HiDreamO1 { result.c_vinput_mask = sd::Tensor(vinput_mask_shape, std::move(vinput_mask)); result.c_image_embeds.reserve(vlm_images.size()); for (const auto& vlm_image : vlm_images) { - auto image_embed = vision_runner->compute(n_threads, vlm_image.second, false, true, true); + auto image_embed = vision_runner->compute(n_threads, vlm_image.second, false); if (image_embed.empty()) { LOG_ERROR("hidream_o1 conditioner: encode VLM image failed"); return SDCondition(); diff --git a/otherarch/sdcpp/src/model/diffusion/hunyuan.hpp b/otherarch/sdcpp/src/model/diffusion/hunyuan.hpp index f81085639..c111674f5 100644 --- a/otherarch/sdcpp/src/model/diffusion/hunyuan.hpp +++ b/otherarch/sdcpp/src/model/diffusion/hunyuan.hpp @@ -1,6 +1,7 @@ #ifndef __SD_MODEL_DIFFUSION_HUNYUAN_HPP__ #define __SD_MODEL_DIFFUSION_HUNYUAN_HPP__ +#include #include #include "model/common/block.hpp" @@ -266,16 +267,16 @@ namespace Hunyuan { GGML_ASSERT(config.hidden_size / config.num_heads == config.axes_dim_sum); if (inferred) { - LOG_DEBUG("hunyuan video: depth = %d, single depth = %d, in_channels = %" PRId64 ", out_channels = %" PRId64 ", hidden_size = %" PRId64 ", context_in_dim = %" PRId64 ", patch_size = %dx%dx%d", - config.depth, - config.depth_single_blocks, - config.in_channels, - config.out_channels, - config.hidden_size, - config.context_in_dim, - std::get<0>(config.patch_size), - std::get<1>(config.patch_size), - std::get<2>(config.patch_size)); + LOG_VERBOSE("hunyuan video: depth = %d, single depth = %d, in_channels = %" PRId64 ", out_channels = %" PRId64 ", hidden_size = %" PRId64 ", context_in_dim = %" PRId64 ", patch_size = %dx%dx%d", + config.depth, + config.depth_single_blocks, + config.in_channels, + config.out_channels, + config.hidden_size, + config.context_in_dim, + std::get<0>(config.patch_size), + std::get<1>(config.patch_size), + std::get<2>(config.patch_size)); } return config; } @@ -615,7 +616,7 @@ namespace Hunyuan { config.theta, config.axes_dim); int64_t pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); - // LOG_DEBUG("pos_len %d", pos_len); + // LOG_VERBOSE("pos_len %d", pos_len); auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len); // pe->data = pe_vec.data(); // print_ggml_tensor(pe, true, "pe"); @@ -654,7 +655,7 @@ namespace Hunyuan { return build_graph(x, timesteps, context, c_concat, y, guidance, byt5, vision, timestep_r); }; - return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), x.dim()); + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim()); } sd::Tensor compute(int n_threads, diff --git a/otherarch/sdcpp/src/model/diffusion/ideogram4.hpp b/otherarch/sdcpp/src/model/diffusion/ideogram4.hpp index 6ce8e1fab..9e2aa3612 100644 --- a/otherarch/sdcpp/src/model/diffusion/ideogram4.hpp +++ b/otherarch/sdcpp/src/model/diffusion/ideogram4.hpp @@ -2,14 +2,18 @@ #define __SD_MODEL_DIFFUSION_IDEOGRAM4_HPP__ #include +#include #include #include #include #include #include -#include "core/ggml_extend.hpp" +#include "core/ggml_extend.h" #include "core/ggml_graph_cut.h" +#include "core/ggml_runner.h" +#include "core/util.h" +#include "model/common/ggml_block.hpp" #include "model/common/rope.hpp" #include "model/diffusion/model.hpp" @@ -58,11 +62,11 @@ namespace Ideogram4 { } if (detected_layers > 0) { config.num_layers = detected_layers; - LOG_DEBUG("ideogram4: num_layers = %" PRId64 ", emb_dim = %" PRId64 ", num_heads = %" PRId64 ", intermediate_size = %" PRId64, - config.num_layers, - config.emb_dim, - config.num_heads, - config.intermediate_size); + LOG_VERBOSE("ideogram4: num_layers = %" PRId64 ", emb_dim = %" PRId64 ", num_heads = %" PRId64 ", intermediate_size = %" PRId64, + config.num_layers, + config.emb_dim, + config.num_heads, + config.intermediate_size); } return config; } @@ -465,7 +469,7 @@ namespace Ideogram4 { } } if (has_uncond_model) { - LOG_DEBUG("using uncond model"); + LOG_VERBOSE("using uncond model"); uncond_model = Ideogram4Transformer(config); uncond_model.init(params_ctx, tensor_storage_map, uncond_prefix); } @@ -537,7 +541,7 @@ namespace Ideogram4 { auto get_graph = [&]() -> ggml_cgraph* { return build_graph(x, timesteps, context, use_uncond_model); }; - return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), x.dim()); + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim()); } sd::Tensor compute(int n_threads, diff --git a/otherarch/sdcpp/src/model/diffusion/krea2.hpp b/otherarch/sdcpp/src/model/diffusion/krea2.hpp index b3947b71d..41760581d 100644 --- a/otherarch/sdcpp/src/model/diffusion/krea2.hpp +++ b/otherarch/sdcpp/src/model/diffusion/krea2.hpp @@ -12,8 +12,11 @@ #include #include -#include "core/ggml_extend.hpp" +#include "core/ggml_extend.h" #include "core/ggml_graph_cut.h" +#include "core/ggml_runner.h" +#include "core/util.h" +#include "model/common/ggml_block.hpp" #include "model/common/rope.hpp" #include "model/diffusion/dit.hpp" #include "model/diffusion/flux.hpp" @@ -143,16 +146,16 @@ namespace Krea2 { } config.update_axes_dim(); - LOG_DEBUG("krea2: layers=%" PRId64 ", features=%" PRId64 ", heads=%" PRId64 ", kv_heads=%" PRId64 ", text_dim=%" PRId64 ", text_layers=%" PRId64 ", text_heads=%" PRId64 ", text_kv_heads=%" PRId64 ", channels=%" PRId64, - config.layers, - config.features, - config.heads, - config.kv_heads, - config.text_dim, - config.text_layers, - config.text_heads, - config.text_kv_heads, - config.in_channels); + LOG_VERBOSE("krea2: layers=%" PRId64 ", features=%" PRId64 ", heads=%" PRId64 ", kv_heads=%" PRId64 ", text_dim=%" PRId64 ", text_layers=%" PRId64 ", text_heads=%" PRId64 ", text_kv_heads=%" PRId64 ", channels=%" PRId64, + config.layers, + config.features, + config.heads, + config.kv_heads, + config.text_dim, + config.text_layers, + config.text_heads, + config.text_kv_heads, + config.in_channels); return config; } }; @@ -775,7 +778,7 @@ namespace Krea2 { auto get_graph = [&]() -> ggml_cgraph* { return build_graph(x, timesteps, context, ref_latents, ref_image_params); }; - return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), x.dim()); + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim()); } sd::Tensor compute(int n_threads, diff --git a/otherarch/sdcpp/src/model/diffusion/lens.hpp b/otherarch/sdcpp/src/model/diffusion/lens.hpp index 931a8527c..35b4b517d 100644 --- a/otherarch/sdcpp/src/model/diffusion/lens.hpp +++ b/otherarch/sdcpp/src/model/diffusion/lens.hpp @@ -1,6 +1,7 @@ #ifndef __SD_MODEL_DIFFUSION_LENS_HPP__ #define __SD_MODEL_DIFFUSION_LENS_HPP__ +#include #include #include @@ -66,14 +67,14 @@ namespace Lens { for (int axis_dim : config.axes_dim) { config.axes_dim_sum += axis_dim; } - LOG_DEBUG("lens: num_layers = %d, selected_layer_count = %d, hidden_size = %" PRId64 ", num_attention_heads = %" PRId64 ", attention_head_dim = %" PRId64 ", in_channels = %" PRId64 ", out_channels = %" PRId64, - config.num_layers, - config.selected_layer_count, - config.num_attention_heads * config.attention_head_dim, - config.num_attention_heads, - config.attention_head_dim, - config.in_channels, - config.out_channels); + LOG_VERBOSE("lens: num_layers = %d, selected_layer_count = %d, hidden_size = %" PRId64 ", num_attention_heads = %" PRId64 ", attention_head_dim = %" PRId64 ", in_channels = %" PRId64 ", out_channels = %" PRId64, + config.num_layers, + config.selected_layer_count, + config.num_attention_heads * config.attention_head_dim, + config.num_attention_heads, + config.attention_head_dim, + config.in_channels, + config.out_channels); return config; } }; @@ -408,7 +409,7 @@ namespace Lens { auto get_graph = [&]() -> ggml_cgraph* { return build_graph(x, timesteps, context); }; - return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), x.dim()); + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim()); } sd::Tensor compute(int n_threads, diff --git a/otherarch/sdcpp/src/model/diffusion/lingbot_video.hpp b/otherarch/sdcpp/src/model/diffusion/lingbot_video.hpp index 1868daa05..1a93b458e 100644 --- a/otherarch/sdcpp/src/model/diffusion/lingbot_video.hpp +++ b/otherarch/sdcpp/src/model/diffusion/lingbot_video.hpp @@ -127,17 +127,17 @@ namespace LingBotVideo { config.topk_group = 2; config.routed_scaling_factor = 2.5f; } - LOG_DEBUG("lingbot_video: depth = %" PRId64 ", hidden_size = %" PRId64 ", heads = %" PRId64 ", text_dim = %" PRId64 ", experts = %" PRId64 ", experts_per_tok = %" PRId64 ", n_group = %" PRId64 ", topk_group = %" PRId64 ", route_scale = %.2f, sparse_layers = %zu", - config.depth, - config.hidden_size, - config.num_attention_heads, - config.text_dim, - config.num_experts, - config.num_experts_per_tok, - config.n_group, - config.topk_group, - config.routed_scaling_factor, - config.sparse_layers.size()); + LOG_VERBOSE("lingbot_video: depth = %" PRId64 ", hidden_size = %" PRId64 ", heads = %" PRId64 ", text_dim = %" PRId64 ", experts = %" PRId64 ", experts_per_tok = %" PRId64 ", n_group = %" PRId64 ", topk_group = %" PRId64 ", route_scale = %.2f, sparse_layers = %zu", + config.depth, + config.hidden_size, + config.num_attention_heads, + config.text_dim, + config.num_experts, + config.num_experts_per_tok, + config.n_group, + config.topk_group, + config.routed_scaling_factor, + config.sparse_layers.size()); return config; } }; @@ -674,7 +674,7 @@ namespace LingBotVideo { auto get_graph = [&]() -> ggml_cgraph* { return build_graph(x, timesteps, context); }; - return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), x.dim()); + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim()); } sd::Tensor compute(int n_threads, diff --git a/otherarch/sdcpp/src/model/diffusion/ltxv.hpp b/otherarch/sdcpp/src/model/diffusion/ltxv.hpp index b75c9f7fa..cddfc2d3b 100644 --- a/otherarch/sdcpp/src/model/diffusion/ltxv.hpp +++ b/otherarch/sdcpp/src/model/diffusion/ltxv.hpp @@ -2,12 +2,15 @@ #define __SD_MODEL_DIFFUSION_LTXV_HPP__ #include +#include #include #include #include #include #include #include +#include "core/ggml_extend_backend.h" +#include "core/ggml_tensor_utils.h" #include "model/common/block.hpp" #include "model/common/rope.hpp" @@ -274,12 +277,12 @@ namespace LTXV { config.audio_connector_apply_gated_attention = true; } } - LOG_DEBUG("ltxav: num_layers = %" PRId64 ", hidden_size = %" PRId64 ", num_attention_heads = %" PRId64 ", audio_hidden_size = %" PRId64 ", audio_num_attention_heads = %" PRId64, - config.num_layers, - config.hidden_size, - config.num_attention_heads, - config.audio_hidden_size, - config.audio_num_attention_heads); + LOG_VERBOSE("ltxav: num_layers = %" PRId64 ", hidden_size = %" PRId64 ", num_attention_heads = %" PRId64 ", audio_hidden_size = %" PRId64 ", audio_num_attention_heads = %" PRId64, + config.num_layers, + config.hidden_size, + config.num_attention_heads, + config.audio_hidden_size, + config.audio_num_attention_heads); return config; } }; @@ -1998,7 +2001,7 @@ namespace LTXV { auto get_graph = [&]() -> ggml_cgraph* { return build_graph(x, timesteps, context, audio_x, audio_timesteps, audio_length, frame_rate, video_positions); }; - auto out = restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), x.dim()); + auto out = restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim()); return out; } @@ -2070,7 +2073,7 @@ namespace LTXV { GGML_ASSERT(!out_opt.empty()); print_sd_tensor(out_opt, false, "ltxav_out"); - LOG_DEBUG("ltxav test done in %lldms", t1 - t0); + LOG_VERBOSE("ltxav test done in %lldms", t1 - t0); } static void load_from_file_and_test(const std::string& model_path, diff --git a/otherarch/sdcpp/src/model/diffusion/mage_flow.hpp b/otherarch/sdcpp/src/model/diffusion/mage_flow.hpp index 6ac2d2aa5..89df0f482 100644 --- a/otherarch/sdcpp/src/model/diffusion/mage_flow.hpp +++ b/otherarch/sdcpp/src/model/diffusion/mage_flow.hpp @@ -142,7 +142,7 @@ namespace MageFlow { auto get_graph = [&]() -> ggml_cgraph* { return build_graph(x, timesteps, context, ref_latents); }; - return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), x.dim()); + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim()); } sd::Tensor compute(int n_threads, diff --git a/otherarch/sdcpp/src/model/diffusion/minimax_h3.hpp b/otherarch/sdcpp/src/model/diffusion/minimax_h3.hpp index a0a6c7aeb..3e2b69db0 100644 --- a/otherarch/sdcpp/src/model/diffusion/minimax_h3.hpp +++ b/otherarch/sdcpp/src/model/diffusion/minimax_h3.hpp @@ -2,12 +2,14 @@ #define __SD_MODEL_DIFFUSION_MINIMAX_H3_HPP__ #include +#include #include #include #include #include #include #include +#include "core/ggml_tensor_utils.h" #include "core/ggml_graph_cut.h" #include "model/diffusion/dit.hpp" @@ -106,14 +108,14 @@ namespace MiniMaxH3 { config.rope_inv_freq_len = inv_freq->ne[0]; } - LOG_DEBUG("minimax_h3: layers=%" PRId64 ", hidden=%" PRId64 ", heads=%" PRId64 - ", head_dim=%" PRId64 ", ffn=%" PRId64 ", adaln_curve=%" PRId64, - config.num_layers, - config.hidden_size, - config.num_attention_heads, - config.attention_head_dim, - config.ffn_hidden_size, - config.adaln_curve_grid); + LOG_VERBOSE("minimax_h3: layers=%" PRId64 ", hidden=%" PRId64 ", heads=%" PRId64 + ", head_dim=%" PRId64 ", ffn=%" PRId64 ", adaln_curve=%" PRId64, + config.num_layers, + config.hidden_size, + config.num_attention_heads, + config.attention_head_dim, + config.ffn_hidden_size, + config.adaln_curve_grid); return config; } }; @@ -1166,11 +1168,9 @@ namespace MiniMaxH3 { extra->video_sigma_shift, extra->audio_sigma_shift); }; - return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, - n_threads, - false, - false, - false), + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, + n_threads, + false), params.x->dim()); } }; diff --git a/otherarch/sdcpp/src/model/diffusion/minit2i.hpp b/otherarch/sdcpp/src/model/diffusion/minit2i.hpp index 284661054..06c4db5a5 100644 --- a/otherarch/sdcpp/src/model/diffusion/minit2i.hpp +++ b/otherarch/sdcpp/src/model/diffusion/minit2i.hpp @@ -2,6 +2,7 @@ #define __SD_MODEL_DIFFUSION_MINIT2I_HPP__ #include +#include #include #include #include @@ -9,7 +10,10 @@ #include #include -#include "core/ggml_extend.hpp" +#include "core/ggml_extend.h" +#include "core/ggml_runner.h" +#include "core/util.h" +#include "model/common/ggml_block.hpp" #include "model/common/rope.hpp" #include "model/diffusion/dit.hpp" #include "model/diffusion/model.hpp" @@ -108,15 +112,15 @@ namespace MiniT2I { config.head_dim = config.hidden_size == 1248 ? 52 : 64; config.num_heads = config.hidden_size / config.head_dim; } - LOG_DEBUG("minit2i: hidden_size=%" PRId64 ", txt_hidden_size=%" PRId64 ", heads=%" PRId64 ", head_dim=%" PRId64 ", double_blocks=%" PRId64 ", txt_blocks=%" PRId64 ", patch=%" PRId64 ", in_channels=%" PRId64, - config.hidden_size, - config.txt_hidden_size, - config.num_heads, - config.head_dim, - config.depth_double, - config.txt_preamble_depth, - config.patch_size, - config.in_channels); + LOG_VERBOSE("minit2i: hidden_size=%" PRId64 ", txt_hidden_size=%" PRId64 ", heads=%" PRId64 ", head_dim=%" PRId64 ", double_blocks=%" PRId64 ", txt_blocks=%" PRId64 ", patch=%" PRId64 ", in_channels=%" PRId64, + config.hidden_size, + config.txt_hidden_size, + config.num_heads, + config.head_dim, + config.depth_double, + config.txt_preamble_depth, + config.patch_size, + config.in_channels); return config; } }; @@ -589,7 +593,7 @@ namespace MiniT2I { auto get_graph = [&]() -> ggml_cgraph* { return build_graph(x, timesteps, context, mask); }; - return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), x.dim()); + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim()); } sd::Tensor compute(int n_threads, diff --git a/otherarch/sdcpp/src/model/diffusion/mmdit.hpp b/otherarch/sdcpp/src/model/diffusion/mmdit.hpp index 6731b5fb5..9063b65d1 100644 --- a/otherarch/sdcpp/src/model/diffusion/mmdit.hpp +++ b/otherarch/sdcpp/src/model/diffusion/mmdit.hpp @@ -2,12 +2,18 @@ #define __SD_MODEL_DIFFUSION_MMDIT_HPP__ #include +#include #include #include #include -#include "core/ggml_extend.hpp" +#include "core/ggml_extend.h" +#include "core/ggml_extend_backend.h" +#include "core/ggml_runner.h" +#include "core/ggml_tensor_utils.h" +#include "core/util.h" #include "model/common/block.hpp" +#include "model/common/ggml_block.hpp" #include "model/diffusion/model.hpp" #include "model_loader.h" @@ -120,16 +126,16 @@ struct MMDiTConfig { } if (has_weight_config) { - LOG_DEBUG("mmdit: num_layers = %" PRId64 ", num_mmdit_x_layers = %" PRId64 ", hidden_size = %" PRId64 ", patch_size = %d, in_channels = %" PRId64 ", out_channels = %" PRId64 ", context_size = %" PRId64 ", adm_in_channels = %" PRId64 ", qk_norm = %s", - config.depth, - config.d_self + 1, - config.hidden_size, - config.patch_size, - config.in_channels, - config.out_channels, - config.context_size, - config.adm_in_channels, - config.qk_norm.empty() ? "none" : config.qk_norm.c_str()); + LOG_VERBOSE("mmdit: num_layers = %" PRId64 ", num_mmdit_x_layers = %" PRId64 ", hidden_size = %" PRId64 ", patch_size = %d, in_channels = %" PRId64 ", out_channels = %" PRId64 ", context_size = %" PRId64 ", adm_in_channels = %" PRId64 ", qk_norm = %s", + config.depth, + config.d_self + 1, + config.hidden_size, + config.patch_size, + config.in_channels, + config.out_channels, + config.context_size, + config.adm_in_channels, + config.qk_norm.empty() ? "none" : config.qk_norm.c_str()); } return config; } @@ -987,7 +993,7 @@ struct MMDiTRunner : public DiffusionModelRunner { return build_graph(x, timesteps, context, y, skip_layers); }; - return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), x.dim()); + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim()); } sd::Tensor compute(int n_threads, @@ -1045,7 +1051,7 @@ struct MMDiTRunner : public DiffusionModelRunner { GGML_ASSERT(!out_opt.empty()); out = std::move(out_opt); print_sd_tensor(out); - LOG_DEBUG("mmdit test done in %lldms", t1 - t0); + LOG_VERBOSE("mmdit test done in %lldms", t1 - t0); } } diff --git a/otherarch/sdcpp/src/model/diffusion/model.hpp b/otherarch/sdcpp/src/model/diffusion/model.hpp index 070ca53d4..0e447cd06 100644 --- a/otherarch/sdcpp/src/model/diffusion/model.hpp +++ b/otherarch/sdcpp/src/model/diffusion/model.hpp @@ -5,7 +5,7 @@ #include #include -#include "core/ggml_extend.hpp" +#include "core/ggml_runner.h" #include "core/tensor_ggml.hpp" #include "model/common/rope.hpp" #include "model_manager.h" diff --git a/otherarch/sdcpp/src/model/diffusion/pid.hpp b/otherarch/sdcpp/src/model/diffusion/pid.hpp index 2a698e0a6..59eac2950 100644 --- a/otherarch/sdcpp/src/model/diffusion/pid.hpp +++ b/otherarch/sdcpp/src/model/diffusion/pid.hpp @@ -1,13 +1,18 @@ #ifndef __SD_MODEL_DIFFUSION_PID_HPP__ #define __SD_MODEL_DIFFUSION_PID_HPP__ +#include #include #include #include #include #include -#include "core/ggml_extend.hpp" +#include "core/ggml_extend.h" +#include "core/ggml_runner.h" +#include "core/ggml_tensor_utils.h" +#include "core/util.h" +#include "model/common/ggml_block.hpp" #include "model/common/rope.hpp" #include "model/diffusion/dit.hpp" #include "model/diffusion/mmdit.hpp" @@ -109,16 +114,16 @@ namespace Pid { config.lq_latent_channels = latent_proj_in_channels; config.lq_latent_down_factor = latent_proj_in_channels >= 64 ? 16 : 8; } - LOG_DEBUG("pid: version = %s, patch_depth = %" PRId64 ", pixel_depth = %" PRId64 ", patch_mlp_hidden_dim = %" PRId64 ", lq_latent_channels = %" PRId64 ", lq_hidden_dim = %" PRId64 ", lq_latent_down_factor = %" PRId64 ", lq_latent_unpatchify_factor = %" PRId64 ", lq_interval = %" PRId64, - config.pit_lq_inject ? "1.5" : "1", - config.patch_depth, - config.pixel_depth, - config.patch_mlp_hidden_dim, - config.lq_latent_channels, - config.lq_hidden_dim, - config.lq_latent_down_factor, - config.lq_latent_unpatchify_factor, - config.lq_interval); + LOG_VERBOSE("pid: version = %s, patch_depth = %" PRId64 ", pixel_depth = %" PRId64 ", patch_mlp_hidden_dim = %" PRId64 ", lq_latent_channels = %" PRId64 ", lq_hidden_dim = %" PRId64 ", lq_latent_down_factor = %" PRId64 ", lq_latent_unpatchify_factor = %" PRId64 ", lq_interval = %" PRId64, + config.pit_lq_inject ? "1.5" : "1", + config.patch_depth, + config.pixel_depth, + config.patch_mlp_hidden_dim, + config.lq_latent_channels, + config.lq_hidden_dim, + config.lq_latent_down_factor, + config.lq_latent_unpatchify_factor, + config.lq_interval); return config; } }; @@ -938,7 +943,7 @@ namespace Pid { auto get_graph = [&]() -> ggml_cgraph* { return build_graph(x, timesteps, context, lq_latent, degrade_sigma); }; - return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), x.dim()); + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim()); } sd::Tensor compute(int n_threads, diff --git a/otherarch/sdcpp/src/model/diffusion/qwen_image.hpp b/otherarch/sdcpp/src/model/diffusion/qwen_image.hpp index 6f08d6180..a9dfdb838 100644 --- a/otherarch/sdcpp/src/model/diffusion/qwen_image.hpp +++ b/otherarch/sdcpp/src/model/diffusion/qwen_image.hpp @@ -2,6 +2,8 @@ #define __SD_MODEL_DIFFUSION_QWEN_IMAGE_HPP__ #include +#include "core/ggml_extend_backend.h" +#include "core/ggml_tensor_utils.h" #include "core/util.h" #include "model/common/block.hpp" @@ -49,9 +51,9 @@ namespace Qwen { } } } - LOG_DEBUG("qwen_image: num_layers = %d, zero_cond_t = %s", - config.num_layers, - config.zero_cond_t ? "true" : "false"); + LOG_VERBOSE("qwen_image: num_layers = %d, zero_cond_t = %s", + config.num_layers, + config.zero_cond_t ? "true" : "false"); return config; } }; @@ -646,7 +648,7 @@ namespace Qwen { circular_x_enabled, config.axes_dim); int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); - // LOG_DEBUG("pos_len %d", pos_len); + // LOG_VERBOSE("pos_len %d", pos_len); auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len); // pe->data = pe_vec.data(); // print_ggml_tensor(pe, true, "pe"); @@ -707,7 +709,7 @@ namespace Qwen { return build_graph(x, timesteps, context, ref_latents, ref_index_mode); }; - return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), x.dim()); + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim()); } sd::Tensor compute(int n_threads, @@ -760,7 +762,7 @@ namespace Qwen { GGML_ASSERT(!out_opt.empty()); out = std::move(out_opt); print_sd_tensor(out); - LOG_DEBUG("qwen_image test done in %lldms", t1 - t0); + LOG_VERBOSE("qwen_image test done in %lldms", t1 - t0); } } diff --git a/otherarch/sdcpp/src/model/diffusion/sefi_image.hpp b/otherarch/sdcpp/src/model/diffusion/sefi_image.hpp index 271919882..6782b68b4 100644 --- a/otherarch/sdcpp/src/model/diffusion/sefi_image.hpp +++ b/otherarch/sdcpp/src/model/diffusion/sefi_image.hpp @@ -1,6 +1,7 @@ #ifndef __SD_MODEL_DIFFUSION_SEFI_IMAGE_HPP__ #define __SD_MODEL_DIFFUSION_SEFI_IMAGE_HPP__ +#include #include #include "model/common/block.hpp" @@ -34,10 +35,10 @@ namespace SefiImage { config.hidden_size = tensor_storage.ne[1] * 2; } } - LOG_DEBUG("sefi_image: semantic_channels = %" PRId64 ", texture_latent_channels = %" PRId64 ", hidden_size = %" PRId64, - config.semantic_channels, - config.texture_latent_channels, - config.hidden_size); + LOG_VERBOSE("sefi_image: semantic_channels = %" PRId64 ", texture_latent_channels = %" PRId64 ", hidden_size = %" PRId64, + config.semantic_channels, + config.texture_latent_channels, + config.hidden_size); return config; } }; diff --git a/otherarch/sdcpp/src/model/diffusion/unet.hpp b/otherarch/sdcpp/src/model/diffusion/unet.hpp index 567573049..1c8ddf4e6 100644 --- a/otherarch/sdcpp/src/model/diffusion/unet.hpp +++ b/otherarch/sdcpp/src/model/diffusion/unet.hpp @@ -3,6 +3,7 @@ #include #include +#include "core/ggml_tensor_utils.h" #include "model.h" #include "model/common/block.hpp" @@ -128,15 +129,15 @@ struct UNetConfig { } } - LOG_DEBUG("unet: in_channels = %d, out_channels = %d, model_channels = %d, time_embed_dim = %d, context_dim = %d, adm_in_channels = %d, num_res_blocks = %d, tiny_unet = %s", - config.in_channels, - config.out_channels, - config.model_channels, - config.time_embed_dim, - config.context_dim, - config.adm_in_channels, - config.num_res_blocks, - config.tiny_unet ? "true" : "false"); + LOG_VERBOSE("unet: in_channels = %d, out_channels = %d, model_channels = %d, time_embed_dim = %d, context_dim = %d, adm_in_channels = %d, num_res_blocks = %d, tiny_unet = %s", + config.in_channels, + config.out_channels, + config.model_channels, + config.time_embed_dim, + config.context_dim, + config.adm_in_channels, + config.num_res_blocks, + config.tiny_unet ? "true" : "false"); return config; } }; @@ -835,7 +836,7 @@ struct UNetModelRunner : public DiffusionModelRunner { return build_graph(x, timesteps, context, c_concat, y, num_video_frames, controls, control_strength, ip_context, ip_scale); }; - return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), x.dim()); + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim()); } sd::Tensor compute(int n_threads, @@ -904,7 +905,7 @@ struct UNetModelRunner : public DiffusionModelRunner { GGML_ASSERT(!out_opt.empty()); out = std::move(out_opt); print_sd_tensor(out); - LOG_DEBUG("unet test done in %lldms", t1 - t0); + LOG_VERBOSE("unet test done in %lldms", t1 - t0); } } }; diff --git a/otherarch/sdcpp/src/model/diffusion/wan.hpp b/otherarch/sdcpp/src/model/diffusion/wan.hpp index 9a907dcfa..1884be5c7 100644 --- a/otherarch/sdcpp/src/model/diffusion/wan.hpp +++ b/otherarch/sdcpp/src/model/diffusion/wan.hpp @@ -1,9 +1,12 @@ #ifndef __SD_MODEL_DIFFUSION_WAN_HPP__ #define __SD_MODEL_DIFFUSION_WAN_HPP__ +#include #include #include #include +#include "core/ggml_extend_backend.h" +#include "core/ggml_tensor_utils.h" #include "model/common/block.hpp" #include "model/common/rope.hpp" @@ -75,13 +78,13 @@ namespace WAN { config.flf_pos_embed_token_number = 514; } } - LOG_DEBUG("wan: model_type = %s, num_layers = %d, vace_layers = %d, dim = %" PRId64 ", ffn_dim = %" PRId64 ", num_heads = %" PRId64, - config.model_type.c_str(), - config.num_layers, - config.vace_layers, - config.dim, - config.ffn_dim, - config.num_heads); + LOG_VERBOSE("wan: model_type = %s, num_layers = %d, vace_layers = %d, dim = %" PRId64 ", ffn_dim = %" PRId64 ", num_heads = %" PRId64, + config.model_type.c_str(), + config.num_layers, + config.vace_layers, + config.dim, + config.ffn_dim, + config.num_heads); return config; } }; @@ -909,7 +912,7 @@ namespace WAN { config.theta, config.axes_dim); int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); - // LOG_DEBUG("pos_len %d", pos_len); + // LOG_VERBOSE("pos_len %d", pos_len); auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len); // pe->data = pe_vec.data(); // print_ggml_tensor(pe); @@ -950,7 +953,7 @@ namespace WAN { return build_graph(x, timesteps, context, clip_fea, c_concat, time_dim_concat, vace_context, vace_strength); }; - return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), x.dim()); + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim()); } sd::Tensor compute(int n_threads, @@ -1007,7 +1010,7 @@ namespace WAN { GGML_ASSERT(!out_opt.empty()); out = std::move(out_opt); print_sd_tensor(out); - LOG_DEBUG("wan test done in %lldms", t1 - t0); + LOG_VERBOSE("wan test done in %lldms", t1 - t0); } } diff --git a/otherarch/sdcpp/src/model/diffusion/z_image.hpp b/otherarch/sdcpp/src/model/diffusion/z_image.hpp index 176adfc18..b29e75457 100644 --- a/otherarch/sdcpp/src/model/diffusion/z_image.hpp +++ b/otherarch/sdcpp/src/model/diffusion/z_image.hpp @@ -2,8 +2,14 @@ #define __SD_MODEL_DIFFUSION_Z_IMAGE_HPP__ #include +#include -#include "core/ggml_extend.hpp" +#include "core/ggml_extend.h" +#include "core/ggml_extend_backend.h" +#include "core/ggml_runner.h" +#include "core/ggml_tensor_utils.h" +#include "core/util.h" +#include "model/common/ggml_block.hpp" #include "model/diffusion/flux.hpp" #include "model/diffusion/mmdit.hpp" #include "model/diffusion/model.hpp" @@ -107,14 +113,14 @@ namespace ZImage { config.num_kv_heads = std::max(1, (qkv_heads - config.num_heads) / 2); } } - LOG_DEBUG("z_image: num_layers = %" PRId64 ", num_refiner_layers = %" PRId64 ", hidden_size = %" PRId64 ", num_heads = %" PRId64 ", num_kv_heads = %" PRId64 ", in_channels = %" PRId64 ", out_channels = %" PRId64, - config.num_layers, - config.num_refiner_layers, - config.hidden_size, - config.num_heads, - config.num_kv_heads, - config.in_channels, - config.out_channels); + LOG_VERBOSE("z_image: num_layers = %" PRId64 ", num_refiner_layers = %" PRId64 ", hidden_size = %" PRId64 ", num_heads = %" PRId64 ", num_kv_heads = %" PRId64 ", in_channels = %" PRId64 ", out_channels = %" PRId64, + config.num_layers, + config.num_refiner_layers, + config.hidden_size, + config.num_heads, + config.num_kv_heads, + config.in_channels, + config.out_channels); return config; } }; @@ -603,7 +609,7 @@ namespace ZImage { circular_x_enabled, config.axes_dim); int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); - // LOG_DEBUG("pos_len %d", pos_len); + // LOG_VERBOSE("pos_len %d", pos_len); auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len); // pe->data = pe_vec.data(); // print_ggml_tensor(pe, true, "pe"); @@ -636,7 +642,7 @@ namespace ZImage { return build_graph(x, timesteps, context, ref_latents, ref_index_mode); }; - return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), x.dim()); + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim()); } sd::Tensor compute(int n_threads, @@ -689,7 +695,7 @@ namespace ZImage { GGML_ASSERT(!out_opt.empty()); out = std::move(out_opt); print_sd_tensor(out); - LOG_DEBUG("z_image test done in %lldms", t1 - t0); + LOG_VERBOSE("z_image test done in %lldms", t1 - t0); } } diff --git a/otherarch/sdcpp/src/model/te/clip.hpp b/otherarch/sdcpp/src/model/te/clip.hpp index 2fde3de79..441a26c2a 100644 --- a/otherarch/sdcpp/src/model/te/clip.hpp +++ b/otherarch/sdcpp/src/model/te/clip.hpp @@ -1,8 +1,11 @@ #ifndef __SD_MODEL_TE_CLIP_HPP__ #define __SD_MODEL_TE_CLIP_HPP__ -#include "core/ggml_extend.hpp" +#include "core/ggml_extend.h" +#include "core/ggml_runner.h" +#include "core/util.h" #include "model.h" +#include "model/common/ggml_block.hpp" #include "tokenizers/clip_tokenizer.h" /*================================================ FrozenCLIPEmbedder ================================================*/ @@ -100,13 +103,13 @@ public: const std::string& graph_cut_prefix = "") { // x: [N, n_token, d_model] int layer_idx = n_layer - 1; - // LOG_DEBUG("clip_skip %d", clip_skip); + // LOG_VERBOSE("clip_skip %d", clip_skip); if (clip_skip > 0) { layer_idx = n_layer - clip_skip; } for (int i = 0; i < n_layer; i++) { - // LOG_DEBUG("layer %d", i); + // LOG_VERBOSE("layer %d", i); if (i == layer_idx + 1) { break; } @@ -116,7 +119,7 @@ public: if (!graph_cut_prefix.empty()) { sd::ggml_graph_cut::mark_graph_cut(x, graph_cut_prefix + ".layers." + std::to_string(i), "x"); } - // LOG_DEBUG("layer %d", i); + // LOG_VERBOSE("layer %d", i); } return x; } @@ -320,7 +323,7 @@ public: if (text_projection != nullptr) { pooled = ggml_ext_linear(ctx->ggml_ctx, pooled, text_projection, nullptr); } else { - LOG_DEBUG("identity projection"); + LOG_VERBOSE("identity projection"); } return pooled; // [hidden_size, 1, 1] } @@ -568,13 +571,11 @@ struct CLIPTextModelRunner : public GGMLRunner { size_t max_token_idx, bool return_pooled, int clip_skip, - bool auto_free = true, - bool free_compute_buffer = true, - bool free_compute_params = true) { + bool auto_runner_end = true) { auto get_graph = [&]() -> ggml_cgraph* { return build_graph(input_ids, num_custom_embeddings, custom_embeddings_data, max_token_idx, return_pooled, clip_skip); }; - auto result = GGMLRunner::compute(get_graph, n_threads, auto_free, free_compute_buffer, free_compute_params); + auto result = GGMLRunner::compute(get_graph, n_threads, auto_runner_end); if (return_pooled) { return take_or_empty(std::move(result)); } diff --git a/otherarch/sdcpp/src/model/te/llm.hpp b/otherarch/sdcpp/src/model/te/llm.hpp index f1a057d7e..d381ac5bb 100644 --- a/otherarch/sdcpp/src/model/te/llm.hpp +++ b/otherarch/sdcpp/src/model/te/llm.hpp @@ -3,6 +3,7 @@ #include #include +#include #include #include #include @@ -18,8 +19,13 @@ #include #include -#include "core/ggml_extend.hpp" +#include "core/ggml_extend.h" +#include "core/ggml_extend_backend.h" +#include "core/ggml_runner.h" +#include "core/ggml_tensor_utils.h" +#include "core/util.h" #include "json.hpp" +#include "model/common/ggml_block.hpp" #include "model/common/rope.hpp" #include "model_loader.h" #include "model_manager.h" @@ -319,11 +325,11 @@ namespace LLM { config.vision.deepstack_visual_indexes = {8, 16, 24}; } } - LOG_DEBUG("llm: num_layers = %" PRId64 ", vocab_size = %" PRId64 ", hidden_size = %" PRId64 ", intermediate_size = %" PRId64, - config.num_layers, - config.vocab_size, - config.hidden_size, - config.intermediate_size); + LOG_VERBOSE("llm: num_layers = %" PRId64 ", vocab_size = %" PRId64 ", hidden_size = %" PRId64 ", intermediate_size = %" PRId64, + config.num_layers, + config.vocab_size, + config.hidden_size, + config.intermediate_size); return config; } }; @@ -1887,9 +1893,9 @@ namespace LLM { enable_vision = false; } if (enable_vision) { - LOG_DEBUG("enable llm vision"); + LOG_VERBOSE("enable llm vision"); if (config.llama_cpp_style) { - LOG_DEBUG("llama.cpp style vision weight"); + LOG_VERBOSE("llama.cpp style vision weight"); } } model = LLM(config, enable_vision, config.llama_cpp_style); @@ -2079,9 +2085,7 @@ namespace LLM { const ImageEmbeds& image_embeds, std::set out_layers, bool return_all_hidden_states = false, - bool auto_free = true, - bool free_compute_buffer = true, - bool free_compute_params = true, + bool auto_runner_end = true, const DeepStackImageEmbeds& deepstack_image_embeds = {}, const std::vector& image_grids = {}) { auto get_graph = [&]() -> ggml_cgraph* { @@ -2093,7 +2097,7 @@ namespace LLM { out_layers, return_all_hidden_states); }; - return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, auto_free, free_compute_buffer, free_compute_params), + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, auto_runner_end), input_ids.dim() + 1); } @@ -2173,13 +2177,11 @@ namespace LLM { sd::Tensor encode_image(const int n_threads, const sd::Tensor& image, - bool auto_free = false, - bool free_compute_buffer = false, - bool free_compute_params = false) { + bool auto_runner_end = false) { auto get_graph = [&]() -> ggml_cgraph* { return build_encode_image_graph(image); }; - return take_or_empty(GGMLRunner::compute(get_graph, n_threads, auto_free, free_compute_buffer, free_compute_params)); + return take_or_empty(GGMLRunner::compute(get_graph, n_threads, auto_runner_end)); } ggml_cgraph* build_encode_image_outputs_graph(const sd::Tensor& image_tensor) { @@ -2287,13 +2289,11 @@ namespace LLM { std::vector> encode_image_outputs(const int n_threads, const sd::Tensor& image, - bool auto_free = false, - bool free_compute_buffer = false, - bool free_compute_params = false) { + bool auto_runner_end = false) { auto get_graph = [&]() -> ggml_cgraph* { return build_encode_image_outputs_graph(image); }; - auto combined = take_or_empty(GGMLRunner::compute(get_graph, n_threads, auto_free, free_compute_buffer, free_compute_params)); + auto combined = take_or_empty(GGMLRunner::compute(get_graph, n_threads, auto_runner_end)); if (combined.empty()) { return {}; } @@ -2312,20 +2312,14 @@ namespace LLM { std::vector> encode_video_block_outputs(const int n_threads, const sd::Tensor& frames, - bool auto_free = false, - bool free_compute_buffer = false, - bool free_compute_params = false) { + bool auto_runner_end = false) { int grid_h = static_cast(frames.shape()[1] / config.vision.patch_size); int grid_w = static_cast(frames.shape()[0] / config.vision.patch_size); auto pixel_values = process_video_block_tensor(frames, config.vision); auto get_graph = [&]() -> ggml_cgraph* { return build_encode_video_block_outputs_graph(pixel_values, grid_h, grid_w); }; - auto combined = take_or_empty(GGMLRunner::compute(get_graph, - n_threads, - auto_free, - free_compute_buffer, - free_compute_params)); + auto combined = take_or_empty(GGMLRunner::compute(get_graph, n_threads, auto_runner_end)); if (combined.empty()) { return {}; } @@ -2387,7 +2381,7 @@ namespace LLM { ss << "['" << item.first << "', " << item.second << "], "; } ss << "]"; - LOG_DEBUG("parse '%s' to %s", text.c_str(), ss.str().c_str()); + LOG_VERBOSE("parse '%s' to %s", text.c_str(), ss.str().c_str()); } std::vector tokens; @@ -2438,7 +2432,7 @@ namespace LLM { out = std::move(out_opt); print_sd_tensor(out, false, "image_embed"); image_embed = out; - LOG_DEBUG("llm encode_image test done in %lldms", t1 - t0); + LOG_VERBOSE("llm encode_image test done in %lldms", t1 - t0); } std::string placeholder = "<|image_pad|>"; @@ -2478,7 +2472,7 @@ namespace LLM { GGML_ASSERT(!out_opt.empty()); out = std::move(out_opt); print_sd_tensor(out); - LOG_DEBUG("llm test done in %lldms", t1 - t0); + LOG_VERBOSE("llm test done in %lldms", t1 - t0); } else if (test_vit) { // auto image = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 280, 280, 3); // ggml_set_f32(image, 0.f); @@ -2497,7 +2491,7 @@ namespace LLM { // auto ref_out = load_tensor_from_file(ctx, "qwen2vl.bin"); // ggml_ext_tensor_diff(ref_out, out, 0.01f); - LOG_DEBUG("llm test done in %lldms", t1 - t0); + LOG_VERBOSE("llm test done in %lldms", t1 - t0); } else if (test_mistral) { std::pair prompt_attn_range; std::string text = "[SYSTEM_PROMPT]You are an AI that reasons about image descriptions. You give structured responses focusing on object relationships, object\nattribution and actions without speculation.[/SYSTEM_PROMPT][INST]"; @@ -2522,7 +2516,7 @@ namespace LLM { GGML_ASSERT(!out_opt.empty()); out = std::move(out_opt); print_sd_tensor(out); - LOG_DEBUG("llm test done in %lldms", t1 - t0); + LOG_VERBOSE("llm test done in %lldms", t1 - t0); } else if (test_qwen3) { std::pair prompt_attn_range; std::string text = "<|im_start|>user\n"; @@ -2547,7 +2541,7 @@ namespace LLM { GGML_ASSERT(!out_opt.empty()); out = std::move(out_opt); print_sd_tensor(out); - LOG_DEBUG("llm test done in %lldms", t1 - t0); + LOG_VERBOSE("llm test done in %lldms", t1 - t0); } else { std::pair prompt_attn_range; std::string text = "<|im_start|>system\nDescribe the image by detailing the color, shape, size, texture, quantity, text, spatial relationships of the objects and background:<|im_end|>\n<|im_start|>user\n"; @@ -2572,7 +2566,7 @@ namespace LLM { GGML_ASSERT(!out_opt.empty()); out = std::move(out_opt); print_sd_tensor(out); - LOG_DEBUG("llm test done in %lldms", t1 - t0); + LOG_VERBOSE("llm test done in %lldms", t1 - t0); } } diff --git a/otherarch/sdcpp/src/model/te/t5.hpp b/otherarch/sdcpp/src/model/te/t5.hpp index 90c781e1d..4c86117cc 100644 --- a/otherarch/sdcpp/src/model/te/t5.hpp +++ b/otherarch/sdcpp/src/model/te/t5.hpp @@ -10,7 +10,12 @@ #include #include -#include "core/ggml_extend.hpp" +#include "core/ggml_extend.h" +#include "core/ggml_extend_backend.h" +#include "core/ggml_runner.h" +#include "core/ggml_tensor_utils.h" +#include "core/util.h" +#include "model/common/ggml_block.hpp" #include "model_loader.h" #include "model_manager.h" #include "tokenizers/t5_unigram_tokenizer.h" @@ -451,13 +456,11 @@ struct T5Runner : public GGMLRunner { sd::Tensor compute(const int n_threads, const sd::Tensor& input_ids, const sd::Tensor& attention_mask, - bool auto_free = true, - bool free_compute_buffer = true, - bool free_compute_params = true) { + bool auto_runner_end = true) { auto get_graph = [&]() -> ggml_cgraph* { return build_graph(input_ids, attention_mask); }; - return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, auto_free, free_compute_buffer, free_compute_params), 3); + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, auto_runner_end), 3); } static std::vector _relative_position_bucket(const std::vector& relative_position, @@ -556,7 +559,7 @@ struct T5Embedder { ss << "['" << item.first << "', " << item.second << "], "; } ss << "]"; - LOG_DEBUG("parse '%s' to %s", text.c_str(), ss.str().c_str()); + LOG_VERBOSE("parse '%s' to %s", text.c_str(), ss.str().c_str()); } std::vector tokens; @@ -614,7 +617,7 @@ struct T5Embedder { GGML_ASSERT(!out_opt.empty()); out = std::move(out_opt); print_sd_tensor(out); - LOG_DEBUG("t5 test done in %lldms", t1 - t0); + LOG_VERBOSE("t5 test done in %lldms", t1 - t0); } } diff --git a/otherarch/sdcpp/src/model/upscaler/esrgan.hpp b/otherarch/sdcpp/src/model/upscaler/esrgan.hpp index 21c977124..66029f5de 100644 --- a/otherarch/sdcpp/src/model/upscaler/esrgan.hpp +++ b/otherarch/sdcpp/src/model/upscaler/esrgan.hpp @@ -7,8 +7,10 @@ #include #include -#include "core/ggml_extend.hpp" +#include "core/ggml_extend.h" +#include "core/ggml_runner.h" #include "core/util.h" +#include "model/common/ggml_block.hpp" /* =================================== ESRGAN =================================== @@ -74,13 +76,13 @@ struct ESRGANConfig { } if (has_model_tensor || has_conv_up1 || has_conv_up2) { - LOG_DEBUG("esrgan: scale = %d, num_block = %d, num_in_ch = %d, num_out_ch = %d, num_feat = %d, num_grow_ch = %d", - config.scale, - config.num_block, - config.num_in_ch, - config.num_out_ch, - config.num_feat, - config.num_grow_ch); + LOG_VERBOSE("esrgan: scale = %d, num_block = %d, num_in_ch = %d, num_out_ch = %d, num_feat = %d, num_grow_ch = %d", + config.scale, + config.num_block, + config.num_in_ch, + config.num_out_ch, + config.num_feat, + config.num_grow_ch); } return config; } @@ -265,7 +267,7 @@ struct ESRGAN : public GGMLRunner { sd::Tensor compute(const int n_threads, const sd::Tensor& x) { auto get_graph = [&]() -> ggml_cgraph* { return build_graph(x); }; - auto result = restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), x.dim()); + auto result = restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim()); return result; } }; diff --git a/otherarch/sdcpp/src/model/upscaler/ltx_latent_upscaler.hpp b/otherarch/sdcpp/src/model/upscaler/ltx_latent_upscaler.hpp index b70e16136..88c6b74ad 100644 --- a/otherarch/sdcpp/src/model/upscaler/ltx_latent_upscaler.hpp +++ b/otherarch/sdcpp/src/model/upscaler/ltx_latent_upscaler.hpp @@ -11,9 +11,11 @@ #include #include -#include "core/ggml_extend.hpp" +#include "core/ggml_extend.h" #include "core/ggml_graph_cut.h" +#include "core/ggml_runner.h" #include "core/util.h" +#include "model/common/ggml_block.hpp" #include "model/diffusion/dit.hpp" #include "model_loader.h" @@ -115,13 +117,13 @@ namespace LTXVUpsampler { } if (inferred) { - LOG_DEBUG("ltx latent upsampler: in_channels = %" PRId64 ", mid_channels = %" PRId64 ", num_blocks_per_stage = %d, spatial_scale = %.3f, temporal_up_factor = %d, rational_resampler = %d", - config.in_channels, - config.mid_channels, - config.num_blocks_per_stage, - config.spatial_scale, - config.temporal_up_factor, - config.rational_resampler); + LOG_VERBOSE("ltx latent upsampler: in_channels = %" PRId64 ", mid_channels = %" PRId64 ", num_blocks_per_stage = %d, spatial_scale = %.3f, temporal_up_factor = %d, rational_resampler = %d", + config.in_channels, + config.mid_channels, + config.num_blocks_per_stage, + config.spatial_scale, + config.temporal_up_factor, + config.rational_resampler); } return config; } @@ -499,7 +501,7 @@ namespace LTXVUpsampler { } size_t expected_dim = static_cast(x.dim()); auto get_graph = [&]() -> ggml_cgraph* { return build_graph(x); }; - return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), expected_dim); + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), expected_dim); } }; diff --git a/otherarch/sdcpp/src/model/vae/audio_vae.hpp b/otherarch/sdcpp/src/model/vae/audio_vae.hpp index bec929097..2fa0f1105 100644 --- a/otherarch/sdcpp/src/model/vae/audio_vae.hpp +++ b/otherarch/sdcpp/src/model/vae/audio_vae.hpp @@ -1,7 +1,8 @@ #ifndef __SD_MODEL_VAE_AUDIO_VAE_HPP__ #define __SD_MODEL_VAE_AUDIO_VAE_HPP__ -#include "core/ggml_extend.hpp" +#include "core/ggml_runner.h" +#include "core/util.h" struct AudioVAERunner : public GGMLRunner { AudioVAERunner(ggml_backend_t backend, diff --git a/otherarch/sdcpp/src/model/vae/auto_encoder_kl.hpp b/otherarch/sdcpp/src/model/vae/auto_encoder_kl.hpp index 604347d77..116a83219 100644 --- a/otherarch/sdcpp/src/model/vae/auto_encoder_kl.hpp +++ b/otherarch/sdcpp/src/model/vae/auto_encoder_kl.hpp @@ -1,6 +1,8 @@ #ifndef __SD_MODEL_VAE_AUTO_ENCODER_KL_HPP__ #define __SD_MODEL_VAE_AUTO_ENCODER_KL_HPP__ +#include +#include "core/ggml_tensor_utils.h" #include "model/vae/vae.hpp" /*================================================== AutoEncoderKL ===================================================*/ @@ -744,7 +746,7 @@ struct AutoEncoderKL : public VAE { auto get_graph = [&]() -> ggml_cgraph* { return build_graph(z, decode_graph); }; - return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), z.dim()); + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), z.dim()); } sd::Tensor gaussian_latent_sample(const sd::Tensor& moments, std::shared_ptr rng) { @@ -864,7 +866,7 @@ struct AutoEncoderKL : public VAE { GGML_ASSERT(!out_opt.empty()); out = std::move(out_opt); print_sd_tensor(out); - LOG_DEBUG("encode test done in %lldms", t1 - t0); + LOG_VERBOSE("encode test done in %lldms", t1 - t0); } if (false) { @@ -884,7 +886,7 @@ struct AutoEncoderKL : public VAE { GGML_ASSERT(!out_opt.empty()); out = std::move(out_opt); print_sd_tensor(out); - LOG_DEBUG("decode test done in %lldms", t1 - t0); + LOG_VERBOSE("decode test done in %lldms", t1 - t0); } }; }; diff --git a/otherarch/sdcpp/src/model/vae/hunyuan_vae.hpp b/otherarch/sdcpp/src/model/vae/hunyuan_vae.hpp index 938f5ffe8..099a15b7e 100644 --- a/otherarch/sdcpp/src/model/vae/hunyuan_vae.hpp +++ b/otherarch/sdcpp/src/model/vae/hunyuan_vae.hpp @@ -825,11 +825,9 @@ namespace Hunyuan { auto get_graph = [&]() -> ggml_cgraph* { return build_graph(graph_input, decode_graph); }; - auto output = restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, - n_threads, - true, - true, - true), + auto output = restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, + n_threads, + false), graph_input.dim()); if (!output.empty() && input.dim() == 4) { output.squeeze_(2); diff --git a/otherarch/sdcpp/src/model/vae/ltx_audio_vae.hpp b/otherarch/sdcpp/src/model/vae/ltx_audio_vae.hpp index 65edd9dda..822f33343 100644 --- a/otherarch/sdcpp/src/model/vae/ltx_audio_vae.hpp +++ b/otherarch/sdcpp/src/model/vae/ltx_audio_vae.hpp @@ -8,7 +8,12 @@ #include #include -#include "core/ggml_extend.hpp" +#include "core/ggml_extend.h" +#include "core/ggml_extend_backend.h" +#include "core/ggml_runner.h" +#include "core/ggml_tensor_utils.h" +#include "core/util.h" +#include "model/common/ggml_block.hpp" #include "model/vae/audio_vae.hpp" #include "model_loader.h" #include "model_manager.h" @@ -212,13 +217,13 @@ namespace LTXV { if (config.audio_channels != 2 || config.latent_channels != 8 || config.mel_bins != 64) { return config; } - LOG_DEBUG("ltx_audio_vae: sample_rate = %d, mel_bins = %d, latent_channels = %d, latent_frequency_bins = %d, has_encoder = %s, has_bwe = %s", - config.sample_rate, - config.mel_bins, - config.latent_channels, - config.latent_frequency_bins, - config.has_encoder ? "true" : "false", - config.has_bwe ? "true" : "false"); + LOG_VERBOSE("ltx_audio_vae: sample_rate = %d, mel_bins = %d, latent_channels = %d, latent_frequency_bins = %d, has_encoder = %s, has_bwe = %s", + config.sample_rate, + config.mel_bins, + config.latent_channels, + config.latent_frequency_bins, + config.has_encoder ? "true" : "false", + config.has_bwe ? "true" : "false"); return config; } }; @@ -1282,7 +1287,7 @@ namespace LTXV { ggml_build_forward_expand(gf, waveform); return gf; }; - auto result = restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), 4); + auto result = restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), 4); int64_t t1 = ggml_time_ms(); LOG_INFO("ltx audio vae decode completed, taking %.2fs", (t1 - t0) * 1.0f / 1000); return result; @@ -1305,7 +1310,7 @@ namespace LTXV { ggml_build_forward_expand(gf, latent); return gf; }; - auto result = restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), 4); + auto result = restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), 4); int64_t t1 = ggml_time_ms(); LOG_INFO("ltx audio vae encode completed, taking %.2fs", (t1 - t0) * 1.0f / 1000); return result; @@ -1326,7 +1331,7 @@ namespace LTXV { GGML_ASSERT(!out.empty()); print_sd_tensor(out, false, "ltx_audio_vae_out"); - LOG_DEBUG("ltx audio vae test done in %lldms", t1 - t0); + LOG_VERBOSE("ltx audio vae test done in %lldms", t1 - t0); } static void load_from_file_and_test(const std::string& model_path, diff --git a/otherarch/sdcpp/src/model/vae/ltx_vae.hpp b/otherarch/sdcpp/src/model/vae/ltx_vae.hpp index 5629b939b..93f6b760e 100644 --- a/otherarch/sdcpp/src/model/vae/ltx_vae.hpp +++ b/otherarch/sdcpp/src/model/vae/ltx_vae.hpp @@ -8,6 +8,8 @@ #include #include #include +#include "core/ggml_extend_backend.h" +#include "core/ggml_tensor_utils.h" #include "model/diffusion/ltxv.hpp" #include "model/vae/vae.hpp" @@ -1126,11 +1128,11 @@ namespace LTXVAE { overlap, window); overlap = window - 1; } - LOG_DEBUG("Using temporal tiling: temporal_tile_frames = %d, temporal_tile_overlap = %d, total frames = %d, resulting in %d tiles", - window, - overlap, - (int)T, - (T + window - overlap - 1) / (window - overlap)); + LOG_VERBOSE("Using temporal tiling: temporal_tile_frames = %d, temporal_tile_overlap = %d, total frames = %d, resulting in %d tiles", + window, + overlap, + (int)T, + (T + window - overlap - 1) / (window - overlap)); ggml_tensor* out = nullptr; for (int i = 0; i < (int)T - overlap; i += (window - overlap)) { int feat_idx = 0; @@ -1327,34 +1329,32 @@ struct LTXVideoVAE : public VAE { const int64_t total_frames = input.shape()[2]; auto plan = make_vae_temporal_tile_plan(total_frames, config); - LOG_DEBUG("Using streaming temporal tiling: temporal_tile_frames=%d, temporal_tile_overlap=%d, total latent frames=%lld, resulting in %d tiles", - plan.tile_frames, - plan.overlap, - (long long)total_frames, - (int)plan.tiles.size()); + LOG_VERBOSE("Using streaming temporal tiling: temporal_tile_frames=%d, temporal_tile_overlap=%d, total latent frames=%lld, resulting in %d tiles", + plan.tile_frames, + plan.overlap, + (long long)total_frames, + (int)plan.tiles.size()); free_cache_ctx_and_buffer(); - cache_tensor_map.clear(); auto output = process_vae_temporal_tiles(input, plan, [&](const sd::Tensor& z_chunk, const VAETemporalTile& tile) { - LOG_DEBUG("LTX VAE temporal tile %lld/%d: latent frames [%lld, %lld), overlap=%d", - (long long)tile.index + 1, - (int)plan.tiles.size(), - (long long)tile.start, - (long long)tile.end, - tile.overlap); + LOG_VERBOSE("LTX VAE temporal tile %lld/%d: latent frames [%lld, %lld), overlap=%d", + (long long)tile.index + 1, + (int)plan.tiles.size(), + (long long)tile.start, + (long long)tile.end, + tile.overlap); auto get_graph = [&]() -> ggml_cgraph* { return build_temporal_tile_graph(z_chunk, static_cast(tile.start), tile.overlap); }; - return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, true, true, true), + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), expected_dim); }); free_cache_ctx_and_buffer(); - cache_tensor_map.clear(); return output; } @@ -1407,7 +1407,7 @@ struct LTXVideoVAE : public VAE { auto get_graph = [&]() -> ggml_cgraph* { return build_graph(input, decode_graph); }; - auto result = restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), expected_dim); + auto result = restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), expected_dim); if (result.empty()) { return {}; } @@ -1420,7 +1420,7 @@ struct LTXVideoVAE : public VAE { auto get_graph = [&]() -> ggml_cgraph* { return build_latent_statistics_graph(z, normalize); }; - return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), static_cast(z.dim())); } @@ -1467,7 +1467,7 @@ struct LTXVideoVAE : public VAE { GGML_ASSERT(!out.empty()); print_sd_tensor(out, false, "ltx_vae_out"); - LOG_DEBUG("ltx vae test done in %lldms", t1 - t0); + LOG_VERBOSE("ltx vae test done in %lldms", t1 - t0); } static void load_from_file_and_test(const std::string& model_path, diff --git a/otherarch/sdcpp/src/model/vae/mage_vae.hpp b/otherarch/sdcpp/src/model/vae/mage_vae.hpp index 39075a836..5567aec22 100644 --- a/otherarch/sdcpp/src/model/vae/mage_vae.hpp +++ b/otherarch/sdcpp/src/model/vae/mage_vae.hpp @@ -490,7 +490,7 @@ namespace MageVAE { auto get_graph = [&]() -> ggml_cgraph* { return build_graph(input, decode_graph); }; - return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), input.dim()); + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), input.dim()); } int get_encoder_output_channels(int input_channels) override { diff --git a/otherarch/sdcpp/src/model/vae/minimax_h3_audio_vae.hpp b/otherarch/sdcpp/src/model/vae/minimax_h3_audio_vae.hpp index 21f6ce167..926d2fbe5 100644 --- a/otherarch/sdcpp/src/model/vae/minimax_h3_audio_vae.hpp +++ b/otherarch/sdcpp/src/model/vae/minimax_h3_audio_vae.hpp @@ -480,7 +480,7 @@ namespace MiniMaxH3 { return graph; }; auto result = restore_trailing_singleton_dims( - GGMLRunner::compute(get_graph, n_threads, false, false, false), + GGMLRunner::compute(get_graph, n_threads, false), 4); int64_t t1 = ggml_time_ms(); LOG_INFO("MiniMax-H3 audio VAE encode completed, taking %.2fs", @@ -500,7 +500,7 @@ namespace MiniMaxH3 { return graph; }; auto result = restore_trailing_singleton_dims( - GGMLRunner::compute(get_graph, n_threads, false, false, false), + GGMLRunner::compute(get_graph, n_threads, false), 4); int64_t t1 = ggml_time_ms(); LOG_INFO("MiniMax-H3 audio VAE decode completed, taking %.2fs", diff --git a/otherarch/sdcpp/src/model/vae/minimax_h3_vae.hpp b/otherarch/sdcpp/src/model/vae/minimax_h3_vae.hpp index a8aa6b5f1..ec65fa37e 100644 --- a/otherarch/sdcpp/src/model/vae/minimax_h3_vae.hpp +++ b/otherarch/sdcpp/src/model/vae/minimax_h3_vae.hpp @@ -791,11 +791,9 @@ namespace MiniMaxH3VAE { return graph; }; return restore_trailing_singleton_dims( - GGMLRunner::compute(get_graph, - n_threads, - false, - false, - false), + GGMLRunner::compute(get_graph, + n_threads, + false), 5); } }; diff --git a/otherarch/sdcpp/src/model/vae/tae.hpp b/otherarch/sdcpp/src/model/vae/tae.hpp index d291bb785..95bc0b271 100644 --- a/otherarch/sdcpp/src/model/vae/tae.hpp +++ b/otherarch/sdcpp/src/model/vae/tae.hpp @@ -1,8 +1,12 @@ #ifndef __SD_MODEL_VAE_TAE_HPP__ #define __SD_MODEL_VAE_TAE_HPP__ -#include "core/ggml_extend.hpp" +#include "core/ggml_extend.h" +#include "core/ggml_runner.h" +#include "core/rng.hpp" +#include "core/util.h" #include "model.h" +#include "model/common/ggml_block.hpp" /* =================================== TinyAutoEncoder =================================== @@ -65,7 +69,7 @@ public: if (n_in != n_out) { auto skip = std::dynamic_pointer_cast(blocks["skip"]); - LOG_DEBUG("skip"); + LOG_VERBOSE("skip"); x = skip->forward(ctx, x); } @@ -787,7 +791,7 @@ struct TinyImageAutoEncoder : public VAE { return build_graph(z_tensor, decode_graph); }; - return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), z_tensor.dim()); + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), z_tensor.dim()); } }; @@ -872,7 +876,7 @@ struct TinyVideoAutoEncoder : public VAE { return build_graph(z_tensor, decode_graph); }; - return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), z_tensor.dim()); + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), z_tensor.dim()); } }; diff --git a/otherarch/sdcpp/src/model/vae/vae.hpp b/otherarch/sdcpp/src/model/vae/vae.hpp index f3adb0ccc..03cfd5e2e 100644 --- a/otherarch/sdcpp/src/model/vae/vae.hpp +++ b/otherarch/sdcpp/src/model/vae/vae.hpp @@ -5,6 +5,7 @@ #include "model/common/block.hpp" #include "model/vae/vae_tiling.hpp" #include "model_manager.h" +#include "runtime/tiling.h" struct VAE : public GGMLRunner { protected: @@ -54,23 +55,23 @@ protected: } auto plan = make_vae_temporal_tile_plan(input.shape()[2], resolved_config); - LOG_DEBUG("%s temporal tiling: tile_frames=%d, overlap=%d, total_frames=%lld, tiles=%d", - get_desc().c_str(), - plan.tile_frames, - plan.overlap, - (long long)input.shape()[2], - (int)plan.tiles.size()); + LOG_VERBOSE("%s temporal tiling: tile_frames=%d, overlap=%d, total_frames=%lld, tiles=%d", + get_desc().c_str(), + plan.tile_frames, + plan.overlap, + (long long)input.shape()[2], + (int)plan.tiles.size()); return process_vae_temporal_tiles_blended( input, plan, output_scale, [&](const sd::Tensor& input_tile, const VAETemporalTile& tile) { - LOG_DEBUG("%s temporal tile %d/%d: input frames [%lld, %lld)", - get_desc().c_str(), - tile.index + 1, - (int)plan.tiles.size(), - (long long)tile.start, - (long long)tile.end); + LOG_VERBOSE("%s temporal tile %d/%d: input frames [%lld, %lld)", + get_desc().c_str(), + tile.index + 1, + (int)plan.tiles.size(), + (long long)tile.start, + (long long)tile.end); return _compute(n_threads, input_tile, true); }); } @@ -230,7 +231,7 @@ public: const float encode_tile_factor = sd_version_is_minimax_h3(version) ? 1.f : (sd_version_is_wan(version) || sd_version_is_hunyuan_video(version) || sd_version_is_ltxav(version)) ? 1.30539f : 2.0f; get_tile_sizes(tile_size_x, tile_size_y, tile_overlap, tiling_params, W, H, encode_tile_factor); - LOG_DEBUG("VAE Tile size: %dx%d", tile_size_x, tile_size_y); + LOG_VERBOSE("VAE Tile size: %dx%d", tile_size_x, tile_size_y); output = tiled_compute(input, n_threads, static_cast(W), @@ -251,14 +252,14 @@ public: tiling_params); } - runner_done(); + runner_end(); if (output.empty()) { LOG_ERROR("vae encode compute failed"); return {}; } int64_t t1 = ggml_time_ms(); - LOG_DEBUG("computing vae encode graph completed, taking %.2fs", (t1 - t0) * 1.0f / 1000); + LOG_VERBOSE("computing vae encode graph completed, taking %.2fs", (t1 - t0) * 1.0f / 1000); return std::move(output); } @@ -281,7 +282,7 @@ public: int tile_size_x, tile_size_y; get_tile_sizes(tile_size_x, tile_size_y, tile_overlap, tiling_params, input.shape()[0], input.shape()[1]); if (!silent) { - LOG_DEBUG("VAE Tile size: %dx%d", tile_size_x, tile_size_y); + LOG_VERBOSE("VAE Tile size: %dx%d", tile_size_x, tile_size_y); } output = tiled_compute( input, @@ -305,7 +306,7 @@ public: tiling_params); } - runner_done(); + runner_end(); if (output.empty()) { LOG_ERROR("vae decode compute failed"); @@ -315,7 +316,7 @@ public: scale_tensor_to_0_1(&output); } int64_t t1 = ggml_time_ms(); - LOG_DEBUG("computing vae decode graph completed, taking %.2fs", (t1 - t0) * 1.0f / 1000); + LOG_VERBOSE("computing vae decode graph completed, taking %.2fs", (t1 - t0) * 1.0f / 1000); return std::move(output); } diff --git a/otherarch/sdcpp/src/model/vae/wan_vae.hpp b/otherarch/sdcpp/src/model/vae/wan_vae.hpp index 422a57823..bfdb4b7d5 100644 --- a/otherarch/sdcpp/src/model/vae/wan_vae.hpp +++ b/otherarch/sdcpp/src/model/vae/wan_vae.hpp @@ -4,6 +4,8 @@ #include #include #include +#include "core/ggml_extend_backend.h" +#include "core/ggml_tensor_utils.h" #include "model/common/block.hpp" #include "model/vae/vae.hpp" @@ -1278,7 +1280,7 @@ namespace WAN { } } if (is_2D) { - LOG_DEBUG("USING 2D VAE"); + LOG_VERBOSE("USING 2D VAE"); } ae = WanVAE(decode_only, version, is_2D); ae.init(params_ctx, tensor_storage_map, prefix); @@ -1409,31 +1411,29 @@ namespace WAN { stateful_config.overlap = 0; auto plan = make_vae_temporal_tile_plan(input.shape()[2], stateful_config); - LOG_DEBUG("Wan VAE stateful temporal tiling: tile_frames=%d, total latent frames=%lld, tiles=%d", - plan.tile_frames, - (long long)input.shape()[2], - (int)plan.tiles.size()); + LOG_VERBOSE("Wan VAE stateful temporal tiling: tile_frames=%d, total latent frames=%lld, tiles=%d", + plan.tile_frames, + (long long)input.shape()[2], + (int)plan.tiles.size()); free_cache_ctx_and_buffer(); - cache_tensor_map.clear(); ae.clear_cache(); auto output = process_vae_temporal_tiles(input, plan, [&](const sd::Tensor& input_tile, const VAETemporalTile& tile) { - LOG_DEBUG("Wan VAE temporal tile %d/%d: latent frames [%lld, %lld)", - tile.index + 1, - (int)plan.tiles.size(), - (long long)tile.start, - (long long)tile.end); + LOG_VERBOSE("Wan VAE temporal tile %d/%d: latent frames [%lld, %lld)", + tile.index + 1, + (int)plan.tiles.size(), + (long long)tile.start, + (long long)tile.end); auto get_graph = [&]() -> ggml_cgraph* { return build_temporal_tile_graph(input_tile, static_cast(tile.start)); }; return restore_trailing_singleton_dims( - GGMLRunner::compute(get_graph, n_threads, true, true, true), + GGMLRunner::compute(get_graph, n_threads, false), static_cast(input.dim())); }); free_cache_ctx_and_buffer(); - cache_tensor_map.clear(); ae.clear_cache(); return output; } @@ -1448,7 +1448,7 @@ namespace WAN { auto get_graph = [&]() -> ggml_cgraph* { return build_graph(input.empty() ? z : input, decode_graph); }; - auto result = restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, true, true, true), + auto result = restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), input.empty() ? z.dim() : input.dim()); if (!result.empty() && z.dim() == 4) { result.squeeze_(2); @@ -1481,7 +1481,7 @@ namespace WAN { GGML_ASSERT(!out_opt.empty()); out = std::move(out_opt); print_sd_tensor(out); - LOG_DEBUG("decode test done in %ldms", t1 - t0); + LOG_VERBOSE("decode test done in %ldms", t1 - t0); } }; diff --git a/otherarch/sdcpp/src/model_io/gguf_reader_ext.h b/otherarch/sdcpp/src/model_io/gguf_reader_ext.h index d4da419db..0a321818c 100644 --- a/otherarch/sdcpp/src/model_io/gguf_reader_ext.h +++ b/otherarch/sdcpp/src/model_io/gguf_reader_ext.h @@ -77,7 +77,7 @@ private: if (align_val != 0 && (align_val & (align_val - 1)) == 0) { alignment_ = align_val; - LOG_DEBUG("Found alignment: %zu", alignment_); + LOG_VERBOSE("Found alignment: %zu", alignment_); } else { LOG_ERROR("Invalid alignment value %u, fallback to default %zu", align_val, alignment_); } @@ -197,8 +197,8 @@ public: if (!safe_read(fin, metadata_kv_count)) return false; - LOG_DEBUG("GGUF v%u, tensor_count=%llu, metadata_kv_count=%llu", - version, (unsigned long long)tensor_count, (unsigned long long)metadata_kv_count); + LOG_VERBOSE("GGUF v%u, tensor_count=%llu, metadata_kv_count=%llu", + version, (unsigned long long)tensor_count, (unsigned long long)metadata_kv_count); // --- Read Metadata --- for (uint64_t i = 0; i < metadata_kv_count; i++) { diff --git a/otherarch/sdcpp/src/model_io/safetensors_io.cpp b/otherarch/sdcpp/src/model_io/safetensors_io.cpp index ef2f34154..10670260d 100644 --- a/otherarch/sdcpp/src/model_io/safetensors_io.cpp +++ b/otherarch/sdcpp/src/model_io/safetensors_io.cpp @@ -254,7 +254,7 @@ bool read_safetensors_file(const std::string& file_path, for (auto& item : header_.items()) { std::string name = item.key(); nlohmann::json tensor_info = item.value(); - // LOG_DEBUG("%s %s\n", name.c_str(), tensor_info.dump().c_str()); + // LOG_VERBOSE("%s %s\n", name.c_str(), tensor_info.dump().c_str()); if (name == "__metadata__") { continue; @@ -403,7 +403,7 @@ bool read_safetensors_file(const std::string& file_path, tensor_storages.push_back(tensor_storage); - // LOG_DEBUG("%s %s", tensor_storage.to_string().c_str(), dtype.c_str()); + // LOG_VERBOSE("%s %s", tensor_storage.to_string().c_str(), dtype.c_str()); } return true; diff --git a/otherarch/sdcpp/src/model_loader.cpp b/otherarch/sdcpp/src/model_loader.cpp index 130a8645a..03ec613f9 100644 --- a/otherarch/sdcpp/src/model_loader.cpp +++ b/otherarch/sdcpp/src/model_loader.cpp @@ -242,7 +242,7 @@ void ModelLoader::add_tensor_storage(const TensorStorage& tensor_storage) { void ModelLoader::set_n_threads(int n_threads) { n_threads_ = n_threads > 0 ? n_threads : sd_get_num_physical_cores(); - LOG_DEBUG("using %d threads for model loading", n_threads_); + LOG_VERBOSE("using %d threads for model loading", n_threads_); } bool ModelLoader::init_from_file(const std::string& file_path, const std::string& prefix) { @@ -282,7 +282,7 @@ void ModelLoader::convert_tensors_name() { for (auto& [_, tensor_storage] : tensor_storage_map) { auto new_name = convert_tensor_name(tensor_storage.name, version); - // LOG_DEBUG("%s -> %s", tensor_storage.name.c_str(), new_name.c_str()); + // LOG_VERBOSE("%s -> %s", tensor_storage.name.c_str(), new_name.c_str()); tensor_storage.name = new_name; new_map[new_name] = std::move(tensor_storage); } @@ -304,7 +304,7 @@ bool ModelLoader::init_from_file_and_convert_name(const std::string& file_path, /*================================================= GGUFModelLoader ==================================================*/ bool ModelLoader::init_from_gguf_file(const std::string& file_path, const std::string& prefix) { - LOG_DEBUG("init from '%s'", file_path.c_str()); + LOG_VERBOSE("init from '%s'", file_path.c_str()); std::vector tensor_storages; std::string error; @@ -316,7 +316,7 @@ bool ModelLoader::init_from_gguf_file(const std::string& file_path, const std::s size_t file_index = add_file_path(file_path); for (auto& tensor_storage : tensor_storages) { - // LOG_DEBUG("%s", tensor_storage.name.c_str()); + // LOG_VERBOSE("%s", tensor_storage.name.c_str()); if (!starts_with(tensor_storage.name, prefix)) { tensor_storage.name = prefix + tensor_storage.name; @@ -332,7 +332,7 @@ bool ModelLoader::init_from_gguf_file(const std::string& file_path, const std::s /*================================================= SafeTensorsModelLoader ==================================================*/ bool ModelLoader::init_from_safetensors_file(const std::string& file_path, const std::string& prefix) { - LOG_DEBUG("init from '%s', prefix = '%s'", file_path.c_str(), prefix.c_str()); + LOG_VERBOSE("init from '%s', prefix = '%s'", file_path.c_str(), prefix.c_str()); std::vector tensor_storages; std::string error; @@ -357,14 +357,14 @@ bool ModelLoader::init_from_safetensors_file(const std::string& file_path, const add_tensor_storage(tensor_storage); - // LOG_DEBUG("%s", tensor_storage.to_string().c_str()); + // LOG_VERBOSE("%s", tensor_storage.to_string().c_str()); } return true; } bool ModelLoader::init_from_safetensors_index_file(const std::string& file_path, const std::string& prefix) { - LOG_DEBUG("init from safetensors index '%s', prefix = '%s'", file_path.c_str(), prefix.c_str()); + LOG_VERBOSE("init from safetensors index '%s', prefix = '%s'", file_path.c_str(), prefix.c_str()); std::vector shard_paths; std::string error; @@ -385,7 +385,7 @@ bool ModelLoader::init_from_safetensors_index_file(const std::string& file_path, /*================================================= TorchLegacyModelLoader ==================================================*/ bool ModelLoader::init_from_torch_legacy_file(const std::string& file_path, const std::string& prefix) { - LOG_DEBUG("init from torch legacy '%s'", file_path.c_str()); + LOG_VERBOSE("init from torch legacy '%s'", file_path.c_str()); std::vector tensor_storages; std::string error; @@ -417,7 +417,7 @@ bool ModelLoader::init_from_torch_legacy_file(const std::string& file_path, cons /*================================================= TorchZipModelLoader ==================================================*/ bool ModelLoader::init_from_torch_zip_file(const std::string& file_path, const std::string& prefix) { - LOG_DEBUG("init from '%s'", file_path.c_str()); + LOG_VERBOSE("init from '%s'", file_path.c_str()); std::vector tensor_storages; std::string error; @@ -436,7 +436,7 @@ bool ModelLoader::init_from_torch_zip_file(const std::string& file_path, const s add_tensor_storage(tensor_storage); - // LOG_DEBUG("%s", tensor_storage.to_string().c_str()); + // LOG_VERBOSE("%s", tensor_storage.to_string().c_str()); } return true; @@ -473,7 +473,7 @@ bool ModelLoader::init_from_diffusers_file(const std::string& file_path, const s // return false; } if (!init_from_safetensors_file(clip_g_path, "te.1.")) { - LOG_DEBUG("Couldn't find working second text encoder in %s", file_path.c_str()); + LOG_VERBOSE("Couldn't find working second text encoder in %s", file_path.c_str()); } return true; } @@ -637,7 +637,7 @@ SDVersion ModelLoader::get_sd_version() { } } if (is_wan) { - LOG_DEBUG("patch_embedding_channels %d", patch_embedding_channels); + LOG_VERBOSE("patch_embedding_channels %d", patch_embedding_channels); if (patch_embedding_channels == 184320 && !has_img_emb) { return VERSION_WAN2_2_I2V; } @@ -894,7 +894,7 @@ void ModelLoader::process_model_files(bool enable_mmap, bool writable_mmap) { fdata.tensors = std::move(file_tensors); if (enable_mmap && !is_zip) { - LOG_DEBUG("using mmap for I/O"); + LOG_VERBOSE("using mmap for I/O"); std::unique_ptr mmapped = MmapWrapper::create(file_path, writable_mmap); if (mmapped) { uint8_t* mmap_data = static_cast(mmapped->writable_data()); @@ -926,7 +926,7 @@ std::vector ModelLoader::mmap_tensors(std::map& tensors, std::mutex tensor_names_mutex; auto on_new_tensor_cb = [&](const TensorStorage& tensor_storage, ggml_tensor** dst_tensor) -> bool { const std::string& name = tensor_storage.name; - // LOG_DEBUG("%s", tensor_storage.to_string().c_str()); + // LOG_VERBOSE("%s", tensor_storage.to_string().c_str()); { std::lock_guard lock(tensor_names_mutex); tensor_names_in_file.insert(name); diff --git a/otherarch/sdcpp/src/model_manager.cpp b/otherarch/sdcpp/src/model_manager.cpp index 8db303dda..166030a5c 100644 --- a/otherarch/sdcpp/src/model_manager.cpp +++ b/otherarch/sdcpp/src/model_manager.cpp @@ -97,6 +97,7 @@ static bool device_supports_param_op(ggml_backend_dev_t device, ModelManager::~ModelManager() { release_all(); + release_prefetch(); } void ModelManager::set_common_ignore_tensors(std::set ignore_tensors) { @@ -142,7 +143,7 @@ size_t estimate_tensors_size(const std::map& tensors) return size; } -void ModelManager::set_split_buffer_type(ggml_backend_t compute_backend, ggml_backend_buffer_type_t split_buft) { +void ModelManager::set_split_buffer_type(ggml_backend_t compute_backend, ggml_backend_buffer_type_t split_buft, const std::vector>& device_limits) { if (compute_backend == nullptr) { return; } @@ -151,6 +152,7 @@ void ModelManager::set_split_buffer_type(ggml_backend_t compute_backend, ggml_ba return; } split_buffer_types_[compute_backend] = split_buft; + split_buffer_devices_[split_buft] = device_limits; } bool ModelManager::tensor_shape_supports_split_buffer(const ggml_tensor* tensor) { @@ -163,11 +165,10 @@ bool ModelManager::tensor_shape_supports_split_buffer(const ggml_tensor* tensor) } ggml_backend_buffer_type_t ModelManager::split_buffer_type_for(const TensorState& state) const { - if (!state.allow_split_buffer || !tensor_shape_supports_split_buffer(state.tensor)) { + if (!tensor_shape_supports_split_buffer(state.tensor)) { return nullptr; } - auto it = split_buffer_types_.find(state.compute_backend); - return it != split_buffer_types_.end() ? it->second : nullptr; + return state.split_buffer_type; } bool ModelManager::register_param_tensors(const std::string& desc, @@ -202,14 +203,17 @@ bool ModelManager::register_param_tensors(const std::string& desc, } ggml_set_name(tensor, name.c_str()); - auto state = std::make_unique(); - state->name = name; - state->tensor = tensor; - state->desc = desc; - state->residency_mode = residency_mode; - state->compute_backend = compute_backend; - state->params_backend = params_backend; - state->allow_split_buffer = allow_split_buffer; + auto state = std::make_unique(); + state->name = name; + state->tensor = tensor; + state->desc = desc; + state->residency_mode = residency_mode; + state->compute_backend = compute_backend; + state->params_backend = params_backend; + auto split_buffer = split_buffer_types_.find(compute_backend); + if (allow_split_buffer && split_buffer != split_buffer_types_.end()) { + state->split_buffer_type = split_buffer->second; + } state->params_follow_compute_backend = params_follow_compute_backend; if (tensor_ops != nullptr) { auto op_it = tensor_ops->find(tensor); @@ -239,7 +243,7 @@ bool ModelManager::unregister_param_tensors(const std::string& desc, size_t* reg if (state == nullptr || state->desc != desc) { continue; } - if (state->active_prepare_count > 0) { + if (state->pin_count > 0) { LOG_ERROR("model manager cannot unregister active %s tensor '%s'", desc.c_str(), state->name.c_str()); @@ -255,6 +259,7 @@ bool ModelManager::unregister_param_tensors(const std::string& desc, size_t* reg return true; } + clear_all_prefetched_params(); release_compute_staging_blocks(false); std::vector storage_blocks_to_release; @@ -285,7 +290,7 @@ bool ModelManager::unregister_param_tensors(const std::string& desc, size_t* reg if (state == nullptr) { continue; } - if (state->active_prepare_count > 0 || state->staged_to_compute_backend) { + if (state->pin_count > 0 || state->staged_to_compute_backend) { LOG_ERROR("model manager cannot unregister %s while tensor '%s' is active", desc.c_str(), state->name.c_str()); @@ -401,14 +406,27 @@ bool ModelManager::load_tensors_to_params_backend(const std::vector prepared; for (ParamsStorageBlock* block : created_storage_blocks) { if (block != nullptr && block->buffer != nullptr) { - LOG_DEBUG("model manager prepared params backend buffer (%6.2f MB, %zu tensors, %s)", - ggml_backend_buffer_get_size(block->buffer) / (1024.f * 1024.f), - block->states.size(), - ggml_backend_buffer_is_host(block->buffer) ? "RAM" : "VRAM"); + auto& stats = prepared[ggml_backend_buffer_get_type(block->buffer)]; + stats.bytes += ggml_backend_buffer_get_size(block->buffer); + stats.tensors += block->states.size(); + ++stats.blocks; } } + for (const auto& entry : prepared) { + LOG_VERBOSE("model manager prepared params backend buffers (%6.2f MB, %zu tensors, %zu blocks, %s) on %s", + entry.second.bytes / (1024.f * 1024.f), + entry.second.tensors, entry.second.blocks, + ggml_backend_buft_is_host(entry.first) ? "RAM" : "VRAM", + ggml_backend_buft_name(entry.first)); + } return true; } @@ -442,68 +460,99 @@ bool ModelManager::stage_tensors_to_compute_backend(const std::vector& states = pair.second; - if (states.empty()) { + ggml_backend_t compute_backend = pair.first.first; + ggml_backend_buffer_type_t staging_buft = pair.first.second; + const std::vector& target_states = pair.second; + if (target_states.empty()) { continue; } - int64_t t0 = ggml_time_ms(); - - ggml_init_params init_params; - init_params.mem_size = std::max(1, states.size()) * ggml_tensor_overhead(); - init_params.mem_buffer = nullptr; - init_params.no_alloc = true; - - ggml_context* staging_ctx = ggml_init(init_params); - GGML_ASSERT(staging_ctx != nullptr); - - std::vector> staged_tensors; - staged_tensors.reserve(states.size()); - for (TensorState* state : states) { - ggml_tensor* staging_tensor = ggml_dup_tensor(staging_ctx, state->tensor); - ggml_set_name(staging_tensor, state->tensor->name); - staged_tensors.push_back({state, staging_tensor}); + const size_t alignment = ggml_backend_buft_get_alignment(staging_buft); + size_t backend_limit = ggml_backend_buft_get_max_size(staging_buft); + if (!ggml_backend_buft_is_host(staging_buft) && + (backend_limit == 0 || backend_limit > MAX_RESIDENCY_BLOCK_BYTES)) { + backend_limit = MAX_RESIDENCY_BLOCK_BYTES; } - ggml_backend_buffer_t compute_buffer = ggml_backend_alloc_ctx_tensors_from_buft(staging_ctx, staging_buft); - if (compute_buffer == nullptr) { - LOG_ERROR("model manager alloc compute params backend buffer failed, num_tensors = %zu", - staged_tensors.size()); - ggml_free(staging_ctx); + const int64_t t0 = ggml_time_ms(); + size_t staged_bytes = 0; + size_t staged_blocks = 0; + auto stage_chunk = [&](const std::vector& chunk) -> bool { + if (chunk.empty()) { + return true; + } + ggml_init_params init_params; + init_params.mem_size = std::max(1, chunk.size()) * ggml_tensor_overhead(); + init_params.mem_buffer = nullptr; + init_params.no_alloc = true; + + ggml_context* staging_ctx = ggml_init(init_params); + GGML_ASSERT(staging_ctx != nullptr); + std::vector> staged_tensors; + staged_tensors.reserve(chunk.size()); + for (TensorState* state : chunk) { + ggml_tensor* staging_tensor = ggml_dup_tensor(staging_ctx, state->tensor); + ggml_set_name(staging_tensor, state->tensor->name); + staged_tensors.push_back({state, staging_tensor}); + } + + ggml_backend_buffer_t compute_buffer = + ggml_backend_alloc_ctx_tensors_from_buft(staging_ctx, staging_buft); + if (compute_buffer == nullptr) { + LOG_ERROR("model manager alloc compute params backend buffer failed, num_tensors = %zu", + staged_tensors.size()); + ggml_free(staging_ctx); + return false; + } + ggml_backend_buffer_set_usage(compute_buffer, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); + for (auto& staged_tensor : staged_tensors) { + TensorState* state = staged_tensor.first; + ggml_tensor* managed_tensor = state->tensor; + ggml_tensor* staging_tensor = staged_tensor.second; + ggml_backend_tensor_copy(managed_tensor, staging_tensor); + std::swap(managed_tensor->buffer, staging_tensor->buffer); + std::swap(managed_tensor->data, staging_tensor->data); + std::swap(managed_tensor->extra, staging_tensor->extra); + state->staged_to_compute_backend = true; + } + ggml_backend_synchronize(compute_backend); + + auto block = std::make_unique(); + block->compute_backend = compute_backend; + block->buffer = compute_buffer; + block->staging_ctx = staging_ctx; + block->staged_tensors = std::move(staged_tensors); + staged_bytes += ggml_backend_buffer_get_size(compute_buffer); + ++staged_blocks; + compute_staging_blocks_.push_back(std::move(block)); + return true; + }; + + std::vector chunk; + size_t chunk_size = 0; + for (TensorState* state : target_states) { + const size_t tensor_size = GGML_PAD( + ggml_backend_buft_get_alloc_size(staging_buft, state->tensor), alignment); + if (!chunk.empty() && backend_limit > 0 && + tensor_size > backend_limit - std::min(chunk_size, backend_limit)) { + if (!stage_chunk(chunk)) { + return false; + } + chunk.clear(); + chunk_size = 0; + } + chunk.push_back(state); + chunk_size = tensor_size > SIZE_MAX - chunk_size ? SIZE_MAX : chunk_size + tensor_size; + } + if (!stage_chunk(chunk)) { return false; } - ggml_backend_buffer_set_usage(compute_buffer, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); - - for (auto& staged_tensor : staged_tensors) { - TensorState* state = staged_tensor.first; - ggml_tensor* managed_tensor = state->tensor; - ggml_tensor* staging_tensor = staged_tensor.second; - ggml_backend_tensor_copy(managed_tensor, staging_tensor); - std::swap(managed_tensor->buffer, staging_tensor->buffer); - std::swap(managed_tensor->data, staging_tensor->data); - std::swap(managed_tensor->extra, staging_tensor->extra); - } - ggml_backend_synchronize(compute_backend); - - auto block = std::make_unique(); - block->compute_backend = compute_backend; - block->buffer = compute_buffer; - block->staging_ctx = staging_ctx; - block->staged_tensors = std::move(staged_tensors); - for (auto& staged_tensor : block->staged_tensors) { - TensorState* state = staged_tensor.first; - state->staged_to_compute_backend = true; - } - compute_staging_blocks_.push_back(std::move(block)); - - int64_t t1 = ggml_time_ms(); - LOG_DEBUG("model manager staged compute params (%6.2f MB, %zu tensors) to %s, taking %.2fs", - ggml_backend_buffer_get_size(compute_buffer) / (1024.f * 1024.f), - states.size(), - ggml_backend_name(compute_backend), - (t1 - t0) * 1.0f / 1000); + LOG_VERBOSE("model manager staged compute params (%6.2f MB, %zu tensors, %zu blocks) to %s, taking %.2fs", + staged_bytes / (1024.f * 1024.f), + target_states.size(), + staged_blocks, + ggml_backend_name(compute_backend), + (ggml_time_ms() - t0) / 1000.f); } return true; @@ -608,6 +657,7 @@ bool ModelManager::apply_loras_to_params(const std::vector& states } void ModelManager::reset_lora_applied_params() { + clear_all_prefetched_params(); release_compute_staging_blocks(true); release_params_storage_blocks(true); for (auto& state : tensor_states_) { @@ -726,6 +776,10 @@ bool ModelManager::alloc_params_buffers(const std::vector& states, const std::vector& states = pair.second; size_t alignment = ggml_backend_buft_get_alignment(params_buft); size_t max_size = ggml_backend_buft_get_max_size(params_buft); + if (!ggml_backend_buft_is_host(params_buft) && + (max_size == 0 || max_size > MAX_RESIDENCY_BLOCK_BYTES)) { + max_size = MAX_RESIDENCY_BLOCK_BYTES; + } auto alloc_chunk = [&](const std::vector& chunk, size_t chunk_size) -> bool { if (chunk.empty() || chunk_size == 0) { @@ -755,10 +809,10 @@ bool ModelManager::alloc_params_buffers(const std::vector& states, initialized->data = nullptr; initialized->extra = nullptr; } - LOG_DEBUG("model manager releasing params backend buffer (%6.2f MB, %zu tensors, %s)", - ggml_backend_buffer_get_size(buffer) / (1024.f * 1024.f), - initialized_tensors.size(), - ggml_backend_buffer_is_host(buffer) ? "RAM" : "VRAM"); + LOG_VERBOSE("model manager releasing params backend buffer (%6.2f MB, %zu tensors, %s)", + ggml_backend_buffer_get_size(buffer) / (1024.f * 1024.f), + initialized_tensors.size(), + ggml_backend_buffer_is_host(buffer) ? "RAM" : "VRAM"); ggml_backend_buffer_free(buffer); return false; } @@ -932,10 +986,6 @@ void ModelManager::free_compute_staging_block(ComputeStagingBlock& block) { } if (block.buffer != nullptr) { - LOG_DEBUG("model manager releasing compute params (%6.2f MB, %zu tensors) from %s", - ggml_backend_buffer_get_size(block.buffer) / (1024.f * 1024.f), - block.staged_tensors.size(), - block.compute_backend != nullptr ? ggml_backend_name(block.compute_backend) : "unknown"); ggml_backend_buffer_free(block.buffer); block.buffer = nullptr; } @@ -948,6 +998,12 @@ void ModelManager::free_compute_staging_block(ComputeStagingBlock& block) { void ModelManager::release_compute_staging_blocks(bool force, const std::unordered_set* target_states) { + struct ReleaseStats { + size_t bytes = 0; + size_t tensors = 0; + size_t blocks = 0; + }; + std::map released; for (auto it = compute_staging_blocks_.begin(); it != compute_staging_blocks_.end();) { ComputeStagingBlock* block = it->get(); bool can_release = force; @@ -963,25 +1019,33 @@ void ModelManager::release_compute_staging_blocks(bool force, target_states->find(state) == target_states->end()) { return false; } - return state->active_prepare_count == 0; + return state->pin_count == 0; }); } if (can_release) { + if (block->buffer != nullptr) { + auto& stats = released[block->compute_backend]; + stats.bytes += ggml_backend_buffer_get_size(block->buffer); + stats.tensors += block->staged_tensors.size(); + ++stats.blocks; + } free_compute_staging_block(*block); it = compute_staging_blocks_.erase(it); } else { ++it; } } + for (const auto& entry : released) { + LOG_DEBUG("model manager released compute params (%6.2f MB, %zu tensors, %zu blocks) from %s", + entry.second.bytes / (1024.f * 1024.f), + entry.second.tensors, entry.second.blocks, + entry.first != nullptr ? ggml_backend_name(entry.first) : "unknown"); + } } void ModelManager::free_params_storage_block(ParamsStorageBlock& block) { if (block.buffer != nullptr) { - LOG_DEBUG("model manager releasing params backend buffer (%6.2f MB, %zu tensors, %s)", - ggml_backend_buffer_get_size(block.buffer) / (1024.f * 1024.f), - block.states.size(), - ggml_backend_buffer_is_host(block.buffer) ? "RAM" : "VRAM"); ggml_backend_buffer_free(block.buffer); block.buffer = nullptr; } @@ -1003,6 +1067,12 @@ void ModelManager::free_params_storage_block(ParamsStorageBlock& block) { void ModelManager::release_params_storage_blocks(bool force, const std::unordered_set* target_states) { + struct ReleaseStats { + size_t bytes = 0; + size_t tensors = 0; + size_t blocks = 0; + }; + std::map released; for (auto it = params_storage_blocks_.begin(); it != params_storage_blocks_.end();) { ParamsStorageBlock* block = it->get(); bool can_release = force; @@ -1017,19 +1087,32 @@ void ModelManager::release_params_storage_blocks(bool force, target_states->find(state) == target_states->end()) { return false; } - return state->active_prepare_count == 0 && + return state->pin_count == 0 && !state->staged_to_compute_backend && state->residency_mode == ResidencyMode::Disk; }); } if (can_release) { + if (block->buffer != nullptr) { + auto& stats = released[ggml_backend_buffer_get_type(block->buffer)]; + stats.bytes += ggml_backend_buffer_get_size(block->buffer); + stats.tensors += block->states.size(); + ++stats.blocks; + } free_params_storage_block(*block); it = params_storage_blocks_.erase(it); } else { ++it; } } + for (const auto& entry : released) { + LOG_VERBOSE("model manager released params backend buffers (%6.2f MB, %zu tensors, %zu blocks, %s) from %s", + entry.second.bytes / (1024.f * 1024.f), + entry.second.tensors, entry.second.blocks, + ggml_backend_buft_is_host(entry.first) ? "RAM" : "VRAM", + ggml_backend_buft_name(entry.first)); + } } void ModelManager::erase_params_storage_block(ParamsStorageBlock* block) { @@ -1044,16 +1127,20 @@ void ModelManager::erase_params_storage_block(ParamsStorageBlock* block) { } void ModelManager::release_all() { + clear_all_prefetched_params(); + runtime_residencies_.clear(); + workspace_reclaimers_.clear(); for (auto& state : tensor_states_) { - state->active_prepare_count = 0; - state->applied_lora_epoch = UINT64_MAX; + state->pin_count = 0; + state->applied_lora_epoch = UINT64_MAX; } release_compute_staging_blocks(true); release_params_storage_blocks(true); } bool ModelManager::resolve_required_tensor_states(const std::vector& tensors, - std::vector& required_states) const { + std::vector& required_states, + ggml_backend_t compute_backend) const { required_states.clear(); std::unordered_set seen; for (ggml_tensor* tensor : tensors) { @@ -1075,7 +1162,9 @@ bool ModelManager::resolve_required_tensor_states(const std::vectorcompute_backend == nullptr || + state->compute_backend == compute_backend) && + seen.insert(state).second) { required_states.push_back(state); } } @@ -1097,6 +1186,7 @@ bool ModelManager::assign_compute_backend(const std::vector& tenso return false; } + clear_all_prefetched_params(); for (TensorState* state : required_states) { if (state == nullptr || state->tensor == nullptr) { continue; @@ -1110,7 +1200,7 @@ bool ModelManager::assign_compute_backend(const std::vector& tenso continue; } - if (state->active_prepare_count > 0 || state->staged_to_compute_backend) { + if (state->pin_count > 0 || state->staged_to_compute_backend) { LOG_ERROR("model manager cannot move active tensor '%s' to another compute backend", state->name.c_str()); return false; @@ -1130,6 +1220,131 @@ bool ModelManager::assign_compute_backend(const std::vector& tenso return true; } +size_t ModelManager::compute_backend_alloc_size(const std::vector& states, + bool missing_only) const { + size_t total_size = 0; + std::unordered_set seen; + for (TensorState* state : states) { + if (state == nullptr || state->tensor == nullptr || !seen.insert(state).second || + should_ignore(*state) || is_optional_missing_tensor(state->name)) { + continue; + } + const bool compute_resident = + state->compute_backend == state->params_backend + ? state->loaded_to_params_backend + : state->staged_to_compute_backend; + if (missing_only && compute_resident) { + continue; + } + + ggml_backend_buffer_type_t buffer_type = nullptr; + if (state->compute_backend == state->params_backend) { + buffer_type = params_buffer_type_for(*state); + } else { + buffer_type = split_buffer_type_for(*state); + if (buffer_type == nullptr && state->compute_backend != nullptr) { + buffer_type = ggml_backend_get_default_buffer_type(state->compute_backend); + } + } + if (buffer_type == nullptr) { + continue; + } + const size_t alignment = ggml_backend_buft_get_alignment(buffer_type); + const size_t tensor_size = ggml_backend_buft_get_alloc_size(buffer_type, state->tensor); + const size_t alloc_size = GGML_PAD(tensor_size, alignment); + if (alloc_size > SIZE_MAX - total_size) { + return SIZE_MAX; + } + total_size += alloc_size; + } + return total_size; +} + +size_t ModelManager::compute_backend_resident_bytes(ggml_backend_t compute_backend) const { + if (compute_backend == nullptr) { + return 0; + } + ggml_backend_dev_t compute_device = ggml_backend_get_device(compute_backend); + if (compute_device == nullptr) { + return 0; + } + + size_t total_size = 0; + auto add_buffer = [&](ggml_backend_buffer_t buffer) { + if (buffer == nullptr || ggml_backend_buffer_is_host(buffer)) { + return; + } + ggml_backend_buffer_type_t buffer_type = ggml_backend_buffer_get_type(buffer); + auto split_devices = split_buffer_devices_.find(buffer_type); + const bool on_device = split_devices == split_buffer_devices_.end() + ? buffer_type != nullptr && ggml_backend_buft_get_device(buffer_type) == compute_device + : std::any_of(split_devices->second.begin(), split_devices->second.end(), [&](const auto& entry) { + return ggml_backend_get_device(entry.first) == compute_device; + }); + if (!on_device) { + return; + } + const size_t buffer_size = ggml_backend_buffer_get_size(buffer); + total_size = buffer_size > SIZE_MAX - total_size ? SIZE_MAX : total_size + buffer_size; + }; + + for (const auto& block : params_storage_blocks_) { + if (block != nullptr) { + add_buffer(block->buffer); + } + } + for (const auto& block : compute_staging_blocks_) { + if (block != nullptr) { + add_buffer(block->buffer); + } + } + for (const auto& entry : prefetch_blocks_) { + if (entry.second != nullptr) { + for (const auto& block : entry.second->staging_blocks) { + if (block != nullptr) { + add_buffer(block->buffer); + } + } + } + } + return total_size; +} + +void ModelManager::update_runtime_residency(uintptr_t owner_id, + ggml_backend_t compute_backend, + size_t resident_bytes) { + if (owner_id == 0) { + return; + } + if (compute_backend == nullptr || resident_bytes == 0) { + runtime_residencies_.erase({owner_id, compute_backend}); + return; + } + runtime_residencies_[{owner_id, compute_backend}] = {compute_backend, resident_bytes}; +} + +size_t ModelManager::other_runtime_resident_bytes(uintptr_t owner_id, + ggml_backend_t compute_backend) const { + if (compute_backend == nullptr) { + return 0; + } + ggml_backend_dev_t compute_device = ggml_backend_get_device(compute_backend); + if (compute_device == nullptr) { + return 0; + } + size_t total_size = 0; + for (const auto& entry : runtime_residencies_) { + if (entry.first.first == owner_id || entry.second.compute_backend == nullptr || + ggml_backend_get_device(entry.second.compute_backend) != compute_device) { + continue; + } + total_size = entry.second.resident_bytes > SIZE_MAX - total_size + ? SIZE_MAX + : total_size + entry.second.resident_bytes; + } + return total_size; +} + bool ModelManager::prepare_params(const std::vector& tensors) { if (tensors.empty()) { return true; @@ -1150,18 +1365,20 @@ bool ModelManager::prepare_params(const std::vector& tensors) { return false; } + // LoRA execution may reclaim other residency blocks while these weights are in use. + const uint64_t use_epoch = ++residency_epoch_; + for (TensorState* state : required_states) { + if (state != nullptr) { + state->pin_count++; + state->last_use_epoch = use_epoch; + } + } if (!apply_loras_to_params(required_states)) { + finish_compute_backend_usage(required_states); release_compute_staging_blocks(false); release_params_storage_blocks(false); return false; } - - for (TensorState* state : required_states) { - if (state == nullptr) { - continue; - } - state->active_prepare_count++; - } return true; } @@ -1175,11 +1392,10 @@ void ModelManager::finish_compute_backend_usage(const std::vector& if (state == nullptr || !target_states.insert(state).second) { continue; } - if (state->active_prepare_count > 0) { - state->active_prepare_count--; + if (state->pin_count > 0) { + state->pin_count--; } } - release_compute_staging_blocks(false, &target_states); } void ModelManager::release_compute_backend_params(const std::vector& tensors) { @@ -1193,7 +1409,7 @@ void ModelManager::release_compute_backend_params(const std::vector& tensors) { +void ModelManager::evict_compute_backend_params(const std::vector& tensors) { if (tensors.empty()) { return; } @@ -1201,9 +1417,285 @@ void ModelManager::release_params_backend_params(const std::vector if (!resolve_required_tensor_states(tensors, required_states)) { return; } - if (required_states.empty()) { - return; - } std::unordered_set target_states(required_states.begin(), required_states.end()); + + for (const auto& block : compute_staging_blocks_) { + const bool intersects = std::any_of( + block->staged_tensors.begin(), + block->staged_tensors.end(), + [&](const std::pair& pair) { + return pair.first != nullptr && target_states.count(pair.first) > 0; + }); + const bool fully_evictable = std::all_of( + block->staged_tensors.begin(), + block->staged_tensors.end(), + [](const std::pair& pair) { + return pair.first == nullptr || pair.first->pin_count == 0; + }); + if (intersects && fully_evictable) { + for (const auto& pair : block->staged_tensors) { + if (pair.first != nullptr) { + target_states.insert(pair.first); + } + } + } + } + release_compute_staging_blocks(false, &target_states); + + for (const auto& block : params_storage_blocks_) { + const bool intersects = std::any_of( + block->states.begin(), + block->states.end(), + [&](TensorState* state) { + return state != nullptr && target_states.count(state) > 0; + }); + const bool fully_evictable = std::all_of( + block->states.begin(), + block->states.end(), + [](TensorState* state) { + return state == nullptr || + (state->pin_count == 0 && !state->staged_to_compute_backend && + state->residency_mode == ResidencyMode::Disk); + }); + if (intersects && fully_evictable) { + target_states.insert(block->states.begin(), block->states.end()); + } + } release_params_storage_blocks(false, &target_states); } +WeightResidencyInfo ModelManager::inspect_compute_backend_params( + const std::vector& tensors) const { + WeightResidencyInfo info; + std::vector states; + if (!resolve_required_tensor_states(tensors, states)) { + return info; + } + + ggml_backend_t prefetch_compute_backend = nullptr; + bool has_missing_params = false; + bool prefetch_candidate = true; + for (TensorState* state : states) { + if (state == nullptr || should_ignore(*state) || + is_optional_missing_tensor(state->name)) { + continue; + } + const bool compute_resident = + state->compute_backend == state->params_backend + ? state->loaded_to_params_backend + : state->staged_to_compute_backend; + if (compute_resident) { + continue; + } + has_missing_params = true; + if (split_buffer_type_for(*state) != nullptr) { + prefetch_candidate = false; + } + if (state->compute_backend == state->params_backend || + state->compute_backend == nullptr || sd_backend_is_cpu(state->compute_backend)) { + prefetch_candidate = false; + continue; + } + if (prefetch_compute_backend == nullptr) { + prefetch_compute_backend = state->compute_backend; + } else if (prefetch_compute_backend != state->compute_backend) { + prefetch_candidate = false; + } + } + info.missing_bytes = compute_backend_alloc_size(states, true); + if (has_missing_params && prefetch_candidate && prefetch_compute_backend != nullptr) { + ggml_backend_dev_t device = ggml_backend_get_device(prefetch_compute_backend); + if (device != nullptr) { + ggml_backend_dev_props props{}; + ggml_backend_dev_get_props(device, &props); + info.async_prefetch_supported = props.caps.async; + } + } + return info; +} + +void ModelManager::set_workspace_reclaimer(uintptr_t owner_id, std::function reclaim) { + workspace_reclaimers_[owner_id] = std::move(reclaim); +} + +void ModelManager::remove_runtime_owner(uintptr_t owner_id) { + workspace_reclaimers_.erase(owner_id); + for (auto it = runtime_residencies_.begin(); it != runtime_residencies_.end();) { + if (it->first.first == owner_id) { + it = runtime_residencies_.erase(it); + } else { + ++it; + } + } +} + +ModelManager::CapacityCheck ModelManager::check_capacity( + const DeviceMemoryRequest& request, + const std::vector& states) const { + CapacityCheck result; + if (request.compute_backend == nullptr || sd_backend_is_cpu(request.compute_backend)) { + return result; + } + auto add = [](size_t a, size_t b) { return b > SIZE_MAX - a ? SIZE_MAX : a + b; }; + const size_t missing = compute_backend_alloc_size(states, true); + result.required_device_bytes = add(request.pending_allocation_bytes, missing); + result.required_budget_bytes = add(request.runtime_peak_bytes(), missing); + auto device = ggml_backend_get_device(request.compute_backend); + if (device != nullptr) { + size_t free_bytes = 0, total_bytes = 0; + ggml_backend_dev_memory(device, &free_bytes, &total_bytes); + if (free_bytes != 0 || total_bytes != 0) { + result.available_device_bytes = free_bytes; + } + } + if (request.max_backend_bytes > 0) { + const size_t resident = add(compute_backend_resident_bytes(request.compute_backend), + other_runtime_resident_bytes(request.owner_id, request.compute_backend)); + result.available_budget_bytes = resident < request.max_backend_bytes + ? request.max_backend_bytes - resident + : 0; + } + std::map split_devices; + for (auto state : states) { + auto placement = split_buffer_devices_.find(split_buffer_type_for(*state)); + if (placement != split_buffer_devices_.end()) { + for (const auto& entry : placement->second) { + auto inserted = split_devices.emplace(entry); + if (!inserted.second && entry.second > 0) { + auto& limit = inserted.first->second; + limit = limit == 0 ? entry.second : std::min(limit, entry.second); + } + } + } + } + // GGML exposes only a split buffer's total size, not per-device allocations. + // Charge that upper bound on every participant instead of undercounting a shard. + for (const auto& entry : split_devices) { + size_t free_bytes = 0, total_bytes = 0; + ggml_backend_dev_memory(ggml_backend_get_device(entry.first), &free_bytes, &total_bytes); + if (free_bytes != 0 || total_bytes != 0) { + result.available_device_bytes = std::min(result.available_device_bytes, free_bytes); + } + if (entry.second > 0) { + const size_t resident = add(compute_backend_resident_bytes(entry.first), + other_runtime_resident_bytes(request.owner_id, entry.first)); + const size_t available = resident < entry.second ? entry.second - resident : 0; + result.available_budget_bytes = std::min(result.available_budget_bytes, available); + } + } + return result; +} + +bool ModelManager::fits_compute_backend_capacity( + const DeviceMemoryRequest& request, + const std::vector& required_params) const { + std::vector states; + return resolve_required_tensor_states(required_params, states, request.compute_backend) && + check_capacity(request, states).fits(); +} + +bool ModelManager::ensure_compute_backend_capacity( + const DeviceMemoryRequest& request, + const std::vector& required_params, + const std::vector>& preferred_eviction_order, + const std::vector& protected_params) { + std::vector required_states; + if (!resolve_required_tensor_states(required_params, required_states, request.compute_backend)) { + return false; + } + + ggml_backend_t compute_backend = request.compute_backend; + if (compute_backend == nullptr) { + LOG_ERROR("model manager cannot reclaim memory for a null compute backend"); + return false; + } + if (sd_backend_is_cpu(compute_backend)) { + return true; + } + + auto fits = [&]() { return check_capacity(request, required_states).fits(); }; + if (fits()) { + return true; + } + for (const auto& entry : workspace_reclaimers_) { + if (entry.first != request.owner_id) { + entry.second(); + if (fits()) { + return true; + } + } + } + + std::unordered_set protected_states; + std::vector resolved_protected; + if (!resolve_required_tensor_states(protected_params, resolved_protected)) { + return false; + } + protected_states.insert(resolved_protected.begin(), resolved_protected.end()); + for (const auto& entry : prefetch_blocks_) { + if (entry.second != nullptr) { + protected_states.insert(entry.second->states.begin(), entry.second->states.end()); + } + } + + std::unordered_set eviction_states; + auto add_evictable_state = [&](TensorState* state) { + if (state == nullptr || state->compute_backend != compute_backend || + state->pin_count > 0 || protected_states.find(state) != protected_states.end()) { + return; + } + const bool reloadable = state->residency_mode == ResidencyMode::Disk || + state->compute_backend != state->params_backend; + const bool resident = state->compute_backend == state->params_backend + ? state->loaded_to_params_backend + : state->staged_to_compute_backend; + if (reloadable && resident) { + eviction_states.insert(state); + } + }; + auto release_eviction_states = [&]() { + release_compute_staging_blocks(false, &eviction_states); + release_params_storage_blocks(false, &eviction_states); + return fits(); + }; + + for (const auto& candidate_params : preferred_eviction_order) { + std::vector candidate_states; + if (!resolve_required_tensor_states(candidate_params, candidate_states)) { + return false; + } + for (TensorState* state : candidate_states) { + add_evictable_state(state); + } + if (release_eviction_states()) { + return true; + } + } + + std::vector global_candidates; + global_candidates.reserve(tensor_states_.size()); + for (const auto& state : tensor_states_) { + if (state != nullptr && eviction_states.find(state.get()) == eviction_states.end()) { + global_candidates.push_back(state.get()); + } + } + std::stable_sort(global_candidates.begin(), + global_candidates.end(), + [](const TensorState* lhs, const TensorState* rhs) { + return lhs->last_use_epoch < rhs->last_use_epoch; + }); + for (TensorState* state : global_candidates) { + add_evictable_state(state); + if (release_eviction_states()) { + return true; + } + } + + const auto capacity = check_capacity(request, required_states); + LOG_WARN("model manager cannot make enough memory available on %s: need %.2f MB device / %.2f MB budget, available %.2f MB device / %.2f MB budget", + ggml_backend_name(compute_backend), + capacity.required_device_bytes / (1024.0 * 1024.0), + capacity.required_budget_bytes / (1024.0 * 1024.0), + capacity.available_device_bytes / (1024.0 * 1024.0), + capacity.available_budget_bytes / (1024.0 * 1024.0)); + return false; +} diff --git a/otherarch/sdcpp/src/model_manager.h b/otherarch/sdcpp/src/model_manager.h index 3dc633864..b23f0db63 100644 --- a/otherarch/sdcpp/src/model_manager.h +++ b/otherarch/sdcpp/src/model_manager.h @@ -9,10 +9,10 @@ #include #include +#include "device_residency_manager.h" #include "model_loader.h" -#include "weight_manager.h" -class ModelManager : public RunnerWeightManager { +class ModelManager : public DeviceResidencyManager { public: enum class ResidencyMode { Disk, @@ -28,24 +28,27 @@ public: }; private: + static constexpr size_t MAX_RESIDENCY_BLOCK_BYTES = 64ULL * 1024ULL * 1024ULL; + struct TensorState { std::string name; ggml_tensor* tensor = nullptr; std::string desc; - ResidencyMode residency_mode = ResidencyMode::ParamBackend; - ggml_backend_t compute_backend = nullptr; - ggml_backend_t params_backend = nullptr; - bool allow_split_buffer = false; - bool params_follow_compute_backend = false; - bool metadata_validated = false; - enum ggml_op usage_op = GGML_OP_NONE; + ResidencyMode residency_mode = ResidencyMode::ParamBackend; + ggml_backend_t compute_backend = nullptr; + ggml_backend_t params_backend = nullptr; + ggml_backend_buffer_type_t split_buffer_type = nullptr; + bool params_follow_compute_backend = false; + bool metadata_validated = false; + enum ggml_op usage_op = GGML_OP_NONE; - int active_prepare_count = 0; + int pin_count = 0; bool loaded_to_params_backend = false; bool staged_to_compute_backend = false; uint64_t applied_lora_epoch = UINT64_MAX; + uint64_t last_use_epoch = 0; }; struct ParamsStorageBlock { @@ -61,26 +64,55 @@ private: std::vector> staged_tensors; }; + struct PrefetchBlock { + std::vector states; + ggml_backend_t compute_backend = nullptr; + ggml_backend_t transfer_backend = nullptr; + ggml_backend_event_t event = nullptr; + std::vector> staging_blocks; + }; + + struct RuntimeResidency { + ggml_backend_t compute_backend = nullptr; + size_t resident_bytes = 0; + }; + ModelLoader model_loader_; std::vector> tensor_states_; std::map tensor_states_by_name_; std::vector> params_storage_blocks_; std::vector> compute_staging_blocks_; std::map split_buffer_types_; + std::map>> split_buffer_devices_; + std::map> prefetch_blocks_; + std::map prefetch_backends_; + std::map, RuntimeResidency> runtime_residencies_; + std::map> workspace_reclaimers_; bool warned_split_lora_skip_ = false; std::set common_ignore_tensors_; std::vector loras_; - SDVersion lora_version_ = VERSION_COUNT; - uint64_t current_lora_epoch_ = 0; - int n_threads_ = 0; - bool enable_mmap_ = false; - bool writable_mmap_ = false; + SDVersion lora_version_ = VERSION_COUNT; + uint64_t current_lora_epoch_ = 0; + uint64_t residency_epoch_ = 0; + int n_threads_ = 0; + bool enable_mmap_ = false; + bool writable_mmap_ = false; + bool segmented_compute_disabled_ = false; + bool prefetch_disabled_ = false; void finish_compute_backend_usage(const std::vector& states); void release_all(); + ggml_backend_t prefetch_backend_for(ggml_backend_t compute_backend); + bool populate_prefetch_block(PrefetchBlock& block); + void synchronize_prefetch_block(PrefetchBlock& block); + void free_prefetch_block(PrefetchBlock& block); + void clear_all_prefetched_params(); + void release_prefetch(); + bool resolve_required_tensor_states(const std::vector& tensors, - std::vector& required_states) const; + std::vector& required_states, + ggml_backend_t compute_backend = nullptr) const; bool should_ignore(const TensorState& state) const; bool is_optional_missing_tensor(const std::string& name) const; bool validate_tensor(const TensorState& state) const; @@ -94,6 +126,21 @@ private: std::vector& created_storage_blocks); bool load_tensors(const std::vector& states); bool stage_tensors_to_compute_backend(const std::vector& states); + size_t compute_backend_alloc_size(const std::vector& states, + bool missing_only) const; + size_t compute_backend_resident_bytes(ggml_backend_t compute_backend) const; + struct CapacityCheck { + size_t required_device_bytes = 0; + size_t required_budget_bytes = 0; + size_t available_device_bytes = SIZE_MAX; + size_t available_budget_bytes = SIZE_MAX; + bool fits() const { + return required_device_bytes <= available_device_bytes && + required_budget_bytes <= available_budget_bytes; + } + }; + CapacityCheck check_capacity(const DeviceMemoryRequest& request, + const std::vector& states) const; ggml_backend_buffer_type_t params_buffer_type_for(const TensorState& state) const; ggml_backend_buffer_type_t split_buffer_type_for(const TensorState& state) const; @@ -105,6 +152,8 @@ private: void free_params_storage_block(ParamsStorageBlock& block); void erase_params_storage_block(ParamsStorageBlock* block); void reset_lora_applied_params(); + size_t other_runtime_resident_bytes(uintptr_t owner_id, + ggml_backend_t compute_backend) const; public: ~ModelManager() override; @@ -116,11 +165,15 @@ public: n_threads_ = n_threads; model_loader_.set_n_threads(n_threads); } + void set_segmented_compute_disabled(bool disabled) { + segmented_compute_disabled_ = disabled; + } + void set_prefetch_disabled(bool disabled) { prefetch_disabled_ = disabled; } void set_enable_mmap(bool enable_mmap) { enable_mmap_ = enable_mmap; } void set_writable_mmap(bool writable_mmap) { writable_mmap_ = writable_mmap; } void set_common_ignore_tensors(std::set ignore_tensors); void set_loras(std::vector loras, SDVersion version); - void set_split_buffer_type(ggml_backend_t compute_backend, ggml_backend_buffer_type_t split_buft); + void set_split_buffer_type(ggml_backend_t compute_backend, ggml_backend_buffer_type_t split_buft, const std::vector>& device_limits); static bool tensor_shape_supports_split_buffer(const ggml_tensor* tensor); @@ -180,8 +233,30 @@ public: bool assign_compute_backend(const std::vector& tensors, ggml_backend_t compute_backend) override; bool prepare_params(const std::vector& tensors) override; + void set_workspace_reclaimer(uintptr_t owner_id, std::function reclaim) override; + void remove_runtime_owner(uintptr_t owner_id) override; + bool fits_compute_backend_capacity(const DeviceMemoryRequest& request, + const std::vector& required_params) const override; + bool segmented_compute_enabled() const override { return !segmented_compute_disabled_; } + bool prefetch_enabled() const override { return !prefetch_disabled_; } void release_compute_backend_params(const std::vector& tensors) override; - void release_params_backend_params(const std::vector& tensors) override; + void evict_compute_backend_params(const std::vector& tensors) override; + WeightResidencyInfo inspect_compute_backend_params( + const std::vector& tensors) const override; + void update_runtime_residency(uintptr_t owner_id, + ggml_backend_t compute_backend, + size_t resident_bytes) override; + bool ensure_compute_backend_capacity( + const DeviceMemoryRequest& request, + const std::vector& required_params, + const std::vector>& preferred_eviction_order, + const std::vector& protected_params) override; + WeightPrefetchResult prefetch_params( + uintptr_t owner_id, + const std::vector& tensors) override; + bool activate_prefetched_params(uintptr_t owner_id, + const std::vector& tensors) override; + void clear_prefetched_params(uintptr_t owner_id) override; }; #endif // __MODEL_MANAGER_H__ diff --git a/otherarch/sdcpp/src/model_manager_prefetch.cpp b/otherarch/sdcpp/src/model_manager_prefetch.cpp new file mode 100644 index 000000000..1d3041938 --- /dev/null +++ b/otherarch/sdcpp/src/model_manager_prefetch.cpp @@ -0,0 +1,353 @@ +#include "model_manager.h" + +#include +#include + +#include "core/ggml_extend_backend.h" +#include "core/util.h" + +ggml_backend_t ModelManager::prefetch_backend_for(ggml_backend_t compute_backend) { + auto existing = prefetch_backends_.find(compute_backend); + if (existing != prefetch_backends_.end()) { + return existing->second; + } + if (compute_backend == nullptr) { + return nullptr; + } + + ggml_backend_dev_t device = ggml_backend_get_device(compute_backend); + if (device == nullptr || ggml_backend_dev_type(device) == GGML_BACKEND_DEVICE_TYPE_CPU) { + return nullptr; + } + ggml_backend_t transfer_backend = ggml_backend_dev_init(device, nullptr); + if (transfer_backend == nullptr) { + LOG_WARN("model manager failed to create a prefetch backend for %s", + ggml_backend_name(compute_backend)); + } + prefetch_backends_[compute_backend] = transfer_backend; + return transfer_backend; +} + +void ModelManager::synchronize_prefetch_block(PrefetchBlock& block) { + if (block.event != nullptr) { + ggml_backend_event_synchronize(block.event); + ggml_backend_event_free(block.event); + block.event = nullptr; + } else if (block.transfer_backend != nullptr) { + ggml_backend_synchronize(block.transfer_backend); + } + block.transfer_backend = nullptr; +} + +void ModelManager::free_prefetch_block(PrefetchBlock& block) { + synchronize_prefetch_block(block); + for (auto& staging_block : block.staging_blocks) { + if (staging_block == nullptr) { + continue; + } + staging_block->staged_tensors.clear(); + if (staging_block->buffer != nullptr) { + ggml_backend_buffer_free(staging_block->buffer); + staging_block->buffer = nullptr; + } + if (staging_block->staging_ctx != nullptr) { + ggml_free(staging_block->staging_ctx); + staging_block->staging_ctx = nullptr; + } + } + block.staging_blocks.clear(); +} + +bool ModelManager::populate_prefetch_block(PrefetchBlock& block) { + if (block.states.empty() || block.compute_backend == nullptr) { + return false; + } + + block.transfer_backend = prefetch_backend_for(block.compute_backend); + if (block.transfer_backend == nullptr) { + return false; + } + + for (TensorState* state : block.states) { + if (state == nullptr || state->tensor == nullptr || + state->tensor->buffer == nullptr || state->tensor->data == nullptr || + state->params_backend == nullptr || state->staged_to_compute_backend || + state->pin_count > 0) { + return false; + } + } + + ggml_backend_buffer_type_t buffer_type = + ggml_backend_get_default_buffer_type(block.compute_backend); + if (buffer_type == nullptr) { + return false; + } + const size_t alignment = ggml_backend_buft_get_alignment(buffer_type); + size_t backend_limit = ggml_backend_buft_get_max_size(buffer_type); + if (!ggml_backend_buft_is_host(buffer_type) && + (backend_limit == 0 || backend_limit > MAX_RESIDENCY_BLOCK_BYTES)) { + backend_limit = MAX_RESIDENCY_BLOCK_BYTES; + } + + auto enqueue_chunk = [&](const std::vector& chunk) -> bool { + if (chunk.empty()) { + return true; + } + ggml_init_params init_params; + init_params.mem_size = chunk.size() * ggml_tensor_overhead(); + init_params.mem_buffer = nullptr; + init_params.no_alloc = true; + ggml_context* staging_ctx = ggml_init(init_params); + if (staging_ctx == nullptr) { + return false; + } + + auto staging_block = std::make_unique(); + staging_block->compute_backend = block.compute_backend; + staging_block->staging_ctx = staging_ctx; + staging_block->staged_tensors.reserve(chunk.size()); + for (TensorState* state : chunk) { + ggml_tensor* staging_tensor = ggml_dup_tensor(staging_ctx, state->tensor); + ggml_set_name(staging_tensor, state->tensor->name); + if (ggml_backend_buffer_is_host(state->tensor->buffer) && + (!ggml_is_contiguous(state->tensor) || !ggml_is_contiguous(staging_tensor) || + ggml_nbytes(state->tensor) != ggml_nbytes(staging_tensor))) { + ggml_free(staging_ctx); + staging_block->staging_ctx = nullptr; + return false; + } + staging_block->staged_tensors.push_back({state, staging_tensor}); + } + staging_block->buffer = + ggml_backend_alloc_ctx_tensors_from_buft(staging_ctx, buffer_type); + if (staging_block->buffer == nullptr) { + ggml_free(staging_ctx); + staging_block->staging_ctx = nullptr; + return false; + } + ggml_backend_buffer_set_usage(staging_block->buffer, + GGML_BACKEND_BUFFER_USAGE_WEIGHTS); + + for (const auto& pair : staging_block->staged_tensors) { + TensorState* state = pair.first; + ggml_tensor* staging_tensor = pair.second; + const bool host_source = ggml_backend_buffer_is_host(state->tensor->buffer); + if (host_source) { + ggml_backend_tensor_set_async(block.transfer_backend, + staging_tensor, + state->tensor->data, + 0, + ggml_nbytes(state->tensor)); + } else { + ggml_backend_tensor_copy_async(state->params_backend, + block.transfer_backend, + state->tensor, + staging_tensor); + } + } + block.staging_blocks.push_back(std::move(staging_block)); + return true; + }; + + std::vector chunk; + size_t chunk_size = 0; + for (TensorState* state : block.states) { + const size_t tensor_size = GGML_PAD( + ggml_backend_buft_get_alloc_size(buffer_type, state->tensor), alignment); + if (!chunk.empty() && backend_limit > 0 && + tensor_size > backend_limit - std::min(chunk_size, backend_limit)) { + if (!enqueue_chunk(chunk)) { + return false; + } + chunk.clear(); + chunk_size = 0; + } + chunk.push_back(state); + chunk_size = tensor_size > SIZE_MAX - chunk_size ? SIZE_MAX : chunk_size + tensor_size; + } + if (!enqueue_chunk(chunk)) { + return false; + } + + ggml_backend_dev_t device = ggml_backend_get_device(block.transfer_backend); + block.event = ggml_backend_event_new(device); + if (block.event != nullptr) { + ggml_backend_event_record(block.event, block.transfer_backend); + } + + size_t total_size = 0; + for (const auto& staging_block : block.staging_blocks) { + if (staging_block != nullptr && staging_block->buffer != nullptr) { + const size_t buffer_size = ggml_backend_buffer_get_size(staging_block->buffer); + total_size = buffer_size > SIZE_MAX - total_size ? SIZE_MAX : total_size + buffer_size; + } + } + LOG_DEBUG("model manager queued segment prefetch (%6.2f MB, %zu tensors) to %s", + total_size / (1024.f * 1024.f), + block.states.size(), + ggml_backend_name(block.compute_backend)); + return true; +} + +WeightPrefetchResult ModelManager::prefetch_params( + uintptr_t owner_id, + const std::vector& tensors) { + if (tensors.empty()) { + return WeightPrefetchResult::AlreadyResident; + } + + std::vector required_states; + if (!resolve_required_tensor_states(tensors, required_states)) { + return WeightPrefetchResult::Failed; + } + + std::vector states; + states.reserve(required_states.size()); + ggml_backend_t compute_backend = nullptr; + bool needs_synchronous_load = false; + for (TensorState* state : required_states) { + if (state == nullptr || should_ignore(*state) || + is_optional_missing_tensor(state->name)) { + continue; + } + if (state->compute_backend == state->params_backend) { + needs_synchronous_load = needs_synchronous_load || + !state->loaded_to_params_backend; + continue; + } + if (state->staged_to_compute_backend || state->pin_count > 0) { + continue; + } + // Split buffers cannot use the primary device's asynchronous upload path. + if (split_buffer_type_for(*state) != nullptr) { + return WeightPrefetchResult::Unsupported; + } + if (compute_backend == nullptr) { + compute_backend = state->compute_backend; + } else if (compute_backend != state->compute_backend) { + return WeightPrefetchResult::Failed; + } + states.push_back(state); + } + if (states.empty()) { + return needs_synchronous_load ? WeightPrefetchResult::Unsupported + : WeightPrefetchResult::AlreadyResident; + } + if (compute_backend == nullptr || sd_backend_is_cpu(compute_backend)) { + return WeightPrefetchResult::Unsupported; + } + ggml_backend_dev_t compute_device = ggml_backend_get_device(compute_backend); + ggml_backend_dev_props compute_props{}; + if (compute_device == nullptr) { + return WeightPrefetchResult::Unsupported; + } + ggml_backend_dev_get_props(compute_device, &compute_props); + if (!compute_props.caps.async) { + return WeightPrefetchResult::Unsupported; + } + clear_prefetched_params(owner_id); + if (!load_tensors_to_params_backend(states)) { + return WeightPrefetchResult::Failed; + } + + auto block = std::make_unique(); + block->states = std::move(states); + block->compute_backend = compute_backend; + if (!populate_prefetch_block(*block)) { + free_prefetch_block(*block); + return WeightPrefetchResult::Failed; + } + prefetch_blocks_[owner_id] = std::move(block); + return WeightPrefetchResult::Scheduled; +} + +bool ModelManager::activate_prefetched_params( + uintptr_t owner_id, + const std::vector& tensors) { + std::vector required_states; + if (!resolve_required_tensor_states(tensors, required_states)) { + return false; + } + + const bool already_staged = std::all_of( + required_states.begin(), + required_states.end(), + [&](TensorState* state) { + return state == nullptr || should_ignore(*state) || + is_optional_missing_tensor(state->name) || + state->compute_backend == state->params_backend || + state->staged_to_compute_backend; + }); + if (already_staged) { + clear_prefetched_params(owner_id); + return true; + } + + auto existing = prefetch_blocks_.find(owner_id); + if (existing == prefetch_blocks_.end()) { + return false; + } + std::unique_ptr block = std::move(existing->second); + prefetch_blocks_.erase(existing); + synchronize_prefetch_block(*block); + + for (const auto& staging_block : block->staging_blocks) { + if (staging_block == nullptr) { + continue; + } + for (const auto& pair : staging_block->staged_tensors) { + TensorState* state = pair.first; + ggml_tensor* staging_tensor = pair.second; + if (state == nullptr || state->tensor == nullptr || staging_tensor == nullptr || + state->staged_to_compute_backend || state->pin_count > 0) { + free_prefetch_block(*block); + return false; + } + } + } + const uint64_t use_epoch = ++residency_epoch_; + for (auto& staging_block : block->staging_blocks) { + if (staging_block == nullptr) { + continue; + } + for (auto& pair : staging_block->staged_tensors) { + TensorState* state = pair.first; + ggml_tensor* staging_tensor = pair.second; + std::swap(state->tensor->buffer, staging_tensor->buffer); + std::swap(state->tensor->data, staging_tensor->data); + std::swap(state->tensor->extra, staging_tensor->extra); + state->staged_to_compute_backend = true; + state->last_use_epoch = use_epoch; + } + compute_staging_blocks_.push_back(std::move(staging_block)); + } + block->staging_blocks.clear(); + return true; +} + +void ModelManager::clear_prefetched_params(uintptr_t owner_id) { + auto existing = prefetch_blocks_.find(owner_id); + if (existing == prefetch_blocks_.end()) { + return; + } + std::unique_ptr block = std::move(existing->second); + prefetch_blocks_.erase(existing); + free_prefetch_block(*block); +} + +void ModelManager::clear_all_prefetched_params() { + for (auto& entry : prefetch_blocks_) { + free_prefetch_block(*entry.second); + } + prefetch_blocks_.clear(); +} + +void ModelManager::release_prefetch() { + clear_all_prefetched_params(); + for (auto& entry : prefetch_backends_) { + if (entry.second != nullptr) { + ggml_backend_free(entry.second); + } + } + prefetch_backends_.clear(); +} diff --git a/otherarch/sdcpp/src/name_conversion.cpp b/otherarch/sdcpp/src/name_conversion.cpp index 7e528c77c..d0ff45c8e 100644 --- a/otherarch/sdcpp/src/name_conversion.cpp +++ b/otherarch/sdcpp/src/name_conversion.cpp @@ -1448,7 +1448,7 @@ std::string convert_tensor_name(std::string name, SDVersion version) { } } - // LOG_DEBUG("name %s %d", name.c_str(), version); + // LOG_VERBOSE("name %s %d", name.c_str(), version); if (sd_version_is_unet(version) || is_underline || is_lycoris_underline) { name = convert_sep_to_dot(name); diff --git a/otherarch/sdcpp/src/runtime/cache_dit.hpp b/otherarch/sdcpp/src/runtime/cache_dit.hpp index bec6e8115..273aa0ef3 100644 --- a/otherarch/sdcpp/src/runtime/cache_dit.hpp +++ b/otherarch/sdcpp/src/runtime/cache_dit.hpp @@ -8,8 +8,8 @@ #include #include -#include "core/ggml_extend.hpp" #include "core/tensor.hpp" +#include "core/util.h" #include "runtime/condition_cache_utils.hpp" struct DBCacheConfig { diff --git a/otherarch/sdcpp/src/runtime/denoiser.hpp b/otherarch/sdcpp/src/runtime/denoiser.hpp index b6f1843ec..981669c2e 100644 --- a/otherarch/sdcpp/src/runtime/denoiser.hpp +++ b/otherarch/sdcpp/src/runtime/denoiser.hpp @@ -11,8 +11,10 @@ #include #include -#include "core/ggml_extend.hpp" +#include "core/rng.hpp" #include "core/tensor.hpp" +#include "core/util.h" +#include "model.h" #include "runtime/gits_noise.h" #include "runtime/guidance.h" @@ -311,7 +313,7 @@ struct BetaScheduler : SigmaScheduler { explicit BetaScheduler(const char* extra_sample_args = nullptr) { parse_extra_sample_args(extra_sample_args); - LOG_DEBUG("Beta scheduler: alpha=%.4f, beta=%.4f", alpha, beta); + LOG_VERBOSE("Beta scheduler: alpha=%.4f, beta=%.4f", alpha, beta); } void parse_extra_sample_args(const char* extra_sample_args) { @@ -692,7 +694,7 @@ struct LTX2Scheduler : SigmaScheduler { float exp_shift = std::exp(sigma_shift); float target_terminal = std::clamp(terminal, 0.0f, 0.99f); - LOG_DEBUG("LTX2 scheduler: tokens=%d, shift=%.4f, stretch=%d, terminal=%.4f", token_count, sigma_shift, stretch ? 1 : 0, target_terminal); + LOG_VERBOSE("LTX2 scheduler: tokens=%d, shift=%.4f, stretch=%d, terminal=%.4f", token_count, sigma_shift, stretch ? 1 : 0, target_terminal); sigmas.reserve(n + 1); for (uint32_t i = 0; i <= n; ++i) { @@ -760,7 +762,7 @@ struct FluxScheduler : SigmaScheduler { sigmas.reserve(n + 1); float mu = compute_mu(); - LOG_DEBUG("Flux scheduler: image_seq_len=%d, steps=%u, mu=%.3f", image_seq_len, n, mu); + LOG_VERBOSE("Flux scheduler: image_seq_len=%d, steps=%u, mu=%.3f", image_seq_len, n, mu); if (n == 0) { sigmas.push_back(1.0f); @@ -811,7 +813,7 @@ struct Flux2Scheduler : SigmaScheduler { sigmas.reserve(n + 1); float mu = compute_empirical_mu(image_seq_len, n); - LOG_DEBUG("Flux2 scheduler: image_seq_len=%d, steps=%u, mu=%.3f", image_seq_len, n, mu); + LOG_VERBOSE("Flux2 scheduler: image_seq_len=%d, steps=%u, mu=%.3f", image_seq_len, n, mu); if (n == 0) { sigmas.push_back(1.0f); @@ -1413,8 +1415,8 @@ struct SefiFlowDenoiser : public FluxFlowDenoiser { sem_sigmas.push_back(sigma_sem); tex_sigmas.push_back(sigma_tex); } - LOG_DEBUG("SefiFlowDenoiser: built %u-step dual schedule (alpha=%.2f delta_t=%.2f)", - n, timestep_shift_alpha, delta_t); + LOG_VERBOSE("SefiFlowDenoiser: built %u-step dual schedule (alpha=%.2f delta_t=%.2f)", + n, timestep_shift_alpha, delta_t); return tex_sigmas; } }; @@ -2690,7 +2692,7 @@ static sd::Tensor sample_lms(denoise_cb_t model, int steps = static_cast(sigmas.size()) - 1; max_order = std::min(max_order, steps); // history can not be larger than steps - LOG_DEBUG("linear multi-step sampler: lms_max_order = %i, lms_shift = %i, lms_divisions = %i", max_order, shift, divisions); + LOG_VERBOSE("linear multi-step sampler: lms_max_order = %i, lms_shift = %i, lms_divisions = %i", max_order, shift, divisions); std::vector lms_coeff(max_order); std::vector> hist = {}; @@ -2793,7 +2795,7 @@ static sd::Tensor sample_gradient_estimation(denoise_cb_t model, LOG_WARN("ignoring invalid euler_ge extra sample arg '%s=%s'", key.c_str(), value.c_str()); continue; } - LOG_DEBUG("setting euler_ge gamma to %.2f", parsed); + LOG_VERBOSE("setting euler_ge gamma to %.2f", parsed); ge_gamma = parsed; } } diff --git a/otherarch/sdcpp/src/runtime/easycache.hpp b/otherarch/sdcpp/src/runtime/easycache.hpp index 75ae3ddad..bae9f305d 100644 --- a/otherarch/sdcpp/src/runtime/easycache.hpp +++ b/otherarch/sdcpp/src/runtime/easycache.hpp @@ -6,8 +6,8 @@ #include #include -#include "core/ggml_extend.hpp" #include "core/tensor.hpp" +#include "core/util.h" #include "runtime/condition_cache_utils.hpp" #include "runtime/denoiser.hpp" diff --git a/otherarch/sdcpp/src/runtime/preprocessing.hpp b/otherarch/sdcpp/src/runtime/preprocessing.hpp index b39a9a9d0..1990d15f4 100644 --- a/otherarch/sdcpp/src/runtime/preprocessing.hpp +++ b/otherarch/sdcpp/src/runtime/preprocessing.hpp @@ -4,7 +4,7 @@ #include #include -#include "core/ggml_extend.hpp" +#include "core/ggml_tensor_utils.h" #define M_PI_ 3.14159265358979323846f diff --git a/otherarch/sdcpp/src/runtime/spectrum.hpp b/otherarch/sdcpp/src/runtime/spectrum.hpp index 7246d6f09..e30f69308 100644 --- a/otherarch/sdcpp/src/runtime/spectrum.hpp +++ b/otherarch/sdcpp/src/runtime/spectrum.hpp @@ -5,8 +5,8 @@ #include #include -#include "core/ggml_extend.hpp" #include "core/tensor.hpp" +#include "core/util.h" struct SpectrumConfig { float w = 0.40f; diff --git a/otherarch/sdcpp/src/runtime/tiling.cpp b/otherarch/sdcpp/src/runtime/tiling.cpp new file mode 100644 index 000000000..f0a9dbc73 --- /dev/null +++ b/otherarch/sdcpp/src/runtime/tiling.cpp @@ -0,0 +1,268 @@ +#include "runtime/tiling.h" + +#include +#include +#include +#include + +#include "core/util.h" +#include "ggml.h" + +static void sd_tiling_calc_tiles(int& num_tiles_dim, + float& tile_overlap_factor_dim, + int small_dim, + int tile_size, + const float tile_overlap_factor, + bool circular) { + int tile_overlap = static_cast(tile_size * tile_overlap_factor); + int non_tile_overlap = tile_size - tile_overlap; + + if (circular) { + // circular means the last and first tile are overlapping (wraping around) + num_tiles_dim = small_dim / non_tile_overlap; + + if (num_tiles_dim < 1) { + num_tiles_dim = 1; + } + + tile_overlap_factor_dim = (tile_size - small_dim / num_tiles_dim) / (float)tile_size; + + // if single tile and tile_overlap_factor is not 0, add one to ensure we have at least two overlapping tiles + if (num_tiles_dim == 1 && tile_overlap_factor_dim > 0) { + num_tiles_dim++; + tile_overlap_factor_dim = 0.5; + } + + return; + } + // else, non-circular means the last and first tile are not overlapping + + num_tiles_dim = (small_dim - tile_overlap) / non_tile_overlap; + int overshoot_dim = ((num_tiles_dim + 1) * non_tile_overlap + tile_overlap) % small_dim; + + if ((overshoot_dim != non_tile_overlap) && (overshoot_dim <= num_tiles_dim * (tile_size / 2 - tile_overlap))) { + // if tiles don't fit perfectly using the desired overlap + // and there is enough room to squeeze an extra tile without overlap becoming >0.5 + num_tiles_dim++; + } + + tile_overlap_factor_dim = (float)(tile_size * num_tiles_dim - small_dim) / (float)(tile_size * (num_tiles_dim - 1)); + if (num_tiles_dim <= 2) { + if (small_dim <= tile_size) { + num_tiles_dim = 1; + tile_overlap_factor_dim = 0; + } else { + num_tiles_dim = 2; + tile_overlap_factor_dim = (2 * tile_size - small_dim) / (float)tile_size; + } + } +} + +static int64_t sd_tensor_plane_size(const sd::Tensor& tensor) { + GGML_ASSERT(tensor.dim() >= 2); + return tensor.shape()[0] * tensor.shape()[1]; +} + +static sd::Tensor sd_tensor_split_2d(const sd::Tensor& input, int width, int height, int x, int y) { + GGML_ASSERT(input.dim() >= 4); + std::vector output_shape = input.shape(); + output_shape[0] = width; + output_shape[1] = height; + sd::Tensor output(std::move(output_shape)); + int64_t input_width = input.shape()[0]; + int64_t input_height = input.shape()[1]; + int64_t input_plane = sd_tensor_plane_size(input); + int64_t output_plane = sd_tensor_plane_size(output); + int64_t plane_count = input.numel() / input_plane; + for (int iy = 0; iy < height; iy++) { + for (int ix = 0; ix < width; ix++) { + int64_t src_xy = (ix + x) % input_width + input_width * ((iy + y) % input_height); + int64_t dst_xy = ix + width * iy; + for (int64_t plane = 0; plane < plane_count; ++plane) { + output[plane * output_plane + dst_xy] = input[plane * input_plane + src_xy]; + } + } + } + return output; +} + +static void sd_tensor_merge_2d(const sd::Tensor& input, + sd::Tensor* output, + int x, + int y, + int overlap_x, + int overlap_y, + bool circular_x, + bool circular_y, + int x_skip, + int y_skip) { + GGML_ASSERT(output != nullptr); + int64_t width = input.shape()[0]; + int64_t height = input.shape()[1]; + int64_t img_width = output->shape()[0]; + int64_t img_height = output->shape()[1]; + int64_t input_plane = sd_tensor_plane_size(input); + int64_t output_plane = sd_tensor_plane_size(*output); + int64_t plane_count = input.numel() / input_plane; + GGML_ASSERT(output->numel() / output_plane == plane_count); + + // unclamped -> expects x in the range [0-1] + auto smootherstep_f32 = [](const float x) -> float { + GGML_ASSERT(x >= 0.f && x <= 1.f); + return x * x * x * (x * (6.0f * x - 15.0f) + 10.0f); + }; + + for (int iy = y_skip; iy < height; iy++) { + for (int ix = x_skip; ix < width; ix++) { + int64_t src_xy = ix + width * iy; + int64_t ox = (x + ix) % img_width; + int64_t oy = (y + iy) % img_height; + int64_t dst_xy = ox + img_width * oy; + for (int64_t plane = 0; plane < plane_count; ++plane) { + float new_value = input[plane * input_plane + src_xy]; + if (overlap_x > 0 || overlap_y > 0) { + float old_value = (*output)[plane * output_plane + dst_xy]; + const float x_f_0 = (circular_x || (overlap_x > 0 && x > 0)) ? (ix - x_skip) / float(overlap_x) : 1.f; + const float x_f_1 = (circular_x || (overlap_x > 0 && x < (img_width - width))) ? (width - ix) / float(overlap_x) : 1.f; + const float y_f_0 = (circular_y || (overlap_y > 0 && y > 0)) ? (iy - y_skip) / float(overlap_y) : 1.f; + const float y_f_1 = (circular_y || (overlap_y > 0 && y < (img_height - height))) ? (height - iy) / float(overlap_y) : 1.f; + const float x_f = std::min(std::min(x_f_0, x_f_1), 1.f); + const float y_f = std::min(std::min(y_f_0, y_f_1), 1.f); + (*output)[plane * output_plane + dst_xy] = + old_value + new_value * smootherstep_f32(y_f) * smootherstep_f32(x_f); + } else { + (*output)[plane * output_plane + dst_xy] = new_value; + } + } + } + } +} + +sd::Tensor process_tiles_2d(const sd::Tensor& input, + int output_width, + int output_height, + int scale, + int p_tile_size_x, + int p_tile_size_y, + float tile_overlap_factor, + bool circular_x, + bool circular_y, + const TileProcessCallback& on_processing, + bool silent) { + sd::Tensor output; + int input_width = static_cast(input.shape()[0]); + int input_height = static_cast(input.shape()[1]); + + GGML_ASSERT(((input_width / output_width) == (input_height / output_height)) && + ((output_width / input_width) == (output_height / input_height))); + GGML_ASSERT(((input_width / output_width) == scale) || + ((output_width / input_width) == scale)); + + int small_width = output_width; + int small_height = output_height; + bool decode = output_width > input_width; + if (decode) { + small_width = input_width; + small_height = input_height; + } + + int num_tiles_x; + float tile_overlap_factor_x; + sd_tiling_calc_tiles(num_tiles_x, tile_overlap_factor_x, small_width, p_tile_size_x, tile_overlap_factor, circular_x); + + int num_tiles_y; + float tile_overlap_factor_y; + sd_tiling_calc_tiles(num_tiles_y, tile_overlap_factor_y, small_height, p_tile_size_y, tile_overlap_factor, circular_y); + + int tile_overlap_x = static_cast(p_tile_size_x * tile_overlap_factor_x); + int non_tile_overlap_x = p_tile_size_x - tile_overlap_x; + int tile_overlap_y = static_cast(p_tile_size_y * tile_overlap_factor_y); + int non_tile_overlap_y = p_tile_size_y - tile_overlap_y; + int tile_size_x = p_tile_size_x < small_width ? p_tile_size_x : small_width; + int tile_size_y = p_tile_size_y < small_height ? p_tile_size_y : small_height; + int input_tile_size_x = tile_size_x; + int input_tile_size_y = tile_size_y; + int output_tile_size_x = tile_size_x; + int output_tile_size_y = tile_size_y; + if (decode) { + output_tile_size_x *= scale; + output_tile_size_y *= scale; + } else { + input_tile_size_x *= scale; + input_tile_size_y *= scale; + } + + int num_tiles = num_tiles_x * num_tiles_y; + int tile_count = 1; + bool last_y = false; + bool last_x = false; + float last_time = 0.0f; + if (!silent) { + LOG_VERBOSE("num tiles : %d, %d ", num_tiles_x, num_tiles_y); + LOG_VERBOSE("optimal overlap : %f, %f (targeting %f)", tile_overlap_factor_x, tile_overlap_factor_y, tile_overlap_factor); + LOG_VERBOSE("processing %i tiles", num_tiles); + pretty_progress(0, num_tiles, 0.0f); + } + for (int y = 0; y < small_height && !last_y; y += non_tile_overlap_y) { + int dy = 0; + if (!circular_y && y + tile_size_y >= small_height) { + int original_y = y; + y = small_height - tile_size_y; + dy = original_y - y; + if (decode) { + dy *= scale; + } + last_y = true; + } + for (int x = 0; x < small_width && !last_x; x += non_tile_overlap_x) { + int dx = 0; + if (!circular_x && x + tile_size_x >= small_width) { + int original_x = x; + x = small_width - tile_size_x; + dx = original_x - x; + if (decode) { + dx *= scale; + } + last_x = true; + } + + int x_in = decode ? x : scale * x; + int y_in = decode ? y : scale * y; + int x_out = decode ? x * scale : x; + int y_out = decode ? y * scale : y; + + int overlap_x_out = decode ? tile_overlap_x * scale : tile_overlap_x; + int overlap_y_out = decode ? tile_overlap_y * scale : tile_overlap_y; + + int64_t t1 = ggml_time_ms(); + auto input_tile = sd_tensor_split_2d(input, input_tile_size_x, input_tile_size_y, x_in, y_in); + auto output_tile = on_processing(input_tile); + if (output_tile.empty()) { + return {}; + } + GGML_ASSERT(output_tile.shape()[0] == output_tile_size_x && output_tile.shape()[1] == output_tile_size_y); + if (output.empty()) { + std::vector output_shape = output_tile.shape(); + output_shape[0] = output_width; + output_shape[1] = output_height; + output = sd::Tensor::zeros(std::move(output_shape)); + } + sd_tensor_merge_2d(output_tile, &output, x_out, y_out, overlap_x_out, overlap_y_out, circular_x, circular_y, dx, dy); + + if (!silent) { + int64_t t2 = ggml_time_ms(); + last_time = (t2 - t1) / 1000.0f; + pretty_progress(tile_count, num_tiles, last_time); + } + tile_count++; + } + last_x = false; + } + if (!silent && tile_count < num_tiles) { + pretty_progress(num_tiles, num_tiles, last_time); + } + if (output.empty()) { + return {}; + } + return output; +} diff --git a/otherarch/sdcpp/src/runtime/tiling.h b/otherarch/sdcpp/src/runtime/tiling.h new file mode 100644 index 000000000..a832a644e --- /dev/null +++ b/otherarch/sdcpp/src/runtime/tiling.h @@ -0,0 +1,22 @@ +#ifndef __SD_RUNTIME_TILING_H__ +#define __SD_RUNTIME_TILING_H__ + +#include + +#include "core/tensor.hpp" + +using TileProcessCallback = std::function(const sd::Tensor&)>; + +sd::Tensor process_tiles_2d(const sd::Tensor& input, + int output_width, + int output_height, + int scale, + int p_tile_size_x, + int p_tile_size_y, + float tile_overlap_factor, + bool circular_x, + bool circular_y, + const TileProcessCallback& on_processing, + bool silent = false); + +#endif // __SD_RUNTIME_TILING_H__ diff --git a/otherarch/sdcpp/src/runtime/ucache.hpp b/otherarch/sdcpp/src/runtime/ucache.hpp index 187e1e789..7f3eac865 100644 --- a/otherarch/sdcpp/src/runtime/ucache.hpp +++ b/otherarch/sdcpp/src/runtime/ucache.hpp @@ -6,8 +6,8 @@ #include #include -#include "core/ggml_extend.hpp" #include "core/tensor.hpp" +#include "core/util.h" #include "runtime/condition_cache_utils.hpp" #include "runtime/denoiser.hpp" diff --git a/otherarch/sdcpp/src/stable-diffusion.cpp b/otherarch/sdcpp/src/stable-diffusion.cpp index a8a7427c4..ead676b22 100644 --- a/otherarch/sdcpp/src/stable-diffusion.cpp +++ b/otherarch/sdcpp/src/stable-diffusion.cpp @@ -1,4 +1,5 @@ #include +#include #include #include #include @@ -9,9 +10,12 @@ #include #include -#include "core/ggml_extend.hpp" +#include "core/ggml_extend_backend.h" #include "core/ggml_graph_cut.h" +#include "core/ggml_runner.h" +#include "core/ggml_tensor_utils.h" #include "core/layer_split_partition.h" +#include "model.h" #include "core/rng.hpp" #include "core/rng_mt19937.hpp" @@ -252,8 +256,9 @@ public: sd_tiling_params_t vae_tiling_params = {false, false, 0, 0, 0.5f, 0, 0, nullptr}; bool enable_mmap = false; sd::ggml_graph_cut::MaxVramAssignment max_vram_assignment; - bool stream_layers = false; - bool eager_load = false; + bool disable_prefetch = false; + bool disable_segmented_compute = false; + bool eager_load = false; std::string backend_spec; std::string params_backend_spec; std::string split_mode_spec; @@ -438,7 +443,11 @@ public: if (split_buft == nullptr) { return fall_back_to_layer_split("backend has no split buffer type"); } - model_manager->set_split_buffer_type(main_backend, split_buft); + std::vector> split_device_limits; + for (auto backend : module_backends) { + split_device_limits.emplace_back(backend, max_vram_assignment.bytes_for_backend(backend)); + } + model_manager->set_split_buffer_type(main_backend, split_buft, split_device_limits); std::map split_tensors; if constexpr (std::is_base_of_v) { @@ -602,6 +611,8 @@ public: version, "", model_manager); + control_net->set_max_graph_vram_bytes( + max_graph_vram_bytes_for_module(SDBackendModule::CONTROL_NET)); if (diffusion_conv_direct) { LOG_INFO("Using Conv2d direct in the control net"); control_net->set_conv2d_direct_enabled(true); @@ -706,7 +717,7 @@ public: } file_alphas_cumprod = std::move(loaded_alphas); - LOG_DEBUG("loaded alphas_cumprod from model file"); + LOG_VERBOSE("loaded alphas_cumprod from model file"); } bool init_model_loader(ModelLoader& model_loader, @@ -1113,14 +1124,15 @@ public: } bool init(const sd_ctx_params_t* sd_ctx_params) { - n_threads = sd_ctx_params->n_threads; - enable_mmap = sd_ctx_params->enable_mmap; - stream_layers = sd_ctx_params->stream_layers; - eager_load = sd_ctx_params->eager_load; - backend_spec = SAFE_STR(sd_ctx_params->backend); - params_backend_spec = SAFE_STR(sd_ctx_params->params_backend); - split_mode_spec = SAFE_STR(sd_ctx_params->split_mode); - auto_fit_enabled = sd_ctx_params->auto_fit; + n_threads = sd_ctx_params->n_threads; + enable_mmap = sd_ctx_params->enable_mmap; + disable_prefetch = sd_ctx_params->disable_prefetch; + disable_segmented_compute = sd_ctx_params->disable_segmented_compute; + eager_load = sd_ctx_params->eager_load; + backend_spec = SAFE_STR(sd_ctx_params->backend); + params_backend_spec = SAFE_STR(sd_ctx_params->params_backend); + split_mode_spec = SAFE_STR(sd_ctx_params->split_mode); + auto_fit_enabled = sd_ctx_params->auto_fit && backend_spec.empty() && params_backend_spec.empty(); max_vram_assignment.reset(0.f); { std::string error; @@ -1144,11 +1156,13 @@ public: sampler_rng = rng; } - ggml_log_set(ggml_log_callback_default, nullptr); + ggml_log_set(sd_ggml_log_callback, nullptr); model_manager = std::make_shared(); model_manager->set_n_threads(n_threads); model_manager->set_enable_mmap(enable_mmap); + model_manager->set_segmented_compute_disabled(disable_segmented_compute); + model_manager->set_prefetch_disabled(disable_prefetch); ModelLoader& model_loader = model_manager->loader(); if (!init_model_loader(model_loader, sd_ctx_params, use_tae, use_audio_vae, use_control_net)) { @@ -1183,10 +1197,6 @@ public: return false; } } - if (stream_layers && !backend_manager.params_backend_is_cpu(SDBackendModule::DIFFUSION)) { - LOG_WARN("--stream-layers has no effect unless diffusion params backend is cpu; ignoring"); - stream_layers = false; - } if (eager_load && graph_cut_layer_split_active()) { LOG_WARN("--eager-load is not supported with graph-cut layer split; weights will be prepared lazily"); eager_load = false; @@ -1216,7 +1226,7 @@ public: LOG_INFO("Diffusion model weight type stat: %s", wtype_stat_to_str(diffusion_model_wtype_stat).c_str()); LOG_INFO("VAE weight type stat: %s", wtype_stat_to_str(vae_wtype_stat).c_str()); - LOG_DEBUG("ggml tensor size = %d bytes", (int)sizeof(ggml_tensor)); + LOG_VERBOSE("ggml tensor size = %d bytes", (int)sizeof(ggml_tensor)); bool have_int8_tensorwise = false; for (const auto& [_, tensor_storage] : model_loader.get_tensor_storage_map()) { @@ -1236,7 +1246,8 @@ public: } // Avoid full-model LoRA merge buffers on constrained setups. const bool params_offloaded = params_backend_for(SDBackendModule::DIFFUSION) != backend_for(SDBackendModule::DIFFUSION); - const bool streaming_constrained = stream_layers || params_offloaded; + const bool streaming_constrained = params_offloaded || + backend_manager.params_backend_is_disk(SDBackendModule::DIFFUSION); if (have_quantized_weight || streaming_constrained || row_split_active()) { apply_lora_immediately = false; } else { @@ -1609,7 +1620,6 @@ public: } diffusion_model->set_max_graph_vram_bytes(max_graph_vram_bytes_for_module(SDBackendModule::DIFFUSION)); - diffusion_model->set_stream_layers_enabled(stream_layers); if (!register_runner_params("Diffusion model", diffusion_model, SDBackendModule::DIFFUSION, @@ -1619,7 +1629,6 @@ public: if (high_noise_diffusion_model) { high_noise_diffusion_model->set_max_graph_vram_bytes(max_graph_vram_bytes_for_module(SDBackendModule::DIFFUSION)); - high_noise_diffusion_model->set_stream_layers_enabled(stream_layers); if (!register_runner_params("High noise diffusion model", high_noise_diffusion_model, SDBackendModule::DIFFUSION, @@ -1824,6 +1833,8 @@ public: version, "", model_manager); + control_net->set_max_graph_vram_bytes( + max_graph_vram_bytes_for_module(SDBackendModule::CONTROL_NET)); if (sd_ctx_params->diffusion_conv_direct) { LOG_INFO("Using Conv2d direct in the control net"); control_net->set_conv2d_direct_enabled(true); @@ -1900,7 +1911,7 @@ public: } } - LOG_DEBUG("validating model metadata"); + LOG_VERBOSE("validating model metadata"); std::set ignore_tensors; if (use_tae && !tae_preview_only) { @@ -1957,9 +1968,9 @@ public: LOG_ERROR("model params eager load failed"); return false; } - LOG_DEBUG("model metadata validated; weights pre-loaded to params backend"); + LOG_VERBOSE("model metadata validated; weights pre-loaded to params backend"); } else { - LOG_DEBUG("model metadata validated; weights will be prepared lazily"); + LOG_VERBOSE("model metadata validated; weights will be prepared lazily"); } { @@ -2155,15 +2166,15 @@ public: } bool is_using_v_parameterization_for_sd2(bool is_inpaint = false) { - struct RunnerDoneOnExit { + struct RunnerEndOnExit { GGMLRunner* runner = nullptr; - ~RunnerDoneOnExit() { + ~RunnerEndOnExit() { if (runner != nullptr) { - runner->runner_done(); + runner->runner_end(); } } }; - RunnerDoneOnExit diffusion_runner_done{diffusion_model.get()}; + RunnerEndOnExit diffusion_runner_end{diffusion_model.get()}; sd::Tensor x_t = sd::full({8, 8, 4, 1}, 0.5f); sd::Tensor c = sd::full({1024, 2, 1, 1}, 0.5f); @@ -2189,7 +2200,7 @@ public: double result = static_cast((out - x_t).mean()); int64_t t1 = ggml_time_ms(); - LOG_DEBUG("check is_using_v_parameterization_for_sd2, taking %.2fs", (t1 - t0) * 1.0f / 1000); + LOG_VERBOSE("check is_using_v_parameterization_for_sd2, taking %.2fs", (t1 - t0) * 1.0f / 1000); return result < -1; } @@ -2204,7 +2215,7 @@ public: return nullptr; } if (lora_spec.is_high_noise) { - LOG_DEBUG("high noise lora: %s", lora_spec.path.c_str()); + LOG_VERBOSE("high noise lora: %s", lora_spec.path.c_str()); } // kcpp // first check the cache @@ -2397,7 +2408,7 @@ public: if (loras[i].is_high_noise) { lora_id = "|high_noise|" + lora_id; } - LOG_DEBUG("lora %s:%.2f", lora_id.c_str(), loras[i].multiplier); + LOG_VERBOSE("lora %s:%.2f", lora_id.c_str(), loras[i].multiplier); } for (auto& extension : generation_extensions) { @@ -2693,7 +2704,7 @@ public: float shifted_t_float = t * (float(shifted_timestep) / float(TIMESTEPS)); int64_t shifted_t = static_cast(roundf(shifted_t_float)); shifted_t = std::max((int64_t)0, std::min((int64_t)(TIMESTEPS - 1), shifted_t)); - LOG_DEBUG("shifting timestep from %.2f to %" PRId64 " (sigma: %.4f)", t, shifted_t, sigma); + LOG_VERBOSE("shifting timestep from %.2f to %" PRId64 " (sigma: %.4f)", t, shifted_t, sigma); return std::vector{(float)shifted_t}; } if (sd_version_is_anima(version)) { @@ -2813,17 +2824,17 @@ public: const sd_cache_params_t* cache_params, bool preview_final_step, const sd::Tensor& video_positions = {}) { - struct RunnerDoneOnExit { + struct RunnerEndOnExit { GGMLRunner* runner = nullptr; - ~RunnerDoneOnExit() { + ~RunnerEndOnExit() { if (runner != nullptr) { - runner->runner_done(); + runner->runner_end(); } } }; - RunnerDoneOnExit sample_diffusion_runner_done{work_diffusion_model.get()}; + RunnerEndOnExit sample_diffusion_runner_end{work_diffusion_model.get()}; - RunnerDoneOnExit sample_control_runner_done{!control_image.empty() && control_net != nullptr ? control_net.get() : nullptr}; + RunnerEndOnExit sample_control_runner_end{!control_image.empty() && control_net != nullptr ? control_net.get() : nullptr}; std::vector skip_layers(guidance.slg.layers, guidance.slg.layers + guidance.slg.layer_count); float cfg_scale = guidance.txt_cfg; @@ -2854,7 +2865,7 @@ public: } } schedule_str += "]"; - LOG_DEBUG("using guidance schedule: %s", schedule_str.c_str()); + LOG_VERBOSE("using guidance schedule: %s", schedule_str.c_str()); } sd_sample::SampleCacheRuntime cache_runtime = sd_sample::init_sample_cache_runtime(version, @@ -2911,7 +2922,7 @@ public: auto denoise = [&](const sd::Tensor& x, float sigma, int step) -> sd::guidance::GuiderOutput { if (get_cancel_flag() == SD_CANCEL_ALL) { - LOG_DEBUG("cancelling generation"); + LOG_VERBOSE("cancelling generation"); return {}; } @@ -3117,7 +3128,7 @@ public: } const std::vector* uncond_skip_layers = nullptr; if (is_skiplayer_step && slg_uncond) { - LOG_DEBUG("Skipping layers at uncond step %d\n", step); + LOG_VERBOSE("Skipping layers at uncond step %d\n", step); uncond_skip_layers = &skip_layer_guidance.layers(); } uncond_out = run_condition(uncond, @@ -3152,7 +3163,7 @@ public: } if (is_skiplayer_step && slg_scale != 0.0f) { - LOG_DEBUG("Skipping layers at step %d\n", step); + LOG_VERBOSE("Skipping layers at step %d\n", step); if (!step_cache.is_step_skipped()) { guidance_input.predict_skip_layer = [&]() -> sd::Tensor { return run_condition(cond, @@ -3195,10 +3206,6 @@ public: LOG_ERROR("Diffusion model sampling failed"); if (control_net) { control_net->free_control_ctx(); - control_net->free_compute_buffer(); - } - if (work_diffusion_model) { - work_diffusion_model->free_compute_buffer(); } return {}; } @@ -3212,10 +3219,6 @@ public: if (control_net) { control_net->free_control_ctx(); - control_net->free_compute_buffer(); - } - if (work_diffusion_model) { - work_diffusion_model->free_compute_buffer(); } return x0; } @@ -3364,7 +3367,6 @@ public: while (decoded.empty() && auto_fit_enabled && sd::backend_fit::prepare_vae_decode_retry_tiling(vae_tiling_params, prefer_temporal_tiling)) { - first_stage_model->free_compute_buffer(); decoded = first_stage_model->decode(n_threads, latents, vae_tiling_params, decode_video, circular_x, circular_y); } return decoded; @@ -3839,26 +3841,27 @@ void sd_hires_params_init(sd_hires_params_t* hires_params) { } void sd_ctx_params_init(sd_ctx_params_t* sd_ctx_params) { - *sd_ctx_params = {}; - sd_ctx_params->n_threads = sd_get_num_physical_cores(); - sd_ctx_params->wtype = SD_TYPE_COUNT; - sd_ctx_params->rng_type = CUDA_RNG; - sd_ctx_params->sampler_rng_type = RNG_TYPE_COUNT; - sd_ctx_params->prediction = PREDICTION_COUNT; - sd_ctx_params->lora_apply_mode = LORA_APPLY_AUTO; - sd_ctx_params->max_vram = nullptr; - sd_ctx_params->stream_layers = false; - sd_ctx_params->eager_load = false; - sd_ctx_params->enable_mmap = false; - sd_ctx_params->diffusion_flash_attn = false; - sd_ctx_params->vae_format = SD_VAE_FORMAT_AUTO; - sd_ctx_params->backend = nullptr; - sd_ctx_params->params_backend = nullptr; - sd_ctx_params->split_mode = nullptr; - sd_ctx_params->auto_fit = false; - sd_ctx_params->rpc_servers = nullptr; - sd_ctx_params->model_args = nullptr; - sd_ctx_params->pulid_weights_path = nullptr; + *sd_ctx_params = {}; + sd_ctx_params->n_threads = sd_get_num_physical_cores(); + sd_ctx_params->wtype = SD_TYPE_COUNT; + sd_ctx_params->rng_type = CUDA_RNG; + sd_ctx_params->sampler_rng_type = RNG_TYPE_COUNT; + sd_ctx_params->prediction = PREDICTION_COUNT; + sd_ctx_params->lora_apply_mode = LORA_APPLY_AUTO; + sd_ctx_params->max_vram = nullptr; + sd_ctx_params->disable_prefetch = false; + sd_ctx_params->disable_segmented_compute = false; + sd_ctx_params->eager_load = false; + sd_ctx_params->enable_mmap = false; + sd_ctx_params->diffusion_flash_attn = false; + sd_ctx_params->vae_format = SD_VAE_FORMAT_AUTO; + sd_ctx_params->backend = nullptr; + sd_ctx_params->params_backend = nullptr; + sd_ctx_params->split_mode = nullptr; + sd_ctx_params->auto_fit = true; + sd_ctx_params->rpc_servers = nullptr; + sd_ctx_params->model_args = nullptr; + sd_ctx_params->pulid_weights_path = nullptr; } char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) { @@ -3892,7 +3895,8 @@ char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) { "sampler_rng_type: %s\n" "prediction: %s\n" "max_vram: %s\n" - "stream_layers: %s\n" + "disable_prefetch: %s\n" + "disable_segmented_compute: %s\n" "eager_load: %s\n" "backend: %s\n" "params_backend: %s\n" @@ -3926,7 +3930,8 @@ char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) { sd_rng_type_name(sd_ctx_params->sampler_rng_type), sd_prediction_name(sd_ctx_params->prediction), SAFE_STR(sd_ctx_params->max_vram), - BOOL_STR(sd_ctx_params->stream_layers), + BOOL_STR(sd_ctx_params->disable_prefetch), + BOOL_STR(sd_ctx_params->disable_segmented_compute), BOOL_STR(sd_ctx_params->eager_load), SAFE_STR(sd_ctx_params->backend), SAFE_STR(sd_ctx_params->params_backend), @@ -4739,7 +4744,7 @@ struct SamplePlan { break; } } - LOG_DEBUG("switching from high noise model at step %d", high_noise_sample_steps); + LOG_VERBOSE("switching from high noise model at step %d", high_noise_sample_steps); } LOG_INFO("sampling using %s method", sampling_methods_str[sample_method]); @@ -5151,11 +5156,11 @@ struct ImageGenerationEmbeds { SDCondition img_uncond; }; -struct ConditionerRunnerDoneOnExit { +struct ConditionerRunnerEndOnExit { Conditioner* conditioner = nullptr; - ~ConditionerRunnerDoneOnExit() { + ~ConditionerRunnerEndOnExit() { if (conditioner != nullptr) { - conditioner->runner_done(); + conditioner->runner_end(); } } }; @@ -5317,7 +5322,7 @@ static std::optional prepare_image_generation_latents(sd t_enc--; } } else { - LOG_DEBUG("Interpreting denoise strength as relative noise level"); + LOG_VERBOSE("Interpreting denoise strength as relative noise level"); // assume x_noised = K * (x * (1-noise_level) + noise * noise_level) = K * lerp(x, noise, noise_level) // K = 1, noise_level = sigma for flow models // K = 1+sigma, noise_level=sigma/(1+sigma) for diffusion models @@ -5341,7 +5346,7 @@ static std::optional prepare_image_generation_latents(sd sigma_sched.assign(plan->sigmas.begin() + plan->sample_steps - t_enc - 1, plan->sigmas.end()); if (target_sigma > 0 && force_first_sigma && strength_as_noise_level) { - LOG_DEBUG("force_first_sigma to %.4f (from %.4f)", target_sigma, sigma_sched[0]); + LOG_VERBOSE("force_first_sigma to %.4f (from %.4f)", target_sigma, sigma_sched[0]); sigma_sched[0] = target_sigma; } @@ -5443,7 +5448,7 @@ static std::optional prepare_image_generation_latents(sd } sd::Tensor ref_latent; if (ref_image_params.resize_before_vae && !sd_version_is_pid(sd_ctx->sd->version)) { - LOG_DEBUG("auto resize ref images"); + LOG_VERBOSE("auto resize ref images"); double vae_width; double vae_height; if (ref_image_params.resize_vae_to_target) { @@ -5466,12 +5471,12 @@ static std::optional prepare_image_generation_latents(sd ref_images[i].shape()[2], ref_images[i].shape()[3]}); - LOG_DEBUG("resize vae ref image %d from %" PRId64 "x%" PRId64 " to %" PRId64 "x%" PRId64, - static_cast(i), - ref_images[i].shape()[1], - ref_images[i].shape()[0], - resized_ref_img.shape()[1], - resized_ref_img.shape()[0]); + LOG_VERBOSE("resize vae ref image %d from %" PRId64 "x%" PRId64 " to %" PRId64 "x%" PRId64, + static_cast(i), + ref_images[i].shape()[1], + ref_images[i].shape()[0], + resized_ref_img.shape()[1], + resized_ref_img.shape()[0]); ref_latent = sd_ctx->sd->encode_first_stage(resized_ref_img); } else { @@ -5581,7 +5586,7 @@ static std::optional prepare_image_generation_embeds(sd_c SamplePlan* plan, ImageGenerationLatents* latents, const RefImageParams& ref_image_params) { - ConditionerRunnerDoneOnExit conditioner_runner_done{sd_ctx->sd->cond_stage_model.get()}; + ConditionerRunnerEndOnExit conditioner_runner_end{sd_ctx->sd->cond_stage_model.get()}; ConditionerParams condition_params; condition_params.text = request->prompt; @@ -6251,6 +6256,15 @@ static std::optional prepare_video_generation_latents(sd sd::Tensor start_image; sd::Tensor end_image; + struct RunnerEndGuard { + GGMLRunner* runner = nullptr; + ~RunnerEndGuard() { + if (runner) { + runner->runner_end(); + } + } + } runner_guard{sd_ctx->sd->audio_vae_model.get()}; + if (sd_vid_gen_params->init_image.data) { start_image = sd_image_to_tensor(sd_vid_gen_params->init_image, request->width, request->height); } @@ -6917,10 +6931,6 @@ static std::optional prepare_video_generation_latents(sd latents.init_latent = pack_ltxav_audio_and_video_latents(latents.init_latent, latents.audio_latent); } - if (sd_ctx->sd->audio_vae_model != nullptr) { - sd_ctx->sd->audio_vae_model->runner_done(); - } - return latents; } @@ -6928,7 +6938,7 @@ static ImageGenerationEmbeds prepare_video_generation_embeds(sd_ctx_t* sd_ctx, const sd_vid_gen_params_t* sd_vid_gen_params, const GenerationRequest& request, const ImageGenerationLatents& latents) { - ConditionerRunnerDoneOnExit conditioner_runner_done{sd_ctx->sd->cond_stage_model.get()}; + ConditionerRunnerEndOnExit conditioner_runner_end{sd_ctx->sd->cond_stage_model.get()}; ImageGenerationEmbeds embeds; ConditionerParams condition_params; @@ -6993,11 +7003,11 @@ static sd_image_t* decode_video_outputs(sd_ctx_t* sd_ctx, video_latent.shape()[3] > sd_ctx->sd->get_latent_channel()) { video_latent = sd::ops::slice(video_latent, 3, 0, sd_ctx->sd->get_latent_channel()); } - LOG_DEBUG("decode_video_outputs latent %dx%dx%dx%d", - (int)video_latent.shape()[0], - (int)video_latent.shape()[1], - (int)video_latent.shape()[2], - (int)video_latent.shape()[3]); + LOG_VERBOSE("decode_video_outputs latent %dx%dx%dx%d", + (int)video_latent.shape()[0], + (int)video_latent.shape()[1], + (int)video_latent.shape()[2], + (int)video_latent.shape()[3]); // auto z = sd::load_tensor_from_file_as_tensor("ltx_vae_z.bin"); int64_t t4 = ggml_time_ms(); sd::Tensor vid = sd_ctx->sd->decode_first_stage(video_latent, true); @@ -7007,11 +7017,11 @@ static sd_image_t* decode_video_outputs(sd_ctx_t* sd_ctx, LOG_ERROR("decode_first_stage failed for video"); return nullptr; } - LOG_DEBUG("decode_video_outputs decoded %dx%dx%dx%d", - (int)vid.shape()[0], - (int)vid.shape()[1], - (int)vid.shape()[2], - (int)vid.shape()[3]); + LOG_VERBOSE("decode_video_outputs decoded %dx%dx%dx%d", + (int)vid.shape()[0], + (int)vid.shape()[1], + (int)vid.shape()[2], + (int)vid.shape()[3]); if (request.frames > 0 && vid.shape()[2] > request.frames) { vid = sd::ops::slice(vid, 2, 0, request.frames); @@ -7361,7 +7371,7 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx, LOG_ERROR("cancelling generation before high-noise sampling"); return false; } - LOG_DEBUG("sample(high noise) %dx%dx%d", W, H, T); + LOG_VERBOSE("sample(high noise) %dx%dx%d", W, H, T); int64_t sampling_start = ggml_time_ms(); std::vector high_noise_sigmas(plan.sigmas.begin(), plan.sigmas.begin() + plan.high_noise_sample_steps + 1); @@ -7408,7 +7418,7 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx, LOG_ERROR("cancelling generation before sampling"); return false; } - LOG_DEBUG("sample %dx%dx%d", W, H, T); + LOG_VERBOSE("sample %dx%dx%d", W, H, T); int64_t sampling_start = ggml_time_ms(); sd::Tensor final_latent = sd_ctx->sd->sample(sd_ctx->sd->diffusion_model, true, @@ -7551,7 +7561,7 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx, sd_vid_gen_params->sample_params.eta, hires_sample_method); - LOG_DEBUG("sample(latent upscale) %dx%dx%d", W, H, T); + LOG_VERBOSE("sample(latent upscale) %dx%dx%d", W, H, T); LOG_INFO("LTX latent spatial upscale refine: scheduler_steps=%d, denoising_strength=%.2f, sampler=%s, sigma_sched_size=%zu%s", hires_scheduler_steps, request.hires.denoising_strength, @@ -7599,6 +7609,15 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx, int64_t latent_end = ggml_time_ms(); LOG_INFO("generating latent video completed, taking %.2fs", (latent_end - latent_start) * 1.0f / 1000); + struct RunnerEndGuard { + GGMLRunner* runner = nullptr; + ~RunnerEndGuard() { + if (runner) { + runner->runner_end(); + } + } + } runner_guard{sd_ctx->sd->audio_vae_model.get()}; + sd_audio_t* generated_audio = nullptr; if ((sd_version_is_ltxav(sd_ctx->sd->version) || sd_version_is_minimax_h3(sd_ctx->sd->version)) && has_input_audio) { @@ -7620,11 +7639,11 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx, latents.audio_length, sd_ctx->sd->get_latent_channel()); if (!audio_latent.empty()) { - LOG_DEBUG("decode audio latent %dx%dx%dx%d", - (int)audio_latent.shape()[0], - (int)audio_latent.shape()[1], - (int)audio_latent.shape()[2], - (int)audio_latent.shape()[3]); + LOG_VERBOSE("decode audio latent %dx%dx%dx%d", + (int)audio_latent.shape()[0], + (int)audio_latent.shape()[1], + (int)audio_latent.shape()[2], + (int)audio_latent.shape()[3]); auto waveform = sd_ctx->sd->decode_ltx_audio_latent(audio_latent); if (!waveform.empty()) { generated_audio = waveform_to_sd_audio(sd_ctx->sd, waveform); @@ -7636,10 +7655,6 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx, LOG_INFO("decoding audio latent completed, taking %.2fs", (audio_latent_decode_end - audio_latent_decode_start) * 1.0f / 1000); } - if (sd_ctx->sd->audio_vae_model != nullptr) { - sd_ctx->sd->audio_vae_model->runner_done(); - } - if (latents.video_conditioning_frame_count > 0) { int64_t target_frames = latents.video_target_frame_count > 0 ? latents.video_target_frame_count : final_latent.shape()[2] - latents.video_conditioning_frame_count; diff --git a/otherarch/sdcpp/src/tokenizers/bpe_tokenizer.cpp b/otherarch/sdcpp/src/tokenizers/bpe_tokenizer.cpp index 2858e346f..a4001b5cf 100644 --- a/otherarch/sdcpp/src/tokenizers/bpe_tokenizer.cpp +++ b/otherarch/sdcpp/src/tokenizers/bpe_tokenizer.cpp @@ -205,7 +205,7 @@ std::vector BPETokenizer::encode(const std::string& text, on_new_token_cb_t ss << "\"" << token << "\", "; } ss << "]"; - LOG_DEBUG("split prompt \"%s\" to %zu tokens %s", text.c_str(), bpe_tokens.size(), ss.str().c_str()); + LOG_VERBOSE("split prompt \"%s\" to %zu tokens %s", text.c_str(), bpe_tokens.size(), ss.str().c_str()); return bpe_tokens; } diff --git a/otherarch/sdcpp/src/tokenizers/clip_tokenizer.cpp b/otherarch/sdcpp/src/tokenizers/clip_tokenizer.cpp index d51eadec4..ceabb40ee 100644 --- a/otherarch/sdcpp/src/tokenizers/clip_tokenizer.cpp +++ b/otherarch/sdcpp/src/tokenizers/clip_tokenizer.cpp @@ -63,7 +63,7 @@ void CLIPTokenizer::load_from_merges(const std::string& merges_utf8_str) { } vocab.push_back(utf8_to_utf32("<|startoftext|>")); vocab.push_back(utf8_to_utf32("<|endoftext|>")); - LOG_DEBUG("vocab size: %zu", vocab.size()); + LOG_VERBOSE("vocab size: %zu", vocab.size()); int i = 0; for (const auto& token : vocab) { encoder[token] = i; diff --git a/otherarch/sdcpp/src/tokenizers/gemma_tokenizer.cpp b/otherarch/sdcpp/src/tokenizers/gemma_tokenizer.cpp index a7b67ef14..8838bf1e0 100644 --- a/otherarch/sdcpp/src/tokenizers/gemma_tokenizer.cpp +++ b/otherarch/sdcpp/src/tokenizers/gemma_tokenizer.cpp @@ -29,7 +29,7 @@ void GemmaTokenizer::load_from_merges(const std::string& merges_utf8_str, const decoder[i] = token; } encoder_len = static_cast(vocab.size()); - LOG_DEBUG("vocab size: %d", encoder_len); + LOG_VERBOSE("vocab size: %d", encoder_len); std::vector merges = split_utf32(merges_utf8_str); std::vector> merge_pairs; @@ -37,7 +37,7 @@ void GemmaTokenizer::load_from_merges(const std::string& merges_utf8_str, const size_t space_pos = merge.find(' '); merge_pairs.emplace_back(merge.substr(0, space_pos), merge.substr(space_pos + 1)); } - LOG_DEBUG("merges size %zu", merge_pairs.size()); + LOG_VERBOSE("merges size %zu", merge_pairs.size()); int rank = 0; for (const auto& merge : merge_pairs) { @@ -214,7 +214,7 @@ void Gemma2Tokenizer::load_from_merges(const std::string& merges_utf8_str, const decoder[i] = token; } encoder_len = static_cast(vocab.size()); - LOG_DEBUG("vocab size: %d", encoder_len); + LOG_VERBOSE("vocab size: %d", encoder_len); std::vector merges = split_utf32(merges_utf8_str); std::vector> merge_pairs; @@ -222,7 +222,7 @@ void Gemma2Tokenizer::load_from_merges(const std::string& merges_utf8_str, const size_t space_pos = merge.find(' '); merge_pairs.emplace_back(merge.substr(0, space_pos), merge.substr(space_pos + 1)); } - LOG_DEBUG("merges size %zu", merge_pairs.size()); + LOG_VERBOSE("merges size %zu", merge_pairs.size()); int rank = 0; for (const auto& merge : merge_pairs) { diff --git a/otherarch/sdcpp/src/tokenizers/gpt_oss_tokenizer.cpp b/otherarch/sdcpp/src/tokenizers/gpt_oss_tokenizer.cpp index 9779734ce..bb11839f2 100644 --- a/otherarch/sdcpp/src/tokenizers/gpt_oss_tokenizer.cpp +++ b/otherarch/sdcpp/src/tokenizers/gpt_oss_tokenizer.cpp @@ -31,7 +31,7 @@ void GPTOSSTokenizer::load_from_merges(const std::string& merges_utf8_str, const encoder_len++; } encoder_len = static_cast(encoder.size()); - LOG_DEBUG("vocab size: %d", encoder_len); + LOG_VERBOSE("vocab size: %d", encoder_len); std::vector merges = split_utf32(merges_utf8_str); std::vector> merge_pairs; @@ -39,7 +39,7 @@ void GPTOSSTokenizer::load_from_merges(const std::string& merges_utf8_str, const size_t space_pos = merge.find(' '); merge_pairs.emplace_back(merge.substr(0, space_pos), merge.substr(space_pos + 1)); } - LOG_DEBUG("merges size %zu", merge_pairs.size()); + LOG_VERBOSE("merges size %zu", merge_pairs.size()); int rank = 0; for (const auto& merge : merge_pairs) { diff --git a/otherarch/sdcpp/src/tokenizers/mistral_tokenizer.cpp b/otherarch/sdcpp/src/tokenizers/mistral_tokenizer.cpp index cc418710b..13e361cab 100644 --- a/otherarch/sdcpp/src/tokenizers/mistral_tokenizer.cpp +++ b/otherarch/sdcpp/src/tokenizers/mistral_tokenizer.cpp @@ -20,7 +20,7 @@ void MistralTokenizer::load_from_merges(const std::string& merges_utf8_str, cons decoder[i] = token; } encoder_len = static_cast(vocab.size()); - LOG_DEBUG("vocab size: %d", encoder_len); + LOG_VERBOSE("vocab size: %d", encoder_len); auto byte_unicode_pairs = bytes_to_unicode(); byte_encoder = std::map(byte_unicode_pairs.begin(), byte_unicode_pairs.end()); @@ -28,7 +28,7 @@ void MistralTokenizer::load_from_merges(const std::string& merges_utf8_str, cons byte_decoder[pair.second] = pair.first; } std::vector merges = split_utf32(merges_utf8_str); - LOG_DEBUG("merges size %zu", merges.size()); + LOG_VERBOSE("merges size %zu", merges.size()); std::vector> merge_pairs; for (const auto& merge : merges) { size_t space_pos = merge.find(' '); diff --git a/otherarch/sdcpp/src/tokenizers/qwen2_tokenizer.cpp b/otherarch/sdcpp/src/tokenizers/qwen2_tokenizer.cpp index 79e683e7b..a0c2e3bdb 100644 --- a/otherarch/sdcpp/src/tokenizers/qwen2_tokenizer.cpp +++ b/otherarch/sdcpp/src/tokenizers/qwen2_tokenizer.cpp @@ -11,7 +11,7 @@ void Qwen2Tokenizer::load_from_merges(const std::string& merges_utf8_str) { } std::vector merges = split_utf32(merges_utf8_str); - LOG_DEBUG("merges size %zu", merges.size()); + LOG_VERBOSE("merges size %zu", merges.size()); std::vector> merge_pairs; for (const auto& merge : merges) { size_t space_pos = merge.find(' '); @@ -36,7 +36,7 @@ void Qwen2Tokenizer::load_from_merges(const std::string& merges_utf8_str) { i++; } encoder_len = i; - LOG_DEBUG("vocab size: %d", encoder_len); + LOG_VERBOSE("vocab size: %d", encoder_len); int rank = 0; for (const auto& merge : merge_pairs) { diff --git a/otherarch/sdcpp/src/tokenizers/t5_unigram_tokenizer.cpp b/otherarch/sdcpp/src/tokenizers/t5_unigram_tokenizer.cpp index 64e9e0712..7ea6fd1e3 100644 --- a/otherarch/sdcpp/src/tokenizers/t5_unigram_tokenizer.cpp +++ b/otherarch/sdcpp/src/tokenizers/t5_unigram_tokenizer.cpp @@ -333,7 +333,7 @@ std::vector T5UniGramTokenizer::encode(const std::string& input, on_new_tok ss << "\"" << token_str << "\", "; } ss << "]"; - LOG_DEBUG("split prompt \"%s\" to tokens %s", input.c_str(), ss.str().c_str()); + LOG_VERBOSE("split prompt \"%s\" to tokens %s", input.c_str(), ss.str().c_str()); return tokens; } diff --git a/otherarch/sdcpp/src/upscaler.cpp b/otherarch/sdcpp/src/upscaler.cpp index dbb99af36..46d4e4f3a 100644 --- a/otherarch/sdcpp/src/upscaler.cpp +++ b/otherarch/sdcpp/src/upscaler.cpp @@ -1,7 +1,8 @@ #include "upscaler.h" -#include "core/ggml_extend.hpp" +#include "core/ggml_extend_backend.h" #include "core/util.h" #include "model_loader.h" +#include "runtime/tiling.h" #include "stable-diffusion.h" #include @@ -32,16 +33,9 @@ void UpscalerGGML::set_max_graph_vram_bytes(size_t max_vram_bytes) { } } -void UpscalerGGML::set_stream_layers_enabled(bool enabled) { - stream_layers_enabled = enabled; - if (esrgan_upscaler) { - esrgan_upscaler->set_stream_layers_enabled(enabled); - } -} - bool UpscalerGGML::load_from_file(const std::string& esrgan_path, int n_threads) { - ggml_log_set(ggml_log_callback_default, nullptr); + ggml_log_set(sd_ggml_log_callback, nullptr); std::string error; if (!backend_manager.init(backend_spec.c_str(), @@ -94,7 +88,6 @@ bool UpscalerGGML::load_from_file(const std::string& esrgan_path, return false; } esrgan_upscaler->set_max_graph_vram_bytes(max_graph_vram_bytes); - esrgan_upscaler->set_stream_layers_enabled(stream_layers_enabled); if (direct) { esrgan_upscaler->set_conv2d_direct_enabled(true); } @@ -139,7 +132,7 @@ sd::Tensor UpscalerGGML::upscale_tensor(const sd::Tensor& input_te false, on_processing); } - esrgan_upscaler->free_compute_buffer(); + esrgan_upscaler->runner_end(); if (upscaled.empty()) { LOG_ERROR("esrgan compute failed"); return {}; diff --git a/otherarch/sdcpp/src/upscaler.h b/otherarch/sdcpp/src/upscaler.h index 38150f59f..867f64440 100644 --- a/otherarch/sdcpp/src/upscaler.h +++ b/otherarch/sdcpp/src/upscaler.h @@ -20,7 +20,6 @@ struct UpscalerGGML { bool direct = false; int tile_size = 128; size_t max_graph_vram_bytes = 0; - bool stream_layers_enabled = false; std::string backend_spec; std::string params_backend_spec; @@ -34,7 +33,6 @@ struct UpscalerGGML { bool load_from_file(const std::string& esrgan_path, int n_threads); void set_max_graph_vram_bytes(size_t max_vram_bytes); - void set_stream_layers_enabled(bool enabled); sd::Tensor upscale_tensor(const sd::Tensor& input_tensor); sd_image_t upscale(sd_image_t input_image, uint32_t upscale_factor); }; diff --git a/otherarch/sdcpp/src/weight_manager.h b/otherarch/sdcpp/src/weight_manager.h deleted file mode 100644 index 82f6d03e4..000000000 --- a/otherarch/sdcpp/src/weight_manager.h +++ /dev/null @@ -1,19 +0,0 @@ -#ifndef __WEIGHT_MANAGER_H__ -#define __WEIGHT_MANAGER_H__ - -#include - -#include "ggml-backend.h" - -struct ggml_tensor; - -struct RunnerWeightManager { - virtual ~RunnerWeightManager() = default; - virtual bool assign_compute_backend(const std::vector& tensors, - ggml_backend_t compute_backend) = 0; - virtual bool prepare_params(const std::vector& tensors) = 0; - virtual void release_compute_backend_params(const std::vector& tensors) = 0; - virtual void release_params_backend_params(const std::vector& tensors) = 0; -}; - -#endif // __WEIGHT_MANAGER_H__