cudagraph.py

import torch

def cudagraph(f):
    _graphs = {}
    def f_(*args):
        key = hash(tuple(tuple(a.shape) for a in args))
        if key in _graphs:
            wrapped, *_ = _graphs[key]
            return wrapped(*args)
        g = torch.cuda.CUDAGraph()
        in_tensors = [a.clone() for a in args]
        f(*in_tensors) # stream warmup
        with torch.cuda.graph(g):
            out_tensors = f(*in_tensors)
        def wrapped(*args):
            [a.copy_(b) for a, b in zip(in_tensors, args)]
            g.replay()
            return [o.clone() for o in out_tensors]
        _graphs[key] = (wrapped, g, in_tensors, out_tensors)
        return wrapped(*args)
    return f_