Safetensors
GGUF
Turkish
llama
Llama-3
instruct
finetune
chatml
gpt4
synthetic data
distillation
function calling
json mode
axolotl
roleplaying
chat
Instructions to use tda45/TdAI with libraries, inference providers, notebooks, and local apps. Follow these links to get started.
- Libraries
- llama-cpp-python
How to use tda45/TdAI with llama-cpp-python:
# !pip install llama-cpp-python from llama_cpp import Llama llm = Llama.from_pretrained( repo_id="tda45/TdAI", filename="llama.cpp/models/ggml-vocab-aquila.gguf", )
output = llm( "Once upon a time,", max_tokens=512, echo=True ) print(output)
- Notebooks
- Google Colab
- Kaggle
- Local Apps Settings
- llama.cpp
How to use tda45/TdAI with llama.cpp:
Install (macOS, Linux)
curl -LsSf https://llama.app/install.sh | sh # Start a local OpenAI-compatible server with a web UI: llama serve -hf tda45/TdAI # Run inference directly in the terminal: llama cli -hf tda45/TdAI
Install from WinGet (Windows)
winget install llama.cpp # Start a local OpenAI-compatible server with a web UI: llama serve -hf tda45/TdAI # Run inference directly in the terminal: llama cli -hf tda45/TdAI
Use pre-built binary
# Download pre-built binary from: # https://github.com/ggerganov/llama.cpp/releases # Start a local OpenAI-compatible server with a web UI: ./llama-server -hf tda45/TdAI # Run inference directly in the terminal: ./llama-cli -hf tda45/TdAI
Build from source code
git clone https://github.com/ggerganov/llama.cpp.git cd llama.cpp cmake -B build cmake --build build -j --target llama-server llama-cli # Start a local OpenAI-compatible server with a web UI: ./build/bin/llama-server -hf tda45/TdAI # Run inference directly in the terminal: ./build/bin/llama-cli -hf tda45/TdAI
Use Docker
docker model run hf.co/tda45/TdAI
- LM Studio
- Jan
- Ollama
How to use tda45/TdAI with Ollama:
ollama run hf.co/tda45/TdAI
- Unsloth Studio
How to use tda45/TdAI with Unsloth Studio:
Install Unsloth Studio (macOS, Linux, WSL)
curl -fsSL https://unsloth.ai/install.sh | sh # Run unsloth studio unsloth studio -H 0.0.0.0 -p 8888 # Then open http://localhost:8888 in your browser # Search for tda45/TdAI to start chatting
Install Unsloth Studio (Windows)
irm https://unsloth.ai/install.ps1 | iex # Run unsloth studio unsloth studio -H 0.0.0.0 -p 8888 # Then open http://localhost:8888 in your browser # Search for tda45/TdAI to start chatting
Using HuggingFace Spaces for Unsloth
# No setup required # Open https://huggingface.co/spaces/unsloth/studio in your browser # Search for tda45/TdAI to start chatting
- Atomic Chat new
- Docker Model Runner
How to use tda45/TdAI with Docker Model Runner:
docker model run hf.co/tda45/TdAI
- Lemonade
How to use tda45/TdAI with Lemonade:
Pull the model
# Download Lemonade from https://lemonade-server.ai/ lemonade pull tda45/TdAI
Run and chat with the model
lemonade run user.TdAI-{{QUANT_TAG}}List all available models
lemonade list
| struct ggml_backend_meta_device; | |
| struct ggml_backend_meta_buffer_type; | |
| struct ggml_backend_meta_buffer; | |
| struct ggml_backend_meta; | |
| const char * ggml_backend_meta_split_axis_name(enum ggml_backend_meta_split_axis split_axis) { | |
| switch (split_axis) { | |
| case GGML_BACKEND_SPLIT_AXIS_0: | |
| return "0"; | |
| case GGML_BACKEND_SPLIT_AXIS_1: | |
| return "1"; | |
| case GGML_BACKEND_SPLIT_AXIS_2: | |
| return "2"; | |
| case GGML_BACKEND_SPLIT_AXIS_3: | |
| return "3"; | |
| case GGML_BACKEND_SPLIT_AXIS_MIRRORED: | |
| return "MIRRORED"; | |
| case GGML_BACKEND_SPLIT_AXIS_PARTIAL: | |
| return "PARTIAL"; | |
| case GGML_BACKEND_SPLIT_AXIS_NONE: | |
| return "NONE"; | |
| case GGML_BACKEND_SPLIT_AXIS_UNKNOWN: | |
| return "UNKNOWN"; | |
| default: | |
| GGML_ABORT("fatal error"); | |
| } | |
| } | |
| // | |
| // meta backend device | |
| // | |
| struct ggml_backend_meta_device_context { | |
| std::vector<ggml_backend_dev_t> simple_devs; | |
| ggml_backend_meta_get_split_state_t get_split_state; | |
| void * get_split_state_ud; | |
| std::string name; | |
| std::string description; | |
| ggml_backend_meta_device_context( | |
| std::vector<ggml_backend_dev_t> simple_devs, ggml_backend_meta_get_split_state_t get_split_state, void * get_split_state_ud) : | |
| simple_devs(std::move(simple_devs)), get_split_state(get_split_state), get_split_state_ud(get_split_state_ud) { | |
| name = std::string("Meta("); | |
| description = std::string("Meta("); | |
| for (size_t i = 0; i < simple_devs.size(); i++) { | |
| if (i > 0) { | |
| name += ","; | |
| description += ","; | |
| } | |
| name += ggml_backend_dev_name (simple_devs[i]); | |
| description += ggml_backend_dev_description(simple_devs[i]); | |
| } | |
| name += ")"; | |
| description += ")"; | |
| } | |
| bool operator<(const ggml_backend_meta_device_context & other) const { | |
| return std::tie(simple_devs, get_split_state, get_split_state_ud) | |
| < std::tie(other.simple_devs, other.get_split_state, other.get_split_state_ud); | |
| } | |
| }; | |
| static bool ggml_backend_dev_is_meta(ggml_backend_dev_t dev); | |
| static const char * ggml_backend_meta_device_get_name(ggml_backend_dev_t dev) { | |
| GGML_ASSERT(ggml_backend_dev_is_meta(dev)); | |
| const ggml_backend_meta_device_context * meta_dev_ctx = (const ggml_backend_meta_device_context *) dev->context; | |
| return meta_dev_ctx->name.c_str(); | |
| } | |
| static const char * ggml_backend_meta_device_get_description(ggml_backend_dev_t dev) { | |
| GGML_ASSERT(ggml_backend_dev_is_meta(dev)); | |
| const ggml_backend_meta_device_context * meta_dev_ctx = (const ggml_backend_meta_device_context *) dev->context; | |
| return meta_dev_ctx->description.c_str(); | |
| } | |
| static void ggml_backend_meta_device_get_memory(ggml_backend_dev_t dev, size_t * free, size_t * total) { | |
| GGML_ASSERT(ggml_backend_dev_is_meta(dev)); | |
| const ggml_backend_meta_device_context * meta_dev_ctx = (const ggml_backend_meta_device_context *) dev->context; | |
| *free = 0; | |
| *total = 0; | |
| for (ggml_backend_dev_t dev : meta_dev_ctx->simple_devs) { | |
| size_t tmp_free, tmp_total; | |
| ggml_backend_dev_memory(dev, &tmp_free, &tmp_total); | |
| *free += tmp_free; | |
| *total += tmp_total; | |
| } | |
| } | |
| static enum ggml_backend_dev_type ggml_backend_meta_device_get_type(ggml_backend_dev_t dev) { | |
| return GGML_BACKEND_DEVICE_TYPE_META; | |
| GGML_UNUSED(dev); | |
| } | |
| static void ggml_backend_meta_device_get_props(ggml_backend_dev_t dev, ggml_backend_dev_props * props) { | |
| GGML_ASSERT(ggml_backend_dev_is_meta(dev)); | |
| const ggml_backend_meta_device_context * meta_dev_ctx = (const ggml_backend_meta_device_context *) dev->context; | |
| // TODO replace placeholders | |
| props->name = ggml_backend_meta_device_get_name(dev); | |
| props->description = ggml_backend_meta_device_get_description(dev); | |
| props->type = ggml_backend_meta_device_get_type(dev); | |
| props->device_id = 0; | |
| ggml_backend_meta_device_get_memory(dev, &props->memory_free, &props->memory_total); | |
| props->caps = { | |
| /* .async = */ true, | |
| /* .host_buffer = */ false, // Not implemented. | |
| /* .buffer_from_host_ptr = */ false, // Not implemented. | |
| /* .events = */ false, // Not implemented. | |
| }; | |
| for (ggml_backend_dev_t simple_dev : meta_dev_ctx->simple_devs) { | |
| ggml_backend_dev_props tmp_props; | |
| ggml_backend_dev_get_props(simple_dev, &tmp_props); | |
| props->caps.async = props->caps.async && tmp_props.caps.async; | |
| props->caps.host_buffer = props->caps.host_buffer && tmp_props.caps.host_buffer; | |
| props->caps.buffer_from_host_ptr = props->caps.buffer_from_host_ptr && tmp_props.caps.buffer_from_host_ptr; | |
| props->caps.events = props->caps.events && tmp_props.caps.events; | |
| } | |
| } | |
| static ggml_backend_t ggml_backend_meta_device_init_backend(ggml_backend_dev_t dev, const char * params); | |
| static ggml_backend_buffer_type_t ggml_backend_meta_device_get_buffer_type(ggml_backend_dev_t dev); | |
| static ggml_backend_buffer_type_t ggml_backend_meta_device_get_host_buffer_type(ggml_backend_dev_t dev); | |
| static bool ggml_backend_meta_device_supports_op(ggml_backend_dev_t dev, const ggml_tensor * op) { | |
| GGML_ASSERT(ggml_backend_dev_is_meta(dev)); | |
| const ggml_backend_meta_device_context * meta_dev_ctx = (const ggml_backend_meta_device_context *) dev->context; | |
| return std::all_of(meta_dev_ctx->simple_devs.begin(), meta_dev_ctx->simple_devs.end(), | |
| [op](ggml_backend_dev_t simple_dev) { return ggml_backend_dev_supports_op(simple_dev, op); }); | |
| } | |
| static bool ggml_backend_meta_device_supports_buft(ggml_backend_dev_t dev, ggml_backend_buffer_type_t buft) { | |
| GGML_ASSERT(ggml_backend_dev_is_meta(dev)); | |
| ggml_backend_dev_t dev_buft = ggml_backend_buft_get_device(buft); | |
| if (!ggml_backend_dev_is_meta(dev_buft)) { | |
| return false; | |
| } | |
| const ggml_backend_meta_device_context * meta_dev_ctx = (const ggml_backend_meta_device_context *) dev->context; | |
| const ggml_backend_meta_device_context * meta_buft_dev_ctx = (const ggml_backend_meta_device_context *) dev_buft->context; | |
| if (meta_dev_ctx->simple_devs.size() != meta_buft_dev_ctx->simple_devs.size()) { | |
| return false; | |
| } | |
| for (size_t i = 0; i < meta_dev_ctx->simple_devs.size(); i++) { | |
| if (meta_dev_ctx->simple_devs[i] != meta_buft_dev_ctx->simple_devs[i]) { | |
| return false; | |
| } | |
| } | |
| return true; | |
| } | |
| static const ggml_backend_device_i ggml_backend_meta_device_iface = { | |
| /* .get_name = */ ggml_backend_meta_device_get_name, | |
| /* .get_description = */ ggml_backend_meta_device_get_description, | |
| /* .get_memory = */ ggml_backend_meta_device_get_memory, | |
| /* .get_type = */ ggml_backend_meta_device_get_type, | |
| /* .get_props = */ ggml_backend_meta_device_get_props, | |
| /* .init_backend = */ ggml_backend_meta_device_init_backend, | |
| /* .get_buffer_type = */ ggml_backend_meta_device_get_buffer_type, | |
| /* .get_host_buffer_type = */ ggml_backend_meta_device_get_host_buffer_type, | |
| /* .buffer_from_host_ptr = */ nullptr, | |
| /* .supports_op = */ ggml_backend_meta_device_supports_op, | |
| /* .supports_buft = */ ggml_backend_meta_device_supports_buft, | |
| /* .offload_op = */ nullptr, | |
| /* .event_new = */ nullptr, | |
| /* .event_free = */ nullptr, | |
| /* .event_synchronize = */ nullptr, | |
| }; | |
| static bool ggml_backend_dev_is_meta(ggml_backend_dev_t dev) { | |
| return dev != nullptr && dev->iface.get_name == ggml_backend_meta_device_iface.get_name; | |
| } | |
| static size_t ggml_backend_meta_dev_n_devs(ggml_backend_dev_t meta_dev) { | |
| GGML_ASSERT(ggml_backend_dev_is_meta(meta_dev)); | |
| const ggml_backend_meta_device_context * meta_dev_ctx = (const ggml_backend_meta_device_context *) meta_dev->context; | |
| return meta_dev_ctx->simple_devs.size(); | |
| } | |
| static ggml_backend_dev_t ggml_backend_meta_dev_simple_dev(ggml_backend_dev_t meta_dev, size_t index) { | |
| GGML_ASSERT(ggml_backend_dev_is_meta(meta_dev)); | |
| const ggml_backend_meta_device_context * meta_dev_ctx = (const ggml_backend_meta_device_context *) meta_dev->context; | |
| GGML_ASSERT(index < meta_dev_ctx->simple_devs.size()); | |
| return meta_dev_ctx->simple_devs[index]; | |
| } | |
| ggml_backend_dev_t ggml_backend_meta_device( | |
| ggml_backend_dev_t * devs, size_t n_devs, ggml_backend_meta_get_split_state_t get_split_state, void * get_split_state_ud) { | |
| GGML_ASSERT(n_devs <= GGML_BACKEND_META_MAX_DEVICES); | |
| // TODO: this is not thread-safe - needs to be fixed | |
| static std::vector<std::unique_ptr<ggml_backend_meta_device_context>> ctxs; | |
| static std::map<ggml_backend_meta_device_context, struct ggml_backend_device> meta_devs; | |
| std::vector<ggml_backend_dev_t> simple_devs; | |
| simple_devs.reserve(n_devs); | |
| for (size_t i = 0; i < n_devs; i++) { | |
| simple_devs.push_back(devs[i]); | |
| } | |
| ggml_backend_meta_device_context ctx(simple_devs, get_split_state, get_split_state_ud); | |
| { | |
| auto it = meta_devs.find(ctx); | |
| if (it != meta_devs.end()) { | |
| return &it->second; | |
| } | |
| } | |
| ctxs.push_back(std::make_unique<ggml_backend_meta_device_context>(ctx)); | |
| struct ggml_backend_device meta_dev = { | |
| /*iface =*/ ggml_backend_meta_device_iface, | |
| /*reg =*/ nullptr, | |
| /*ctx =*/ ctxs.back().get(), | |
| }; | |
| auto result = meta_devs.emplace(*ctxs.back(), meta_dev); | |
| return &result.first->second; | |
| } | |
| // | |
| // meta backend buffer type | |
| // | |
| struct ggml_backend_meta_buffer_type_context { | |
| std::vector<ggml_backend_buffer_type_t> simple_bufts; | |
| std::string name; | |
| ggml_backend_meta_buffer_type_context(std::vector<ggml_backend_buffer_type_t> simple_bufts) : simple_bufts(std::move(simple_bufts)) { | |
| name = "Meta("; | |
| for (size_t i = 0; i < simple_bufts.size(); i++) { | |
| if (i > 0) { | |
| name += ","; | |
| } | |
| name += ggml_backend_buft_name(simple_bufts[i]); | |
| } | |
| name += ")"; | |
| } | |
| bool operator<(const ggml_backend_meta_buffer_type_context & other) const { | |
| return simple_bufts < other.simple_bufts; | |
| } | |
| }; | |
| static size_t ggml_backend_meta_buft_n_bufts(ggml_backend_buffer_type_t meta_buft) { | |
| GGML_ASSERT(ggml_backend_buft_is_meta(meta_buft)); | |
| const ggml_backend_meta_buffer_type_context * meta_buft_ctx = (const ggml_backend_meta_buffer_type_context *) meta_buft->context; | |
| return meta_buft_ctx->simple_bufts.size(); | |
| } | |
| static const char * ggml_backend_meta_buffer_type_get_name(ggml_backend_buffer_type_t buft) { | |
| GGML_ASSERT(ggml_backend_buft_is_meta(buft)); | |
| const ggml_backend_meta_buffer_type_context * meta_buft_ctx = (const ggml_backend_meta_buffer_type_context *) buft->context; | |
| return meta_buft_ctx->name.c_str(); | |
| } | |
| static ggml_backend_buffer_type_t ggml_backend_meta_buft_simple_buft(ggml_backend_buffer_type_t meta_buft, size_t index) { | |
| GGML_ASSERT(ggml_backend_buft_is_meta(meta_buft)); | |
| const ggml_backend_meta_buffer_type_context * meta_buft_ctx = (const ggml_backend_meta_buffer_type_context *) meta_buft->context; | |
| GGML_ASSERT(index < meta_buft_ctx->simple_bufts.size()); | |
| return meta_buft_ctx->simple_bufts[index]; | |
| } | |
| static ggml_backend_buffer_t ggml_backend_meta_buffer_type_alloc_buffer(ggml_backend_buffer_type_t buft, size_t size); | |
| static size_t ggml_backend_meta_buffer_type_get_alignment(ggml_backend_buffer_type_t buft) { | |
| const size_t n_simple_bufts = ggml_backend_meta_buft_n_bufts(buft); | |
| size_t max_alignment = 1; | |
| for (size_t i = 0; i < n_simple_bufts; i++) { | |
| const size_t alignment = ggml_backend_buft_get_alignment(ggml_backend_meta_buft_simple_buft(buft, i)); | |
| max_alignment = std::max(max_alignment, alignment); | |
| GGML_ASSERT(max_alignment % alignment == 0); | |
| } | |
| return max_alignment; | |
| } | |
| static size_t ggml_backend_meta_buffer_type_get_max_size(ggml_backend_buffer_type_t buft) { | |
| const size_t n_simple_bufts = ggml_backend_meta_buft_n_bufts(buft); | |
| size_t max_size = SIZE_MAX; | |
| for (size_t i = 0; i < n_simple_bufts; i++) { | |
| max_size = std::min(max_size, ggml_backend_buft_get_max_size(ggml_backend_meta_buft_simple_buft(buft, i))); | |
| } | |
| return max_size; | |
| } | |
| static size_t ggml_backend_meta_buffer_type_get_alloc_size(ggml_backend_buffer_type_t buft, const ggml_tensor * tensor) { | |
| const size_t n_simple_bufts = ggml_backend_meta_buft_n_bufts(buft); | |
| size_t max_alloc_size = 0; | |
| for (size_t i = 0; i < n_simple_bufts; i++) { | |
| const size_t alloc_size = ggml_backend_buft_get_alloc_size(ggml_backend_meta_buft_simple_buft(buft, i), tensor); | |
| max_alloc_size = std::max(max_alloc_size, alloc_size); | |
| } | |
| return max_alloc_size; | |
| } | |
| static bool ggml_backend_meta_buffer_type_is_host(ggml_backend_buffer_type_t buft) { | |
| const size_t n_simple_bufts = ggml_backend_meta_buft_n_bufts(buft); | |
| for (size_t i = 0; i < n_simple_bufts; i++) { | |
| if (!ggml_backend_buft_is_host(ggml_backend_meta_buft_simple_buft(buft, i))) { | |
| return false; | |
| } | |
| } | |
| return true; | |
| } | |
| static const struct ggml_backend_buffer_type_i ggml_backend_meta_buffer_type_iface = { | |
| /* .get_name = */ ggml_backend_meta_buffer_type_get_name, | |
| /* .alloc_buffer = */ ggml_backend_meta_buffer_type_alloc_buffer, | |
| /* .get_alignment = */ ggml_backend_meta_buffer_type_get_alignment, | |
| /* .get_max_size = */ ggml_backend_meta_buffer_type_get_max_size, | |
| /* .get_alloc_size = */ ggml_backend_meta_buffer_type_get_alloc_size, | |
| /* .is_host = */ ggml_backend_meta_buffer_type_is_host, | |
| }; | |
| bool ggml_backend_buft_is_meta(ggml_backend_buffer_type_t buft) { | |
| return buft != nullptr && buft->iface.get_name == ggml_backend_meta_buffer_type_iface.get_name; | |
| } | |
| static ggml_backend_buffer_type_t ggml_backend_meta_device_get_buffer_type(ggml_backend_dev_t dev) { | |
| static std::map<ggml_backend_dev_t, struct ggml_backend_buffer_type> meta_bufts; | |
| GGML_ASSERT(ggml_backend_dev_is_meta(dev)); | |
| { | |
| auto it = meta_bufts.find(dev); | |
| if (it != meta_bufts.end()) { | |
| return &it->second; | |
| } | |
| } | |
| const size_t n_devs = ggml_backend_meta_dev_n_devs(dev); | |
| std::vector<ggml_backend_buffer_type_t> simple_bufts; | |
| simple_bufts.reserve(n_devs); | |
| for (size_t i = 0; i < n_devs; i++) { | |
| simple_bufts.push_back(ggml_backend_dev_buffer_type(ggml_backend_meta_dev_simple_dev(dev, i))); | |
| } | |
| ggml_backend_meta_buffer_type_context * buft_ctx = new ggml_backend_meta_buffer_type_context(simple_bufts); | |
| struct ggml_backend_buffer_type meta_buft = { | |
| /*iface =*/ ggml_backend_meta_buffer_type_iface, | |
| /*device =*/ dev, | |
| /*ctx =*/ buft_ctx, | |
| }; | |
| auto result = meta_bufts.emplace(dev, meta_buft); | |
| return &result.first->second; | |
| } | |
| static ggml_backend_buffer_type_t ggml_backend_meta_device_get_host_buffer_type(ggml_backend_dev_t dev) { | |
| GGML_ASSERT(ggml_backend_dev_is_meta(dev)); | |
| const ggml_backend_meta_device_context * meta_dev_ctx = (const ggml_backend_meta_device_context *) dev->context; | |
| ggml_backend_buffer_type_t host_buft = nullptr; | |
| for (ggml_backend_dev_t simple_dev : meta_dev_ctx->simple_devs) { | |
| ggml_backend_buffer_type_t simple_host_buft = ggml_backend_dev_host_buffer_type(simple_dev); | |
| if (simple_host_buft == nullptr) { | |
| return nullptr; | |
| } | |
| if (host_buft == nullptr) { | |
| host_buft = simple_host_buft; | |
| } else if (host_buft != simple_host_buft) { | |
| // if different simple devices have different host buffer types, | |
| // we cannot provide a single host buffer type for the meta device | |
| return nullptr; | |
| } | |
| } | |
| return host_buft; | |
| } | |
| // | |
| // meta backend buffer | |
| // | |
| // Container to hold the tensor slices per simple ggml backend buffer. | |
| struct ggml_backend_meta_simple_tensor_container { | |
| std::vector<ggml_context_ptr> ctxs; | |
| std::map<const ggml_tensor *, std::vector<ggml_tensor *>> simple_tensors; | |
| ggml_backend_meta_simple_tensor_container(const ggml_init_params & params, const int n_simple) { | |
| ctxs.reserve(n_simple); | |
| for (int i = 0; i < n_simple; i++) { | |
| ctxs.emplace_back(ggml_init(params)); | |
| } | |
| } | |
| ggml_backend_meta_simple_tensor_container() {} | |
| }; | |
| struct ggml_backend_meta_buffer_context { | |
| // FIXME | |
| // Most tensors can simply be stored statically in their own buffer. | |
| // Externally created views however also need a mapping to simple tensors but they use the buffer of the view source. | |
| // If external views are simply using that buffer they will slowly deplete its memory. | |
| // Current solution: rotating set of 2 "compute" containers to hold external views, works correctly for llama.cpp. | |
| // Long-term: tie the lifetime of external views to the meta backend executing the graph instead, | |
| // currently not possible due to graph-external operations in the backend scheduler. | |
| ggml_backend_meta_simple_tensor_container stc_static; | |
| ggml_backend_meta_simple_tensor_container stc_compute[2]; | |
| int stc_compute_index = 0; | |
| int stc_compute_index_next = 0; | |
| std::vector<ggml_backend_buffer_ptr> bufs; | |
| // FIXME | |
| // The size of the split state cache is unbounded and can theoretically grow infinitely large. | |
| // However, it is also expensive to build and clearing it on every rebuild in ggml_backend_meta_graph_compute is too expensive. | |
| static constexpr size_t nbtc = GGML_TENSOR_SIZE - sizeof(ggml_tensor::padding); | |
| std::map<std::pair<const ggml_tensor *, bool>, std::pair<ggml_backend_meta_split_state, char[nbtc]>> split_state_cache; | |
| int debug; | |
| ggml_backend_meta_buffer_context( | |
| ggml_backend_meta_simple_tensor_container & stc_static, | |
| ggml_backend_meta_simple_tensor_container & stc_compute_0, | |
| ggml_backend_meta_simple_tensor_container & stc_compute_1, | |
| const std::vector<ggml_backend_buffer_t> & bufs) | |
| : stc_static(std::move(stc_static)), stc_compute{std::move(stc_compute_0), std::move(stc_compute_1)} { | |
| this->bufs.reserve(bufs.size()); | |
| for (ggml_backend_buffer_t buf : bufs) { | |
| this->bufs.emplace_back(buf); | |
| } | |
| const char * GGML_META_DEBUG = getenv("GGML_META_DEBUG"); | |
| debug = GGML_META_DEBUG ? atoi(GGML_META_DEBUG) : 0; | |
| } | |
| ggml_backend_meta_simple_tensor_container & get_simple_tensor_container(const ggml_tensor * tensor) { | |
| if (stc_static.simple_tensors.find(tensor) != stc_static.simple_tensors.end()) { | |
| return stc_static; | |
| } | |
| return stc_compute[stc_compute_index]; | |
| } | |
| }; | |
| static void ggml_backend_meta_buffer_free_buffer(ggml_backend_buffer_t buffer) { | |
| GGML_ASSERT(ggml_backend_buffer_is_meta(buffer)); | |
| ggml_backend_meta_buffer_context * buf_ctx = (ggml_backend_meta_buffer_context *) buffer->context; | |
| delete buf_ctx; | |
| } | |
| static size_t ggml_backend_meta_buffer_n_bufs(ggml_backend_buffer_t meta_buf) { | |
| GGML_ASSERT(ggml_backend_buffer_is_meta(meta_buf)); | |
| ggml_backend_meta_buffer_context * buf_ctx = (ggml_backend_meta_buffer_context *) meta_buf->context; | |
| return buf_ctx->bufs.size(); | |
| } | |
| static ggml_backend_buffer_t ggml_backend_meta_buffer_simple_buffer(ggml_backend_buffer_t meta_buf, size_t index) { | |
| GGML_ASSERT(ggml_backend_buffer_is_meta(meta_buf)); | |
| ggml_backend_meta_buffer_context * buf_ctx = (ggml_backend_meta_buffer_context *) meta_buf->context; | |
| GGML_ASSERT(index < buf_ctx->bufs.size()); | |
| return buf_ctx->bufs[index].get(); | |
| } | |
| static struct ggml_tensor * ggml_backend_meta_buffer_simple_tensor(const struct ggml_tensor * tensor, size_t index) { | |
| GGML_ASSERT(ggml_backend_buffer_is_meta(tensor->buffer)); | |
| ggml_backend_meta_buffer_context * buf_ctx = (ggml_backend_meta_buffer_context *) tensor->buffer->context; | |
| GGML_ASSERT(index < buf_ctx->bufs.size()); | |
| ggml_backend_meta_simple_tensor_container & stc = buf_ctx->get_simple_tensor_container(tensor); | |
| auto it = stc.simple_tensors.find(tensor); | |
| if (it == stc.simple_tensors.end()) { | |
| return nullptr; | |
| } | |
| return it->second[index]; | |
| } | |
| static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(const struct ggml_tensor * tensor, bool assume_sync); | |
| static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( | |
| ggml_backend_meta_simple_tensor_container & stc, const struct ggml_tensor * tensor, bool assume_sync) { | |
| // FIXME Currently this function preserves/erases the information in n_segments and nr in an inconsistent way. | |
| // Since the operations in question are developed specifically for llama.cpp this currently does not manifest as a bug there. | |
| // However, in a broader ggml context with arbitrary ggml graphs this can lead to unexpected results. | |
| const size_t n_bufs = ggml_backend_meta_buffer_n_bufs(tensor->buffer); | |
| ggml_backend_meta_buffer_context * buf_ctx = (ggml_backend_meta_buffer_context *) tensor->buffer->context; | |
| auto split_states_equal = [&](const ggml_backend_meta_split_state & a, const ggml_backend_meta_split_state & b) -> bool { | |
| if (a.axis != b.axis) { | |
| return false; | |
| } | |
| for (size_t j = 0; j < n_bufs; j++) { | |
| int64_t sum_a = 0; | |
| for (size_t s = 0; s < a.n_segments; s++) { | |
| sum_a += a.ne[s*n_bufs + j] * a.nr[s]; | |
| } | |
| int64_t sum_b = 0; | |
| for (size_t s = 0; s < b.n_segments; s++) { | |
| sum_b += b.ne[s*n_bufs + j] * b.nr[s]; | |
| } | |
| if (sum_a != sum_b) { | |
| return false; | |
| } | |
| } | |
| return true; | |
| }; | |
| auto handle_generic = [&](const std::vector<ggml_backend_meta_split_state> & src_ss, bool scalar_only) -> ggml_backend_meta_split_state { | |
| ggml_backend_meta_split_state ret = {GGML_BACKEND_SPLIT_AXIS_NONE, {0}, {1}, 1}; | |
| for (size_t i = 0; i < GGML_MAX_SRC; i++) { | |
| if (tensor->src[i] == nullptr || tensor->src[i] == tensor) { | |
| continue; | |
| } | |
| if (ret.axis == GGML_BACKEND_SPLIT_AXIS_NONE) { | |
| ret = src_ss[i]; | |
| } else if (!split_states_equal(src_ss[i], ret)) { | |
| ret = {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, {1}, 1}; | |
| break; | |
| } | |
| } | |
| if (ret.axis == GGML_BACKEND_SPLIT_AXIS_NONE) { | |
| ret = {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, {1}, 1}; | |
| } | |
| if (scalar_only && ret.axis >= 0 && ret.axis < GGML_MAX_DIMS) { | |
| ret = {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, {1}, 1}; | |
| } | |
| GGML_ASSERT(ret.axis != GGML_BACKEND_SPLIT_AXIS_UNKNOWN); | |
| return ret; | |
| }; | |
| // Some ops process data on a per-row bases: | |
| auto handle_per_row = [&](const std::vector<ggml_backend_meta_split_state> & src_ss) -> ggml_backend_meta_split_state { | |
| GGML_ASSERT(src_ss[0].axis != GGML_BACKEND_SPLIT_AXIS_0); | |
| return src_ss[0]; | |
| }; | |
| // Some ops broadcast the src1 data across src0: | |
| auto handle_bin_bcast = [&](const std::vector<ggml_backend_meta_split_state> & src_ss) -> ggml_backend_meta_split_state { | |
| if (src_ss[0].axis >= 0 && src_ss[0].axis < GGML_MAX_DIMS && | |
| tensor->src[1]->ne[src_ss[0].axis] == 1 && src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED) { | |
| return src_ss[0]; | |
| } | |
| if (src_ss[2].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED && (src_ss[0].axis == src_ss[1].axis || | |
| (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED && (src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_PARTIAL)))) { | |
| return src_ss[0]; // GGML_OP_ADD_ID | |
| } | |
| GGML_ASSERT(tensor->src[2] == nullptr || src_ss[2].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED); | |
| return handle_generic(src_ss, /*scalar_only =*/ false); | |
| }; | |
| auto handle_concat = [&](const std::vector<ggml_backend_meta_split_state> & src_ss) -> ggml_backend_meta_split_state { | |
| const ggml_backend_meta_split_axis concat_axis = ggml_backend_meta_split_axis(ggml_get_op_params_i32(tensor, 0)); | |
| if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED && src_ss[1].axis >= 0 && src_ss[1].axis < GGML_MAX_DIMS) { | |
| GGML_ASSERT(concat_axis != src_ss[1].axis); | |
| return src_ss[1]; | |
| } | |
| if (src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED && src_ss[0].axis >= 0 && src_ss[0].axis < GGML_MAX_DIMS) { | |
| GGML_ASSERT(concat_axis != src_ss[0].axis); | |
| return src_ss[0]; | |
| } | |
| if (src_ss[0].axis == src_ss[1].axis && src_ss[0].axis != concat_axis) { | |
| return src_ss[0]; | |
| } | |
| return handle_generic(src_ss, /*scalar_only =*/ true); | |
| }; | |
| auto handle_mul_mat = [&](const std::vector<ggml_backend_meta_split_state> & src_ss) -> ggml_backend_meta_split_state { | |
| if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED && src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED) { | |
| return {GGML_BACKEND_SPLIT_AXIS_MIRRORED, {0}, {1}, 1}; | |
| } | |
| if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_1 && src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED) { | |
| ggml_backend_meta_split_state ret = src_ss[0]; | |
| ret.axis = GGML_BACKEND_SPLIT_AXIS_0; | |
| ret.nr[0] = 1; | |
| ret.n_segments = 1; | |
| return ret; | |
| } | |
| if (src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_1 && src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED) { | |
| return src_ss[1]; | |
| } | |
| if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_0 && src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_0) { | |
| GGML_ASSERT(split_states_equal(src_ss[0], src_ss[1])); | |
| return {assume_sync ? GGML_BACKEND_SPLIT_AXIS_MIRRORED : GGML_BACKEND_SPLIT_AXIS_PARTIAL, {0}, {1}, 1}; | |
| } | |
| GGML_ABORT("fatal error"); | |
| //return {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, {1}, 1}; | |
| }; | |
| auto handle_reshape = [&](const std::vector<ggml_backend_meta_split_state> & src_ss) -> ggml_backend_meta_split_state { | |
| switch (src_ss[0].axis) { | |
| case GGML_BACKEND_SPLIT_AXIS_0: | |
| case GGML_BACKEND_SPLIT_AXIS_1: | |
| case GGML_BACKEND_SPLIT_AXIS_2: | |
| case GGML_BACKEND_SPLIT_AXIS_3: { | |
| GGML_ASSERT(src_ss[0].n_segments == 1); | |
| if (src_ss[0].axis == ggml_n_dims(tensor->src[0]) - 1 && src_ss[0].nr[0] == 1) { | |
| return {ggml_backend_meta_split_axis(ggml_n_dims(tensor) - 1), {0}, {1}, 1}; | |
| } | |
| int64_t base_ne_in = tensor->src[0]->ne[0]; | |
| for (int dim = 1; dim <= src_ss[0].axis; dim++) { | |
| base_ne_in *= tensor->src[0]->ne[dim]; | |
| } | |
| base_ne_in /= src_ss[0].nr[0]; | |
| int64_t base_ne_out = 1; | |
| for (int dim = 0; dim < GGML_MAX_DIMS; dim++) { | |
| const int64_t base_ne_out_next = base_ne_out *= tensor->ne[dim]; | |
| if (base_ne_out_next % base_ne_in == 0) { | |
| return {ggml_backend_meta_split_axis(dim), {0}, {uint32_t(base_ne_out_next/base_ne_in)}, 1}; | |
| } | |
| if (base_ne_out_next > base_ne_in) { | |
| GGML_ASSERT(src_ss[0].n_segments == 1); | |
| GGML_ASSERT(src_ss[0].nr[0] == 1); | |
| return {ggml_backend_meta_split_axis(dim), {0}, {1}, 1}; | |
| } | |
| base_ne_out = base_ne_out_next; | |
| } | |
| GGML_ABORT("shape mismatch for %s", ggml_op_name(tensor->op)); | |
| } | |
| case GGML_BACKEND_SPLIT_AXIS_MIRRORED: | |
| case GGML_BACKEND_SPLIT_AXIS_PARTIAL: { | |
| return src_ss[0]; | |
| } | |
| default: { | |
| GGML_ABORT("fatal error"); | |
| //return {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, {1}, 1}; | |
| } | |
| } | |
| }; | |
| auto handle_cpy = [&](const std::vector<ggml_backend_meta_split_state> & src_ss) -> ggml_backend_meta_split_state { | |
| if (src_ss[0].axis >= 0 && src_ss[0].axis < GGML_MAX_DIMS) { | |
| return handle_reshape(src_ss); | |
| } | |
| return handle_generic(src_ss, /*scalar_only =*/ false); | |
| }; | |
| auto handle_view = [&](const std::vector<ggml_backend_meta_split_state> & src_ss) -> ggml_backend_meta_split_state { | |
| if (ggml_is_contiguous(tensor) && ggml_is_contiguous(tensor->src[0])) { | |
| return handle_reshape(src_ss); | |
| } | |
| const int axis = src_ss[0].axis; | |
| { | |
| bool all_strides_the_same = true; | |
| for (int dim = 0; dim < GGML_MAX_DIMS; dim++) { | |
| if (tensor->ne[dim] == 1 && tensor->src[0]->ne[dim] == 1) { | |
| continue; | |
| } | |
| if (tensor->nb[dim] != tensor->src[0]->nb[dim]) { | |
| all_strides_the_same = false; | |
| break; | |
| } | |
| } | |
| if (all_strides_the_same) { | |
| return src_ss[0]; | |
| } | |
| } | |
| if (!ggml_is_permuted(tensor) && !ggml_is_permuted(tensor->src[0]) && axis >= 0 && axis < GGML_MAX_DIMS-1) { | |
| for (int dim = 0; dim < GGML_MAX_DIMS-1; dim++) { | |
| if (tensor->nb[dim+1] == tensor->src[0]->nb[axis+1]) { | |
| return {ggml_backend_meta_split_axis(dim), {0}, {1}, 1}; | |
| } | |
| } | |
| GGML_ABORT("fatal error"); | |
| } | |
| if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED || src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_PARTIAL) { | |
| return src_ss[0]; | |
| } | |
| GGML_ABORT("view of permuted tensor not implemented"); | |
| //return {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, {1}, 1}; | |
| }; | |
| auto handle_permute = [&](const std::vector<ggml_backend_meta_split_state> & src_ss) -> ggml_backend_meta_split_state { | |
| switch (src_ss[0].axis) { | |
| case GGML_BACKEND_SPLIT_AXIS_0: | |
| case GGML_BACKEND_SPLIT_AXIS_1: | |
| case GGML_BACKEND_SPLIT_AXIS_2: | |
| case GGML_BACKEND_SPLIT_AXIS_3: { | |
| GGML_ASSERT(src_ss[0].n_segments == 1 || src_ss[0].nr[0] == 1); | |
| return {ggml_backend_meta_split_axis(tensor->op_params[src_ss[0].axis]), {0}, {src_ss[0].nr[0]}, 1}; | |
| } | |
| case GGML_BACKEND_SPLIT_AXIS_MIRRORED: | |
| case GGML_BACKEND_SPLIT_AXIS_PARTIAL: { | |
| return src_ss[0]; | |
| } | |
| default: { | |
| GGML_ABORT("fatal error"); | |
| //return {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, {1}, 1}; | |
| } | |
| } | |
| }; | |
| auto handle_transpose = [&](const std::vector<ggml_backend_meta_split_state> & src_ss) -> ggml_backend_meta_split_state { | |
| switch (src_ss[0].axis) { | |
| case GGML_BACKEND_SPLIT_AXIS_0: | |
| case GGML_BACKEND_SPLIT_AXIS_1: { | |
| GGML_ASSERT(src_ss[0].n_segments == 1 || src_ss[0].nr[0] == 1); | |
| return {ggml_backend_meta_split_axis(int(src_ss[0].axis) ^ 1), {0}, {src_ss[0].nr[0]}, 1}; | |
| } | |
| case GGML_BACKEND_SPLIT_AXIS_2: | |
| case GGML_BACKEND_SPLIT_AXIS_3: | |
| case GGML_BACKEND_SPLIT_AXIS_MIRRORED: | |
| case GGML_BACKEND_SPLIT_AXIS_PARTIAL: { | |
| return src_ss[0]; | |
| } | |
| default: { | |
| GGML_ABORT("fatal error"); | |
| //return {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, {1}, 1}; | |
| } | |
| } | |
| }; | |
| auto handle_get_rows = [&](const std::vector<ggml_backend_meta_split_state> & src_ss) -> ggml_backend_meta_split_state { | |
| if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_0 && src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED) { | |
| return src_ss[0]; | |
| } | |
| return handle_generic(src_ss, /*scalar_only =*/ true); | |
| }; | |
| auto handle_set_rows = [&](const std::vector<ggml_backend_meta_split_state> & src_ss) -> ggml_backend_meta_split_state { | |
| GGML_ASSERT(src_ss[0].axis != GGML_BACKEND_SPLIT_AXIS_1); | |
| GGML_ASSERT(src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED); | |
| GGML_ASSERT(split_states_equal(src_ss[0], src_ss[2])); | |
| return src_ss[0]; | |
| }; | |
| auto handle_rope = [&](const std::vector<ggml_backend_meta_split_state> & src_ss) -> ggml_backend_meta_split_state { | |
| GGML_ASSERT(src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED); | |
| return src_ss[0]; | |
| }; | |
| auto handle_pad = [&](const std::vector<ggml_backend_meta_split_state> & src_ss) -> ggml_backend_meta_split_state { | |
| if (src_ss[0].axis >= 0 && src_ss[0].axis < GGML_MAX_DIMS) { | |
| GGML_ASSERT(tensor->op_params[2*src_ss[0].axis + 0] == 0); | |
| GGML_ASSERT(tensor->op_params[2*src_ss[0].axis + 1] == 0); | |
| } | |
| return src_ss[0]; | |
| }; | |
| auto handle_flash_attn_ext = [&](const std::vector<ggml_backend_meta_split_state> & src_ss) -> ggml_backend_meta_split_state { | |
| GGML_ASSERT( src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_2); | |
| GGML_ASSERT( src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_2); | |
| GGML_ASSERT( src_ss[2].axis == GGML_BACKEND_SPLIT_AXIS_2); | |
| GGML_ASSERT(tensor->src[4] == nullptr || src_ss[3].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED); | |
| GGML_ASSERT(tensor->src[4] == nullptr || src_ss[4].axis == GGML_BACKEND_SPLIT_AXIS_0); | |
| return {GGML_BACKEND_SPLIT_AXIS_1, {0}, {1}, 1}; | |
| }; | |
| auto handle_ssm_conv = [&](const std::vector<ggml_backend_meta_split_state> & src_ss) -> ggml_backend_meta_split_state { | |
| if (src_ss[0].axis == src_ss[1].axis) { | |
| if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_0) { | |
| return {GGML_BACKEND_SPLIT_AXIS_1, {0}, {1}, 1}; | |
| } | |
| if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_1) { | |
| return {GGML_BACKEND_SPLIT_AXIS_0, {0}, {1}, 1}; | |
| } | |
| } | |
| return handle_generic(src_ss, /*scalar_only =*/ false); | |
| }; | |
| auto handle_gated_delta_net = [&](const std::vector<ggml_backend_meta_split_state> & src_ss) -> ggml_backend_meta_split_state { | |
| if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED && src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED && | |
| src_ss[2].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED && src_ss[3].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED && | |
| src_ss[4].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED && src_ss[5].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED) { | |
| return src_ss[0]; | |
| } | |
| GGML_ASSERT(src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_1); | |
| GGML_ASSERT(src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_1); | |
| GGML_ASSERT(src_ss[2].axis == GGML_BACKEND_SPLIT_AXIS_1); | |
| GGML_ASSERT(src_ss[3].axis == GGML_BACKEND_SPLIT_AXIS_1); | |
| GGML_ASSERT(src_ss[4].axis == GGML_BACKEND_SPLIT_AXIS_1); | |
| // state shape is [S_v, S_v, H_v, n_seqs] (s0 only); the heads dim is its own axis 2, | |
| // so a head-aligned split on the input cache lands on axis 2 here. | |
| GGML_ASSERT(src_ss[5].axis == GGML_BACKEND_SPLIT_AXIS_2 || src_ss[5].axis == GGML_BACKEND_SPLIT_AXIS_1 || src_ss[5].axis == GGML_BACKEND_SPLIT_AXIS_0); | |
| return {GGML_BACKEND_SPLIT_AXIS_0, {0}, {1}, 1}; | |
| }; | |
| auto calculate_split_state = [&]() -> ggml_backend_meta_split_state { | |
| if (ggml_nelements(tensor) == 0) { | |
| return {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, {1}, 1}; | |
| } | |
| if (ggml_backend_buffer_get_usage(tensor->buffer) != GGML_BACKEND_BUFFER_USAGE_COMPUTE && tensor->view_src == nullptr) { | |
| ggml_backend_dev_t dev = ggml_backend_buft_get_device(ggml_backend_buffer_get_type(tensor->buffer)); | |
| const ggml_backend_meta_device_context * dev_ctx = (const ggml_backend_meta_device_context *) dev->context; | |
| ggml_backend_meta_split_state ret = dev_ctx->get_split_state(tensor, dev_ctx->get_split_state_ud); | |
| if (ret.axis >= 0 && ret.axis <= GGML_MAX_DIMS) { | |
| const int64_t granularity = ret.axis == GGML_BACKEND_SPLIT_AXIS_0 ? ggml_blck_size(tensor->type) : 1; | |
| int64_t ne_sum = 0; | |
| for (size_t s = 0; s < ret.n_segments; s++) { | |
| for (size_t j = 0; j < n_bufs; j++) { | |
| GGML_ASSERT(ret.ne[s*n_bufs + j] % granularity == 0); | |
| ne_sum += ret.ne[s*n_bufs + j] * ret.nr[s]; | |
| } | |
| } | |
| GGML_ASSERT(ne_sum == tensor->ne[ret.axis]); | |
| } | |
| return ret; | |
| } | |
| std::vector<ggml_backend_meta_split_state> src_ss(GGML_MAX_SRC, {GGML_BACKEND_SPLIT_AXIS_NONE, {0}, {1}, 1}); | |
| for (size_t i = 0; i < GGML_MAX_SRC; i++) { | |
| if (tensor->src[i] == nullptr || tensor->src[i] == tensor) { | |
| src_ss[i] = {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, {1}, 1}; | |
| continue; | |
| } | |
| src_ss[i] = ggml_backend_meta_get_split_state(stc, tensor->src[i], /*assume_sync =*/ true); | |
| GGML_ASSERT(src_ss[i].axis != GGML_BACKEND_SPLIT_AXIS_UNKNOWN); | |
| } | |
| ggml_backend_meta_split_state split_state; | |
| switch (tensor->op) { | |
| case GGML_OP_NONE: { | |
| split_state = {GGML_BACKEND_SPLIT_AXIS_MIRRORED, {0}, {1}, 1}; | |
| } break; | |
| case GGML_OP_DUP: { | |
| split_state = handle_generic(src_ss, /*scalar_only =*/ true); | |
| } break; | |
| case GGML_OP_ADD: | |
| case GGML_OP_ADD_ID: { | |
| split_state = handle_bin_bcast(src_ss); | |
| } break; | |
| case GGML_OP_ADD1: | |
| case GGML_OP_ACC: { | |
| split_state = handle_generic(src_ss, /*scalar_only =*/ true); | |
| } break; | |
| case GGML_OP_SUB: | |
| case GGML_OP_MUL: | |
| case GGML_OP_DIV: { | |
| split_state = handle_bin_bcast(src_ss); | |
| } break; | |
| case GGML_OP_SQR: | |
| case GGML_OP_SQRT: | |
| case GGML_OP_LOG: | |
| case GGML_OP_SIN: | |
| case GGML_OP_COS: { | |
| split_state = handle_generic(src_ss, /*scalar_only =*/ false); | |
| } break; | |
| case GGML_OP_SUM: { | |
| split_state = handle_generic(src_ss, /*scalar_only =*/ true); | |
| } break; | |
| case GGML_OP_SUM_ROWS: | |
| case GGML_OP_CUMSUM: | |
| case GGML_OP_MEAN: | |
| case GGML_OP_ARGMAX: | |
| case GGML_OP_COUNT_EQUAL: { | |
| split_state = handle_per_row(src_ss); | |
| } break; | |
| case GGML_OP_REPEAT: | |
| case GGML_OP_REPEAT_BACK: { | |
| split_state = handle_generic(src_ss, /*scalar_only =*/ false); | |
| } break; | |
| case GGML_OP_CONCAT: { | |
| split_state = handle_concat(src_ss); | |
| } break; | |
| case GGML_OP_SILU_BACK: { | |
| split_state = handle_generic(src_ss, /*scalar_only =*/ false); | |
| } break; | |
| case GGML_OP_NORM: | |
| case GGML_OP_RMS_NORM: | |
| case GGML_OP_RMS_NORM_BACK: | |
| case GGML_OP_GROUP_NORM: | |
| case GGML_OP_L2_NORM: { | |
| split_state = handle_per_row(src_ss); | |
| } break; | |
| case GGML_OP_MUL_MAT: | |
| case GGML_OP_MUL_MAT_ID: { | |
| split_state = handle_mul_mat(src_ss); | |
| } break; | |
| case GGML_OP_OUT_PROD: { | |
| split_state = handle_generic(src_ss, /*scalar_only =*/ true); | |
| } break; | |
| case GGML_OP_SCALE: { | |
| split_state = handle_generic(src_ss, /*scalar_only =*/ false); | |
| } break; | |
| case GGML_OP_SET: { | |
| split_state = handle_generic(src_ss, /*scalar_only =*/ true); | |
| } break; | |
| case GGML_OP_CPY: { | |
| split_state = handle_cpy(src_ss); | |
| } break; | |
| case GGML_OP_CONT: | |
| case GGML_OP_RESHAPE: { | |
| split_state = handle_reshape(src_ss); | |
| } break; | |
| case GGML_OP_VIEW: { | |
| split_state = handle_view(src_ss); | |
| } break; | |
| case GGML_OP_PERMUTE: { | |
| split_state = handle_permute(src_ss); | |
| } break; | |
| case GGML_OP_TRANSPOSE: { | |
| split_state = handle_transpose(src_ss); | |
| } break; | |
| case GGML_OP_GET_ROWS: { | |
| split_state = handle_get_rows(src_ss); | |
| } break; | |
| case GGML_OP_GET_ROWS_BACK: { | |
| split_state = handle_generic(src_ss, /*scalar_only =*/ true); | |
| } break; | |
| case GGML_OP_SET_ROWS: { | |
| split_state = handle_set_rows(src_ss); | |
| } break; | |
| case GGML_OP_DIAG: | |
| case GGML_OP_DIAG_MASK_INF: | |
| case GGML_OP_DIAG_MASK_ZERO: { | |
| split_state = handle_generic(src_ss, /*scalar_only =*/ true); | |
| } break; | |
| case GGML_OP_SOFT_MAX: | |
| case GGML_OP_SOFT_MAX_BACK: { | |
| split_state = handle_generic(src_ss, /*scalar_only =*/ false); | |
| } break; | |
| case GGML_OP_ROPE: { | |
| split_state = handle_rope(src_ss); | |
| } break; | |
| case GGML_OP_ROPE_BACK: { | |
| split_state = handle_generic(src_ss, /*scalar_only =*/ true); | |
| } break; | |
| case GGML_OP_CLAMP: { | |
| split_state = handle_generic(src_ss, /*scalar_only =*/ false); | |
| } break; | |
| case GGML_OP_CONV_TRANSPOSE_1D: | |
| case GGML_OP_IM2COL: | |
| case GGML_OP_IM2COL_BACK: | |
| case GGML_OP_IM2COL_3D: | |
| case GGML_OP_CONV_2D: | |
| case GGML_OP_CONV_3D: | |
| case GGML_OP_CONV_2D_DW: | |
| case GGML_OP_CONV_TRANSPOSE_2D: | |
| case GGML_OP_POOL_1D: | |
| case GGML_OP_POOL_2D: | |
| case GGML_OP_POOL_2D_BACK: | |
| case GGML_OP_UPSCALE: { | |
| split_state = handle_generic(src_ss, /*scalar_only =*/ true); | |
| } break; | |
| case GGML_OP_PAD: { | |
| split_state = handle_pad(src_ss); | |
| } break; | |
| case GGML_OP_PAD_REFLECT_1D: | |
| case GGML_OP_ROLL: | |
| case GGML_OP_ARANGE: | |
| case GGML_OP_TIMESTEP_EMBEDDING: { | |
| split_state = handle_generic(src_ss, /*scalar_only =*/ true); | |
| } break; | |
| case GGML_OP_ARGSORT: | |
| case GGML_OP_TOP_K: { | |
| split_state = handle_per_row(src_ss); | |
| } break; | |
| case GGML_OP_LEAKY_RELU: { | |
| split_state = handle_generic(src_ss, /*scalar_only =*/ false); | |
| } break; | |
| case GGML_OP_TRI: { | |
| split_state = handle_generic(src_ss, /*scalar_only =*/ true); | |
| } break; | |
| case GGML_OP_FILL: { | |
| split_state = handle_generic(src_ss, /*scalar_only =*/ false); | |
| } break; | |
| case GGML_OP_FLASH_ATTN_EXT: { | |
| split_state = handle_flash_attn_ext(src_ss); | |
| } break; | |
| case GGML_OP_FLASH_ATTN_BACK: { | |
| split_state = handle_generic(src_ss, /*scalar_only =*/ true); | |
| } break; | |
| case GGML_OP_SSM_CONV: { | |
| split_state = handle_ssm_conv(src_ss); | |
| } break; | |
| case GGML_OP_SSM_SCAN: | |
| case GGML_OP_WIN_PART: | |
| case GGML_OP_WIN_UNPART: | |
| case GGML_OP_GET_REL_POS: | |
| case GGML_OP_ADD_REL_POS: | |
| case GGML_OP_RWKV_WKV6: | |
| case GGML_OP_GATED_LINEAR_ATTN: | |
| case GGML_OP_RWKV_WKV7: | |
| case GGML_OP_SOLVE_TRI: { | |
| split_state = handle_generic(src_ss, /*scalar_only =*/ true); | |
| } break; | |
| case GGML_OP_GATED_DELTA_NET: { | |
| split_state = handle_gated_delta_net(src_ss); | |
| } break; | |
| case GGML_OP_UNARY: { | |
| split_state = handle_generic(src_ss, /*scalar_only =*/ false); | |
| } break; | |
| case GGML_OP_MAP_CUSTOM1: | |
| case GGML_OP_MAP_CUSTOM2: | |
| case GGML_OP_MAP_CUSTOM3: | |
| case GGML_OP_CUSTOM: { | |
| split_state = handle_generic(src_ss, /*scalar_only =*/ true); | |
| } break; | |
| case GGML_OP_CROSS_ENTROPY_LOSS: | |
| case GGML_OP_CROSS_ENTROPY_LOSS_BACK: { | |
| split_state = handle_per_row(src_ss); | |
| } break; | |
| case GGML_OP_OPT_STEP_ADAMW: | |
| case GGML_OP_OPT_STEP_SGD: | |
| case GGML_OP_GLU: { | |
| split_state = handle_generic(src_ss, /*scalar_only =*/ false); | |
| } break; | |
| default: { | |
| GGML_ABORT("ggml op not implemented: %s", ggml_op_name(tensor->op)); | |
| split_state = {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, {1}, 1}; | |
| } break; | |
| } | |
| if (split_state.axis >= 0 && split_state.axis < GGML_MAX_DIMS) { | |
| bool first_src_split_by_axis = true; | |
| const size_t n_bufs = ggml_backend_meta_buffer_n_bufs(tensor->buffer); | |
| for (size_t i = 0; i < GGML_MAX_SRC; i++) { | |
| if (tensor->src[i] == nullptr || src_ss[i].axis < 0 || src_ss[i].axis >= GGML_MAX_DIMS) { | |
| continue; | |
| } | |
| if (first_src_split_by_axis) { | |
| for (size_t j = 0; j < n_bufs; j++) { | |
| // Take over ratio from src: | |
| for (size_t s = 0; s < src_ss[i].n_segments; s++) { | |
| split_state.ne[s*n_bufs + j] = 0; | |
| } | |
| for (size_t s = 0; s < src_ss[i].n_segments; s++) { | |
| split_state.ne[j] += src_ss[i].ne[s*n_bufs + j] * src_ss[i].nr[s]; | |
| } | |
| split_state.ne[j] *= tensor->ne[split_state.axis]; | |
| if (split_state.ne[j] != 0 || tensor->src[i]->ne[src_ss[i].axis] != 0) { | |
| const int64_t div = tensor->src[i]->ne[src_ss[i].axis] * split_state.nr[0]; | |
| GGML_ASSERT(split_state.ne[j] % div == 0); | |
| split_state.ne[j] /= div; | |
| } | |
| } | |
| } else { | |
| GGML_ASSERT(split_state.n_segments == 1); | |
| for (size_t j = 0; j < n_bufs; j++) { | |
| // Assert that ratio is consistent: | |
| int64_t sum = 0; | |
| for (size_t s = 0; s < src_ss[i].n_segments; s++) { | |
| sum += src_ss[i].ne[s*n_bufs + j] * src_ss[i].nr[s]; | |
| } | |
| GGML_ASSERT(split_state.ne[j]*split_state.nr[0] * tensor->src[i]->ne[src_ss[i].axis] | |
| == sum * tensor->ne[split_state.axis]); | |
| } | |
| } | |
| first_src_split_by_axis = false; | |
| } | |
| GGML_ASSERT(!first_src_split_by_axis); | |
| } | |
| return split_state; | |
| }; | |
| const std::pair key = std::make_pair(tensor, assume_sync); | |
| auto it = buf_ctx->split_state_cache.find(key); | |
| if (it != buf_ctx->split_state_cache.end() && memcmp(it->second.second, (const char *) tensor, sizeof(it->second.second)) != 0) { | |
| buf_ctx->split_state_cache.clear(); | |
| it = buf_ctx->split_state_cache.end(); | |
| } | |
| if (it == buf_ctx->split_state_cache.end()) { | |
| buf_ctx->split_state_cache[key].first = calculate_split_state(); | |
| memcpy(buf_ctx->split_state_cache[key].second, tensor, sizeof(buf_ctx->split_state_cache[key].second)); | |
| if (buf_ctx->debug > 0) { | |
| std::string srcs_info; | |
| for (size_t i = 0; i < GGML_MAX_SRC; i++) { | |
| if (tensor->src[i] == nullptr) { | |
| continue; | |
| } | |
| if (!srcs_info.empty()) { | |
| srcs_info += ", "; | |
| } | |
| const ggml_backend_meta_split_state split_state = ggml_backend_meta_get_split_state(tensor->src[0], true); | |
| GGML_ASSERT(split_state.n_segments == 1); | |
| const char * axis_name = ggml_backend_meta_split_axis_name(split_state.axis); | |
| std::string ne_info; | |
| for (size_t j = 0; j < n_bufs; j++) { | |
| if (!ne_info.empty()) { | |
| ne_info += ", "; | |
| } | |
| ne_info += std::to_string(split_state.ne[j]) + "x" + std::to_string(split_state.nr[0]); | |
| } | |
| srcs_info += std::string(tensor->src[i]->name) + "[" + ggml_op_name(tensor->src[i]->op) + ", " + axis_name + ", {" + ne_info + "}]"; | |
| } | |
| std::string ne_info; | |
| for (size_t j = 0; j < n_bufs; j++) { | |
| if (!ne_info.empty()) { | |
| ne_info += ", "; | |
| } | |
| const ggml_backend_meta_split_state & ss = buf_ctx->split_state_cache[key].first; | |
| ne_info += std::to_string(ss.ne[j]) + "x" + std::to_string(ss.nr[0]); | |
| } | |
| GGML_LOG_DEBUG("SPLIT_STATE: {%s} -> %s[%s, %s, {%s}]\n", srcs_info.c_str(), tensor->name, ggml_op_name(tensor->op), | |
| ggml_backend_meta_split_axis_name(buf_ctx->split_state_cache[key].first.axis), ne_info.c_str()); | |
| } | |
| } | |
| ggml_backend_meta_split_state ret = buf_ctx->split_state_cache[key].first; | |
| GGML_ASSERT(ret.axis != GGML_BACKEND_SPLIT_AXIS_NONE); | |
| if (ret.axis >= 0 && ret.axis < GGML_MAX_DIMS) { | |
| int64_t ne_ret = 0; | |
| for (size_t s = 0; s < ret.n_segments; s++) { | |
| for (size_t j = 0; j < n_bufs; j++) { | |
| ne_ret += ret.ne[s*n_bufs + j] * ret.nr[s]; | |
| } | |
| } | |
| assert(ne_ret == tensor->ne[int(ret.axis)]); | |
| } | |
| return ret; | |
| } | |
| static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(const struct ggml_tensor * tensor, bool assume_sync) { | |
| GGML_ASSERT(ggml_backend_buffer_is_meta(tensor->buffer)); | |
| ggml_backend_meta_buffer_context * buf_ctx = (ggml_backend_meta_buffer_context *) tensor->buffer->context; | |
| return ggml_backend_meta_get_split_state(buf_ctx->get_simple_tensor_container(tensor), tensor, assume_sync); | |
| } | |
| static void * ggml_backend_meta_buffer_get_base(ggml_backend_buffer_t buffer) { | |
| GGML_UNUSED(buffer); | |
| return (void *) 0x1000000000000000; // FIXME | |
| } | |
| static enum ggml_status ggml_backend_meta_buffer_init_tensor_impl(ggml_backend_meta_simple_tensor_container & stc, ggml_tensor * tensor) { | |
| GGML_ASSERT(ggml_backend_buffer_is_meta(tensor->buffer)); | |
| ggml_backend_meta_buffer_context * buf_ctx = (ggml_backend_meta_buffer_context *) tensor->buffer->context; | |
| const size_t n_simple_bufs = ggml_backend_meta_buffer_n_bufs(tensor->buffer); | |
| const ggml_backend_meta_split_state split_state = ggml_backend_meta_get_split_state(stc, tensor, /*assume_sync =*/ true); | |
| GGML_ASSERT(ggml_nelements(tensor) == 0 || split_state.axis != GGML_BACKEND_SPLIT_AXIS_UNKNOWN); | |
| GGML_ASSERT(split_state.n_segments <= 16); | |
| int split_dim = split_state.axis; | |
| int64_t ne[GGML_MAX_DIMS]; | |
| size_t nb[GGML_MAX_DIMS]; | |
| for (size_t k = 0; k < GGML_MAX_DIMS; k++) { | |
| ne[k] = tensor->ne[k]; | |
| nb[k] = tensor->nb[k]; | |
| } | |
| std::vector<ggml_tensor *> simple_tensors; | |
| simple_tensors.reserve(n_simple_bufs); | |
| for (size_t j = 0; j < n_simple_bufs; j++) { | |
| ggml_context * simple_ctx = stc.ctxs[j].get(); | |
| ggml_backend_buffer_t simple_buf = buf_ctx->bufs[j].get(); | |
| if (split_dim >= 0 && split_dim < GGML_MAX_DIMS) { | |
| // TODO: the following assert fails for llama-parallel even though the results are correct: | |
| // GGML_ASSERT(ggml_is_contiguously_allocated(tensor)); | |
| ne[split_dim] = 0; | |
| for (size_t s = 0; s < split_state.n_segments; s++) { | |
| ne[split_dim] += split_state.ne[s*n_simple_bufs + j] * split_state.nr[s]; | |
| } | |
| for (int i = 0; i < GGML_MAX_DIMS; i++) { | |
| if (tensor->nb[i] > tensor->nb[split_dim]) { | |
| nb[i] = tensor->nb[i] * ne[split_dim]/tensor->ne[split_dim]; | |
| } | |
| } | |
| } | |
| ggml_tensor * t_ij = ggml_new_tensor(simple_ctx, tensor->type, GGML_MAX_DIMS, ne); | |
| t_ij->op = tensor->op; | |
| for (int i = 0; i < GGML_MAX_DIMS; i++) { | |
| t_ij->nb[i] = nb[i]; | |
| } | |
| t_ij->flags = tensor->flags; | |
| memcpy(t_ij->op_params, tensor->op_params, sizeof(tensor->op_params)); | |
| ggml_set_name(t_ij, tensor->name); | |
| t_ij->buffer = simple_buf; | |
| t_ij->view_src = tensor->view_src; | |
| t_ij->view_offs = tensor->view_offs; | |
| if (t_ij->view_src != nullptr && ggml_backend_buffer_is_meta(t_ij->view_src->buffer)) { | |
| t_ij->view_src = ggml_backend_meta_buffer_simple_tensor(tensor->view_src, j); | |
| if (t_ij->view_offs > 0 && split_dim >= 0 && split_dim < GGML_MAX_DIMS) { | |
| GGML_ASSERT(tensor->ne[split_dim] != 0); | |
| const int split_dim_view_src = ggml_backend_meta_get_split_state(tensor->view_src, /*assume_sync =*/ true).axis; | |
| GGML_ASSERT(split_dim_view_src >= 0 && split_dim_view_src < GGML_MAX_DIMS); | |
| // The offset can be internal to the data split, in those cases the view offset should not be scaled. | |
| // If however, the offset is larger than the data split then it needs to be scaled proportionally. | |
| bool split_internal_offset = t_ij->view_offs <= tensor->view_src->nb[split_dim_view_src]; | |
| for (int i = 0; i < GGML_MAX_DIMS; i++) { | |
| const size_t dim_size = tensor->ne[i] * tensor->nb[i]; | |
| if (tensor->view_offs <= dim_size && dim_size < tensor->nb[split_dim]) { | |
| split_internal_offset = true; | |
| break; | |
| } | |
| } | |
| if (!split_internal_offset) { | |
| t_ij->view_offs = t_ij->view_offs * ne[split_dim]/tensor->ne[split_dim]; | |
| } | |
| } | |
| } | |
| if (t_ij->view_src != nullptr) { | |
| t_ij->data = (char *) t_ij->view_src->data + t_ij->view_offs; | |
| } else if (simple_buf != nullptr) { | |
| t_ij->data = (char *) ggml_backend_buffer_get_base(simple_buf) | |
| + size_t(tensor->data) - size_t(ggml_backend_buffer_get_base(tensor->buffer)); | |
| } | |
| t_ij->extra = tensor->extra; | |
| for (int i = 0; i < GGML_MAX_SRC; i++) { | |
| t_ij->src[i] = tensor->src[i]; | |
| if (tensor->src[i] == tensor) { | |
| t_ij->src[i] = t_ij; | |
| } else if (t_ij->src[i] != nullptr && ggml_backend_buffer_is_meta(t_ij->src[i]->buffer)) { | |
| t_ij->src[i] = ggml_backend_meta_buffer_simple_tensor(tensor->src[i], j); | |
| } | |
| } | |
| simple_tensors.push_back(t_ij); | |
| } | |
| // If one of the sources has a zero-sized slice, disable the computation: | |
| for (int i = 0; i < GGML_MAX_SRC; i++) { | |
| if (tensor->src[i] == nullptr || !ggml_backend_buffer_is_meta(tensor->src[i]->buffer)) { | |
| continue; | |
| } | |
| const ggml_backend_meta_split_state split_state_src = ggml_backend_meta_get_split_state(tensor->src[i], /*assume_sync =*/ true); | |
| if (split_state_src.axis < 0 || split_state_src.axis >= GGML_MAX_DIMS) { | |
| continue; | |
| } | |
| for (size_t j = 0; j < n_simple_bufs; j++) { | |
| int64_t ne_sum = 0; | |
| for (size_t s = 0; s < split_state_src.n_segments; s++) { | |
| ne_sum += split_state_src.ne[s*n_simple_bufs + j] * split_state_src.nr[s]; | |
| } | |
| if (ne_sum == 0) { | |
| simple_tensors[j]->flags &= ~GGML_TENSOR_FLAG_COMPUTE; | |
| } | |
| } | |
| } | |
| stc.simple_tensors[tensor] = simple_tensors; | |
| return GGML_STATUS_SUCCESS; | |
| } | |
| static enum ggml_status ggml_backend_meta_buffer_init_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor) { | |
| GGML_ASSERT(ggml_backend_buffer_is_meta(buffer)); | |
| ggml_backend_meta_buffer_context * buf_ctx = (ggml_backend_meta_buffer_context *) buffer->context; | |
| buf_ctx->stc_compute_index = buf_ctx->stc_compute_index_next; | |
| return ggml_backend_meta_buffer_init_tensor_impl(buf_ctx->get_simple_tensor_container(tensor), tensor); | |
| } | |
| static void ggml_backend_meta_buffer_set_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor, const void * data, size_t offset, size_t size) { | |
| const size_t n_bufs = ggml_backend_meta_buffer_n_bufs(buffer); | |
| GGML_ASSERT(ggml_is_contiguous(tensor)); | |
| const ggml_backend_meta_split_state split_state = ggml_backend_meta_get_split_state(tensor, /*assume_sync =*/ false); | |
| if (split_state.n_segments != 1 || split_state.nr[0] != 1) { | |
| GGML_ASSERT(split_state.axis >= 0 && split_state.axis < GGML_MAX_DIMS); | |
| GGML_ASSERT(split_state.nr[0] != 0); | |
| GGML_ASSERT(tensor->ne[3] == 1); | |
| size_t offset_data = 0; | |
| std::vector<size_t> simple_offsets(n_bufs, 0); | |
| if (split_state.axis == GGML_BACKEND_SPLIT_AXIS_0) { | |
| GGML_ASSERT(tensor->ne[2] == 1); | |
| const size_t row_stride = tensor->nb[1]; | |
| GGML_ASSERT(offset % row_stride == 0); | |
| GGML_ASSERT(size % row_stride == 0); | |
| const int64_t row_start = offset / row_stride; | |
| const int64_t row_count = size / row_stride; | |
| GGML_ASSERT(row_start + row_count <= tensor->ne[1]); | |
| const int64_t blck_size = ggml_blck_size(tensor->type); | |
| for (size_t s = 0; s < split_state.n_segments; s++) { | |
| for (size_t r = 0; r < split_state.nr[s]; r++) { | |
| for (size_t j = 0; j < n_bufs; j++) { | |
| ggml_tensor * simple_tensor = ggml_backend_meta_buffer_simple_tensor(tensor, j); | |
| GGML_ASSERT(split_state.ne[s*n_bufs + j] % blck_size == 0); | |
| const size_t nbytes = split_state.ne[s*n_bufs + j]/blck_size * tensor->nb[0]; | |
| ggml_backend_tensor_set_2d(simple_tensor, (const char *) data + offset_data, | |
| simple_offsets[j] + row_start * simple_tensor->nb[1], nbytes, | |
| row_count, simple_tensor->nb[1], tensor->nb[1]); | |
| offset_data += nbytes; | |
| simple_offsets[j] += nbytes; | |
| } | |
| } | |
| } | |
| GGML_ASSERT(offset_data*row_count == size); | |
| return; | |
| } | |
| GGML_ASSERT(split_state.axis == GGML_BACKEND_SPLIT_AXIS_1); | |
| const size_t row_stride = tensor->nb[2]; | |
| GGML_ASSERT(offset % row_stride == 0); | |
| GGML_ASSERT(size % row_stride == 0); | |
| const int64_t row_start = offset / row_stride; | |
| const int64_t row_count = size / row_stride; | |
| GGML_ASSERT(row_start + row_count <= tensor->ne[2]); | |
| for (size_t s = 0; s < split_state.n_segments; s++) { | |
| for (size_t r = 0; r < split_state.nr[s]; r++) { | |
| for (size_t j = 0; j < n_bufs; j++) { | |
| ggml_tensor * simple_tensor = ggml_backend_meta_buffer_simple_tensor(tensor, j); | |
| const size_t nbytes = split_state.ne[s*n_bufs + j] * tensor->nb[1]; | |
| ggml_backend_tensor_set_2d(simple_tensor, (const char *) data + offset_data, | |
| simple_offsets[j] + row_start * simple_tensor->nb[2], nbytes, | |
| row_count, simple_tensor->nb[2], tensor->nb[2]); | |
| offset_data += nbytes; | |
| simple_offsets[j] += nbytes; | |
| } | |
| } | |
| } | |
| GGML_ASSERT(offset_data*row_count == size); | |
| return; | |
| } | |
| switch (split_state.axis) { | |
| case GGML_BACKEND_SPLIT_AXIS_0: | |
| case GGML_BACKEND_SPLIT_AXIS_1: | |
| case GGML_BACKEND_SPLIT_AXIS_2: { | |
| // Exploit that tensors are contiguous to splice it with simple tensors as "chunks". | |
| const size_t chunk_size_full = tensor->nb[split_state.axis + 1]; | |
| GGML_ASSERT(offset % chunk_size_full == 0); | |
| GGML_ASSERT(size % chunk_size_full == 0); | |
| const int64_t i_start = offset /chunk_size_full; | |
| const int64_t i_stop = (offset + size)/chunk_size_full; | |
| size_t offset_j = 0; | |
| for (size_t j = 0; j < n_bufs; j++) { | |
| ggml_tensor * simple_tensor = ggml_backend_meta_buffer_simple_tensor(tensor, j); | |
| const size_t chunk_size_j = simple_tensor->nb[split_state.axis + 1]; | |
| if (chunk_size_j == 0) { | |
| continue; | |
| } | |
| const size_t simple_offset = i_start * chunk_size_j; | |
| ggml_backend_tensor_set_2d(simple_tensor, (const char *) data + offset_j, simple_offset, chunk_size_j, i_stop - i_start, chunk_size_j, chunk_size_full); | |
| offset_j += chunk_size_j; | |
| } | |
| GGML_ASSERT(offset_j == chunk_size_full); | |
| } break; | |
| case GGML_BACKEND_SPLIT_AXIS_MIRRORED: { | |
| for (size_t j = 0; j < n_bufs; j++) { | |
| ggml_tensor * simple_tensor = ggml_backend_meta_buffer_simple_tensor(tensor, j); | |
| ggml_backend_tensor_set(simple_tensor, data, offset, size); | |
| } | |
| } break; | |
| case GGML_BACKEND_SPLIT_AXIS_PARTIAL: { | |
| GGML_ASSERT(tensor->type == GGML_TYPE_F32); | |
| const int64_t ne = ggml_nelements(tensor); | |
| std::vector<float> tmp; | |
| tmp.reserve(ne); | |
| for (int64_t i = 0; i < ne; i++) { | |
| tmp.push_back(((const float *) data)[i] / n_bufs); | |
| } | |
| for (size_t j = 0; j < n_bufs; j++) { | |
| ggml_tensor * simple_tensor = ggml_backend_meta_buffer_simple_tensor(tensor, j); | |
| ggml_backend_tensor_set(simple_tensor, tmp.data(), offset, size); | |
| } | |
| } break; | |
| default: { | |
| GGML_ABORT("fatal error"); | |
| } | |
| } | |
| } | |
| static void ggml_backend_meta_buffer_get_tensor(ggml_backend_buffer_t buffer, const ggml_tensor * tensor, void * data, size_t offset, size_t size) { | |
| const size_t n_bufs = ggml_backend_meta_buffer_n_bufs(buffer); | |
| GGML_ASSERT(ggml_is_contiguous(tensor)); | |
| const ggml_backend_meta_split_state split_state = ggml_backend_meta_get_split_state(tensor, /*assume_sync =*/ false); | |
| if (split_state.n_segments != 1 || split_state.nr[0] != 1) { | |
| GGML_ASSERT(split_state.axis >= 0 && split_state.axis < GGML_MAX_DIMS); | |
| GGML_ASSERT(split_state.nr[0] != 0); | |
| GGML_ASSERT(tensor->ne[3] == 1); | |
| size_t offset_data = 0; | |
| std::vector<size_t> simple_offsets(n_bufs, 0); | |
| if (split_state.axis == GGML_BACKEND_SPLIT_AXIS_0) { | |
| GGML_ASSERT(tensor->ne[2] == 1); | |
| const size_t row_stride = tensor->nb[1]; | |
| GGML_ASSERT(offset % row_stride == 0); | |
| GGML_ASSERT(size % row_stride == 0); | |
| const int64_t row_start = offset / row_stride; | |
| const int64_t row_count = size / row_stride; | |
| GGML_ASSERT(row_start + row_count <= tensor->ne[1]); | |
| const int64_t blck_size = ggml_blck_size(tensor->type); | |
| for (size_t s = 0; s < split_state.n_segments; s++) { | |
| for (size_t r = 0; r < split_state.nr[s]; r++) { | |
| for (size_t j = 0; j < n_bufs; j++) { | |
| const ggml_tensor * simple_tensor = ggml_backend_meta_buffer_simple_tensor(tensor, j); | |
| GGML_ASSERT(split_state.ne[s*n_bufs + j] % blck_size == 0); | |
| const size_t nbytes = split_state.ne[s*n_bufs + j]/blck_size * tensor->nb[0]; | |
| ggml_backend_tensor_get_2d(simple_tensor, (char *) data + offset_data, | |
| simple_offsets[j] + row_start * simple_tensor->nb[1], nbytes, | |
| row_count, simple_tensor->nb[1], tensor->nb[1]); | |
| offset_data += nbytes; | |
| simple_offsets[j] += nbytes; | |
| } | |
| } | |
| } | |
| GGML_ASSERT(offset_data*row_count == size); | |
| return; | |
| } | |
| GGML_ASSERT(split_state.axis == GGML_BACKEND_SPLIT_AXIS_1); | |
| const size_t row_stride = tensor->nb[2]; | |
| GGML_ASSERT(offset % row_stride == 0); | |
| GGML_ASSERT(size % row_stride == 0); | |
| const int64_t row_start = offset / row_stride; | |
| const int64_t row_count = size / row_stride; | |
| GGML_ASSERT(row_start + row_count <= tensor->ne[2]); | |
| for (size_t s = 0; s < split_state.n_segments; s++) { | |
| for (size_t r = 0; r < split_state.nr[s]; r++) { | |
| for (size_t j = 0; j < n_bufs; j++) { | |
| const ggml_tensor * simple_tensor = ggml_backend_meta_buffer_simple_tensor(tensor, j); | |
| const size_t nbytes = split_state.ne[s*n_bufs + j] * tensor->nb[1]; | |
| ggml_backend_tensor_get_2d(simple_tensor, (char *) data + offset_data, | |
| simple_offsets[j] + row_start * simple_tensor->nb[2], nbytes, | |
| row_count, simple_tensor->nb[2], tensor->nb[2]); | |
| offset_data += nbytes; | |
| simple_offsets[j] += nbytes; | |
| } | |
| } | |
| } | |
| GGML_ASSERT(offset_data*row_count == size); | |
| return; | |
| } | |
| switch (split_state.axis) { | |
| case GGML_BACKEND_SPLIT_AXIS_0: | |
| case GGML_BACKEND_SPLIT_AXIS_1: | |
| case GGML_BACKEND_SPLIT_AXIS_2: { | |
| // Exploit that tensors are contiguous to splice it with simple tensors as "chunks". | |
| const size_t chunk_size_full = tensor->nb[split_state.axis + 1]; | |
| GGML_ASSERT(offset % chunk_size_full == 0); | |
| GGML_ASSERT(size % chunk_size_full == 0); | |
| const int64_t i_start = offset /chunk_size_full; | |
| const int64_t i_stop = (offset + size)/chunk_size_full; | |
| size_t offset_j = 0; | |
| for (size_t j = 0; j < n_bufs; j++){ | |
| const ggml_tensor * simple_tensor = ggml_backend_meta_buffer_simple_tensor(tensor, j); | |
| const size_t chunk_size_j = simple_tensor->nb[split_state.axis + 1]; | |
| if (chunk_size_j == 0) { | |
| continue; | |
| } | |
| const size_t simple_offset = i_start * chunk_size_j; | |
| ggml_backend_tensor_get_2d(simple_tensor, (char *) data + offset_j, simple_offset, chunk_size_j, i_stop - i_start, chunk_size_j, chunk_size_full); | |
| offset_j += chunk_size_j; | |
| } | |
| GGML_ASSERT(offset_j == chunk_size_full); | |
| } break; | |
| case GGML_BACKEND_SPLIT_AXIS_MIRRORED: { | |
| // TODO other simple backend may be better | |
| const ggml_tensor * simple_tensor = ggml_backend_meta_buffer_simple_tensor(tensor, 0); | |
| ggml_backend_tensor_get(simple_tensor, data, offset, size); | |
| } break; | |
| default: { | |
| GGML_ABORT("fatal error"); | |
| } | |
| } | |
| } | |
| static void ggml_backend_meta_buffer_clear(ggml_backend_buffer_t buffer, uint8_t value) { | |
| const size_t n_buffers = ggml_backend_meta_buffer_n_bufs(buffer); | |
| for (size_t i = 0; i < n_buffers; i++) { | |
| ggml_backend_buffer_clear(ggml_backend_meta_buffer_simple_buffer(buffer, i), value); | |
| } | |
| } | |
| static void ggml_backend_meta_buffer_reset(ggml_backend_buffer_t buffer) { | |
| GGML_ASSERT(ggml_backend_buffer_is_meta(buffer)); | |
| ggml_backend_meta_buffer_context * buf_ctx = (ggml_backend_meta_buffer_context *) buffer->context; | |
| for (size_t i = 0; i < buf_ctx->bufs.size(); i++) { | |
| ggml_backend_buffer_reset(ggml_backend_meta_buffer_simple_buffer(buffer, i)); | |
| } | |
| } | |
| static const ggml_backend_buffer_i ggml_backend_meta_buffer_iface = { | |
| /* .free_buffer = */ ggml_backend_meta_buffer_free_buffer, | |
| /* .get_base = */ ggml_backend_meta_buffer_get_base, | |
| /* .init_tensor = */ ggml_backend_meta_buffer_init_tensor, | |
| /* .memset_tensor = */ nullptr, // TODO implement | |
| /* .set_tensor = */ ggml_backend_meta_buffer_set_tensor, | |
| /* .get_tensor = */ ggml_backend_meta_buffer_get_tensor, | |
| /* .set_tensor_2d = */ nullptr, | |
| /* .get_tensor_2d = */ nullptr, | |
| /* .cpy_tensor = */ nullptr, | |
| /* .clear = */ ggml_backend_meta_buffer_clear, | |
| /* .reset = */ ggml_backend_meta_buffer_reset, | |
| }; | |
| bool ggml_backend_buffer_is_meta(ggml_backend_buffer_t buf) { | |
| return buf != nullptr && buf->iface.free_buffer == ggml_backend_meta_buffer_iface.free_buffer; | |
| } | |
| static ggml_backend_buffer_t ggml_backend_meta_buffer_type_alloc_buffer(ggml_backend_buffer_type_t buft, size_t size) { | |
| const size_t n_simple_bufts = ggml_backend_meta_buft_n_bufts(buft); | |
| const ggml_init_params params = { | |
| /*.mem_size =*/ 1024*1024*ggml_tensor_overhead(), // FIXME | |
| /*.mem_buffer =*/ nullptr, | |
| /*.no_alloc =*/ true, | |
| }; | |
| ggml_backend_meta_simple_tensor_container stc_static; | |
| ggml_backend_meta_simple_tensor_container stc_compute_0(params, n_simple_bufts); | |
| ggml_backend_meta_simple_tensor_container stc_compute_1(params, n_simple_bufts); | |
| size_t max_size = 0; | |
| std::vector<ggml_backend_buffer_t> bufs; | |
| bufs.reserve(n_simple_bufts); | |
| for (size_t i = 0; i < n_simple_bufts; i++) { | |
| bufs.push_back(ggml_backend_buft_alloc_buffer(ggml_backend_meta_buft_simple_buft(buft, i), size)); | |
| GGML_ASSERT(bufs.back() != nullptr); | |
| max_size = std::max(max_size, ggml_backend_buffer_get_size(bufs.back())); | |
| } | |
| ggml_backend_meta_buffer_context * buf_ctx = new ggml_backend_meta_buffer_context(stc_static, stc_compute_0, stc_compute_1, bufs); | |
| return ggml_backend_buffer_init(buft, ggml_backend_meta_buffer_iface, buf_ctx, max_size); | |
| } | |
| struct ggml_backend_buffer * ggml_backend_meta_alloc_ctx_tensors_from_buft(struct ggml_context * ctx, ggml_backend_buffer_type_t buft) { | |
| const size_t n_simple_bufts = ggml_backend_meta_buft_n_bufts(buft); | |
| constexpr size_t compute_headroom = 16; // Maximum number of views per statically allocated tensor that can be created between evals. | |
| const ggml_init_params params_static = { | |
| /*.mem_size =*/ ggml_get_mem_size(ctx), | |
| /*.mem_buffer =*/ nullptr, | |
| /*.no_alloc =*/ true, | |
| }; | |
| const ggml_init_params params_compute = { | |
| /*.mem_size =*/ compute_headroom*ggml_get_mem_size(ctx), | |
| /*.mem_buffer =*/ nullptr, | |
| /*.no_alloc =*/ true, | |
| }; | |
| ggml_backend_meta_simple_tensor_container stc_static (params_static, n_simple_bufts); | |
| ggml_backend_meta_simple_tensor_container stc_compute_0(params_compute, n_simple_bufts); | |
| ggml_backend_meta_simple_tensor_container stc_compute_1(params_compute, n_simple_bufts); | |
| std::vector<ggml_backend_buffer_t> bufs(n_simple_bufts, nullptr); | |
| ggml_backend_meta_buffer_context * meta_buf_ctx = new ggml_backend_meta_buffer_context(stc_static, stc_compute_0, stc_compute_1, bufs); | |
| ggml_backend_buffer_t meta_buf = ggml_backend_buffer_init(buft, ggml_backend_meta_buffer_iface, meta_buf_ctx, 0); | |
| for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) { | |
| t->buffer = meta_buf; | |
| ggml_backend_meta_buffer_init_tensor_impl(meta_buf_ctx->stc_static, t); | |
| t->data = (void *) 0x2000000000000000; // FIXME | |
| } | |
| for (size_t i = 0; i < n_simple_bufts; i++) { | |
| ggml_context * ctx = meta_buf_ctx->stc_static.ctxs[i].get(); | |
| ggml_backend_buffer_type_t simple_buft = ggml_backend_meta_buft_simple_buft(buft, i); | |
| // If a ggml_context only has zero-sized tensors, ggml_backend_alloc_ctx_tensors_from_buft returns NULL. | |
| // For those edge cases, allocate a dummy buffer instead. | |
| bool any_nonzero_slice = false; | |
| for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) { | |
| if (ggml_nelements(t) != 0) { | |
| any_nonzero_slice = true; | |
| break; | |
| } | |
| } | |
| if (any_nonzero_slice) { | |
| meta_buf_ctx->bufs[i].reset(ggml_backend_alloc_ctx_tensors_from_buft(ctx, simple_buft)); | |
| } else { | |
| meta_buf_ctx->bufs[i].reset(ggml_backend_buft_alloc_buffer(simple_buft, 0)); | |
| for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) { | |
| t->buffer = meta_buf_ctx->bufs[i].get(); | |
| } | |
| } | |
| GGML_ASSERT(meta_buf_ctx->bufs[i]); | |
| meta_buf->size = std::max(meta_buf->size, ggml_backend_buffer_get_size(meta_buf_ctx->bufs[i].get())); | |
| } | |
| return meta_buf; | |
| } | |
| // | |
| // meta backend | |
| // | |
| static ggml_guid_t ggml_backend_meta_guid() { | |
| static ggml_guid guid = {0xf1, 0x0e, 0x34, 0xcf, 0x9c, 0x6f, 0x43, 0xcb, 0x96, 0x92, 0xbe, 0x8e, 0xbb, 0x71, 0x3f, 0xda}; | |
| return &guid; | |
| } | |
| struct ggml_backend_meta_context { | |
| struct cgraph_config { | |
| ggml_cgraph * cgraph_main = nullptr; | |
| int offset = 0; // Node offset vs. original graph | |
| std::vector<ggml_cgraph *> cgraphs_aux; | |
| }; | |
| struct backend_config { | |
| ggml_backend_t backend; | |
| std::vector<cgraph_config> cgraphs; | |
| std::vector<ggml_tensor *> nodes; | |
| std::vector<ggml_backend_buffer_ptr> bufs; | |
| backend_config(ggml_backend_t backend, const size_t n_reduce_steps) : backend(backend) { | |
| bufs.resize(n_reduce_steps); | |
| } | |
| }; | |
| std::string name; | |
| std::vector<backend_config> backend_configs; | |
| ggml_context_ptr ctx; | |
| std::vector<ggml_cgraph *> cgraphs_aux; | |
| std::vector<ggml_tensor *> nodes_aux; | |
| size_t n_reduce_steps; | |
| int max_nnodes = 0; | |
| size_t max_tmp_size = 0; | |
| size_t max_subgraphs = 0; | |
| size_t n_subgraphs = 0; | |
| uint64_t uid = 0; | |
| void * comm_ctx = nullptr; | |
| ggml_backend_comm_allreduce_tensor_t comm_allreduce = nullptr; | |
| ggml_backend_meta_context(ggml_backend_dev_t meta_dev, const char * params) { | |
| const size_t n_devs = ggml_backend_meta_dev_n_devs(meta_dev); | |
| n_reduce_steps = std::ceil(std::log2(n_devs)); | |
| name = "Meta("; | |
| std::vector<ggml_backend_t> simple_backends; | |
| backend_configs.reserve(n_devs); | |
| simple_backends.reserve(n_devs); | |
| for (size_t i = 0; i < n_devs; i++) { | |
| ggml_backend_dev_t simple_dev = ggml_backend_meta_dev_simple_dev(meta_dev, i); | |
| if (i > 0) { | |
| name += ","; | |
| } | |
| name += ggml_backend_dev_name(simple_dev); | |
| simple_backends.push_back(ggml_backend_dev_init(simple_dev, params)); | |
| backend_configs.emplace_back(simple_backends.back(), n_reduce_steps); | |
| } | |
| name += ")"; | |
| if (n_devs > 1) { | |
| ggml_backend_comm_init_t comm_init = (ggml_backend_comm_init_t) ggml_backend_reg_get_proc_address( | |
| ggml_backend_dev_backend_reg(ggml_backend_get_device(simple_backends[0])), "ggml_backend_comm_init"); | |
| if (comm_init != nullptr) { | |
| comm_ctx = comm_init(simple_backends.data(), simple_backends.size()); | |
| } | |
| } | |
| if (comm_ctx != nullptr) { | |
| comm_allreduce = (ggml_backend_comm_allreduce_tensor_t) | |
| ggml_backend_reg_get_proc_address(ggml_backend_dev_backend_reg( | |
| ggml_backend_get_device(simple_backends[0])), "ggml_backend_comm_allreduce_tensor"); | |
| GGML_ASSERT(comm_allreduce != nullptr); | |
| } | |
| } | |
| ~ggml_backend_meta_context() { | |
| if (comm_ctx != nullptr) { | |
| ggml_backend_comm_free_t comm_free = (ggml_backend_comm_free_t) ggml_backend_reg_get_proc_address( | |
| ggml_backend_dev_backend_reg(ggml_backend_get_device(backend_configs[0].backend)), "ggml_backend_comm_free"); | |
| GGML_ASSERT(comm_free != nullptr); | |
| comm_free(comm_ctx); | |
| } | |
| for (auto & bc : backend_configs) { | |
| ggml_backend_free(bc.backend); | |
| } | |
| } | |
| }; | |
| static const char * ggml_backend_meta_get_name(ggml_backend_t backend) { | |
| GGML_ASSERT(ggml_backend_is_meta(backend)); | |
| const ggml_backend_meta_context * backend_ctx = (const ggml_backend_meta_context *) backend->context; | |
| return backend_ctx->name.c_str(); | |
| } | |
| static void ggml_backend_meta_free(ggml_backend_t backend) { | |
| GGML_ASSERT(ggml_backend_is_meta(backend)); | |
| ggml_backend_meta_context * backend_ctx = (ggml_backend_meta_context *) backend->context; | |
| delete backend_ctx; | |
| delete backend; | |
| } | |
| static void ggml_backend_meta_set_tensor_async(ggml_backend_t backend, ggml_tensor * tensor, const void * data, size_t offset, size_t size) { | |
| const size_t n_backends = ggml_backend_meta_n_backends(backend); | |
| GGML_ASSERT(offset == 0); | |
| GGML_ASSERT(ggml_is_contiguous(tensor)); | |
| const ggml_backend_meta_split_state split_state = ggml_backend_meta_get_split_state(tensor, /*assume_sync =*/ false); | |
| GGML_ASSERT(split_state.n_segments == 1); | |
| GGML_ASSERT(split_state.nr[0] == 1); | |
| switch (split_state.axis) { | |
| case GGML_BACKEND_SPLIT_AXIS_0: | |
| case GGML_BACKEND_SPLIT_AXIS_1: | |
| case GGML_BACKEND_SPLIT_AXIS_2: { | |
| // Exploit that tensors are contiguous to splice it with simple tensors as "chunks". | |
| const size_t chunk_size_full = tensor->nb[split_state.axis + 1]; | |
| GGML_ASSERT(offset % chunk_size_full == 0); | |
| GGML_ASSERT(size % chunk_size_full == 0); | |
| const int64_t i_start = offset /chunk_size_full; | |
| const int64_t i_stop = (offset + size)/chunk_size_full; | |
| size_t offset_j = 0; | |
| for (size_t j = 0; j < n_backends; j++){ | |
| ggml_backend_t simple_backend = ggml_backend_meta_simple_backend(backend, j); | |
| ggml_tensor * simple_tensor = ggml_backend_meta_buffer_simple_tensor(tensor, j); | |
| const size_t chunk_size_j = simple_tensor->nb[split_state.axis + 1]; | |
| if (chunk_size_j == 0) { | |
| continue; | |
| } | |
| ggml_backend_tensor_set_2d_async(simple_backend, simple_tensor, (const char *) data + offset_j, offset, chunk_size_j, | |
| i_stop - i_start, chunk_size_j, chunk_size_full); | |
| offset_j += chunk_size_j; | |
| } | |
| GGML_ASSERT(offset_j == chunk_size_full); | |
| } break; | |
| case GGML_BACKEND_SPLIT_AXIS_MIRRORED: { | |
| for (size_t j = 0; j < n_backends; j++) { | |
| ggml_backend_tensor_set_async( | |
| ggml_backend_meta_simple_backend(backend, j), ggml_backend_meta_buffer_simple_tensor(tensor, j), data, offset, size); | |
| } | |
| } break; | |
| default: { | |
| GGML_ABORT("fatal error"); | |
| } | |
| } | |
| } | |
| static void ggml_backend_meta_get_tensor_async(ggml_backend_t backend, const ggml_tensor * tensor, void * data, size_t offset, size_t size) { | |
| const size_t n_backends = ggml_backend_meta_n_backends(backend); | |
| GGML_ASSERT(offset == 0); | |
| GGML_ASSERT(ggml_is_contiguous(tensor)); | |
| const ggml_backend_meta_split_state split_state = ggml_backend_meta_get_split_state(tensor, /*assume_sync =*/ false); | |
| GGML_ASSERT(split_state.n_segments == 1); | |
| GGML_ASSERT(split_state.nr[0] == 1); | |
| switch (split_state.axis) { | |
| case GGML_BACKEND_SPLIT_AXIS_0: | |
| case GGML_BACKEND_SPLIT_AXIS_1: | |
| case GGML_BACKEND_SPLIT_AXIS_2: { | |
| // Exploit that tensors are contiguous to splice it with simple tensors as "chunks". | |
| const size_t chunk_size_full = tensor->nb[split_state.axis + 1]; | |
| GGML_ASSERT(offset % chunk_size_full == 0); | |
| GGML_ASSERT(size % chunk_size_full == 0); | |
| const int64_t i_start = offset /chunk_size_full; | |
| const int64_t i_stop = (offset + size)/chunk_size_full; | |
| size_t offset_j = 0; | |
| for (size_t j = 0; j < n_backends; j++){ | |
| ggml_backend_t simple_backend = ggml_backend_meta_simple_backend(backend, j); | |
| const ggml_tensor * simple_tensor = ggml_backend_meta_buffer_simple_tensor(tensor, j); | |
| const size_t chunk_size_j = simple_tensor->nb[split_state.axis + 1]; | |
| if (chunk_size_j == 0) { | |
| continue; | |
| } | |
| ggml_backend_tensor_get_2d_async(simple_backend, simple_tensor, (char *) data + offset_j, offset, chunk_size_j, | |
| i_stop - i_start, chunk_size_j, chunk_size_full); | |
| offset_j += chunk_size_j; | |
| } | |
| GGML_ASSERT(offset_j == chunk_size_full); | |
| } break; | |
| case GGML_BACKEND_SPLIT_AXIS_MIRRORED: { | |
| // TODO other simple backend may be better | |
| ggml_backend_t simple_backend = ggml_backend_meta_simple_backend(backend, 0); | |
| const ggml_tensor * simple_tensor = ggml_backend_meta_buffer_simple_tensor(tensor, 0); | |
| ggml_backend_tensor_get_async(simple_backend, simple_tensor, data, offset, size); | |
| } break; | |
| default: { | |
| GGML_ABORT("fatal error"); | |
| } | |
| } | |
| } | |
| static void ggml_backend_meta_synchronize(ggml_backend_t backend) { | |
| const size_t n_backends = ggml_backend_meta_n_backends(backend); | |
| for (size_t i = 0; i < n_backends; i++) { | |
| ggml_backend_synchronize(ggml_backend_meta_simple_backend(backend, i)); | |
| } | |
| } | |
| static enum ggml_status ggml_backend_meta_graph_compute(ggml_backend_t backend, struct ggml_cgraph * cgraph) { | |
| GGML_ASSERT(cgraph->grads == nullptr); | |
| const size_t n_backends = ggml_backend_meta_n_backends(backend); | |
| ggml_backend_meta_context * backend_ctx = (ggml_backend_meta_context *) backend->context; | |
| // If the previous cgraph had a defined UID it can be used to skip rebuilding the subgraphs per simple backend. | |
| const bool needs_rebuild = (cgraph->uid == 0) || (cgraph->uid != backend_ctx->uid); | |
| bool max_nnodes_raised = false; | |
| if (cgraph->n_nodes > backend_ctx->max_nnodes) { | |
| for (size_t j = 0; j < n_backends; j++) { | |
| auto & bcj = backend_ctx->backend_configs[j]; | |
| bcj.nodes.resize(cgraph->n_nodes); | |
| bcj.cgraphs.resize(cgraph->n_nodes); | |
| } | |
| backend_ctx->max_nnodes = cgraph->n_nodes; | |
| max_nnodes_raised = true; | |
| assert(needs_rebuild); | |
| } | |
| if (needs_rebuild) { | |
| std::set<ggml_backend_buffer_t> used_buffers; | |
| for (int i = 0; i < cgraph->n_leafs; i++) { | |
| if (ggml_backend_buffer_is_meta(cgraph->leafs[i]->buffer)) { | |
| used_buffers.emplace(cgraph->leafs[i]->buffer); | |
| } | |
| } | |
| for (int i = 0; i < cgraph->n_nodes; i++) { | |
| if (ggml_backend_buffer_is_meta(cgraph->nodes[i]->buffer)) { | |
| used_buffers.emplace(cgraph->nodes[i]->buffer); | |
| } | |
| } | |
| for (ggml_backend_buffer_t buf : used_buffers) { | |
| ggml_backend_meta_buffer_context * buf_ctx = (ggml_backend_meta_buffer_context *) buf->context; | |
| buf_ctx->stc_compute_index_next = buf_ctx->stc_compute_index ^ 1; | |
| ggml_backend_meta_simple_tensor_container & stc = buf_ctx->stc_compute[buf_ctx->stc_compute_index_next]; | |
| for (ggml_context_ptr & ctx : stc.ctxs) { | |
| ggml_reset(ctx.get()); | |
| } | |
| stc.simple_tensors.clear(); | |
| } | |
| size_t n_subgraphs = 0; | |
| size_t max_tmp_size = 0; | |
| for (size_t j = 0; j < n_backends; j++) { | |
| auto & bcj = backend_ctx->backend_configs[j]; | |
| for (int i = 0; i < cgraph->n_nodes; i++) { | |
| ggml_tensor * node = cgraph->nodes[i]; | |
| if (node->view_src != nullptr && node->view_src->op == GGML_OP_NONE && ggml_backend_buffer_is_host(node->view_src->buffer)) { | |
| // FIXME s_copy_main is on the CPU and its view seems to be incorrectly added to the graph nodes. | |
| // For regular usage this doesn't matter since it's a noop but trying to call ggml_backend_meta_buffer_simple_tensor results in a crash. | |
| bcj.nodes[i] = node; | |
| continue; | |
| } | |
| bcj.nodes[i] = ggml_backend_meta_buffer_simple_tensor(node, j); | |
| GGML_ASSERT(bcj.nodes[i]); | |
| } | |
| } | |
| { | |
| // For MoE models it may make sense to delay the AllReduce in order to reduce I/O: | |
| auto get_i_delayed = [&](const int i) -> int { | |
| int id = i; // i_delayed | |
| int idr = i; // i_delayed return, last safe return value | |
| ggml_tensor * node = cgraph->nodes[id]; | |
| int32_t n_used = ggml_node_get_use_count(cgraph, id); | |
| // Skip MIRRORED nodes that don't consume node | |
| auto skip_unrelated = [&]() { | |
| while (id + 1 < cgraph->n_nodes) { | |
| ggml_tensor * next = cgraph->nodes[id+1]; | |
| if (ggml_backend_meta_get_split_state(next, false).axis != GGML_BACKEND_SPLIT_AXIS_MIRRORED) { | |
| break; | |
| } | |
| bool safe = true; | |
| for (int s = 0; s < GGML_MAX_SRC; s++) { | |
| if (next->src[s] == nullptr) { | |
| continue; | |
| } | |
| if (next->src[s] == node) { | |
| safe = false; | |
| break; | |
| } | |
| if (ggml_backend_meta_get_split_state(next->src[s], false).axis != GGML_BACKEND_SPLIT_AXIS_MIRRORED) { | |
| safe = false; | |
| break; | |
| } | |
| } | |
| if (!safe) { | |
| break; | |
| } | |
| id++; | |
| } | |
| }; | |
| skip_unrelated(); | |
| if (id + 1 >= cgraph->n_nodes) { | |
| return idr; | |
| } | |
| { | |
| ggml_tensor * next = cgraph->nodes[id+1]; | |
| if (next->op == GGML_OP_ADD_ID && next->src[0] == node && | |
| ggml_backend_meta_get_split_state(next->src[1], false).axis == GGML_BACKEND_SPLIT_AXIS_PARTIAL && | |
| ggml_backend_meta_get_split_state(next->src[2], false).axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED) { | |
| node = next; | |
| id++; | |
| idr = id; | |
| n_used = ggml_node_get_use_count(cgraph, id); | |
| } | |
| } | |
| // Chain of MULs with MIRRORED src[1] | |
| while (true) { | |
| skip_unrelated(); | |
| if (id + 1 >= cgraph->n_nodes) { | |
| return idr; | |
| } | |
| ggml_tensor * next = cgraph->nodes[id+1]; | |
| if (next->op == GGML_OP_MUL && next->src[0] == node && | |
| ggml_backend_meta_get_split_state(next->src[1], false).axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED) { | |
| node = next; | |
| id++; | |
| idr = id; | |
| n_used = ggml_node_get_use_count(cgraph, id); | |
| } else { | |
| break; | |
| } | |
| } | |
| if (n_used != node->ne[1] || id + 2*n_used-1 >= cgraph->n_nodes) { | |
| return idr; | |
| } | |
| for (int32_t k = 0; k < n_used; k++) { | |
| ggml_tensor * next = cgraph->nodes[id+1]; | |
| if (next->op != GGML_OP_VIEW || next->view_src != node || next->view_offs != k*node->nb[1] || | |
| next->ne[0] != node->ne[0] || next->ne[1] != node->ne[2] || next->nb[1] != node->nb[2] || | |
| ggml_node_get_use_count(cgraph, id+1) != 1) { | |
| return idr; | |
| } | |
| id++; | |
| } | |
| { | |
| ggml_tensor * next = cgraph->nodes[id+1]; | |
| if (next->op != GGML_OP_ADD || next->src[0] != cgraph->nodes[id - (n_used-1)] || | |
| next->src[1] != cgraph->nodes[id - (n_used-2)] || ggml_node_get_use_count(cgraph, id+1) != 1) { | |
| return idr; | |
| } | |
| id++; | |
| } | |
| for (int32_t k = 0; k < n_used - 2; k++) { | |
| ggml_tensor * next = cgraph->nodes[id+1]; | |
| if (next->op != GGML_OP_ADD || next->src[0] != cgraph->nodes[id] || | |
| next->src[1] != cgraph->nodes[id - (n_used-2)] || ggml_node_get_use_count(cgraph, id+1) != 1) { | |
| return idr; | |
| } | |
| id++; | |
| } | |
| idr = id; | |
| return idr; | |
| }; | |
| int i_start = 0; | |
| for (int i = 0; i < cgraph->n_nodes; i++) { | |
| ggml_tensor * node = cgraph->nodes[i]; | |
| if (node->view_src != nullptr && node->view_src->op == GGML_OP_NONE && ggml_backend_buffer_is_host(node->view_src->buffer)) { | |
| continue; | |
| } | |
| const ggml_backend_meta_split_state split_state = ggml_backend_meta_get_split_state(node, /*assume_sync =*/ false); | |
| if (split_state.axis == GGML_BACKEND_SPLIT_AXIS_PARTIAL) { | |
| max_tmp_size = std::max(max_tmp_size, ggml_nbytes(node)); | |
| } | |
| const bool new_subgraph = i + 1 == cgraph->n_nodes || split_state.axis == GGML_BACKEND_SPLIT_AXIS_PARTIAL; | |
| if (!new_subgraph) { | |
| continue; | |
| } | |
| const int i_delayed = get_i_delayed(i); | |
| // If we can delay the AllReduce we need to consider the interaction with zero-sized tensor slices. | |
| // A backend with such a slice would normally have valid data after participating in the AllReduce with a node that has | |
| // its compute flag disabled and thus gets its data zeroed out. | |
| // If the AllReduce is delayed then the nodes until that point also need to have their compute flag disabled. | |
| if (i_delayed > i) { | |
| for (size_t j = 0; j < n_backends; j++) { | |
| auto & bcj = backend_ctx->backend_configs[j]; | |
| if ((bcj.nodes[i]->flags & GGML_TENSOR_FLAG_COMPUTE) == 0) { | |
| for (int ii = i + 1; ii <= i_delayed; ii++) { | |
| bcj.nodes[ii]->flags &= ~GGML_TENSOR_FLAG_COMPUTE; | |
| } | |
| } | |
| } | |
| } | |
| i = i_delayed; | |
| for (size_t j = 0; j < n_backends; j++) { | |
| auto & bcj = backend_ctx->backend_configs[j]; | |
| bcj.cgraphs[n_subgraphs].offset = i_start; | |
| } | |
| n_subgraphs++; | |
| i_start = i + 1; | |
| } | |
| GGML_ASSERT(i_start == cgraph->n_nodes); | |
| } | |
| backend_ctx->uid = cgraph->uid; | |
| backend_ctx->n_subgraphs = n_subgraphs; | |
| if (max_tmp_size > backend_ctx->max_tmp_size) { | |
| for (size_t j = 0; j < n_backends; j++) { | |
| auto & bcj = backend_ctx->backend_configs[j]; | |
| for (size_t i = 0; i < backend_ctx->n_reduce_steps; i++) { | |
| bcj.bufs[i].reset(ggml_backend_alloc_buffer(bcj.backend, max_tmp_size)); | |
| } | |
| } | |
| backend_ctx->max_tmp_size = max_tmp_size; | |
| } | |
| if (max_nnodes_raised || n_subgraphs > backend_ctx->max_subgraphs) { | |
| backend_ctx->max_subgraphs = std::max(backend_ctx->max_subgraphs, n_subgraphs); | |
| const size_t n_nodes_per_device = 3 * backend_ctx->n_reduce_steps; // tmp + ADD (+zeroing) graph per step and device | |
| const size_t n_cgraphs_per_device = 2 * backend_ctx->n_reduce_steps; // ADD ( + zeroing) graph per step and device | |
| const size_t mem_per_device_graphs_main = backend_ctx->max_subgraphs*ggml_graph_overhead_custom(backend_ctx->max_nnodes, cgraph->grads); | |
| const size_t mem_per_device_graphs_aux = n_cgraphs_per_device*backend_ctx->max_subgraphs*ggml_graph_overhead_custom(1, cgraph->grads); | |
| const size_t mem_per_device_nodes_aux = n_nodes_per_device*backend_ctx->max_subgraphs*ggml_tensor_overhead(); | |
| const ggml_init_params params = { | |
| /*.mem_size =*/ n_backends * (mem_per_device_graphs_main + mem_per_device_graphs_aux + mem_per_device_nodes_aux), | |
| /*.mem_buffer =*/ nullptr, | |
| /*.no_alloc =*/ true, | |
| }; | |
| backend_ctx->ctx.reset(ggml_init(params)); | |
| for (size_t j = 0; j < n_backends; j++) { | |
| auto & bcj = backend_ctx->backend_configs[j]; | |
| for (size_t i = 0; i < n_subgraphs; i++) { | |
| bcj.cgraphs[i].cgraph_main = ggml_new_graph_custom(backend_ctx->ctx.get(), cgraph->n_nodes, /*grads =*/ false); | |
| } | |
| } | |
| backend_ctx->cgraphs_aux.resize(n_backends*n_cgraphs_per_device*backend_ctx->max_subgraphs); | |
| for (size_t k = 0; k < backend_ctx->cgraphs_aux.size(); k++) { | |
| backend_ctx->cgraphs_aux[k] = ggml_new_graph_custom(backend_ctx->ctx.get(), 1, cgraph->grads); | |
| } | |
| backend_ctx->nodes_aux.resize(n_backends*n_nodes_per_device*backend_ctx->max_subgraphs); | |
| for (size_t k = 0; k < backend_ctx->nodes_aux.size(); k++) { | |
| backend_ctx->nodes_aux[k] = ggml_new_tensor_1d(backend_ctx->ctx.get(), GGML_TYPE_F32, 1); | |
| } | |
| } | |
| for (size_t j = 0; j < n_backends; j++) { | |
| auto & bcj = backend_ctx->backend_configs[j]; | |
| for (size_t i_graph = 0; i_graph < n_subgraphs; i_graph++) { | |
| ggml_cgraph * cgraph_ij = bcj.cgraphs[i_graph].cgraph_main; | |
| const size_t i_node_start = bcj.cgraphs[i_graph].offset; | |
| const size_t i_node_stop = i_graph + 1 < n_subgraphs ? bcj.cgraphs[i_graph + 1].offset : cgraph->n_nodes; | |
| cgraph_ij->n_nodes = i_node_stop - i_node_start; | |
| ggml_hash_set_reset(&cgraph_ij->visited_hash_set); | |
| for (size_t i_node = i_node_start; i_node < i_node_stop; i_node++) { | |
| ggml_tensor * node_ij = bcj.nodes[i_node]; | |
| cgraph_ij->nodes[i_node - i_node_start] = node_ij; | |
| const size_t hash_pos_orig = ggml_hash_find(&cgraph->visited_hash_set, cgraph->nodes[i_node]); | |
| const size_t hash_pos_ij = ggml_hash_insert(&cgraph_ij->visited_hash_set, node_ij); | |
| cgraph_ij->use_counts[hash_pos_ij] = cgraph->use_counts[hash_pos_orig]; | |
| } | |
| cgraph_ij->uid = ggml_graph_next_uid(); | |
| } | |
| } | |
| } | |
| size_t iga = 0; // i graph aux | |
| size_t ina = 0; // i node aux | |
| auto get_node_aux = [&](ggml_tensor * t) -> ggml_tensor * { | |
| ggml_tensor * ret = backend_ctx->nodes_aux[ina++]; | |
| memset(ret, 0, sizeof(ggml_tensor)); | |
| ret->op = GGML_OP_NONE; | |
| ret->type = t->type; | |
| for (size_t k = 0; k < GGML_MAX_DIMS; k++) { | |
| ret->ne[k] = t->ne[k]; | |
| ret->nb[k] = t->nb[k]; | |
| } | |
| return ret; | |
| }; | |
| auto set_tmp_data = [&](ggml_tensor * tensor, const size_t j, const size_t i_buf) { | |
| auto & bcj = backend_ctx->backend_configs[j]; | |
| ggml_backend_buffer_ptr & buf_ptr = bcj.bufs[i_buf]; | |
| if (!buf_ptr || ggml_backend_buffer_get_size(buf_ptr.get()) < backend_ctx->max_tmp_size) { | |
| buf_ptr.reset(ggml_backend_alloc_buffer(bcj.backend, backend_ctx->max_tmp_size)); | |
| } | |
| tensor->buffer = buf_ptr.get(); | |
| tensor->data = ggml_backend_buffer_get_base(buf_ptr.get()); | |
| }; | |
| // FIXME usage_counts | |
| auto get_cgraph_aux = [&]() -> ggml_cgraph * { | |
| ggml_cgraph * ret = backend_ctx->cgraphs_aux[iga++]; | |
| return ret; | |
| }; | |
| // Preferentially use backend-specific allreduce_tensor_async (e.g. NCCL for CUDA), use a generic fallback if unavailable: | |
| auto allreduce_fallback = [&](size_t i) -> ggml_status { | |
| std::vector<ggml_cgraph *> step_cgraphs(n_backends, nullptr); | |
| // Zero out nodes that were disabled due to having a zero-sized slice: | |
| for (size_t j = 0; j < n_backends; j++) { | |
| auto & bcj = backend_ctx->backend_configs[j]; | |
| ggml_tensor * node = bcj.cgraphs[i].cgraph_main->nodes[bcj.cgraphs[i].cgraph_main->n_nodes - 1]; | |
| if (node->flags & GGML_TENSOR_FLAG_COMPUTE) { | |
| continue; | |
| } | |
| ggml_tensor * node_zero = get_node_aux(node); | |
| node_zero->op = GGML_OP_SCALE; // FIXME 0.0f * NaN == NaN | |
| node_zero->src[0] = node; | |
| ggml_set_op_params_f32(node_zero, 0, 0.0f); | |
| node_zero->data = node->data; | |
| node_zero->buffer = node->buffer; | |
| node_zero->flags |= GGML_TENSOR_FLAG_COMPUTE; | |
| step_cgraphs[j] = get_cgraph_aux(); | |
| step_cgraphs[j]->nodes[0] = node_zero; | |
| step_cgraphs[j]->n_nodes = 1; | |
| const ggml_status status = ggml_backend_graph_compute_async(bcj.backend, step_cgraphs[j]); | |
| if (status != GGML_STATUS_SUCCESS) { | |
| return status; | |
| } | |
| } | |
| std::fill(step_cgraphs.begin(), step_cgraphs.end(), nullptr); | |
| auto push_data = [&](const size_t j_src, const size_t j_dst, const size_t i_buf) { | |
| assert(step_cgraphs[j_dst] == nullptr); | |
| auto & bcj_src = backend_ctx->backend_configs[j_src]; | |
| auto & bcj_dst = backend_ctx->backend_configs[j_dst]; | |
| ggml_tensor * node_src = bcj_src.cgraphs[i].cgraph_main->nodes[bcj_src.cgraphs[i].cgraph_main->n_nodes - 1]; | |
| ggml_tensor * node_dst = bcj_dst.cgraphs[i].cgraph_main->nodes[bcj_dst.cgraphs[i].cgraph_main->n_nodes - 1]; | |
| GGML_ASSERT(ggml_is_contiguous(node_src)); | |
| GGML_ASSERT(ggml_is_contiguous(node_dst)); | |
| ggml_tensor * node_tmp = get_node_aux(node_dst); | |
| set_tmp_data(node_tmp, j_dst, i_buf); | |
| ggml_backend_tensor_copy_async(bcj_src.backend, bcj_dst.backend, node_src, node_tmp); | |
| ggml_tensor * node_red = get_node_aux(node_dst); | |
| node_red->view_src = node_dst->view_src == nullptr ? node_dst : node_dst->view_src; | |
| node_red->view_offs = node_dst->view_offs; | |
| node_red->op = GGML_OP_ADD; | |
| node_red->src[0] = node_dst; | |
| node_red->src[1] = node_tmp; | |
| node_red->flags |= GGML_TENSOR_FLAG_COMPUTE; | |
| ggml_backend_view_init(node_red); | |
| ggml_cgraph * cgraph_aux = get_cgraph_aux(); | |
| cgraph_aux->nodes[0] = node_red; | |
| cgraph_aux->n_nodes = 1; | |
| step_cgraphs[j_dst] = cgraph_aux; | |
| }; | |
| size_t offset_j = n_backends/2; | |
| while ((offset_j & (offset_j - 1)) != 0) { | |
| offset_j--; | |
| } | |
| const size_t offset_j_max = offset_j; | |
| size_t i_buf = 0; | |
| // If n_backends is not a power of 2, fold in the excess prior to butterfly reduction: | |
| for (size_t j_src = 2*offset_j_max; j_src < n_backends; j_src++) { | |
| const size_t j_dst = j_src - 2*offset_j_max; | |
| push_data(j_src, j_dst, i_buf); | |
| const ggml_status status = ggml_backend_graph_compute_async(backend_ctx->backend_configs[j_dst].backend, step_cgraphs[j_dst]); | |
| if (status != GGML_STATUS_SUCCESS) { | |
| return status; | |
| } | |
| i_buf = 1; | |
| } | |
| // Butterfly reduction: | |
| for (; offset_j >= 1; offset_j /= 2) { | |
| std::fill(step_cgraphs.begin(), step_cgraphs.end(), nullptr); | |
| for (size_t j = 0; j < 2*offset_j_max; j++) { | |
| const size_t j_other = j ^ offset_j; | |
| if (j_other >= n_backends) { | |
| continue; | |
| } | |
| push_data(j, j_other, i_buf); | |
| } | |
| for (size_t j = 0; j < 2*offset_j_max; j++) { | |
| if (step_cgraphs[j] == nullptr) { | |
| continue; | |
| } | |
| auto & bcj = backend_ctx->backend_configs[j]; | |
| const ggml_status status = ggml_backend_graph_compute_async(bcj.backend, step_cgraphs[j]); | |
| if (status != GGML_STATUS_SUCCESS) { | |
| return status; | |
| } | |
| } | |
| i_buf++; | |
| } | |
| assert(i_buf == backend_ctx->n_reduce_steps); | |
| // If n_backends is not a power of 2, copy back the reduced tensors to the excess: | |
| for (size_t j = 2*offset_j_max; j < n_backends; j++) { | |
| auto & bcj_src = backend_ctx->backend_configs[j - 2*offset_j_max]; | |
| auto & bcj_dst = backend_ctx->backend_configs[j]; | |
| ggml_tensor * node_src = bcj_src.cgraphs[i].cgraph_main->nodes[bcj_src.cgraphs[i].cgraph_main->n_nodes - 1]; | |
| ggml_tensor * node_dst = bcj_dst.cgraphs[i].cgraph_main->nodes[bcj_dst.cgraphs[i].cgraph_main->n_nodes - 1]; | |
| ggml_backend_tensor_copy_async(bcj_src.backend, bcj_dst.backend, node_src, node_dst); | |
| } | |
| return GGML_STATUS_SUCCESS; | |
| }; | |
| for (size_t i = 0; i < backend_ctx->n_subgraphs; i++) { | |
| for (size_t j = 0; j < n_backends; j++) { | |
| auto & bcj = backend_ctx->backend_configs[j]; | |
| const ggml_status status = ggml_backend_graph_compute_async(bcj.backend, bcj.cgraphs[i].cgraph_main); | |
| if (status != GGML_STATUS_SUCCESS) { | |
| return status; | |
| } | |
| } | |
| if (n_backends > 1 && i < backend_ctx->n_subgraphs - 1) { | |
| bool backend_allreduce_success = false; | |
| if (backend_ctx->comm_ctx) { | |
| std::vector<ggml_tensor *> nodes; | |
| nodes.reserve(n_backends); | |
| for (size_t j = 0; j < n_backends; j++) { | |
| auto & bcj = backend_ctx->backend_configs[j]; | |
| ggml_cgraph * cgraph_ij = bcj.cgraphs[i].cgraph_main; | |
| nodes.push_back(cgraph_ij->nodes[cgraph_ij->n_nodes-1]); | |
| } | |
| backend_allreduce_success = backend_ctx->comm_allreduce(backend_ctx->comm_ctx, nodes.data()); | |
| } | |
| if (!backend_allreduce_success) { | |
| const ggml_status status = allreduce_fallback(i); | |
| if (status != GGML_STATUS_SUCCESS) { | |
| return status; | |
| } | |
| } | |
| } | |
| } | |
| return GGML_STATUS_SUCCESS; | |
| } | |
| static const ggml_backend_i ggml_backend_meta_i = { | |
| /* .get_name = */ ggml_backend_meta_get_name, | |
| /* .free = */ ggml_backend_meta_free, | |
| /* .set_tensor_async = */ ggml_backend_meta_set_tensor_async, | |
| /* .get_tensor_async = */ ggml_backend_meta_get_tensor_async, | |
| /* .set_tensor_2d_async = */ nullptr, | |
| /* .get_tensor_2d_async = */ nullptr, | |
| /* .cpy_tensor_async = */ nullptr, | |
| /* .synchronize = */ ggml_backend_meta_synchronize, | |
| /* .graph_plan_create = */ nullptr, | |
| /* .graph_plan_free = */ nullptr, | |
| /* .graph_plan_update = */ nullptr, | |
| /* .graph_plan_compute = */ nullptr, | |
| /* .graph_compute = */ ggml_backend_meta_graph_compute, | |
| /* .event_record = */ nullptr, | |
| /* .event_wait = */ nullptr, | |
| /* .graph_optimize = */ nullptr, | |
| }; | |
| bool ggml_backend_is_meta(ggml_backend_t backend) { | |
| return backend != nullptr && backend->iface.get_name == ggml_backend_meta_i.get_name; | |
| } | |
| static ggml_backend_t ggml_backend_meta_device_init_backend(ggml_backend_dev_t dev, const char * params) { | |
| ggml_backend_meta_context * backend_ctx = new ggml_backend_meta_context(dev, params); | |
| ggml_backend_t backend = new struct ggml_backend; | |
| backend->guid = ggml_backend_meta_guid(); | |
| backend->iface = ggml_backend_meta_i; | |
| backend->device = dev; | |
| backend->context = backend_ctx; | |
| return backend; | |
| } | |
| size_t ggml_backend_meta_n_backends(ggml_backend_t meta_backend) { | |
| GGML_ASSERT(ggml_backend_is_meta(meta_backend)); | |
| const ggml_backend_meta_context * backend_ctx = (const ggml_backend_meta_context *) meta_backend->context; | |
| return backend_ctx->backend_configs.size(); | |
| } | |
| ggml_backend_t ggml_backend_meta_simple_backend(ggml_backend_t meta_backend, size_t index) { | |
| GGML_ASSERT(ggml_backend_is_meta(meta_backend)); | |
| const ggml_backend_meta_context * backend_ctx = (const ggml_backend_meta_context *) meta_backend->context; | |
| return backend_ctx->backend_configs[index].backend; | |
| } | |