ROCm / rocmProfileData

MIT License
14 stars 8 forks source link

rpdTraceControl as a context manager, callable in a loop #52

Open jpvillam-amd opened 5 months ago

jpvillam-amd commented 5 months ago

The goal is to allow something like:

import torch
from rpdTracerControl import rpdTracerControl

for step in [1,2,3]:
    trace_file_name = f"{step}_torch.rpd"
    print(f"in {step}!!")
    with rpdTracerControl(trace_file_name, nvtx=True) as p:
        a =torch.rand(step).cuda()
        torch.add(a, 20)
    p.top_totals()