microsoft / microsoft/onnxruntime

[WebGPU EP] Add unit test for SessionBufferPool

Open
#28,970 0 comments 0 reactions 1 assignee Claimed by @qjia7 View on GitHub
ep:WebGPU platform:web
Dominant language
C++
Stars
21.9k
Forks
4.2k
Avg merge
4d 11h
Merged PRs (30d)
184

Description

## Background

PR #28761 introduces `SessionBufferPool` and the `sessionBufferPoolGenerations` provider option for the WebGPU EP, letting a session retain retired per-graph `BufferManager` buffer caches and seed them into newly created generators. End-to-end behavior is validated via `verify_multi_gen.py` and `verify_max_length_change.py` (Phi-4 on graph capture: `storage hits=171 misses=0, uniform hits=296 misses=0` after the first generator donates), but there is no dedicated unit test for the pool itself.

This issue tracks adding one as a follow-up. From [the PR review thread](https://github.com/microsoft/onnxruntime/pull/28761#issuecomment-4663803276):

> A unit test with a mock `IBufferCacheManager` exercising `Donate → SeedInto` roundtrip, `Donate` eviction at capacity, and `Clear` would lock in the contract and make the pool tractable to modify safely later.

## Proposed coverage

A new `session_buffer_pool_test.cc` (alongside the other WebGPU EP unit tests) with a fake `IBufferCacheManager` / `BufferManager` shim. At minimum:

1. **Donate → SeedInto roundtrip** — donate a slot with N storage + M uniform entries, immediately seed into a fresh manager, assert all entries land in the receiving cache with sizes and handles preserved.
2. **FIFO eviction at capacity** — with `max_generations = k`, donate `k + d` slots and assert that the `d` oldest slots' buffers were released (via `wgpuBufferRelease`) while the newest `k` remain pooled. The mock should track release counts to prove no leak.
3. **LIFO consume** — confirm `SeedInto` drains the most-recently-donated slot first; relevant for hot-reuse / shape stability.
4. **`max_generations == 0`** — `Donate` is a no-op (caches still extracted but immediately discarded with proper release on the retiring manager side).
5. **Empty-slot early-exit** — donating a manager whose caches extract to empty does not push a slot.
6. **`Clear` releases everything** — after `Clear`, the mock buffer release count equals the number of donated handles still pending; pool `Size()` returns 0.
7. **Destructor releases** — same property as `Clear` but via `~SessionBufferPool()`.

## Non-goals

- GPU-side fence behavior (covered by the documented invariants on `Donate` plus existing end-to-end tests).
- Threading (`SessionBufferPool` is documented as relying on `InferenceSession::session_mutex_` for serialization; no need to test concurrency).

## Notes

- Should land after the companion GenAI-side `State::~State() → SessionReleaseCapturedGraph` change so the feature is actually exercised in production. There is no urgency on this issue until the pool is in the steady-state hot path.
- The mock `IBufferCacheManager` only needs to implement `ExtractCachedBuffers` and `AbsorbCachedBuffers` — the rest of the cache interface is irrelevant for the pool's contract.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.