stream_ordered_timing.hpp

 
#pragma once
 
#include <memory>
#include <optional>
#include <string>
 
#include <rmm/cuda_stream_view.hpp>
 
#include <rapidsmpf/statistics.hpp>
 
namespace rapidsmpf {
 
class StreamOrderedTiming {
  public:
    StreamOrderedTiming(
        rmm::cuda_stream_view stream, std::shared_ptr<Statistics> statistics
    );
 
    void stop_and_record(
        std::string const& name,
        std::optional<std::string> stream_delay_name = std::nullopt
    );
 
    static void cancel_inflight_timings(Statistics const* statistics);
 
  private:
    std::uintptr_t uid_{0};
    rmm::cuda_stream_view stream_;
    std::shared_ptr<Statistics> statistics_;
};
 
}  // namespace rapidsmpf