[feature request] add memory counter event to pytorch trace result

[FindDefinition]

New Feature

We can use memory instant events in pytorch profiler result to generate a nice gpu memory trace in perfetto:

This memory trace is aligned to python stack trace so it's more useful than pytorch solutions such as export_memory_timeline

I have implemented this by add profile_memory to torch.profiler.profile and parse pytorch trace events, I think viztracer can add this feature easily. we may use --log_torch and --log_torch_memory to enable this feature.

            for ev in pth_trace_dict["traceEvents"]:
                if (ev["ph"] == "I" or ev["ph"] == "i") and "args" in ev:
                    if "Total Allocated" in ev["args"]:
                        total_alloc = ev["args"]["Total Allocated"]
                        total_reserv = ev["args"]["Total Reserved"]
                        if ev["args"]["Device Id"] >= 0:
                            external_events.append({
                                "name": "pth_memory",
                                "ph": "C",
                                "pid": pid,
                                "tid": 0,
                                "ts": ev["ts"] + base_ts_us,
                                "args": {
                                    "Total Allocated": total_alloc,
                                    "Total Reserved": total_reserv,
                                },
                            })

[gaogaotiantian]

Sure, that sounds reasonable. I can work on it later.