Fix usage of stream for memory pattern allocation with arena

[skottmckay]

Description

Use real stream for memory pattern block allocation so that the cleanup of chunks from the arena correctly marks it as unused.

Motivation and Context

#29351

[Copilot]

Pull request overview

This PR updates ORT’s stream-aware memory-pattern block allocation to use the actual per-device stream (including any user-provided override stream) so that stream-aware arena cleanup can correctly mark those chunks as reusable across iterations. This aligns memory-pattern allocations with the streams that will actually consume them, addressing GPU memory growth when stream tagging/cleanup mismatches.

Changes:

• Allocate memory-pattern buffers on the stream associated with the target OrtDevice (preferring a RunOptions override stream when applicable).
• Replace the GetRootStream() dummy/root-stream concept with a GetStreamForDevice(const OrtDevice&) lookup.
• Ensure arena stream-buffer cleanup also runs for the user-provided override stream (since it’s not owned by the collection).

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated no comments.

File	Description
onnxruntime/core/framework/execution_frame.cc	Allocates the memory-pattern backing buffer on the real stream for the target device (or override stream).
onnxruntime/core/framework/device_stream_collection.h	Replaces GetRootStream() with GetStreamForDevice(const OrtDevice&).
onnxruntime/core/framework/device_stream_collection.cc	Removes dummy root stream plumbing, adds device-based stream lookup, and releases stream-aware arena buffers for override streams during cleanup.

[tianleiwu]

Review summary

Solid, well-scoped fix. Replacing the dummy root_stream_ (whose default-constructed OrtDevice never matched a CUDA arena in ReleaseSingleStreamBuffers, so the pattern chunk was never untagged) with GetStreamForDevice() routes the memory-pattern buffer to a real device stream. The stream tag is then cleared at CleanUp, so the ~1 GiB pattern chunk becomes reusable across runs/threads. This matches the root-cause analysis in #29351.

Verified:

• GetStreamForDevice() override-precedence is consistent with GetStream() — because SetStreamOverride requires an exact device match, the first if always returns the override for that device, so the loop never returns the stale overridden device_streams_[i].
• The alloc->Alloc() fallback taken when no device stream matches allocates with stream == nullptr (BFCArena::Alloc → AllocateRawInternal(size, false, nullptr)), i.e. an untagged chunk even on a StreamAwareBFCArena, so it remains safe-to-reuse.

Two non-blocking notes:

1. Cleanup coverage for externally-set (SetDeviceStream) streams — see inline comment.
2. Test coverage. No unit test exercises the new GetStreamForDevice() path or the untag-on-cleanup behavior. Consider extending DeviceStreamCollection.TestOverride in onnxruntime/test/framework/inference_session_test.cc to assert GetStreamForDevice(device) returns (a) the override stream when its device matches, (b) the bound device stream otherwise, and (c) nullptr for an unmapped device.

[tianleiwu]

CleanUp untags owned_streams_ and the override stream, but not streams injected via SetDeviceStream() — i.e. the subgraph parent_stream set in UpdateWithParentStream (utils.cc). GetStreamForDevice() can return that non-owned parent_stream for the memory-pattern allocation, so its chunk isn't untagged by this collection's CleanUp; it relies on the parent collection (which owns parent_stream) untagging it at the parent's run end.

That invariant holds today (main-graph streams are always owned via AddDeviceStream; the only non-null external stream is the subgraph parent), so it isn't a leak — but the dependency is implicit. A one-line comment stating that non-owned streams are untagged by their owning collection would make this explicit and guard against future callers of SetDeviceStream.