Skip to content

fix(linalg): enforce cosine u8 length contract - #8737

Merged
Xuanwo merged 5 commits into
mainfrom
gatekeeper/fix-8638-1
Sep 2, 2026
Merged

fix(linalg): enforce cosine u8 length contract#8737
Xuanwo merged 5 commits into
mainfrom
gatekeeper/fix-8638-1

Conversation

@lance-gatefixer

@lance-gatefixer lance-gatefixer Bot commented Aug 24, 2026

Copy link
Copy Markdown
Contributor

Summary

  • enforce equal input lengths at the scalar and dispatched accumulation boundaries before selecting a SIMD backend
  • retain debug-only assertions inside the unsafe SIMD kernels to match the sibling u8 distance implementations
  • cover both shorter and longer right-hand inputs in the regression test

Root cause

The safe u8 cosine dispatcher selected a backend whose only length guard was a debug assertion. Release builds removed that guard, while SIMD loops bounded loads by the left slice length and loaded from the right slice unchecked. The scalar path also silently truncated mismatched inputs through iterator zipping.

Validation

  • cargo fmt --all -- --check
  • cargo test -p lance-linalg --lib distance::cosine_u8 (8 passed)
  • cargo test -p lance-linalg --lib (250 passed, 1 ignored)
  • cargo test --release -p lance-linalg --lib distance::cosine_u8::tests::rejects_mismatched_lengths (2 passed)
  • cargo clippy --all --tests --benches -- -D warnings
  • cargo doc -p lance-linalg --no-deps

Fixes #8638

@github-actions github-actions Bot added A-index Vector index, linalg, tokenizer bug Something isn't working labels Aug 24, 2026
lance-gatekeeper[bot]

This comment was marked as outdated.

@lance-gatekeeper lance-gatekeeper Bot added the K-approved Latest Gatekeeper recommendation permits acceptance. label Aug 24, 2026
@lance-gatefixer

lance-gatefixer Bot commented Aug 25, 2026

Copy link
Copy Markdown
Contributor Author

Blocked: The current-head linux-build check is blocked by runner memory exhaustion in an unrelated io_uring test.

Remote head 5d0c78a contains the current main tip 5e38587; Rust run 32790844710, job 97631926922 failed only at lance-io::uring::tests::test_read_range because both io_uring workers returned Cannot allocate memory (os error 12). The live logs and branch diff were inspected: this PR changes only rust/lance-linalg/src/distance/cosine_u8.rs, so repairing lance-io here would be an unrelated speculative change. Rerun linux-build on a runner with sufficient memory; alternatively, run the failing lance-io test on an io_uring-enabled host with adequate memory and provide the result.

@LuciferYang

Copy link
Copy Markdown
Contributor

I had an independent fix for #8638 on a branch and opened #8860 before checking this issue's timeline, which was my mistake. I am closing #8860 in favour of this one. Three things came out of reviewing my own version that apply here.

The kernel-level asserts diverge from the two merged siblings. dot_u8 and l2_u8 were converted by #8594 and both kept debug_assert_eq! inside their unsafe kernels (dot_u8.rs:50 and :92, l2_u8.rs:56 and :95), putting the always-on check only at the public entries. This PR replaces the kernel debug_assert_eq!s with assert_equal_lengths, so cosine_u8 becomes the only one of the three that checks inside a #[target_feature] function.

Coverage is not the reason to prefer one or the other, since both close the hole. mod x86 has no pub, so the kernels are unreachable from outside the file, and the dispatcher calls exactly one of the three guarded functions either way. The argument for matching the siblings is that a reader comparing the three files sees one shape. If you want the check adjacent to the dispatch rather than spread across the arms, the spot is cosine_u8_accum, which is the function that calls DISPATCH and so is the structural equivalent of dot_u8.rs:151 and l2_u8.rs:160.

The three # Panics sections leave cosine_u8 the only one of the three u8 files with any. dot_u8 and l2_u8 carry the identical always-on check and have zero # Panics between them, so after this lands a reader comparing them can reasonably conclude dot_u8 tolerates a mismatch. I had these sections in my branch and took them out for that reason. The nearby precedent that looks like it argues for keeping them does not: dot_distance_arrow_batch's # Panics at dot.rs:807 documents a condition guarded by a debug_assert_eq! at :815, which is a panic release does not have. Either all three files get the sections in one change or none do here.

The rstest test is better than what I wrote. Two cases covering both length directions catches more than one direction does, and #[case::{name}(...)] is what rust/AGENTS.md asks for. Worth noting for anyone reading the diff: the reverse direction is not just symmetry. With a shorter, every read is in bounds and the failure is a distance normalized by the norm of a prefix of b, so it is wrong rather than unsafe, which is a different bug from the out-of-bounds read in the other direction.

One measurement that may be useful in the description. Replacing the always-on checks with a no-op takes the new test to failed under --profile ci as well as --release, because cosine_u8_accum_scalar's zip truncates instead of panicking once the check is gone. So this test has teeth in both profiles, which is not true of the debug_assert_eq! shape it replaces.

@lance-gatefixer

Copy link
Copy Markdown
Contributor Author

Addressed in 69ff6fe: moved the always-on length check to the dispatcher boundary, restored debug-only assertions inside the SIMD kernels to match dot_u8 and l2_u8, and removed the cosine-only panic sections. The two-direction rstest remains unchanged.

@lance-gatekeeper lance-gatekeeper Bot removed the K-approved Latest Gatekeeper recommendation permits acceptance. label Aug 29, 2026
lance-gatekeeper[bot]

This comment was marked as outdated.

@lance-gatekeeper lance-gatekeeper Bot added the K-approved Latest Gatekeeper recommendation permits acceptance. label Aug 29, 2026
@lance-gatekeeper lance-gatekeeper Bot removed the K-approved Latest Gatekeeper recommendation permits acceptance. label Aug 29, 2026
lance-gatekeeper[bot]

This comment was marked as outdated.

@lance-gatekeeper lance-gatekeeper Bot added the K-approved Latest Gatekeeper recommendation permits acceptance. label Aug 29, 2026
@lance-gatefixer

Copy link
Copy Markdown
Contributor Author

Blocked: The Ready PR is blocked by two unrelated current-head CI failures.

Remote head ace45a7 contains the current main tip 108f78e. Rust run 33250719598, linux-arm job 99095818271 failed only in the randomized Float32 legacy IVF_PQ multivector test at recall 0.49 versus the 0.50 requirement (3,271 passed, 1 failed); the exact focused test passed locally on this same head. Python run 33250719594, windows job 99095818216 failed only in test_index_cast_centroids because Torch Inductor could not read its generated temporary module (PermissionError, errno 13; 1,363 passed, 1 failed). The live logs and diff were inspected: this PR changes only rust/lance-linalg/src/distance/cosine_u8.rs, not either failing Float32/Torch path, so changing code here would be speculative and unrelated.

Rerun linux-arm and windows on clean runners. Alternatively, run the exact Rust test on ARM and the Python test on Windows with a writable Torch Inductor cache, then provide those results.

@lance-gatekeeper lance-gatekeeper Bot removed the K-approved Latest Gatekeeper recommendation permits acceptance. label Sep 2, 2026

@lance-gatekeeper lance-gatekeeper Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Gate recommendation: approve.

The latest base sync preserves the reviewed cosine_u8 patch exactly. New adjacent changes widen dot/L2 accumulation but do not alter cosine routing or its scalar and dispatched validation boundaries, so the release-mode safety and mismatched-length behavior remain intact.

@lance-gatekeeper lance-gatekeeper Bot added the K-approved Latest Gatekeeper recommendation permits acceptance. label Sep 2, 2026
@Xuanwo
Xuanwo merged commit b2cbefb into main Sep 2, 2026
38 checks passed
@Xuanwo
Xuanwo deleted the gatekeeper/fix-8638-1 branch September 2, 2026 04:58
@LuciferYang

Copy link
Copy Markdown
Contributor

Thank you @Xuanwo

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

A-index Vector index, linalg, tokenizer bug Something isn't working K-approved Latest Gatekeeper recommendation permits acceptance.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

bug: cosine_u8 SIMD kernels read past the end of b in release builds

2 participants