Skip to content

cuda.core: allow Buffer deallocation stream to be updated before teardown #2600

Description

@Andy-Jost

Summary

Add a public method for changing the stream that will order a Buffer's eventual deallocation without closing the buffer:

Buffer.set_deallocation_stream(stream: Stream | GraphBuilder) -> None

This is a follow-up to #2526.

Problem

A Buffer records a deallocation recipe when it acquires an allocation. Buffer.close(stream=...) can replace that recipe, but it immediately closes the Buffer.

Applications may transfer responsibility for a live allocation from work on stream A to work on stream B while retaining the Buffer or relying on automatic/deferred destruction. After establishing the necessary stream dependency, there is currently no public way to make stream B order the eventual deallocation without also closing the buffer.

This matters when the final release occurs through garbage collection or through another retained DevicePtrHandle, such as a graph attachment.

Proposed API

buffer.set_deallocation_stream(stream)

The method should:

Buffer.close(stream=stream) should remain equivalent to updating the deallocation stream and then closing the buffer.

Ordering semantics

This method does not enqueue synchronization, transfer ownership to a CUDA stream, or establish dependencies between streams. The caller must ensure that the new deallocation stream is ordered after allocation and all uses of the memory, using events or other CUDA synchronization mechanisms as appropriate.

Concurrent mutation and destruction of the same buffer remain the caller's responsibility.

Reflection

Do not add a public getter in this initial change.

The stored DeallocationStream is a teardown recipe rather than an ordinary cuda.core.Stream: default-stream recipes include a bound context, and PTDS recipes also retain host-thread metadata. Returning only a Stream would not faithfully represent or round-trip the complete recipe.

A separate immutable diagnostic representation can be considered later if a concrete introspection use case arises.

Acceptance criteria

  • A live buffer can have its deallocation stream changed without being closed.
  • Subsequent close() or automatic destruction uses the replacement recipe.
  • The buffer remains usable after set_deallocation_stream().
  • Context-bound default-stream and PTDS behavior matches cuda.core: capture complete Buffer deallocation recipe at creation #2526.
  • Existing close(stream=...) behavior remains unchanged.

Related

Metadata

Metadata

Assignees

Labels

P1Medium priority - Should docuda.coreEverything related to the cuda.core moduleenhancementAny code-related improvements

Projects

No projects

Relationships

None yet

Development

No branches or pull requests

Issue actions