()
| 130 | torch is None or not torch.cuda.is_available(), reason="Requires torch and CUDA" |
| 131 | ) |
| 132 | def test_torch_graph() -> None: |
| 133 | assert torch is not None |
| 134 | mod = gen_check_stream_mod() |
| 135 | device_id = torch.cuda.current_device() |
| 136 | device = tvm_ffi.device("cuda", device_id) |
| 137 | device_type = device.dlpack_device_type() |
| 138 | graph = torch.cuda.CUDAGraph() |
| 139 | stream = torch.cuda.Stream(device_id) |
| 140 | x = torch.zeros(1, device="cuda") |
| 141 | with tvm_ffi.use_torch_stream(torch.cuda.graph(graph, stream=stream)): |
| 142 | assert torch.cuda.current_stream() == stream |
| 143 | mod.check_stream(device_type, device_id, stream.cuda_stream) |
| 144 | # avoid cuda graph no capture warning |
| 145 | x = x + 1 |
nothing calls this directly
no test coverage detected