Stream

class nvtripy.Stream(priority: int = 0)

Bases: object

Represents a CUDA stream that can be used to manage concurrent operations.

Note

Streams can only be used with compiled functions.

This class is a wrapper around the underlying stream object, allowing management of CUDA streams.

Parameters:

priority (int) – Assign priority for the new stream. Lower number signifies higher priority.

Example: Creating New Streams

stream_a = tp.Stream()
stream_b = tp.Stream()

Local Variables

>>> stream_a
<Stream(id=133553298685472)>

>>> stream_b
<Stream(id=133553298682352)>

Example: Using Streams With Compiled Functions

func = tp.compile(tp.relu, args=[tp.InputInfo((2, 2), dtype=tp.float32)])

# Run the compiled linear function on a custom stream:
stream = tp.Stream()
func.stream = stream

input = tp.ones((2, 2), dtype=tp.float32)
output = func(input)

Local Variables

>>> stream
<Stream(id=133553343611424)>

>>> input
tensor(
    [[1, 1],
     [1, 1]], 
    dtype=float32, loc=gpu:0, shape=(2, 2))

>>> output
tensor(
    [[1, 1],
     [1, 1]], 
    dtype=float32, loc=gpu:0, shape=(2, 2))

synchronize() → None

Synchronize the stream, blocking until all operations in this stream are complete.

Example: Using Synchronize For Benchmarking

import time

func = tp.compile(tp.relu, args=[tp.InputInfo((2, 2), dtype=tp.float32)])

input = tp.ones((2, 2), dtype=tp.float32)

func.stream = tp.Stream()

num_iters = 10
start_time = time.perf_counter()
for _ in range(num_iters):
    _ = func(input)
func.stream.synchronize()
end_time = time.perf_counter()

time = (end_time - start_time) / num_iters
print(f"Execution took {time * 1000} ms")

Output

Execution took 2.704306799569167 ms

Return type:

None

nvtripy.default_stream(device: device = device(kind='gpu', index=0)) → Stream

Provides access to the default Tripy CUDA stream for a given device. There is only one default stream instance per device.

Parameters:

device (device) – The device for which to get the default stream.

Returns:

The default stream for the specified device.

Raises:

TripyException – If the device is not of type ‘gpu’ or if the device index is not 0.

Return type:

Stream

Note

Calling default_stream() with the same device always returns the same Stream instance for that device.

Example: Retrieving The Default Stream

# Get the default stream for the current device.
default = tp.default_stream()

Local Variables

>>> default
<Stream(id=133553300060288)>