(replay_first_step, func, *args, **kwargs)
| 90 | |
| 91 | |
| 92 | def graph_process(replay_first_step, func, *args, **kwargs): |
| 93 | # `func` should only contain operations on the GPU |
| 94 | # Please ensure that the memory address of the data required by 'func' remains constant |
| 95 | if func.__name__ not in graph_cache: |
| 96 | cuda_stream = get_accelerator().Stream() |
| 97 | cuda_stream.wait_stream(get_accelerator().current_stream()) |
| 98 | with get_accelerator().stream(cuda_stream): |
| 99 | func(*args, **kwargs) |
| 100 | get_accelerator().current_stream().wait_stream(cuda_stream) |
| 101 | graph_cache[func.__name__] = get_accelerator().create_graph() |
| 102 | with get_accelerator().capture_to_graph(graph_cache[func.__name__]): |
| 103 | func(*args, **kwargs) |
| 104 | if replay_first_step: |
| 105 | get_accelerator().replay_graph(graph_cache[func.__name__]) |
| 106 | else: |
| 107 | get_accelerator().replay_graph(graph_cache[func.__name__]) |
| 108 | |
| 109 | |
| 110 | def noop_decorator(func): |
no test coverage detected