[Frontend] Fix conv wrapper crash when the input is a view of a lower-rank buffer#294
Open
YWHyuk wants to merge 13 commits into
Open
[Frontend] Fix conv wrapper crash when the input is a view of a lower-rank buffer#294YWHyuk wants to merge 13 commits into
YWHyuk wants to merge 13 commits into
Conversation
Port the C++ MathLog/MathAtanToVCIX patterns to the in-process Python VCIX pass: add math.log (opcode 2, imm 2) and math.atan (opcode 2, imm 3) to _MATH_VIV and MARKERS. Encoding matches the spike rs1 / gem5 VS1 sub-opcodes (sin=0, cos=1, log=2, atan=3). No C++ mlir-opt pass needed.
Add the pointwise op test under tests/ops/elementwise and register it in the test workflow allowlist.
Bump the third-party pins to the releases shipping the new pointwise instructions (spike torchsim_vlog/vatan, gem5 CustomVlog/Vatan + decode).
reduction_init returned "-inf"/"inf" for max/min regardless of dtype, so an
integer max/min reduction fed an invalid inf identity into every consumer:
the accumulator init, the template reduction init, and the masked-DMA tail
fill (torch.tensor(inf, dtype=int64) overflowed -- surfaced compiling swinv2
with batch>1). reduction_combine_vec also hardcoded the float multi_reduction
kinds <maximumf>/<minimumf>, which reject integer vectors.
Return the dtype's representable extreme from reduction_init for integer
max/min (max identity -inf -> iinfo.min, min identity +inf -> iinfo.max;
bool -> 0/1), and emit the signed-integer multi_reduction kinds <maxsi>/<minsi>
for integer element types. sum/prod need no change: the <add>/<mul> combining
kinds are polymorphic and ops.add/ops.mul already select arith.add{i,f} by
dtype (as ops.maximum/minimum already select maxsi vs maximumf).
Verified: int amax/amin/sum/prod reductions compile and match CPU; float
test_reduce (ReduceMax) and test_softmax unchanged. Also clears the swinv2
batch>1 masked-fill overflow (it then hits a separate unsupported
shifted-window view, tracked separately).
…lar DMA index
torch.roll has no Inductor lowering; it decomposes to index_select(fmod(...)) -- a
cyclic-shift gather whose ModularIndexing DMA index the affine-only descriptor
cannot express ("Unlinearized floor/mod in DMA index"). Rewrite it as the two
contiguous halves the shift produces, stitched by cat:
roll(x, s, d) == cat([slice(x, N-s, N), slice(x, 0, N-s)], d)
each half is a plain strided slice with no modulo. Do it as a lowering (removing
roll from the decomp table so it is not decomposed first) rather than a
decomposition, so we can also REALIZE the input with copy_input: roll's producer
is often a reshaped view (e.g. swinv2's window_reverse), and the slice offset would
otherwise fuse into that reshape's modular index, re-creating a shifted (p+c)%m
form. A plain .contiguous()/clone does NOT create the buffer boundary (Inductor
inlines it); copy_input does.
Verified: torch.roll (1-D / multi-dim / arbitrary shift) and roll+reduction match
CPU; swinv2 batch>1 now compiles through the shifted-window attention (previously
blocked by the roll's cyclic shift and then by the fused slice offset).
A composition of aligned reshapes on one iteration variable (e.g. swinv2's window partition fused with a later reshape) leaves a nested index like ModularIndexing(ModularIndexing(p, 1, 64), 1, 8) that neither sympy nor simplify_with_ranges reduces. collect_boundaries then skips its cut points (the inner base is not a bare variable) and the affine-only DMA check rejects it. Add a general, pattern-free canonicalizer (flatten_nested_floormod / _as_digit): any nesting of FloorDiv/ModularIndexing over a single symbol is a digit extractor (v // A) % M and collapses to one level by composing divisors, from four divisibility-guarded algebraic identities. Applied in collect_boundaries and in _fold_with_ranges. Multi-variable inner arguments (e.g. a roll shift v+c) are left untouched. Verified: numeric equivalence on the swinv2 window index; test_reduce, test_layernorm, test_cat, test_indirect_access unaffected.
The wrapper header imported empty_strided but not empty_strided_cpu, so a graph that allocates a CPU-side buffer (e.g. swinv2's shifted-window attention mask, which Inductor keeps on CPU) failed at runtime with "NameError: name 'empty_strided_cpu' is not defined". Add the same binding the default Inductor PythonWrapperCodegen header provides.
…yout A per-lane reduction is lowered as a 2-D [reduction | batch] collapse: multi_reduction reshapes the accumulator to <reduction_numel x red_size> and reduces axis 0, which assumes the reduction axis is the outermost in-lane run and every non-reduction (batch) axis is inner. get_dma_info reorders the tile axes for reduction loads to guarantee that, but only for the 2-D and 3-D cases. The 4-D case mis-tested is_reduction (reduction_depth < 3, false once there are three batch dims) and then raised NotImplementedError, and rank 5+ fell through to the default row-major order -- both leaving the reduction axis innermost, so the reduce collapsed a batch axis instead of the reduction axis. This is hit whenever a non-reduction dim-merge is blocked and an extra batch axis stays in-lane: SwinV2 cosine window attention adds a gathered relative-position bias (table[idx[q,k], head]) before the softmax amax, keeping head separate from query (4-D tile, head bled into the max), and its post-attention LayerNorm var_mean splits the token axis into several loop vars (6-D tile, mean off by ~0.005). Generalize the reorder to any tile rank: a tile that carries the reduction axis places that axis-group outermost via range(r, L) + range(r-1, -1, -1), which reduces to the existing 3-D order for L=3. Adds test_reduce_gather_bias (fails max abs diff ~3.3 before the fix).
SwinV2 batch>1 used to fail codegen with "Unlinearized floor/mod in DMA index" because the shifted-window path (torch.roll composed with window partition/reverse) produced views axis-split could not linearize. With the roll->slice+cat lowering, the nested/shifted mod handling, and the reduction-axis layout fix, the whole backbone now compiles and matches CPU end to end (max diff ~5e-6 for image 64 / window 8 / batch 2). Add tests/models/test_swinv2.py, wire it as a self-hosted CI job next to the other model tests, and list SwinV2 in the README model coverage.
conv_multi_tile_mapping (used for batch > 1 convs) collects every tile that fits SPAD into tile_candidates and returns them sorted, but it guarded the "no mapping" error on max_used_spad_size instead of on tile_candidates. max_used_spad_size is only bumped when a candidate also satisfies max_k_h_w <= k_h, and max_k_h_w is initialized to K_W, so it only ever fires for the full-kernel (k_h == K_H) tile. When that tile overflows SPAD -- e.g. a CLIP ViT-B/32 patch conv (3->768, 32x32 stride 32) at batch > 1 -- the guard raised "Cannot find a valid mapping" even though smaller-k_h tiles fit and were already in tile_candidates. The mapping variable it tracks is never returned. Guard on `not tile_candidates` instead, so the conv is rejected only when no tile fits at all. The returned candidate list is unchanged, so convs that already worked are unaffected. Adds a batched patch-conv case to test_conv2d.py. Fixes issue #252.
With the conv-mapping fix, CLIP's vision transformer runs with batch > 1 and matches CPU end to end (max diff ~1e-5). Add tests/models/test_clip.py (CLIPVisionModel, batch 2, patch 32), wire it as a self-hosted CI job, and list CLIP in the README model coverage.
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.
YWHyuk
force-pushed
the
feature/togsim-cpp-trace
branch
from
ae26ede to
2463f82
Compare
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
3 participants
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.