[Model] Enable ConvNeXt V2 with batch > 1 (issue #255)#295
Open
YWHyuk wants to merge 15 commits into
Open
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.
ConvNextV2 with batch > 1 died in codegen with "Index names length
mismatch: 2 != 3". The scheduler had approved fusing a reduction epilogue
into a GEMM whose 2-D (M rows x N cols) frame cannot express it, and
set_ranges then found more loop ranges than the template's coordinate map
has entries.
It approved it because the eligibility check read the reduction's stride
out of the node's repr:
stride = [i.strip()[:-1].split(",")[-1].strip()
for i in str(node).split("\n") if "r0" in i][1]
The intent (reject a reduction over the contiguous, stride-1 axis, which
the GEMM reduction template cannot codegen) was right, but picking the
second line that mentions "r0" lands on the wrong index expression once
the node has more than two dimensions -- so ConvNextV2's channels-first
LayerNorm mean was let through. Reading the stride off the LoopBody's
index expression instead makes it decline for the right reason. Note a
MemoryDep's index cannot be used here: it is normalized to a flat
contiguous access and no longer carries the per-axis strides.
Fixing that exposed the deeper problem: only a template knows its output
coordinate frame, yet can_fuse_horizontal hardcoded a case per template
kind and re-guessed that frame. It also mutated nodes (revert_group) from
inside a predicate the scheduler calls speculatively for every candidate
pair, carried a dead isinstance(MaxPoolTemplate) special case, and
declined without ever saying why.
So the decision moves to the templates, following how Inductor's own
CUTLASS and CPP backends do it:
- MLIRTemplate.try_fuse_epilogue / try_fuse_prologue own every condition
that depends on the frame -- including the config gates -- and return a
FusionPlan or None. They are pure; a node mutation the fusion needs is
deferred into the plan's remap and applied from MLIRScheduling.fuse(),
the commit hook Inductor calls once it really fuses. The plan is
recomputed there rather than cached, since can_fuse runs many times per
pair and node identities change as fusions land.
- The scheduler keeps only graph facts: which node carries the template,
the direction, and that this is a single-node to single-node fusion.
Case 1 (pointwise) and Case 2 (reduction) collapse into one delegation.
- Every decline logs a reason, like Inductor's WhyNoFuseNames. Declining
is a normal answer: upstream declines reduction epilogues outright.
- A template declares REDUCTION_EPILOGUE_ALIASING, the coordinate map its
reduction-epilogue codegen addresses. render() consumes it, and its
length is exactly what set_ranges asserts against -- so the fusion gate
is that assert, raised as a decline instead of a crash. Templates
without one cannot absorb a reduction epilogue, which retires the
support_reduction_fusion flag. GEMM's map has two entries, BMM's three.
When the rows do not fit, LoopBody.merge_loops() (which returns a new
body, so this stays pure) says whether merging them would; if it would,
the merge becomes the plan's remap.
tests/ops/fusion covers pointwise, reduction, prologue, conv and
attention fusion and is unchanged, as are the kernels each of them fuses.
Adds a matmul whose reduction is over the contiguous axis: it must not
fuse, and wrongly fusing it returns wrong values rather than failing to
compile.
ConvNextV2 batch > 1 no longer hits the assertion; it now stops at an
unrelated SPAD overflow, so issue #255's report is not fully closed.
SpadOverflowError was raised with no arguments and the only diagnostic was a logger.debug line that is off by default, so a failing compile said nothing beyond "SPAD overflow occurred." and gave no way to tell which kernel, which tiling, or by how much. load() already has the measured usage, the budget, the tile size and the kernel's origins in scope. Put them in the exception message.
Spad globals follow a fixed convention: the MLIR memref carries the full tile shape (all lanes), the .spad C header that Spike links declares the per-lane slice, and the gem5 header declares the full tile. codegen_spad_buffer() emits exactly that, tile_numel_per_lane and tile_numel_per_lane * vector_lane. index_expr() had the two headers swapped. The iota scratch buffer it allocates for vectorized index arithmetic declared compute_vec_size * vector_lane entries in the .spad header and compute_vec_size entries in the gem5 header. So the buffer Spike sees is vector_lane times too large, and the one gem5 sees is one element too small: the initializer stores two elements per iteration, so its last iteration writes one slot past compute_vec_size. The oversize side is what breaks compiles. Since the spad guard tightened to spad/2 for double buffering, a kernel with compute_vec_size 64 spends 64 KB of its 64 KB budget on an iota that only ever uses 64 entries. ConvNextV2 with batch > 1 fails there: its channels-first LayerNorm var_mean kernel needs 81984 bytes/lane, of which 65536 is the iota and only 16448 is real data. Declare compute_vec_size + 1 entries per lane, and vector_lane times that for gem5. The var_mean kernel now measures 16968 bytes/lane.
Add tests/models/test_convnextv2.py, a self-hosted CI job for it, and a model coverage row. batch=2 with depths=[1,1,1,1] and image_size=64 now matches CPU end to end, max abs diff 5.0e-06. The test keeps the runner on self-hosted like the other model tests: the depthwise 7x7 conv lowers to one gather kernel per tap, so the graph is wide.
YWHyuk
changed the title
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 #255. ConvNeXt V2 with
batch > 1died at codegen withAssertionError: Index names length mismatch: 2 != 3. Two independent frontend bugs were stacked behind that message; the second only became visible once the first was fixed.batch=2now matches CPU end to end, max abs diff 5.0e-06.1. The fusion gate read the reduction stride out of a repr string
can_fuse_horizontaldecided whether a reduction epilogue could be folded into a GEMM by parsing the node'srepr():The intent was right: a reduction over the output's contiguous (stride-1) axis cannot be expressed in the GEMM template's 2-D reduction-epilogue frame, so it must not fuse. But taking the second line containing
"r0"lands on the wrong index expression once the node has more than two dimensions. ConvNeXt V2's channels-first LayerNorm mean is exactly that shape, so it was let through, andset_ranges(mlir_common.py:584) assertedlen(dim_aliasing) == len(ranges)->2 != 3.The fix: templates own their own fusion conditions
Rather than repair the string parse, the decision layer now follows the pattern both upstream Inductor backends use (
CUDACPPScheduling._can_fuse_epilogue_implfor CUTLASS,template_fusion_with_epilogues_supportedfor CPP):can_fusestays pure, expressibility is judged from the IR, and declining is cheap and logged.MLIRTemplate.try_fuse_epilogue(template_node, existing_epilogues, node_to_fuse)andtry_fuse_prologue(...)own every frame-dependent condition, config gates included, and returnFusionPlan | None. They are pure, so the scheduler can call them speculatively.FusionPlan.remap, a zero-arg thunk applied fromMLIRScheduling.fuse()-- the commit hookScheduler.fuse_nodes_oncealready calls._single_template_epilogue_pair,_single_prologue_template_pair) and delegates. It loses 155 lines.why_no_fuse()and logs a reason at DEBUG onPyTorchSimFrontend.mlir.mlir_template.The stride check now reads the real index expression via
LoopBody.get_all_read_expr(buffer_name)and the real reduction symbol viabody.vars[1]. Note thatMemoryDep.indexcannot be used here: it is normalized to a flat contiguous access and carries no per-axis strides, soindex.coeff(reduce_var)is always 0, which silently disables the check and produces wrong values rather than a crash. That is presumably why the original parsed the repr.REDUCTION_EPILOGUE_ALIASINGThe coordinate map that
render()already used is hoisted to a class attribute. Its length is exactly whatset_rangesasserts against, so the fusion gate is that assert, raised early as a decline instead of late as a crash.Nonemeans the template cannot absorb a reduction epilogue at all, which retires thesupport_reduction_fusionflag. GEMM declares 2 entries (row + reduced N), BMM declares 3 (batch + row + reduced N) -- that length is the only difference between them, so the conditions live in the base class andmlir_gemm_template.pyends up with no fusion logic at all.When the epilogue needs more loop ranges than the frame addresses, the template asks
body.merge_loops()(which is pure, unlikeSchedulerNode.merge_loops()) whether collapsing would make it fit; if so the merge becomes theplan.remap. For ConvNeXt V2 it correctly refuses -- the reduced C axis sits between batch and spatial in the channels-first layout, so the rows are not contiguous -- and the reduction runs as its own kernel.2. The index-expression scratch buffer was sized per tile, not per lane
With the LayerNorm reduction now correctly running as its own kernel, ConvNeXt V2 hit
SpadOverflowError.Spad globals follow a fixed convention, visible in
codegen_spad_buffer()(mlir_codegen_backend.py:1552-1553):memref.globaltile_desc.get_mlir_shape().spadC header (Spike links this)tile_size // vector_lanetile_sizeindex_expr()had the two headers swapped. The iota buffer ([0, 1, ..., compute_vec_size-1]) it allocates for vectorized index arithmetic declaredcompute_vec_size * vector_laneentries in the.spadheader andcompute_vec_sizein the gem5 header. So Spike's buffer isvector_lanetimes too large, and gem5's is one element too small -- the initializer stores two elements per iteration (ops.broadcast(init_iter, 2), kept wide to avoid scalar ops), so its last iteration writes one slot pastcompute_vec_size.Only
[0, compute_vec_size)is ever touched: oneaffine.parallel (%i) = (0) to (compute_vec_size)writes it, oneaffine.vector_load %buf[0] : vector<compute_vec_size x index>reads it.Since the spad guard tightened to
spad/2for double buffering, the oversize side is fatal. ConvNeXt V2'svar_meankernel:buf0..buf3(data tiles)buf4, buf5(mean/var out)index_expr_64_spad[8192]The real data is 16 KB; an iota that uses 64 entries (512 B) exhausts the whole budget. The heuristic mapping path has no tile-shrink fallback --
SpadOverflowErroris only caught insideautotune()-- so this is a hard failure.The fix declares
compute_vec_size + 1entries per lane andvector_lanetimes that for gem5. The+1is the two-wide initializer store, made explicit in a comment; it is the same slack idea asmax(tile_numel_per_lane, 2)incodegen_spad_buffer(). Thevar_meankernel now measures 16968 bytes/lane, and every otherindex_exprkernel shrinks by the same factor (index_expr_16_spad: 2048 entries -> 17).The MLIR
memreftype is left alone: it already carries the full-tile count, matching the convention, and the.lldeclares the symbolexternalso the C header is what defines the storage.3.
SpadOverflowErrornow says which kernel overflowedIt was raised with no arguments and the only diagnostic was an off-by-default
logger.debug, so a failing compile printedSPAD overflow occurred.and nothing else.load()already has the measured usage, the budget, the tile size and the kernel's origins in scope. That is how the kernel above was identified.4. Test, CI, coverage
tests/models/test_convnextv2.py(batch=2,depths=[1,1,1,1],image_size=64), a self-hosted CI job, and a README coverage row.Verification
Test Passed, max abs diff5.0068e-06.conv2d(x, w, b, padding=3, groups=96)), the smallest op that emitsindex_expr: functional mode, rebuilt binaries, max abs diff7.6e-06. This exercises the resized buffer end to end and rules out an alignment regression from the smaller allocation.tests/ops/fusion/:test_matmul_reduction,test_bmm_reduction,test_matmul_activation,test_conv_fusion,test_prologue_fusionall still compile with their expected fusedorigins(e.g.mm, max_1andbmm, max_1). Worth watching on review: BMM also supports reduction fusion, so putting the reduction conditions in the GEMM template alone silently splitstest_bmm_reductioninto two kernels.tests/ops/fusion/test_matmul_reduction.py::test_matmul_reduce_last_dim: a matmul whose reduction is over the contiguous axis must not fuse. Fusing it anyway returns wrong values rather than failing to compile, so the test checks numbers. Every pre-existing test reducesdim=-2(a non-contiguous axis), so this hole existed.The original ConvNeXt V2 crash has no self-contained repro --
matmul+permute+mean,+addmm, and a full channels-first LayerNorm all fail to reproduce it, because the trigger needs a node with more than two dims whose second"r0"line is the wrong one. Rather than fabricate one, the model test guards it.Follow-ups (not in this PR)
try_fuse_epiloguealready takesexisting_epilogues, so enabling it only needs the scheduler's single-node gate relaxed (upstream CUTLASS does this).support_epilogue_fusion/support_prologue_fusionremain as declarative gates inside the template layer.