Pointwise ops on the C++ trace pipeline#291
Merged
Merged
Conversation
The cat template's _calculate_input_tile_sizes handed each input a concat-dim tile of min(extent, remaining_spad_budget). When that budget did not cover an input's full concat extent, the tile width could fail to divide the extent (e.g. extent 64, tile 62 on the 8-lane config). The copy loop steps by that tile over [0, extent), so a non-dividing tile produces a ragged final iteration with no remainder mask: the MVIN over-reads past the input and the MVOUT over-writes past the output, corrupting the result. This surfaced on the 8x8 (wrapper3) CI config, where the per-input spad budget lands at 126 for two 64-wide inputs, truncating the second tile to 62. It scrambled ~45% of the cat output and broke LlamaDecoderLayer (rotary rotate_half cat) and the diffusion UNet skip-connection concat. Larger arrays give a bigger budget so both inputs got the full 64-wide (dividing) tile, which is why 32x32 and 128x128 passed. Snap each input's concat-dim tile down to the largest divisor of its extent that fits the budget, so every loop iteration is full-width and in-bounds. When the budget already covers the full extent the tile is unchanged, so 32x32/128x128 codegen is a no-op. Verified with a minimal cat repro (64+64 -> 128 on 8x8: fails before, diff 0 after) and the real test_llama LlamaDecoderLayer at 8x8 (now passes). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
The vlane index_expr lowering reconstructs each lane element's logical index along the split axis as stride_dim + outer_dim * nr_vector_lane, where outer_dim is the sublane-row index. The next sublane row starts vector_lane * vlane_stride elements later (vector_lane lanes, each holding vlane_stride positions), but the multiplier used nr_vector_lane alone, dropping the vlane_stride factor. Any index tensor (e.g. arange/iota) whose split axis spills into a second sublane row then had every row past the first come out shifted, producing wrong values. This fired on the 8x8 (wrapper3) config: arange(32) at 8 lanes with vlane_stride 2 needs two sublane rows, so positions 16-31 came out as p-8. Larger arrays keep the 32-wide axis in a single row (outer_dim always 0), which is why 32x32/128x128 were unaffected. It broke test_llama's LlamaModel (rotary position ids) at 8x8. Multiply the outer/row term by vector_lane * vlane_stride. This equals the old value whenever vlane_stride == 1, so single-row splits and larger configs are a codegen no-op. Verified with a minimal arange(32) repro at 8x8 (positions 16-31 wrong before, diff 0 after) and the full test_llama LlamaModel at 8x8 (now passes). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
init_tile_size gave a 1-D tile of 2*vlane_stride*vector_lane regardless of the data extent. When the extent is smaller than that tile the reduction/pointwise MVIN reads past the end of the DRAM buffer, and the compute loop folds the out-of-range elements into the result. For a sum reduction that means adjacent-DRAM garbage is squared/added into the sum. This broke test_single_perceptron's Loss on the 128x128 config: the MSE mean reduces 128 elements but the tile was 2*2*128 = 512, so the MVIN over-read 384 elements past the 128-wide input and roughly doubled the squared-error sum (loss diff ~4593). It fired on wide-lane configs (tile 512 > extent 128); 32x32 gives tile 128 == extent and 8x8 gives a tile that divides 128, so both were unaffected. Cap the 1-D tile to min(2*vlane_stride*vector_lane, extent) so the tile never exceeds the data. It is a no-op whenever extent >= tile, so larger reductions and other configs are unchanged. Verified with a minimal mse_loss(1-D) repro at 128x128 (over-reads before -> diff 0 after, across extents 64/100/127/200/300) and the real test_single_perceptron at 128x128 (Loss now passes). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
torch.sort's frontend lowering exposes the sorted VALUES as the sole scheduler-visible template output; the INDICES buffer is written in place by the same kernel (make_inplace) but is not registered as an output. For argsort the values output is dead, so Inductor DCEs the whole sort kernel and the indices buffer is returned uninitialized (all-zeros), silently corrupting any argsort result (e.g. DeepSeek MoE expert routing). Register the in-place indices write as a MutationOutput owned by the sort operation so the scheduler keeps the kernel alive whenever indices is consumed. OperationBuffer.get_outputs() hardcodes [self], so add a small MLIRTemplateBuffer subclass that can own extra MutationOutputs and have MLIRTemplateCaller.output_node build it; the sort lowering attaches the indices mutation via add_mutation_output. For templates with no mutation this is behaviorally identical to the base (get_outputs == [self]). Then gate the values MVOUT on the values output being live: in argsort it is pruned from the kernel args, so emitting its MVOUT referenced an undeclared %YV. The values scratchpad is still produced internally for the bitonic compare; only the dead DRAM store is skipped. Verified: t.argsort() now returns correct indices (was all-zeros); torch.sort with both outputs live is unchanged; test_sort, test_matmul, test_cat, test_indirect_access unaffected by the output_node change. Note: x[argsort] end to end additionally needs the indirect-gather codegen fix (separate). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
convert_indirect_indexing moves an indirect index into an offset spad and
zeros the indirect term in the DRAM base-offset expression, but it only
walked index.args. A BARE indirect index (x[idx], where the load index IS
the symbol so index.args is empty) was left unzeroed, so the DMA base
offset's affine.apply referenced the loop-local index SSA value before it
was defined ("use of undeclared SSA value name"). Any 1-D gather with a
runtime index tensor failed to compile at codegen; it was masked for
argsort only because the sort DCE bug left the indices all-zero (a
degenerate constant load).
Subs any remaining indirect symbol to 0 after the arg loop (mirrors the
store path). The per-position gather is carried by the offset spad, so the
base offset is correctly 0.
Verified: x[idx] (N=32/64) and x[x.argsort()] now compile and match CPU;
test_indirect_access.py (incl. scatter/index_add and multi-dim) and
test_sort.py unchanged.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
_masked_bounds reverse-engineered per-dim padding from the single flat index offset and clamped each padded load to [pad, pad + input_extent). That recovery is ill-posed: the offset mixes the tap shift with the padding, and under channels_last the stride-1 channel axis absorbs the offset remainder, producing an impossible clamp (e.g. [4, 8) on a size-2 channel tile). Spike then zeroed the whole load, making channels_last depthwise conv 99.9 percent wrong while NCHW was fine. The clamp was also unnecessary: at pad positions the consumer already yields the correct value (a compute-side select for a padded gather, or the pad op's own fill), so reading OOB is harmless. Keep only the ragged tail clamp (glo=0, ghi=loop_extent), which is genuinely load-bearing for non-dividing tiles. Verified on 32x32: channels_last depthwise conv now matches CPU; NCHW conv and standalone ragged F.pad still pass. Broader regressions (padded reduction loads, full MobileNet) to be addressed separately.
Implement abs, sign, isnan, isinf, fmod, bitwise shifts, copysign, erfc, hypot, cosh, sinh, acos, asin, atan, atan2, acosh, asinh, atanh, and log/ atan (VCIX) in the Inductor MLIR backend. Float-only transcendentals cast non-float inputs to f32 while leaving native float widths (f16/f64) untouched.
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).
YWHyuk
force-pushed
the
feature/pointwise-ops-trace
branch
from
6ed4863 to
6d83bb1
Compare
YWHyuk
force-pushed
the
feature/togsim-cpp-trace
branch
from
3fd7ae1 to
79b1aec
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.
Re-homes the pointwise ops work onto the C++ trace pipeline (feature/togsim-cpp-trace).
Background: the original pointwise PRs (#289 and the coupled spike/gem5/llvm PRs) were
built against the old architecture where VCIX lowering is a C++ mlir-opt pass
(-test-pytorchsim-to-vcix, llvm-project#10). On the trace branch that pass is an
in-process Python pass (passes/lower_to_vcix.py), so this branch adapts the feature
to the trace architecture instead of rebasing #289 as-is.
What is included:
bitwise shifts, copysign, erfc, hypot, cosh, sinh, acos, asin, atan, atan2, acosh,
asinh, atanh, log/atan (VCIX). Float-only transcendentals cast non-float inputs to
f32 and keep native float widths (f16/f64).
MathLog/MathAtanToVCIX patterns to _MATH_VIV/MARKERS. math.log = (opcode 2, imm 2),
math.atan = (opcode 2, imm 3), matching spike rs1 / gem5 VS1 (sin=0, cos=1, log=2,
atan=3).
Not needed on trace (intentionally omitted):
exists in the trace pipeline.
Coupled hardware (already merged + released): riscv-isa-sim#6 (spike v1.0.6),
gem5#2 (v1.0.2), both adding vsin/vcos/vlog/vatan decode + exec.
Note: gem5 decode also gained the missing vcos (VS1=0x1) case so cos no longer
faults; spike gained e64/double support for vlog/vatan.