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=133895968596368)>

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

Example: Using Streams With Compiled Functions

linear = tp.Linear(2, 3)

compiled_linear = tp.compile(
    linear, args=[tp.InputInfo((2, 2), dtype=tp.float32)]
)

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

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

Local Variables

>>> linear
Linear(
    weight: Parameter = (shape=[3, 2], dtype=float32),
    bias: Parameter = (shape=[3], dtype=float32),
)
>>> linear.state_dict()
{
    weight: tensor(
        [[0.0000, 1.0000],
         [2.0000, 3.0000],
         [4.0000, 5.0000]], 
        dtype=float32, loc=gpu:0, shape=(3, 2)),
    bias: tensor([0.0000, 1.0000, 2.0000], dtype=float32, loc=gpu:0, shape=(3,)),
}

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

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

>>> output
tensor(
    [[1.0000, 6.0000, 11.0000],
     [1.0000, 6.0000, 11.0000]], 
    dtype=float32, loc=gpu:0, shape=(2, 3))

synchronize() → None

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

Example: Using Synchronize For Benchmarking

import time

linear = tp.Linear(2, 3)
compiled_linear = tp.compile(
    linear, args=[tp.InputInfo((2, 2), dtype=tp.float32)]
)

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

compiled_linear.stream = tp.Stream()

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

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

Output

Execution took 0.26191530050709844 ms

Return type:

None

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

Provides access to the default 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=133895969977392)>