InfiniTrain/infini_train/include/profiler.h at 0df5c998fc93790aac6faab36b742e45ae019fe4 · InfiniTensor/InfiniTrain · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
#pragma once

#include <chrono>
#include <functional>
#include <iostream>
#include <map>
#include <mutex>
#include <string>
#include <vector>

#include "glog/logging.h"

#include "infini_train/include/device.h"

namespace infini_train {
namespace core {
class Event;
}

inline thread_local int tls_profiling_depth = 0;

struct ProfileContext {
    std::string name;
    Device::DeviceType device;
};

inline thread_local ProfileContext tls_profile_context;

inline void SetProfileContext(const std::string &name, Device::DeviceType device) {
    if (tls_profiling_depth == 0) {
        tls_profile_context.name = name;
        tls_profile_context.device = device;
    }
}

inline const ProfileContext &GetProfileContext() { return tls_profile_context; }

struct KernelProfileInfo {
    int64_t host_total_us = 0;
    int64_t device_total_us = 0;
    int count = 0;
};

struct KernelCallRecord {
    std::string tag;
    std::string timestamp;
    int64_t rank;
    std::string name;
    std::string device;
    int64_t host_us;
    int64_t device_us;
    int64_t max_device_mem_usage_mb;
};

class Profiler {
public:
    enum class SortBy {
        HostTimeTotal,
        HostTimePercentage,
        HostTimeAverage,
        DeviceTimeTotal,
        DeviceTimePercentage,
        DeviceTimeAverage,
        Count,
        NotSorted
    };

    static Profiler &Instance();

    void StartRecord(const std::string &name, Device::DeviceType device);
    void EndRecord(const std::string &name, Device::DeviceType device);

    void Report(std::ostream &os = std::cout, SortBy sort_by = SortBy::NotSorted) const;
    void Report(const std::string &file_path, SortBy sort_by = SortBy::NotSorted) const;
    void PrintRecords(std::ostream &os = std::cout) const;
    void PrintRecords(const std::string &file_path) const;

    void Reset();
    void SetTag(const std::string &tag);

private:
    void RecordKernel(const std::string &name, const int &rank, const std::string &device, int64_t host_us,
                      int64_t device_us = 0, int64_t max_device_mem_usage_mb = 0);
    void ReportGroupedByRank(std::function<std::ostream &(int64_t)> get_os, SortBy sort_by) const;
    void PrintRecordsGroupedByRank(std::function<std::ostream &(int64_t)> get_os) const;

    mutable std::mutex mtx_;
    std::vector<KernelCallRecord> call_records_;
    std::string current_tag_ = "Untagged";

    // thread-local tracking
    thread_local static inline std::map<std::string, std::chrono::high_resolution_clock::time_point>
        tls_cpu_timing_map_;

    struct EventPair {
        core::Event *start = nullptr;
        core::Event *stop = nullptr;
    };

    thread_local static inline std::map<std::string, EventPair> tls_device_timing_map_;
};
} // namespace infini_train