[Pytorch] torch.autograd.profiler细节

官网使用说明；

Python部分：

核心类class profile：用户侧用with来创建和退出之；self.function_events成员是核心数据；缺点：DataLoader发起的多进程调用，其无法get到其他进程的操作的cuda时间；

    成员函数__enter__: with开始时调用，调用C++底层的torch.autograd._enable_profiler开始统计；

    成员函数__exit__: with结束时调用，调用C++底层的torch.autograd._disable_profiler结束统计；调用parse_cpu_trace将来自C++的“时间点”加工为“线段”，保存在self.function_events数据成员中；

    以下成员函数都是对self.function_events数据进行加工： _check_finish；table；export_chrome_trace；key_averages；total_average；self_cpu_time_total；其中好几个是先调用_check_finish（内部调用self.function_events.populate_cpu_children），建立FunctionEvent彼此之间的父子关系（之所以不在__exit__里调用，可能是lazy策略，这步就可以放在程序最末尾，减少对原程序的时间延迟打扰）

函数parse_cpu_trace：

    将来自C++的“时间点”(push或pop类型)加工为“线段”(FunctionEvent)；

    自带cuda的线段要把cuda起始时间记录下来；(cudaEventRecord相对于本<node_id, device>的投一个cuda event的elapsed time，记作本时间点）（即这里得到的"线段"对应的cuda时间，等于其发射过的"头一个kernel开始执行至最后一个kernel结束"的总时长）（同一个kernel的执行时间，会被所有包含它的“线段”所"重复"统计，以为每个“线段”的开头和结尾都会调用cudaEventRecord）(这里，一个“线段”只能关联一个所谓的"kernel"）

    记录该”线段“使用的总内存和总显存（只要遇到memory_alloc事件，则将所有包含其的”线段“，全部累加上这次的size，即每个”线段“记录的是total size而不是self size)；

类EventList：继承自python的list类，profile类的self.function_events就是该类型，list成员是FunctionEvent(即"线段")；

    成员函数populate_cpu_children：目的是在”线段“之间建立父子包含关系；把"线段"们先按<thread, node_id>分成不同的组，每组分别操作：

        先按线段的<起始时间，-终止时间>做key来排序，目的是确保“母线段”一定在其“子线段”之前；

        维护一个栈，里面所有线段具有”包含关系“，新来的线段违背”包含关系“时，就pop栈知道满足”包含关系“；对所有满足包含关系的相邻线段，建立父子关系；

    成员函数table：调用build_table来格式化成好看的字符串类型表格；

    成员函数export_chrome_trace：导出到json格式的文件，可以被chrome://tracing打开看trace-view；包括：CPU线段，GPU线段，CPU线段和GPU线段之间的联系线；(args字段里可以放很多附加内容）

    成员函数key_averages：统计，以<event.key, event.node_id>做key把线段们aggregation起来；也可把input_shape加到key里；

    成员函数self_cpu_time_total：统计，所有线段的CPU self time全部加起来；

    成员函数total_average：统计，把所有线段的数值全加和起来；

类FunctionEvent：”线段“；CPU侧的线段，里面self.kernels成员可以包含多个"kernel"线段；关键成员字段:

    node_id, name, thread, cpu_interval(包含起始时间点), kernels, cpu_children, input_shapes, 

    cpu_memory_usage, cuda_memory_usage, is_async, is_remote,  sequence_nr

    成员函数self_cpu_memory_usage

    成员函数self_cuda_memory_usage

    成员函数self_cpu_time_total

    （注意：唯独没有self_cuda_time_total）

    成员函数cuda_time_total

    成员函数cpu_time_total

类FunctionEventAvg：”线段“们的平均；

类emit_nvtx：把通过C++函数和cudaEventRecord打点，换成nvtx的push、pop来打点；还是调用C++层的torch.autograd._enable_profiler来开始统计，只是传过去的参数变成了NVTX；使用nvprof启动进程，代码里有with torch.autograd.profiler.emit_nvtx()则生效；（类开头的一大段注释解释sequence_nr的用法，有一些不太懂）

parse_nvprof_trace：读取nvtx的.sqlite文件，把里面的"点“转为"线段"(FunctionEvent), 有kernel的要在成员变量里挂上kernel（一个”线段“可以挂在多个kernel，kernel的所有”祖先“线段都会挂载它）；

NVTX的parse这里做的很粗糙，暂时没有考虑node_id和thread；

类record_function：用户自定义“范围”，用with把一段代码包起来统计之；成员函数_call_end_callbacks_on_future的用法暂时未知；

build_table：把"线段"们格式化成好看的字符串类型表格

 

其他：

is_async：FunctionEvent如果开头和结尾不是同一个thread，就叫做async的；

<node_id, device>: 每个cuda event都有这些属性，可以拿到不同机器的不同卡的cuda event?

<event.key, event.node_id>: 在key_averages里面使用了这个做aggregation的key，即考虑了不同机器分开统计；

sequence_nr: 用于将forward和backward操作关联起来的序号；C++层埋的信息；

 

C++部分：

profiler.h文件：

Event类：打的“点”；字段如下：

  int64_t cpu_ns_ = 0;
  at::StringView name_;
  EventKind kind_;
  uint16_t thread_id_;
  at::RecordFunctionHandle handle_ {0};
  std::vector<std::vector<int64_t>> shapes_;
  int64_t cpu_memory_usage_ = 0;
  int64_t cuda_memory_usage_ = 0;
  int device_ = -1;
  CUDAEventStub cuda_event = nullptr;
  int node_id_ = 0;
  bool is_remote_ = false;
  int64_t cuda_us_ = -1;
  int64_t sequence_nr_ = -1;

node_id_: torch\distributed\rpc\__init__.py里面，_set_profiler_node_id(rank)，即设置rpc的process rank作为node_id_;

device_: 在reportMemoryUsageToProfiler时赋值的，注意就是个整数；

profiler.cpp文件：

mark、pushRange、popRange、reportMemoryUsage，分别是打点4种EventKind；

成员函数enableProfiler：为每个deivce打初始点__cuda_start_event，同时为host打点__start_profile，后续cuda的点就可以和自己device上的初始点计算elapsed time，同理host后续的点也和host初始点计算elapsed time，因为所有device和host都几乎同时打的初始点，所以所有这些elapsed time可以当作“wall-clock time"使用！（cuda上打初始点似乎有overhead，所以人家先warm-up了一把，即__cuda_startup）

profiler.py文件里的parse_cpu_trace，会调用adjusted_time：

return cuda_time_0.cuda_elapsed_us(cuda_record) + start_record.cpu_elapsed_us(cuda_time_0)

其中前半部分计算cuda_record的Event减去cuda_time_0的Event的时间差；后半部分计算cuda_time_0的cpu_ns_减去start_record的cpu_ns_的时间差；假设：cuda_time_0的Event和其cpu_ns_是几乎同时时刻！目的：总体等价于cuda_record的Event在GPU上执行到的时刻相对于start_record的cpu_ns_的时间差；为什么需要后半部分：因为在C++里是一个循环为每个Device挨个打__cuda_start_event，最后再打__start_profile，所以两者之间可能是有微笑时间差的，即__cuda_start_event的时刻不完全等于__start_profile，为了更精确，所以加了后半部分。