Skip to content

StridedMemoryView should expose access to the producer stream if possible #259

@leofang

Description

@leofang

In an offline discussion it was brought to my attention that some use cases need a way to order the producer stream after the consumer one, when the work is done. (Ex: a ping pong exchange between producer/consumer in a tight loop.)

Right now, StridedMemoryView follows both DLPack and CAI and does the hand-shake (stream ordering) only once. But a second ordering can be done if the producer stream is exposed. In general, this is not considered safe (because the producer may not guarantee the stream pointer is always valid during the lifetime of a StridedMemoryView instance). However, we do have one situation that guarantees the pointer lifetime, which is when StridedMemoryView instances are created inside a @args_viewable_as_strided_memory context (WIP in #247). In this case, before the decorated function returns the producer stream is considered always valid, and by exposing it we allow users to perform such ordering before this function returns, e.g.

@args_viewable_as_strided_memory((0,))
def my_func(arr, work_stream):
    # work_stream is a "consumer" stream
    view = arr.view(work_stream.handle)
    assert isinstance(view.producer_stream, Stream)
    # ... do work on work_stream ...
    view.producer_stream.wait(work_stream)
    return

Metadata

Metadata

Assignees

Labels

cuda.coreEverything related to the cuda.core moduleenhancementAny code-related improvementstriageNeeds the team's attention

Type

No type

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions