events.hpp

Go to the documentation of this file.

#ifndef KOKKOS_EXECUTION_GRAPH_EVENTS_HPP
#define KOKKOS_EXECUTION_GRAPH_EVENTS_HPP
 
#include "Kokkos_Macros.hpp"
#ifdef KOKKOS_ENABLE_EXPERIMENTAL_CXX20_MODULES
import kokkos.core;
import kokkos.graph;
import kokkos.graph_impl;
#else
#    include "Kokkos_Core.hpp"
#    include "Kokkos_Graph.hpp"
#endif
 
#if defined(KOKKOS_EXECUTION_ENABLE_EVENT_DISPATCH)
#    include "kokkos-utils/callbacks/Manager.hpp"
#endif
 
#include "kokkos-execution/impl/label.hpp"
#include "kokkos-execution/utils/streamers.hpp"
 
namespace Kokkos::Execution::GraphImpl {
 
#if defined(__cpp_impl_reflection)
using Utils::operator<<;
#endif

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

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

struct KOKKOS_EXECUTION_UTILS_STREAMABLE GraphAddAggregateNodeEvent {
    const void* graph = nullptr;
    std::vector<const void*> predecessors{};
    const void* node = nullptr;
 
    constexpr auto operator<=>(const GraphAddAggregateNodeEvent&) const = default;
 
#if !defined(__cpp_impl_reflection)
    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 << '}';
    }
#endif
};

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

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

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)
    const auto root = graph.root_node();
    Kokkos::utils::callbacks::dispatch(
        GraphCreateEvent{
            .graph = get_graph_impl_ptr(root),
            .root_node = get_node_ptr(root),
            .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 <typename Label, NodeRef Predecessor, NodeRef NodeType>
void graph_add_node_event(Label&& label, const Predecessor& predecessor, const NodeType& node) {
#if defined(KOKKOS_EXECUTION_ENABLE_EVENT_DISPATCH)
    const auto* node_ptr = get_node_ptr(node);
    Kokkos::utils::callbacks::dispatch(
        GraphAddNodeEvent{
            .name = Impl::as_convertible_to_string(std::forward<Label>(label)),
            .graph = get_graph_impl_ptr(predecessor),
            .predecessor = get_node_ptr(predecessor),
            .node = node_ptr,
            .dev_id = Kokkos::Tools::Experimental::device_id(node_ptr->get_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