Gate output-centric image_to_nchw staging to discrete GPUs#21022
Conversation
Summary: The `image_to_nchw` output-staging shader can dispatch one thread per NCHW output element so writes to the host-visible staging buffer are fully coalesced. On a discrete GPU the staging buffer is host-visible but not device-local (PCIe-backed system RAM), where coalescing writes is a large win: ~11x end-to-end on an RTX 4080 laptop for an ESPCN 2x super-resolution model. That output-centric scheme fetches each texel up to 4 times (once per component). On unified-memory (mobile) GPUs the staging buffer is not PCIe-backed, so write coalescing buys nothing and the redundant fetches are a net loss: measured +15-42% `image_to_nchw` per-dispatch on Mali-G715 and +3-6% on Adreno 750. Gate the two strategies on `Adapter::has_unified_memory()`: - discrete (no unified memory) -> `image_to_nchw_coalesced_*` (output-centric) - unified memory (mobile) -> `image_to_nchw_*` texel-centric (one thread per texel, single fetch, contiguous NCHW writes), which is the default. `image_to_nchw.glsl` emits both variants via a `COALESCED_WRITES` codegen flag (default False). `get_tensor_to_nchw_shader` selects the coalesced variant when `!has_unified_memory()`; the global-workgroup-size picker keys off the shader name so the dispatch topology matches the chosen variant. Differential Revision: D112599710 Signed-off-by: Stephen Jia <ssjia@meta.com>
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/21022
Note: Links to docs will display an error until the docs builds have been completed. ❗ 2 Merge Blocking SEVsThere is 2 active merge blocking SEVs. Please view them below:
If you must merge, use ✅ You can merge normally! (1 Unrelated Failure)As of commit 37460f1 with merge base 931ba24 ( FLAKY - The following job failed but was likely due to flakiness present on trunk:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
meta-cla
Bot
added
the
CLA Signed
|
@SS-JIA has exported this pull request. If you are a Meta employee, you can view the originating Diff in D112599710. |
This PR needs a
|
2 participants
Summary:
The
image_to_nchwoutput-staging shader can dispatch one thread per NCHWoutput element so writes to the host-visible staging buffer are fully coalesced.
On a discrete GPU the staging buffer is host-visible but not device-local
(PCIe-backed system RAM), where coalescing writes is a large win: ~11x
end-to-end on an RTX 4080 laptop for an ESPCN 2x super-resolution model.
That output-centric scheme fetches each texel up to 4 times (once per
component). On unified-memory (mobile) GPUs the staging buffer is not
PCIe-backed, so write coalescing buys nothing and the redundant fetches are a
net loss: measured +15-42%
image_to_nchwper-dispatch on Mali-G715 and +3-6%on Adreno 750.
Gate the two strategies on
Adapter::has_unified_memory():image_to_nchw_coalesced_*(output-centric)image_to_nchw_*texel-centric (one thread pertexel, single fetch, contiguous NCHW writes), which is the default.
image_to_nchw.glslemits both variants via aCOALESCED_WRITEScodegen flag(default False).
get_tensor_to_nchw_shaderselects the coalesced variant when!has_unified_memory(); the global-workgroup-size picker keys off the shadername so the dispatch topology matches the chosen variant.
Differential Revision: D112599710