compute_graph.h

#pragma once
#include <map>
#include <functional>
#include <set>
#include <muda/launch/stream.h>
#include <muda/launch/event.h>
#include <muda/mstl/span.h>
#include <muda/graph/graph.h>
#include <muda/graph/graph_viewer.h>
#include <muda/compute_graph/compute_graph_flag.h>
#include <muda/compute_graph/compute_graph_phase.h>
#include <muda/compute_graph/compute_graph_node_type.h>
#include <muda/compute_graph/compute_graph_node_id.h>
#include <muda/compute_graph/compute_graph_closure_id.h>
#include <muda/compute_graph/compute_graph_var_id.h>
#include <muda/compute_graph/compute_graph_var_usage.h>
#include <muda/compute_graph/compute_graph_dependency.h>
#include <muda/compute_graph/graphviz_options.h>
#include <muda/compute_graph/compute_graph_fwd.h>
 
namespace muda
{
namespace details
{
    class LocalVarId : public U64IdWithType
    {
        using U64IdWithType::U64IdWithType;
    };
    class LocalVarInfo
    {
      public:
        LocalVarId           id{};
        ComputeGraphVarBase* var = nullptr;
    };
}  // namespace details
 
class ComputeGraph
{
  public:
    class AddNodeProxy
    {
        ComputeGraph& m_cg;
        std::string   m_node_name;
 
      public:
        AddNodeProxy(ComputeGraph& cg, std::string_view node_name);
        ComputeGraph& operator<<(std::function<void()>&& f) &&;
    };
    // A depends on B : from B to A
    using Dependency = ComputeGraphDependency;
 
    class GraphPhaseGuard
    {
        ComputeGraph& m_cg;
 
      public:
        GraphPhaseGuard(ComputeGraph& cg, ComputeGraphPhase phase);
        ~GraphPhaseGuard();
    };
 
    // delete copy
    ComputeGraph(const ComputeGraph&)            = delete;
    ComputeGraph& operator=(const ComputeGraph&) = delete;
 
    // delete move
    ComputeGraph(ComputeGraph&&)            = delete;
    ComputeGraph& operator=(ComputeGraph&&) = delete;
 
  private:
    //class TempNodeInfo
    //{
    //  public:
    //    std::map<VarId, ComputeGraphVarUsage> var_usage;
    //};
    template <typename T>
    using U = std::unique_ptr<T>;
    template <typename T>
    using S = std::shared_ptr<T>;
 
    friend class ComputeGraphVarBase;
 
    Graph        m_graph;
    S<GraphExec> m_graph_exec{nullptr};
 
    std::unordered_map<NodeId::value_type, cudaGraph_t> m_sub_graphs;
 
    std::vector<std::pair<std::string, ComputeGraphClosure*>> m_closures;
 
    std::map<VarId, details::LocalVarId> m_global_to_local_var_id;
    std::vector<details::LocalVarInfo>   m_related_vars;
    void emplace_related_var(ComputeGraphVarBase* var);
 
 
    std::vector<ComputeGraphNodeBase*>              m_nodes;
    std::vector<std::vector<ComputeGraphNodeBase*>> m_graph_nodes;
    std::vector<Dependency>                         m_deps;
 
    std::vector<int>        m_closure_need_update;
    ComputeGraphVarManager* m_var_manager = nullptr;
 
    friend class ComputeGraphVarManager;
 
    Event                      m_event;
    mutable Event::QueryResult m_event_result = Event::QueryResult::eFinished;
    Flags<GraphInstantiateFlagBit> m_flags;
 
  public:
    ComputeGraph(ComputeGraphVarManager& manager,
                 std::string_view        name = "graph",
                 ComputeGraphFlag        flag = ComputeGraphFlag::HostLaunch);
 
    ~ComputeGraph();
 
    /**************************************************************
    * 
    * Info API
    * 
    ***************************************************************/
 
    std::string_view name() const { return m_name; }
 
    /**************************************************************
    * 
    * GraphNode API
    * 
    ***************************************************************/
 
    AddNodeProxy create_node(std::string_view node_name);
 
 
    /**************************************************************
    * 
    * Graph Launch API
    * 
    ***************************************************************/
 
    void update();
 
    void build();
 
    void launch(bool single_stream, cudaStream_t s = nullptr);
 
    void launch(cudaStream_t s = nullptr) { return launch(false, s); }
 
    /**************************************************************
    * 
    * Graph Event Query API
    * 
    ***************************************************************/
 
    Event::QueryResult query() const;
 
    /**************************************************************
    * 
    * Graph Closure Capture Node API
    * 
    ***************************************************************/
 
    void capture(std::function<void(cudaStream_t)>&& f);
    void capture(std::string_view name, std::function<void(cudaStream_t)>&& f);
 
    /**************************************************************
    * 
    * Graph Visualization API
    * 
    ***************************************************************/
 
    void graphviz(std::ostream& o, const ComputeGraphGraphvizOptions& options = {});
 
    /**************************************************************
    * 
    * Graph Viewer API
    * 
    ***************************************************************/
 
    GraphViewer viewer();
 
    operator GraphViewer() { return viewer(); }
 
  private:  // internal method
    void topo_build();
 
    void cuda_graph_add_deps();
 
    void build_deps();
 
    void serial_launch();
 
    void _update();
 
    void check_vars_valid();
 
    friend class AddNodeProxy;
    ComputeGraph& add_node(std::string&& name, const std::function<void()>& f);
 
    friend class ComputeGraphNodeBase;
    friend class ComputeGraphClosure;
    span<const Dependency> dep_span(size_t begin, size_t count) const;
 
    void set_current_graph_as_this();
 
    static void clear_current_graph();
 
    static Stream& shared_capture_stream();
 
    friend class ComputeGraphBuilder;
    ClosureId current_closure_id() const { return m_current_closure_id; };
 
    NodeId current_node_id() const { return m_current_node_id; };
 
    size_t current_access_index() const { return m_access_graph_index; }
 
    ComputeGraphPhase current_graph_phase() const;
 
  private:  // internal data
    friend class muda::details::ComputeGraphAccessor;
    std::string       m_name;
    bool              m_need_update = false;
    ClosureId         m_current_closure_id;
    NodeId            m_current_node_id;
    ComputeGraphPhase m_current_graph_phase = ComputeGraphPhase::None;
    bool              m_allow_access_graph  = false;
    size_t            m_access_graph_index  = 0;
    bool              m_allow_node_adding   = true;
    // TempNodeInfo      m_temp_node_info;
    cudaStream_t m_current_single_stream = nullptr;
    bool         m_is_capturing          = false;
    // in capture func, we don't allow any var eval()
    bool m_is_in_capture_func = false;
    // if we have already built the topo, we don't do that again
    bool m_is_topo_built = false;
};
}  // namespace muda
 
#include "details/compute_graph.inl"