code: def do_bench_cudagraph(fn, rep=20, grad_to_none=None):
      import torch
  
      """
      Benchmark the runtime of the provided function.
  
      :param fn: Function to benchmark
      :type fn: Callable
      :param rep: Repetition time (in ms)
      :type rep: int
      :param grad_to_none: Reset the gradient of the provided tensor to None
      :type grad_to_none: torch.tensor, optional
      """
      if torch.cuda.current_stream() == torch.cuda.default_stream():
          raise RuntimeError(
              "Cannot capture graph in default stream. Please use side stream in benchmark code."
          )
      # warmup
      fn()
      # step 1 - we estimate the amount of time the kernel call takes
      # NOTE: this estimate isn't super accurate because the GPU isn't warmed up at this point
      #       but it is probably good enough
      if grad_to_none is not None:
          for x in grad_to_none:
              x.detach_()
              x.requires_grad_(True)
              x.grad = None
      g = torch.cuda.CUDAGraph()
      with torch.cuda.graph(g):
          fn()
      torch.cuda.synchronize()
      start_event = torch.cuda.Event(enable_timing=True)
      end_event = torch.cuda.Event(enable_timing=True)
      start_event.record()
      g.replay()
      end_event.record()
      torch.cuda.synchronize()
      estimate_ms = start_event.elapsed_time(end_event)
      n_repeat = max(1, int(rep / estimate_ms))
      # step 2 - construct a cuda graph with `n_repeat` unrolled function calls to minimize
      # host overhead
      g = torch.cuda.CUDAGraph()
      with torch.cuda.graph(g):
          for i in range(n_repeat):
              if grad_to_none is not None:
                  for x in grad_to_none:
                      x.grad = None
              fn()
      torch.cuda.synchronize()
      # measure time and return
      ret = []
      n_retries = 10
      for i in range(n_retries):
          start_event = torch.cuda.Event(enable_timing=True)
          end_event = torch.cuda.Event(enable_timing=True)
          start_event.record()
          g.replay()
          end_event.record()
          torch.cuda.synchronize()
          ret += [start_event.elapsed_time(end_event) / n_repeat]
      return torch.mean(torch.tensor(ret)).item()
  
imports: import functools
  import os
  import subprocess
  import sys
  from contextlib import contextmanager
  from typing import Any, Dict, List
  from .None import language
label_desc: Use image node_id 2 for calling a global function with example usage: do_bench_cudagraph(fn, rep, grad_to_none) and returns: torch