[Feature]: conv3d_implicit — accept and emit NDHWC so chained convs skip the layout round-trip

Suggestion Description

Ask

conv3d_implicit takes NCDHW and calls _ncdhw_to_ndhwc on every invocation (L705). In a network that is a chain of 3D convolutions, every intermediate could have stayed in NDHWC. Please add a layout parameter so a caller can convert twice per network instead of twice per layer:

conv3d_implicit(
    x, weight, bias=None, stride=1, padding=0,
    layout="NCDHW",        # "NCDHW" (default, unchanged) | "NDHWC"
    out_layout=None,       # None -> same as `layout`
    **kwargs,
)

layout="NDHWC" skips the input transpose; out_layout="NDHWC" lets the non-split-K epilogue write (n, do, ho, wo, k), which is a change of output view rather than new kernel work.

Why

On a Wan 3D causal VAE encoder (bf16, input (1, 3, 33, 352, 640), 278 conv calls, MI350X / gfx950 / ROCm 7.1.1), layout movement is 17.4% of the 193 ms encoder — 992 batched_transpose kernels, 33.6 ms. Swapping MIOpen for conv3d_implicit already removes the weight and output round-trips (_prep_weight caches, the splitk==1 epilogue writes NCDHW), and is 1.20x faster overall on these shapes: 108.2 → 89.9 ms measured per-shape. What remains is the per-call input transpose, roughly 20 ms that nothing in the network needs.

The dominant shapes are rank-5 in=[1,96,6,354,642] w=[96,96,3,3,3], in=[1,384,3,46,82] w=[384,384,3,3,3] and rank-4 in=[4,96,353,641] w=[96,96,3,3], all bf16. Happy to test a branch on them.

Related

• #861 [Issue]: Tune conv perf — tuning the conv itself; this request is about the data movement around it.

Operating System

No response

GPU

No response

ROCm Component

No response

Start at conv3d_implicit and its L705 call to _ncdhw_to_ndhwc; trace how input and non-split-K epilogue layouts are selected. Check the existing conv3d behavior and listed shape paths, then verify that default NCDHW behavior is unchanged and NDHWC chaining avoids intermediate transposes.