profiler_hook.hpp

// SPDX-FileCopyrightText: 2017 - 2024 The Ginkgo authors
//
// SPDX-License-Identifier: BSD-3-Clause
 
#ifndef GKO_PUBLIC_CORE_LOG_PROFILER_HOOK_HPP_
#define GKO_PUBLIC_CORE_LOG_PROFILER_HOOK_HPP_
 
 
#include <iostream>
#include <unordered_map>
 
#include <ginkgo/config.hpp>
#include <ginkgo/core/base/timer.hpp>
#include <ginkgo/core/log/logger.hpp>
 
 
namespace gko {
namespace log {
 
 
enum class profile_event_category {
    memory,
    operation,
    object,
    linop,
    factory,
    solver,
    criterion,
    user,
    internal,
};
 
 
class profiling_scope_guard;
 
 
class ProfilerHook : public Logger {
public:
    using hook_function =
        std::function<void(const char*, profile_event_category)>;
 
    void on_allocation_started(const gko::Executor* exec,
                               const gko::size_type&) const override;
 
    void on_allocation_completed(const gko::Executor* exec,
                                 const gko::size_type&,
                                 const gko::uintptr&) const override;
 
    void on_free_started(const gko::Executor* exec,
                         const gko::uintptr&) const override;
 
    void on_free_completed(const gko::Executor* exec,
                           const gko::uintptr&) const override;
 
    void on_copy_started(const gko::Executor* from, const gko::Executor* to,
                         const gko::uintptr&, const gko::uintptr&,
                         const gko::size_type&) const override;
 
    void on_copy_completed(const gko::Executor* from, const gko::Executor* to,
                           const gko::uintptr&, const gko::uintptr&,
                           const gko::size_type&) const override;
 
    /* Operation events */
    void on_operation_launched(const Executor* exec,
                               const Operation* operation) const override;
 
    void on_operation_completed(const Executor* exec,
                                const Operation* operation) const override;
 
    /* PolymorphicObject events */
    void on_polymorphic_object_copy_started(
        const Executor* exec, const PolymorphicObject* from,
        const PolymorphicObject* to) const override;
 
    void on_polymorphic_object_copy_completed(
        const Executor* exec, const PolymorphicObject* from,
        const PolymorphicObject* to) const override;
 
    void on_polymorphic_object_move_started(
        const Executor* exec, const PolymorphicObject* from,
        const PolymorphicObject* to) const override;
 
    void on_polymorphic_object_move_completed(
        const Executor* exec, const PolymorphicObject* from,
        const PolymorphicObject* to) const override;
 
    /* LinOp events */
    void on_linop_apply_started(const LinOp* A, const LinOp* b,
                                const LinOp* x) const override;
 
    void on_linop_apply_completed(const LinOp* A, const LinOp* b,
                                  const LinOp* x) const override;
 
    void on_linop_advanced_apply_started(const LinOp* A, const LinOp* alpha,
                                         const LinOp* b, const LinOp* beta,
                                         const LinOp* x) const override;
 
    void on_linop_advanced_apply_completed(const LinOp* A, const LinOp* alpha,
                                           const LinOp* b, const LinOp* beta,
                                           const LinOp* x) const override;
 
    /* LinOpFactory events */
    void on_linop_factory_generate_started(const LinOpFactory* factory,
                                           const LinOp* input) const override;
 
    void on_linop_factory_generate_completed(
        const LinOpFactory* factory, const LinOp* input,
        const LinOp* output) const override;
 
    /* Criterion events */
    void on_criterion_check_started(const stop::Criterion* criterion,
                                    const size_type& num_iterations,
                                    const LinOp* residual,
                                    const LinOp* residual_norm,
                                    const LinOp* solution,
                                    const uint8& stopping_id,
                                    const bool& set_finalized) const override;
 
    void on_criterion_check_completed(
        const stop::Criterion* criterion, const size_type& num_iterations,
        const LinOp* residual, const LinOp* residual_norm,
        const LinOp* solution, const uint8& stopping_id,
        const bool& set_finalized, const array<stopping_status>* status,
        const bool& one_changed, const bool& all_stopped) const override;
 
    void on_criterion_check_completed(
        const stop::Criterion* criterion, const size_type& num_iterations,
        const LinOp* residual, const LinOp* residual_norm,
        const LinOp* implicit_sq_resnorm, const LinOp* solution,
        const uint8& stopping_id, const bool& set_finalized,
        const array<stopping_status>* status, const bool& one_changed,
        const bool& all_stopped) const override;
 
    /* Internal solver events */
    void on_iteration_complete(
        const LinOp* solver, const LinOp* right_hand_side,
        const LinOp* solution, const size_type& num_iterations,
        const LinOp* residual, const LinOp* residual_norm,
        const LinOp* implicit_sq_residual_norm,
        const array<stopping_status>* status, bool stopped) const override;
 
    GKO_DEPRECATED(
        "Please use the version with the additional stopping "
        "information.")
    void on_iteration_complete(const LinOp* solver,
                               const size_type& num_iterations,
                               const LinOp* residual, const LinOp* solution,
                               const LinOp* residual_norm) const override;
 
    GKO_DEPRECATED(
        "Please use the version with the additional stopping "
        "information.")
    void on_iteration_complete(
        const LinOp* solver, const size_type& num_iterations,
        const LinOp* residual, const LinOp* solution,
        const LinOp* residual_norm,
        const LinOp* implicit_sq_residual_norm) const override;
 
    bool needs_propagation() const override;
 
    void set_object_name(ptr_param<const PolymorphicObject> obj,
                         std::string name);
 
    void set_synchronization(bool synchronize);
 
    profiling_scope_guard user_range(const char* name) const;
 
    constexpr static uint32 color_yellow_argb = 0xFFFFCB05U;
 
    static std::shared_ptr<ProfilerHook> create_tau(bool initialize = true);
 
    static std::shared_ptr<ProfilerHook> create_vtune();
 
    static std::shared_ptr<ProfilerHook> create_nvtx(
        uint32 color_argb = color_yellow_argb);
 
    static std::shared_ptr<ProfilerHook> create_roctx();
 
    static std::shared_ptr<ProfilerHook> create_for_executor(
        std::shared_ptr<const Executor> exec);
 
    struct summary_entry {
        std::string name;
        std::chrono::nanoseconds inclusive{0};
        std::chrono::nanoseconds exclusive{0};
        int64 count{};
    };
 
    struct nested_summary_entry {
        std::string name;
        std::chrono::nanoseconds elapsed{0};
        int64 count{};
        std::vector<nested_summary_entry> children{};
    };
 
    class SummaryWriter {
    public:
        virtual ~SummaryWriter() = default;
 
        virtual void write(const std::vector<summary_entry>& entries,
                           std::chrono::nanoseconds overhead) = 0;
    };
 
    class NestedSummaryWriter {
    public:
        virtual ~NestedSummaryWriter() = default;
 
        virtual void write_nested(const nested_summary_entry& root,
                                  std::chrono::nanoseconds overhead) = 0;
    };
 
    class TableSummaryWriter : public SummaryWriter,
                               public NestedSummaryWriter {
    public:
        TableSummaryWriter(std::ostream& output = std::cerr,
                           std::string header = "Runtime summary");
 
        void write(const std::vector<summary_entry>& entries,
                   std::chrono::nanoseconds overhead) override;
 
        void write_nested(const nested_summary_entry& root,
                          std::chrono::nanoseconds overhead) override;
 
    private:
        std::ostream* output_;
        std::string header_;
    };
 
    static std::shared_ptr<ProfilerHook> create_summary(
        std::shared_ptr<Timer> timer = std::make_shared<CpuTimer>(),
        std::unique_ptr<SummaryWriter> writer =
            std::make_unique<TableSummaryWriter>(),
        bool debug_check_nesting = false);
 
    static std::shared_ptr<ProfilerHook> create_nested_summary(
        std::shared_ptr<Timer> timer = std::make_shared<CpuTimer>(),
        std::unique_ptr<NestedSummaryWriter> writer =
            std::make_unique<TableSummaryWriter>(),
        bool debug_check_nesting = false);
 
    static std::shared_ptr<ProfilerHook> create_custom(hook_function begin,
                                                       hook_function end);
 
private:
    ProfilerHook(hook_function begin, hook_function end);
 
    void maybe_synchronize(const Executor* exec) const;
 
    std::string stringify_object(const PolymorphicObject* obj) const;
 
    std::unordered_map<const PolymorphicObject*, std::string> name_map_;
    bool synchronize_;
    hook_function begin_hook_;
    hook_function end_hook_;
};
 
 
class profiling_scope_guard {
public:
    profiling_scope_guard();
 
    profiling_scope_guard(const char* name, profile_event_category category,
                          ProfilerHook::hook_function begin,
                          ProfilerHook::hook_function end);
 
    ~profiling_scope_guard();
 
    profiling_scope_guard(const profiling_scope_guard&) = delete;
 
    profiling_scope_guard(profiling_scope_guard&& other) = delete;
 
    profiling_scope_guard& operator=(const profiling_scope_guard&) = delete;
 
    profiling_scope_guard& operator=(profiling_scope_guard&&) = delete;
 
private:
    bool empty_;
    const char* name_;
    profile_event_category category_;
    ProfilerHook::hook_function end_;
};
 
 
}  // namespace log
}  // namespace gko
 
 
#endif  // GKO_PUBLIC_CORE_LOG_PROFILER_HOOK_HPP_