Skip to content

[Frontend] Fix conv wrapper crash when the input is a view of a lower-rank buffer#294

Open
YWHyuk wants to merge 2 commits into
feature/togsim-cpp-tracefrom
fix/conv-wrapper-view-input
Open

[Frontend] Fix conv wrapper crash when the input is a view of a lower-rank buffer#294
YWHyuk wants to merge 2 commits into
feature/togsim-cpp-tracefrom
fix/conv-wrapper-view-input

Conversation

@YWHyuk

@YWHyuk YWHyuk commented Jul 9, 2026

Copy link
Copy Markdown
Collaborator

Fixes a conv codegen bug where the generated wrapper crashes with

padded_shape[3] += 2 * 0
IndexError: list index out of range

Root cause

A conv input node can be a ReinterpretView. call_kernel passes template args by
buffer name, so the wrapper is handed the base buffer, whose rank/shape may differ
from the view codegen assumed:

X
codegen (extract_info) X.layout.size = [1, 32, 8, 8] (the 4D view)
runtime (wrapper) X.shape = (1, 64, 32) (the base buffer)

The wrapper then inferred geometry from the runtime tensor (list(X.shape), X.shape[2],
X.shape[3]) and blew up.

It only fires when the view is free: a contiguous (B, N, C) buffer already is the
channels_last layout of (B, C, H, W), so inductor inserts no materializing copy. That is
exactly the transformer (B, N, C) -> (B, C, H, W) pattern — e.g. SegFormer's
efficient-attention spatial-reduction conv (Conv2d(C, C, k=sr, stride=sr), padding 0).
When the layout does not match, inductor inserts a copy kernel, a real 4D buffer arrives,
and the bug is invisible. That is why it went unnoticed.

Fix

Stop inferring geometry from the runtime tensor. Bake the size/stride/offset the codegen
used into the wrapper and rebuild the logical NCHW inputs from it:

X = reinterpret_tensor(X, (1, 32, 8, 8), (0, 1, 256, 32), 0)
W = reinterpret_tensor(W, (16, 32, 2, 2), (128, 1, 64, 32), 0)
X_padding = torch.zeros((1, 32, 8, 8), device=X.device, dtype=X.dtype)
X_padding[:, :, 0:8, 0:8] = X

X.layout describes the data within the storage of the tensor that arrives (both come from
the same node), so the reinterpret is correct — and idempotent — whether the wrapper is
handed the base buffer, an already reinterpreted view, or a materialized copy.

The conv templates were the only place reading .shape at runtime; everything else already
derives geometry from the node layout at codegen time. No change to call_kernel,
mlir_argdefs, the MLIR signature, or arg_attributes
— the call site still passes the
raw base buffer and it now works.

Also fixed for free: the padding buffer was allocated via torch.zeros(shape) (always f32)
then .to(device=...) (CPU round-trip). It now allocates directly on device with X.dtype.

Verification

  • New regression test tests/ops/conv/test_conv_view_input.py: fails with the above
    IndexError on the pre-fix wrapper, passes after (max diff 2.3e-05 / 2e-04 vs CPU).
  • Normal conv path (real 4D buffer, padding=1) unchanged and numerically correct
    (max diff 3.8e-05 vs CPU); reinterpret_tensor is an identity there.

Follow-ups (not in this PR)

  • padding == 0 fast path: shapes are literals now, so the zeros alloc + full copy can be
    skipped for 1x1 / strided convs.
  • mlir_argdefs derives buffer_types numel from the base buffer, so a slicing view
    feeding e.g. gemm can still get a wrong MLIR arg size. Same class of bug, separate fix.
  • def_conv_kernel patches the signature with re.sub(r'(\d+)(?=xf32)', ...), hardcoding
    xf32; an f16 conv would not get its padded input size applied.

Found while checking #254 on this branch: SegFormer no longer hits the originally reported
Not supporting format error, but was blocked here instead.

A conv input node can be a ReinterpretView. call_kernel passes template args
by buffer *name*, so the wrapper receives the base buffer, whose rank/shape may
differ from the view codegen assumed. The wrapper then did

    padded_shape = list(X.shape)
    padded_shape[3] += 2 * PADDING_W        # IndexError on a 3D base buffer

This fires whenever the view is free: a contiguous (B, N, C) buffer already is
the channels_last layout of (B, C, H, W), so inductor inserts no materializing
copy. That is exactly the transformer (B, N, C) -> (B, C, H, W) pattern, e.g.
SegFormer efficient attention.

Stop inferring geometry from the runtime tensor. Bake the size/stride/offset the
codegen used into the wrapper and rebuild the logical NCHW inputs from it, so the
wrapper is correct whether it is handed the base buffer, an already reinterpreted
view, or a materialized copy (reinterpret is idempotent in all three cases). The
conv templates were the only place reading .shape at runtime; everything else
already derives geometry from the node layout at codegen time.

No change to call_kernel, mlir_argdefs, the MLIR signature, or arg_attributes.

Also allocate the padding buffer directly on the input device with the input
dtype: torch.zeros() defaulted to f32 and to(device=...) round-tripped via CPU.
Covers a conv fed by a free ReinterpretView of a lower-rank buffer. Fails with
IndexError on the pre-fix wrapper and passes after it.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants