events.hpp

Go to the documentation of this file.

#ifndef KOKKOS_EXECUTION_GRAPH_EVENTS_HPP
#define KOKKOS_EXECUTION_GRAPH_EVENTS_HPP
 
#include "Kokkos_Core.hpp"
#include "Kokkos_Graph.hpp"
 
#if defined(KOKKOS_EXECUTION_ENABLE_EVENT_DISPATCH)
#    include "kokkos-utils/callbacks/Manager.hpp"
#endif
 
#include "kokkos-execution/utils/streamers.hpp"
 
namespace Kokkos::Execution::GraphImpl {

struct GraphCreateEvent {
    void* graph = nullptr;
    uint32_t dev_id = 0;
    uint64_t event_id = 0;
 
    constexpr auto operator<=>(const GraphCreateEvent&) const = default;
 
    friend std::ostream& operator<<(std::ostream& out, const GraphCreateEvent& event) {
        return out << "GraphCreateEvent: {graph = " << event.graph << ", dev_id = " << event.dev_id
                   << ", event_id = " << event.event_id << '}';
    }
};

struct GraphAddNodeEvent {
    void* graph = nullptr;
    void* predecessor = nullptr;
    void* node = nullptr;
    uint32_t dev_id = 0;
 
    constexpr auto operator<=>(const GraphAddNodeEvent&) const = default;
 
    friend std::ostream& operator<<(std::ostream& out, const GraphAddNodeEvent& event) {
        return out << "GraphAddNodeEvent: {graph = " << event.graph << ", predecessor = " << event.predecessor
                   << ", node = " << event.node << ", dev_id = " << event.dev_id << '}';
    }
};

struct GraphAddAggregateNodeEvent {
    void* graph = nullptr;
    std::vector<void*> predecessors{};
    void* node = nullptr;
 
    constexpr auto operator<=>(const GraphAddAggregateNodeEvent&) const = default;
 
    friend std::ostream& operator<<(std::ostream& out, const GraphAddAggregateNodeEvent& event) {
        using Kokkos::Execution::Utils::operator<<;
        return out << "GraphAddAggregateNodeEvent: {graph = " << event.graph
                   << ", predecessors = " << event.predecessors << ", node = " << event.node << '}';
    }
};

struct GraphInstantiateEvent {
    void* graph = nullptr;
 
    constexpr auto operator<=>(const GraphInstantiateEvent&) const = default;
 
    friend std::ostream& operator<<(std::ostream& out, const GraphInstantiateEvent& event) {
        return out << "GraphInstantiateEvent: {graph = " << event.graph << '}';
    }
};

struct GraphSubmitEvent {
    void* graph = nullptr;
    uint32_t dev_id = 0;
 
    constexpr auto operator<=>(const GraphSubmitEvent&) const = default;
 
    friend std::ostream& operator<<(std::ostream& out, const GraphSubmitEvent& event) {
        return out << "GraphSubmitEvent: {graph = " << event.graph << ", dev_id = " << event.dev_id << '}';
    }
};

template <typename T>
concept NodeRef = Kokkos::Impl::is_specialization_of_v<T, Kokkos::Experimental::GraphNodeRef>;

template <NodeRef NodeType>
auto* get_graph_impl_ptr(const NodeType& node) noexcept {
    return Kokkos::Impl::GraphAccess::get_graph_weak_ptr(node).lock().get();
}

template <NodeRef NodeType>
auto* get_node_ptr(const NodeType& node) noexcept {
    return Kokkos::Impl::GraphAccess::get_node_ptr(node).get();
}

template <Kokkos::ExecutionSpace Exec>
void graph_create_event(const Kokkos::Experimental::Graph<Exec>& graph) {
#if defined(KOKKOS_EXECUTION_ENABLE_EVENT_DISPATCH)
    Kokkos::utils::callbacks::dispatch(
        GraphCreateEvent{
            .graph = get_graph_impl_ptr(graph.root_node()),
            .dev_id = Kokkos::Tools::Experimental::device_id(graph.get_device_handle().m_exec),
            .event_id = Kokkos::utils::callbacks::get_next_event_id()});
#endif
}

template <Kokkos::ExecutionSpace Exec, typename... Args>
auto create_graph(const Kokkos::Impl::DeviceHandle<Exec>& device_handle, Args&&... args) {
    Kokkos::Experimental::Graph<Exec> graph{device_handle, std::forward<Args>(args)...};
    graph_create_event(graph);
    return graph;
}

template <NodeRef Predecessor, NodeRef NodeType, Kokkos::ExecutionSpace Exec>
void graph_add_node_event(
    const Predecessor& predecessor,
    const NodeType& node,
    const Kokkos::Impl::DeviceHandle<Exec>& device_handle) {
#if defined(KOKKOS_EXECUTION_ENABLE_EVENT_DISPATCH)
    const bool is_predecessor_root = predecessor.get_node_kind() == Kokkos::Experimental::GraphNodeKind::Root;
 
    Kokkos::utils::callbacks::dispatch(
        GraphAddNodeEvent{
            .graph = get_graph_impl_ptr(predecessor),
            .predecessor = is_predecessor_root ? nullptr : get_node_ptr(predecessor),
            .node = get_node_ptr(node),
            .dev_id = Kokkos::Tools::Experimental::device_id(device_handle.m_exec)});
#endif
}

template <NodeRef NodeType, NodeRef... Predecessors>
void graph_add_aggregate_node_event(const NodeType& aggregate, const Predecessors&... predecessors) {
#if defined(KOKKOS_EXECUTION_ENABLE_EVENT_DISPATCH)
    Kokkos::utils::callbacks::dispatch(
        GraphAddAggregateNodeEvent{
            .graph = get_graph_impl_ptr(aggregate),
            .predecessors = {get_node_ptr(predecessors)...},
            .node = get_node_ptr(aggregate)});
#endif
}

template <Kokkos::ExecutionSpace Exec>
void graph_instantiate_event(const Kokkos::Experimental::Graph<Exec>& graph) {
#if defined(KOKKOS_EXECUTION_ENABLE_EVENT_DISPATCH)
    Kokkos::utils::callbacks::dispatch(GraphInstantiateEvent{.graph = get_graph_impl_ptr(graph.root_node())});
#endif
}
 

template <Kokkos::ExecutionSpace Exec>
void graph_submit_event(const Kokkos::Experimental::Graph<Exec>& graph, const Exec& exec) {
#if defined(KOKKOS_EXECUTION_ENABLE_EVENT_DISPATCH)
    Kokkos::utils::callbacks::dispatch(
        GraphSubmitEvent{
            .graph = get_graph_impl_ptr(graph.root_node()), .dev_id = Kokkos::Tools::Experimental::device_id(exec)});
#endif
}

template <Kokkos::ExecutionSpace Exec>
void submit_graph(const Kokkos::Experimental::Graph<Exec>& graph, const Exec& exec) {
    graph_submit_event(graph, exec);
    graph.submit(exec);
}
 
} // namespace Kokkos::Execution::GraphImpl
 
#endif // KOKKOS_EXECUTION_GRAPH_EVENTS_HPP