Skip to content

Move weight prepack into exec plan construction - #30

Open
Human9000-bit wants to merge 7 commits into
enthropy7:mainfrom
Human9000-bit:opt-rework
Open

Move weight prepack into exec plan construction#30
Human9000-bit wants to merge 7 commits into
enthropy7:mainfrom
Human9000-bit:opt-rework

Conversation

@Human9000-bit

@Human9000-bit Human9000-bit commented Aug 13, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator

tracking issue: #29

  • tests: cover the four weight-sharing shapes no benchmark model has
  • plan: permute conv weights for the kernel, not in the loader
  • gpu: seed the accelerator env with the packed weight, not the initializer
  • docs: describe where conv weight layout is decided now
  • tests: build the siamese shape instead of looking for one to download

@Human9000-bit
tests: cover the four weight-sharing shapes no benchmark model has
a4a77ce 
The tracker model was the first thing to hit `fe95f5c`'s bug, in production
rather than in CI, because none of yolov8n / yolo11n / resnet-18 /
mobilenet-v3-small contains the structures that break weight prepacking.
Moving the permute into plan construction is about to touch every one of
them, so they get fixtures first.

Four hazards, in `tests/prepack.rs`:

- A depthwise weight read directly and through an `Identity`. Constant
  folding turns the alias into a new name; tracking the permuted layout per
  name loses it there. Reverting `fe95f5c` reproduces the tracker's exact
  error, `Grouped Conv channel mismatch: IC=16, OC=3, group=16` — the kernel
  height read back as the output-channel count.
- A Conv weight shared by two branches, each with its own BatchNormalization.
  This is the first model in the suite to reach `fold_conv_bn`'s shared-weight
  guard (fold_conv_bn.rs:80), which has never been exercised: neutering it
  changes the output at element 0.
- A Conv whose kernel is a graph input, as a Siamese head's cross-correlation
  is. Nothing to prepack, so it must fall through to the runtime repack —
  which agrees bitwise with the prepacked path today. Forcing the layout
  probes true, the mistake a kernel-derived gate would make, fails it with
  `weight I/G=1, input I/G=4`.
- All three permutes in one graph, with the depthwise pair fused, so a change
  that keeps two of them working does not pass.

Each oracle is a second graph that computes the same thing without the
sharing, rather than a hand-computed constant: reusing a tensor, or feeding
it in at run time, must not change what the graph computes.

`all_three_weight_layouts_reach_their_kernel` is coverage, not correctness.
Kernel selection turns out to be layout-independent — `weight_geometry`'s
OIHW branch is the true reading and the KHWC branches are the workaround — so
it cannot catch a missing permute. It exists to stop the fixture beside it
from quietly ceasing to exercise one.

Verified with `scripts/check-ci-local.sh`, which includes the
`cargo test -p yscv-onnx --features gpu --lib` job that AGENTS.md's gate list
omits: 246 there, against 243 before, the two CPU-gated fixtures correctly
not running.
@Human9000-bit
plan: permute conv weights for the kernel, not in the loader
aa1b117 
`load_onnx_model` pre-permuted Conv weights into whichever layout the NHWC
kernels wanted and overwrote the initializer with the result. Nothing in the
model then said what those bytes meant, so a name-keyed side table had to,
and every consumer had to consult it before reading a weight: three sets on
`OnnxModel`, an `ir::WeightLayout` tag to give the passes one place to ask,
inverse permutes in the exporter and the quantizer, and four-way layout
dispatches in shape inference, cost modelling and conv-param resolution.

It was also wrong. A pass that renames a weight — constant-folding an
`Identity` yields an alias — left the tag behind, and the tracker model died
on `Grouped Conv channel mismatch: IC=3, OC=3, group=16`, a `[KH, KW, C, 1]`
weight read as `[O, I, KH, KW]`. `fe95f5c` fixed that by carrying the tag
through the IR, correctly; this removes the reason a tag exists.

The permute moves to `plan::prepack`, keyed by slot id, and the permuted
tensors stay there instead of being written back. `initializers` is
ONNX-native OIHW throughout, so passes read what the model wrote and a
renamed weight is simply re-examined next time the plan is built. `TensorEnv`
serves the packed copy to the kernels and the initializer to everyone else.

That deletes, rather than relocates: `ir/weight_layout.rs` and the
`weight_layouts` map, the three sets on `OnnxModel`, `apply_ir`'s rebuild of
them, `fold_constants`' permuted-operand guard and its `Identity` carve-out,
`rewrite_convtranspose_dts`' layout check, and the un-permute helpers in
`exporter.rs` and `quantize/rewriter.rs`. `fold_conv_bn` and
`fold_conv_const_binary` lose the `channel_of` indirection: with OIHW the
output channel is axis 0 and each owns one contiguous block. Net -375 lines.

One bug found by the model diff and fixed here. Selecting weights to pack by
`op_type == "Conv"` skips every Conv `fuse_conv_relu` has renamed to
`Conv_Relu` — on mobilenet-v3-small, all nine squeeze-excite `fc1` layers,
which silently fell back to repacking each inference. The loader could not hit
this because it ran before any pass; the plan runs after. Selection is now by
`NodeKind`, as `resolve_conv_params` already did, and
`a_conv_renamed_by_fusion_keeps_its_packed_weight` pins it — the numbers do
not, since the fallback is slower rather than wrong.

Held to byte-identical output against main across 32 runs: yolov8n, yolo11n,
resnet-18 and mobilenet-v3-small, optimizer on and off, `YSCV_REORDER_FUSION_OFF`
on and off, at 1 and 4 threads. Thread count is pinned because
mobilenet-v3-small's own reduction order varies with it — the unpinned diff
shows ~7 ULP between two runs of the same binary.

`scripts/check-ci-local.sh` passes, including the `--features gpu` test job.
The accelerator builds still skip the depthwise and grouped permutes; that
build-time gate is now the only reason the layout has to be recorded at all,
and unifying it is left alone here because it would newly activate metal and
wgpu paths this host cannot run.
@Human9000-bit
gpu: seed the accelerator env with the packed weight, not the initial…
f4035d0 
…izer

Moving the conv permute into the plan changed what `model.initializers`
holds, and two accelerator paths were reading it on the old assumption.
Neither is reachable from the Linux gate.

`metal/graph.rs` reverse-permuted KHWC back to OIHW before uploading
constants to MPSGraph. That inverse is now unnecessary — the initializer is
already OIHW — and it read `model.khwc_weights`, which no longer exists, so
the file would not have compiled on macOS at all. The same blind spot as
`ad5b7ca`: `target_os = "macos"` code is invisible to every job this host can
run.

The wgpu paths pre-populate `TensorEnv` from `model.initializers` instead of
letting `get` fall back. That now writes ONNX-native bytes into a slot the
layout table describes as channel-last, and the Conv dispatch believes the
table — `gpu/mod.rs:1672` reads the shape as `[kH, kW, C_in/group, C_out]`
when `is_khwc_weight` says so. Silent corruption rather than a failure.

Both seeds now come from one place. `TensorEnv::insert_model_weights` takes
the packed copy where the plan made one, so the bytes and the layout tag
cannot disagree; the two metal sites that build GPU buffers keep reading the
initializer, since OIHW is what MPSGraph wants.

Still unverified on real hardware: this host has neither a Metal device nor a
GPU adapter, so `--features gpu` gets compiled and unit-tested but its conv
dispatch never executes.
@Human9000-bit
docs: describe where conv weight layout is decided now
506260e 
Four comments still told the reader that `load_onnx_model` permutes conv
weights and records the result in `khwc_weights` and friends. The kernel
contracts in `fused_pw_dw_3x3` named those tables as the source of the packed
form they require, and `plan/build.rs` still described weight prepacking as
the next thing to lift out of it.

Also says, at the top of the fusion scan, which of the two weight views a
reader is looking at: `initializers` is ONNX-native, `conv_weight` is the
packed copy, and code matching on a `Khwc`-shaped comment needs the second.
Getting that wrong is silent — the shapes are the same rank.
@Human9000-bit
tests: build the siamese shape instead of looking for one to download
1a26c7c 
Real Siamese trackers ship as separate backbone and head files. The two-input
graph is assembled by the host code that calls the backbone twice, so no
single exported file contains a weight shared between two Convs or a
parallelizable split — the two structures that made the tracker the model
worth testing against. Constructing it is the only way to have it in CI.

Two towers over separate inputs, twelve nodes each, sharing all six weights,
merged by an Add. That reaches two things nothing else in the suite did.

Weight sharing at scale: every Conv weight is read twice, across all three
packed layouts, rather than the single shared weight the existing fixture
exercises.

And tower-parallel execution. `node_branches` is non-empty here, so
`run_onnx_model_jit` forks the environment and runs the towers concurrently —
a path with no test coverage at all until now, and one the packed-weight
tables are reached through by reference. Pointing `fork`'s
`prepacked_conv_weights` at an empty slice while leaving the layout tags
intact — a fork that keeps the claim and loses the bytes — fails it in
`matmul` with `right.len() >= k * n`, on the worker threads, which is also
what confirms the fork is really being taken.

Oracle is `YSCV_FORCE_TOWER_PARALLEL` against `YSCV_NO_TOWER_PARALLEL`,
bitwise: the towers share only weights, which nobody writes, so running them
concurrently must not move a single bit.
@Human9000-bit
fix precision
e494a62 
it falied on arm target due to floating point error
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@enthropy7 enthropy7 Awaiting requested review from enthropy7

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@Human9000-bit