()
| 61 | |
| 62 | |
| 63 | def test_raw_stream() -> None: |
| 64 | mod = gen_check_stream_mod() |
| 65 | device = tvm_ffi.device("cuda:0") |
| 66 | stream_1 = 123456789 |
| 67 | stream_2 = 987654321 |
| 68 | with tvm_ffi.use_raw_stream(device, stream_1): |
| 69 | mod.check_stream(device.dlpack_device_type(), device.index, stream_1) |
| 70 | assert tvm_ffi.get_raw_stream(device) == stream_1 |
| 71 | |
| 72 | with tvm_ffi.use_raw_stream(device, stream_2): |
| 73 | mod.check_stream(device.dlpack_device_type(), device.index, stream_2) |
| 74 | assert tvm_ffi.get_raw_stream(device) == stream_2 |
| 75 | |
| 76 | mod.check_stream(device.dlpack_device_type(), device.index, stream_1) |
| 77 | assert tvm_ffi.get_raw_stream(device) == stream_1 |
| 78 | |
| 79 | |
| 80 | @pytest.mark.skipif( |
nothing calls this directly
no test coverage detected