()
| 81 | torch is None or not torch.cuda.is_available(), reason="Requires torch and CUDA" |
| 82 | ) |
| 83 | def test_torch_stream() -> None: |
| 84 | assert torch is not None |
| 85 | mod = gen_check_stream_mod() |
| 86 | device_id = torch.cuda.current_device() |
| 87 | device = tvm_ffi.device("cuda", device_id) |
| 88 | device_type = device.dlpack_device_type() |
| 89 | stream_1 = torch.cuda.Stream(device_id) |
| 90 | stream_2 = torch.cuda.Stream(device_id) |
| 91 | with tvm_ffi.use_torch_stream(torch.cuda.stream(stream_1)): |
| 92 | assert torch.cuda.current_stream() == stream_1 |
| 93 | mod.check_stream(device_type, device_id, stream_1.cuda_stream) |
| 94 | |
| 95 | with tvm_ffi.use_torch_stream(torch.cuda.stream(stream_2)): |
| 96 | assert torch.cuda.current_stream() == stream_2 |
| 97 | mod.check_stream(device_type, device_id, stream_2.cuda_stream) |
| 98 | |
| 99 | assert torch.cuda.current_stream() == stream_1 |
| 100 | mod.check_stream(device_type, device_id, stream_1.cuda_stream) |
| 101 | |
| 102 | |
| 103 | @pytest.mark.skipif( |
nothing calls this directly
no test coverage detected