[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
Open
[Frontend] Fix conv wrapper crash when the input is a view of a lower-rank buffer#294YWHyuk wants to merge 2 commits into
YWHyuk wants to merge 2 commits into
Conversation
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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Fixes a conv codegen bug where the generated wrapper crashes with
Root cause
A conv input node can be a
ReinterpretView.call_kernelpasses template args bybuffer name, so the wrapper is handed the base buffer, whose rank/shape may differ
from the view codegen assumed:
extract_info)X.layout.size = [1, 32, 8, 8](the 4D view)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 thechannels_last layout of
(B, C, H, W), so inductor inserts no materializing copy. That isexactly the transformer
(B, N, C) -> (B, C, H, W)pattern — e.g. SegFormer'sefficient-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.layoutdescribes the data within the storage of the tensor that arrives (both come fromthe 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
.shapeat runtime; everything else alreadyderives geometry from the node layout at codegen time. No change to
call_kernel,mlir_argdefs, the MLIR signature, orarg_attributes— the call site still passes theraw 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 withX.dtype.Verification
tests/ops/conv/test_conv_view_input.py: fails with the aboveIndexErroron the pre-fix wrapper, passes after (max diff 2.3e-05 / 2e-04 vs CPU).padding=1) unchanged and numerically correct(max diff 3.8e-05 vs CPU);
reinterpret_tensoris an identity there.Follow-ups (not in this PR)
padding == 0fast path: shapes are literals now, so thezerosalloc + full copy can beskipped for 1x1 / strided convs.
mlir_argdefsderivesbuffer_typesnumel from the base buffer, so a slicing viewfeeding e.g. gemm can still get a wrong MLIR arg size. Same class of bug, separate fix.
def_conv_kernelpatches the signature withre.sub(r'(\d+)(?=xf32)', ...), hardcodingxf32; 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 formaterror, but was blocked here instead.