cyx-6 opened a new pull request, #5:
URL: https://github.com/apache/tvm-ffi/pull/5
This PR adds the stream context into ffi, so that ffi env stream can be
updated. The `tvm_ffi.use_torch_stream` is for wrapping the torch stream/graph
context. And lower-level `tvm_ffi.use_raw_stream` is for creating context with
existing stream handle.
Example for `tvm_ffi.use_torch_stream`:
case with torch stream:
```python
stream = torch.cuda.Stream()
stream_context = torch.cuda.stream(stream)
with tvm_ffi.use_torch_stream(stream_context):
...
```
case with torch cuda graph
```python
graph = torch.cuda.CUDAGraph()
graph_context = torch.cuda.graph(graph)
with tvm_ffi.use_torch_stream(graph_context):
...
```
case with current stream by default
```python
stream = torch.cuda.Stream()
stream_context = torch.cuda.stream(stream)
with torch.cuda.stream(stream):
with tvm_ffi.use_torch_stream():
...
```
Eaxmple for `tvm_ffi.use_raw_stream`:
```python
device = tvm_ffi.device(...)
stream_handle = ...
with tvm_ffi.use_raw_stream(device, stream_handle):
...
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
