[Model] Enable ConvNeXt V2 with batch > 1 (issue #255)

[YWHyuk]

Fixes #255. ConvNeXt V2 with batch > 1 died at codegen with AssertionError: 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=2 now 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_horizontal decided whether a reduction epilogue could be folded into a GEMM by parsing the node's repr():

[i.strip()[:-1].split(",")[-1].strip() for i in str(node).split("\n") if "r0" in i][1]

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, and set_ranges (mlir_common.py:584) asserted len(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_impl for CUTLASS, template_fusion_with_epilogues_supported for CPP): can_fuse stays pure, expressibility is judged from the IR, and declining is cheap and logged.

• MLIRTemplate.try_fuse_epilogue(template_node, existing_epilogues, node_to_fuse) and try_fuse_prologue(...) own every frame-dependent condition, config gates included, and return FusionPlan | None. They are pure, so the scheduler can call them speculatively.
• Mutations (loop merges, group realignment) are deferred into FusionPlan.remap, a zero-arg thunk applied from MLIRScheduling.fuse() -- the commit hook Scheduler.fuse_nodes_once already calls.
• The scheduler keeps only graph facts (_single_template_epilogue_pair, _single_prologue_template_pair) and delegates. It loses 155 lines.
• Every decline goes through why_no_fuse() and logs a reason at DEBUG on PyTorchSimFrontend.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 via body.vars[1]. Note that MemoryDep.index cannot be used here: it is normalized to a flat contiguous access and carries no per-axis strides, so index.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_ALIASING

The coordinate map that render() already used is hoisted to a class attribute. Its length is exactly what set_ranges asserts against, so the fusion gate is that assert, raised early as a decline instead of late as a crash. None means the template cannot absorb a reduction epilogue at all, which retires the support_reduction_fusion flag. 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 and mlir_gemm_template.py ends 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, unlike SchedulerNode.merge_loops()) whether collapsing would make it fit; if so the merge becomes the plan.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):

	value	meaning
MLIR memref.global	tile_desc.get_mlir_shape()	full tile, all lanes
.spad C header (Spike links this)	tile_size // vector_lane	per lane
gem5 C header	tile_size	full tile

index_expr() had the two headers swapped. The iota buffer ([0, 1, ..., compute_vec_size-1]) it allocates for vectorized index arithmetic declared compute_vec_size * vector_lane entries in the .spad header and compute_vec_size in the gem5 header. So Spike's buffer is vector_lane times 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 past compute_vec_size.

Only [0, compute_vec_size) is ever touched: one affine.parallel (%i) = (0) to (compute_vec_size) writes it, one affine.vector_load %buf[0] : vector<compute_vec_size x index> reads it.

Since the spad guard tightened to spad/2 for double buffering, the oversize side is fatal. ConvNeXt V2's var_mean kernel:

symbol	bytes/lane
buf0..buf3 (data tiles)	16384
buf4, buf5 (mean/var out)	64
index_expr_64_spad[8192]	65536
total	81984 (budget 65536)

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 -- SpadOverflowError is only caught inside autotune() -- so this is a hard failure.

The fix declares compute_vec_size + 1 entries per lane and vector_lane times that for gem5. The +1 is the two-wide initializer store, made explicit in a comment; it is the same slack idea as max(tile_numel_per_lane, 2) in codegen_spad_buffer(). The var_mean kernel now measures 16968 bytes/lane, and every other index_expr kernel shrinks by the same factor (index_expr_16_spad: 2048 entries -> 17).

The MLIR memref type is left alone: it already carries the full-tile count, matching the convention, and the .ll declares the symbol external so the C header is what defines the storage.

3. SpadOverflowError now says which kernel overflowed

It was raised with no arguments and the only diagnostic was an off-by-default logger.debug, so a failing compile printed SPAD 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

• ConvNeXt V2 batch=2: functional mode (Spike) vs CPU, Test Passed, max abs diff 5.0068e-06.
• Depthwise 7x7 conv (conv2d(x, w, b, padding=3, groups=96)), the smallest op that emits index_expr: functional mode, rebuilt binaries, max abs diff 7.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_fusion all still compile with their expected fused origins (e.g. mm, max_1 and bmm, max_1). Worth watching on review: BMM also supports reduction fusion, so putting the reduction conditions in the GEMM template alone silently splits test_bmm_reduction into two kernels.
• New regression test 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 reduces dim=-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)

• Chained epilogue fusion is still unsupported. try_fuse_epilogue already takes existing_epilogues, so enabling it only needs the scheduler's single-node gate relaxed (upstream CUTLASS does this).
• support_epilogue_fusion / support_prologue_fusion remain as declarative gates inside the template layer.