()
| 104 | torch is None or not torch.cuda.is_available(), reason="Requires torch and CUDA" |
| 105 | ) |
| 106 | def test_torch_current_stream() -> None: |
| 107 | assert torch is not None |
| 108 | mod = gen_check_stream_mod() |
| 109 | device_id = torch.cuda.current_device() |
| 110 | device = tvm_ffi.device("cuda", device_id) |
| 111 | device_type = device.dlpack_device_type() |
| 112 | stream_1 = torch.cuda.Stream(device_id) |
| 113 | stream_2 = torch.cuda.Stream(device_id) |
| 114 | with torch.cuda.stream(stream_1): |
| 115 | assert torch.cuda.current_stream() == stream_1 |
| 116 | with tvm_ffi.use_torch_stream(): |
| 117 | mod.check_stream(device_type, device_id, stream_1.cuda_stream) |
| 118 | |
| 119 | with torch.cuda.stream(stream_2): |
| 120 | assert torch.cuda.current_stream() == stream_2 |
| 121 | with tvm_ffi.use_torch_stream(): |
| 122 | mod.check_stream(device_type, device_id, stream_2.cuda_stream) |
| 123 | |
| 124 | assert torch.cuda.current_stream() == stream_1 |
| 125 | with tvm_ffi.use_torch_stream(): |
| 126 | mod.check_stream(device_type, device_id, stream_1.cuda_stream) |
| 127 | |
| 128 | |
| 129 | @pytest.mark.skipif( |
nothing calls this directly
no test coverage detected