Add session-owned PreparedGraph execution and scheduler support

CMakeLists.txt

@@ -534,6 +534,7 @@ add_library(engine_core OBJECT
     src/framework/core/deferred_tensor_writer.cpp
     src/framework/core/disk_expert_weight_store.cpp
     src/framework/core/graph_workspace.cpp
+    src/framework/core/prepared_graph.cpp
     src/framework/core/execution_context.cpp
     src/framework/core/host_memory.cpp
     src/framework/core/load_progress.cpp
@@ -3802,6 +3803,10 @@ if (ENGINE_BUILD_TESTS OR ENGINE_BUILD_EXTENDED_TESTS OR ENGINE_BUILD_MODEL_TEST
 
     add_engine_unittest(ctc_host_decoder_test tests/unittests/test_ctc_host_decoder.cpp)
     add_test(NAME ctc_host_decoder_test COMMAND ctc_host_decoder_test)
+    add_engine_unittest(prepared_graph_test tests/unittests/test_prepared_graph.cpp)
+    target_include_directories(prepared_graph_test PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/external/ggml/src)
+    add_test(NAME prepared_graph_test COMMAND prepared_graph_test
+        --log ${CMAKE_CURRENT_BINARY_DIR}/prepared_graph_test.log)
     add_engine_unittest(rope_transformer_encoder_test tests/unittests/test_rope_transformer_encoder.cpp)
     add_test(NAME rope_transformer_encoder_test COMMAND rope_transformer_encoder_test)
     add_engine_unittest(shared_weight_cache_test tests/unittests/test_shared_weight_cache.cpp)
@@ -4669,6 +4674,7 @@ if (ENGINE_BUILD_TESTS OR ENGINE_BUILD_EXTENDED_TESTS OR ENGINE_BUILD_MODEL_TEST
 
         if(fish_drama3 IN_LIST AUDIOCPP_LINKED_MODELS)
             add_engine_unittest(fish_drama3_backbone_allocation_test tests/community_models/fish_drama3/backbone_allocation_test.cpp)
+            target_include_directories(fish_drama3_backbone_allocation_test PRIVATE external/ggml/src)
             add_test(NAME fish_drama3_backbone_allocation_test COMMAND fish_drama3_backbone_allocation_test)
             add_engine_unittest(fish_drama3_codec_source_test tests/community_models/fish_drama3/codec_source_test.cpp)
             add_test(NAME fish_drama3_codec_source_test COMMAND fish_drama3_codec_source_test)

external/ggml/src/ggml-backend.cpp

@@ -1558,7 +1558,9 @@ static bool ggml_backend_sched_alloc_splits(ggml_backend_sched_t sched) {
             ggml_backend_synchronize(sched->backends[i]);
         }
 
-        ggml_gallocr_reserve_n(sched->galloc, &sched->graph, sched->node_backend_ids, sched->leaf_backend_ids);
+        if (!ggml_gallocr_reserve_n(sched->galloc, &sched->graph, sched->node_backend_ids, sched->leaf_backend_ids)) {
+            return false;
+        }
         if (!ggml_gallocr_alloc_graph(sched->galloc, &sched->graph)) {
             GGML_LOG_ERROR("%s: failed to allocate graph\n", __func__);
             return false;

include/engine/framework/core/graph_workspace.h

@@ -8,6 +8,8 @@
 
 namespace engine::core {
 
+class PreparedGraph;
+
 // Shared activation allocation for graphs that execute serially. Graph owners
 // retain their contexts, plans and backend captures; this owns only scratch.
 class GraphWorkspace {
@@ -17,16 +19,22 @@ class GraphWorkspace {
     GraphWorkspace(const GraphWorkspace &) = delete;
     GraphWorkspace & operator=(const GraphWorkspace &) = delete;
     void add(ggml_cgraph * graph);
+    void add(PreparedGraph & graph);
     // Register every graph before preparing; allocation addresses remain fixed
     // afterward. Graphs must stay alive through this call.
     void prepare();
-    // Start a new graph plan without dropping allocation capacity. The caller
-    // must release the previous graphs and their backend resources first.
+    // Start a new graph plan without dropping allocation capacity. Registered
+    // PreparedGraphs are drained, invalidated and detached; reset them before
+    // rebuilding. Raw graph callers must release their graphs/resources first.
     void reset();
     size_t bytes() const;
 private:
+    friend class PreparedGraph;
+    void forget(PreparedGraph & graph);
+    ggml_backend_t backend_;
     std::unique_ptr<std::remove_pointer_t<ggml_gallocr_t>, decltype(&ggml_gallocr_free)> allocator_;
     std::vector<ggml_cgraph *> graphs_;
+    std::vector<PreparedGraph *> prepared_graphs_;
     bool prepared_ = false;
 };
 

include/engine/framework/core/prepared_graph.h

@@ -0,0 +1,108 @@
+#pragma once
+
+#include "engine/framework/core/execution_context.h"
+
+#include <ggml-alloc.h>
+
+#include <memory>
+#include <string>
+#include <vector>
+
+namespace engine::core {
+
+class GraphWorkspace;
+
+struct PreparedGraphConfig {
+    std::string label;
+    size_t metadata_bytes = 0;
+    size_t graph_nodes = 0;
+    size_t persistent_metadata_bytes = 0;
+};
+
+// Session-owned graph set. All graphs share metadata and execute serially using
+// one scratch workspace. Persistent tensors have separate metadata/storage and
+// survive reset(). ExecutionContext must outlive this owner. Graph IDs, tensor
+// handles and build contexts for transient metadata are invalid after reset().
+// Preallocated external tensors can also be referenced; their context/buffer
+// must outlive every graph using them. This owner never frees external storage.
+class PreparedGraph {
+public:
+    using GraphId = size_t;
+
+    PreparedGraph(const ExecutionContext & execution, PreparedGraphConfig config);
+    ~PreparedGraph();
+    PreparedGraph(const PreparedGraph &) = delete;
+    PreparedGraph & operator=(const PreparedGraph &) = delete;
+
+    ModuleBuildContext & build_context() noexcept { return build_context_; }
+    // Graph zero exists at construction/reset; additional graphs share its context.
+    GraphId add_graph(std::string name, size_t graph_nodes = 0);
+    // An execution-only entry point shares the allocation graph's tensors and
+    // layout. Its nodes must be an ordered subset of that allocation graph.
+    // Build both graphs before allocation; only the allocation graph is allocated.
+    GraphId add_entry_point(std::string name, GraphId allocation_graph, size_t graph_nodes = 0);
+    GraphId allocation_graph(GraphId id = 0) const;
+    ggml_cgraph * graph(GraphId id = 0) const;
+    size_t graph_count() const noexcept { return graphs_.size(); }
+    const ExecutionContext & execution() const noexcept { return execution_; }
+    const std::string & label() const noexcept { return config_.label; }
+    bool allocated() const noexcept { return allocated_; }
+    bool has_host_plan(GraphId id = 0) const;
+
+    // Declare persistent inputs/state before allocation. Initialize their values
+    // with the existing tensor-write APIs; allocation never initializes state.
+    // Use these tensors across independent allocation graphs. Entry points of
+    // the same allocation graph may share its transient tensors directly.
+    ModuleBuildContext & persistent_build_context();
+    void allocate_persistent();
+    size_t persistent_bytes() const noexcept;
+    size_t workspace_bytes() const noexcept;
+    // Explicit zero fill, not a model-specific state initializer. Drains pending
+    // work first; graph rebuild/reset never changes persistent tensor contents.
+    void clear_persistent();
+
+    // Finish graph construction before calling once. Prepares every graph and
+    // reserves shared scratch. Structure/addresses stay fixed until reset().
+    // Persistent tensors must already be allocated before preparing the graphs.
+    void allocate();
+    // Optional single-graph scheduler execution. The returned handle is borrowed:
+    // configure placement/evaluation with GGML's APIs before allocate(). Backend
+    // handles (primary first, CPU last) and callback data must outlive this owner. A scheduler
+    // owns its split allocations, so it cannot share a GraphWorkspace or entry points.
+    ggml_backend_sched_t create_scheduler(std::vector<ggml_backend_t> backends);
+    void compute(GraphId id = 0);
+    // Drain work, release plans/captures and rebuild empty transient metadata.
+    // Retains persistent state and owned scratch capacity. For an external
+    // workspace, reset that workspace before resetting this owner.
+    void reset();
+    // Drain work and free all resources, including persistent state and scratch.
+    // Idempotent; call reset() to start building again after release().
+    void release();
+
+private:
+    friend class GraphWorkspace;
+    struct GraphEntry;
+    void validate_graphs() const;
+    void prepare_shared_allocation();
+    void release_graphs();
+    void synchronize();
+
+    const ExecutionContext & execution_;
+    PreparedGraphConfig config_;
+    std::unique_ptr<ggml_context, decltype(&ggml_free)> context_{nullptr, ggml_free};
+    std::unique_ptr<ggml_context, decltype(&ggml_free)> persistent_context_{nullptr, ggml_free};
+    std::unique_ptr<ggml_backend_buffer, decltype(&ggml_backend_buffer_free)> persistent_buffer_{nullptr, ggml_backend_buffer_free};
+    ModuleBuildContext build_context_;
+    ModuleBuildContext persistent_build_context_;
+    std::vector<std::unique_ptr<GraphEntry>> graphs_;
+    std::unique_ptr<GraphWorkspace> owned_workspace_;
+    GraphWorkspace * workspace_ = nullptr;
+    std::vector<ggml_backend_t> scheduler_backends_;
+    std::unique_ptr<ggml_backend_sched, decltype(&ggml_backend_sched_free)> scheduler_{nullptr, ggml_backend_sched_free};
+    std::unique_ptr<ggml_context, decltype(&ggml_free)> cleanup_context_{nullptr, ggml_free};
+    ggml_cgraph * cleanup_graph_ = nullptr;
+    bool allocation_started_ = false;
+    bool allocated_ = false;
+};
+
+}  // namespace engine::core