fix: validate unsigned AVX2 take indices

[connortsui20]

Summary

• Fix unsigned AVX2 take bounds checks with unsigned SIMD comparisons and masked gathers.
• Fall back to scalar when i32 gathers cannot address valid u32 indices.
• Accumulate validity masks in SIMD and reduce once after the loop.
• Add regression coverage and a 1M-index primitive-take benchmark.

This is actually not as large of a change as it looks, and nothing changes in the hot path (there's just a bit more validation). The rest of the diff comes from the regression tests.

[connortsui20]

I wonder if we its better to just always mask it out instead of checking for better perf? do we even allow take indices larger than i32::MAX?

[connortsui20]

so I asked codex to do another pass and it found this:

Details

The unsafe audit is not clean: I found one remaining UB path.

P1: incorrect mask packing for u64 indices and 32-bit values

At vortex-array/src/arrays/primitive/compute/take/avx2.rs:367, each 64-bit comparison mask is represented as two identical
32-bit lanes:

[m0, m0, m1, m1, m2, m2, m3, m3]

But _mm_shuffle_epi32::<0x55> selects dword 1 four times. After unpacking, the gather mask becomes:

[m0, m0, m2, m2]

It should be:

[m0, m1, m2, m3]

_mm256_mask_i64gather_epi32 pairs four i64 offsets with four i32 mask lanes, as confirmed by the official Rust
implementation (https://doc.rust-lang.org/src/core/stdarch/crates/core_arch/src/x86/avx2.rs.html). Therefore, an invalid
index in lane 1 or 3 can inherit a valid neighboring mask and perform an out-of-bounds gather before the post-loop
assertion runs.

The likely correction is 0b11_01_11_01 (0xDD), which selects dwords 1 and 3 from each half. We should add a direct mask-
packing regression so the test fails safely without executing an invalid gather.

Other unsafe boundaries

• vortex-array/src/arrays/primitive/compute/take/avx2.rs:412: size_of::() == size_of::() is only a
  debug_assert!, but the raw stores and set_len rely on it for soundness. The current dispatcher always pairs them
  correctly, but this safe generic function should use a release assertion or encode the relationship in its type/unsafe
  contract.

• vortex-array/src/arrays/primitive/compute/take/avx2.rs:135: an empty source with nonempty indices returns a zeroed output
  instead of panicking. Because the function accepts arbitrary V: Copy, zero is not necessarily a valid V. The empty-source
  branch should only return an empty buffer when indices is also empty.

• vortex-array/src/arrays/primitive/compute/take/avx2.rs:494: tests call the AVX2 function without runtime feature
  detection. Production dispatch is correctly gated, but the tests violate the unsafe contract on x86-64 CPUs without AVX2.

• vortex-array/src/arrays/primitive/compute/take/avx2.rs:107: the unchecked array construction is valid through the
  production call chain because output and validity lengths both equal the indices length. Still, the safe TakeImpl::take
  signature does not encode that invariant; using the safe constructor would remove this proof obligation outside the hot
  loop.

The index-slice transmute, unaligned loads/stores, deferred set_len, and signed gather offsets are otherwise sound once the
mask-packing bug is fixed: their bounds, alignment, initialization, and feature-detection invariants hold.

I guess I need to fix these too?

[connortsui20]

im going to try this on my machine at home that actually has avx2 and make sure everything here is correct

[connortsui20]

Follow-up on reducing the validity mask once after the AVX2 loop. (This is an implementation detail for the new changes, not a comparison to the current buggy develop. I will make a comment with that comparison soon)

Local benchmark on an AMD Ryzen 9 7950X, pinned to CPU 0. Each run used 200 Divan samples; values below are the median of 3 runs.

taskset -c 0 cargo bench -p vortex-array --bench take_primitive -- primitive_take_u32

u32 values / u32 indices	Median per 1M-index take	200 takes	Throughput
Per-vector movemask	401.7 µs	80.3 ms	2.489 Gindices/s
SIMD accumulation	331.6 µs	66.3 ms	3.015 Gindices/s
Change	-17.5%	-17.5%	+21.1%

Optimized LLVM IR diff (abridged; SSA names normalized)

- %all_valid = phi i32 [ -1, %entry ], [ %next_bits, %loop ]
+ %all_valid = phi <4 x i64> [ splat (i64 -1), %entry ], [ %next_mask, %loop ]

  %lane_mask = sext <8 x i1> %valid to <8 x i32>
- %mask_bytes = bitcast <8 x i32> %lane_mask to <32 x i8>
- %sign_bits = icmp slt <32 x i8> %mask_bytes, zeroinitializer
- %mask_bits = bitcast <32 x i1> %sign_bits to i32
- %next_bits = and i32 %all_valid, %mask_bits
+ %mask_vec = bitcast <8 x i32> %lane_mask to <4 x i64>
+ %next_mask = and <4 x i64> %all_valid, %mask_vec

+ ; Reduced once after the loop.
+ %mask_bytes = bitcast <4 x i64> %all_valid.lcssa to <32 x i8>
+ %invalid = icmp sgt <32 x i8> %mask_bytes, splat (i8 -1)
+ %mask_bits = bitcast <32 x i1> %invalid to i32

x86-64 assembly diff (u32/u32 hot loop)

- vpcmpgtd   ymm3, ymm0, ymm3
- vpmovmskb  r9d, ymm3
- and        r9d, ecx
+ vpcmpgtd   ymm4, ymm0, ymm4
+ vpand      ymm1, ymm1, ymm4
  vpgatherdd ymm5, dword ptr [r13 + 4*ymm3], ymm4

- vpcmpgtd   ymm3, ymm0, ymm3
- vpmovmskb  ecx, ymm3
- and        ecx, r9d
+ vpcmpgtd   ymm4, ymm0, ymm4
+ vpand      ymm1, ymm1, ymm4
  vpgatherdd ymm5, dword ptr [r13 + 4*ymm3], ymm4

  jne        .Lloop
+ vpmovmskb  ecx, ymm1
  cmp        ecx, -1

[connortsui20]

Comparison against the actual buggy develop base (f3d67955), rather than the intermediate per-vector-movemask implementation.

Local benchmark on an AMD Ryzen 9 7950X, pinned to CPU 0: 1M u32 indices into 256 u32 values. The binaries were alternated for four runs, with 200 Divan samples per run.

Version	Median per take	Throughput
Buggy develop (f3d67955)	338.4 µs	2.955 Gindices/s
Current PR (f27ef984)	335.4 µs	2.982 Gindices/s
Observed change	-0.9%	+0.9%

Paired deltas ranged from 2.1% faster to 1.0% slower, so this is effectively performance-neutral. The correctness fix is not itself a performance optimization; accumulating the validity mask in SIMD keeps the added safety check from causing a measurable regression.

Optimized LLVM IR diff (u32/u32 hot loop, abridged; SSA names normalized)

- ; Signed inclusive comparison (incorrect for high-bit u32 and index == len).
- %valid = icmp sle <8 x i32> %indices, %len_splat
+ ; Sign-bit bias maps unsigned ordering to AVX2's signed comparison.
+ %biased_indices = xor <8 x i32> %indices, splat (i32 -2147483648)
+ %valid = icmp sgt <8 x i32> %biased_len_splat, %biased_indices
  %mask = sext <8 x i1> %valid to <8 x i32>
+ %mask64 = bitcast <8 x i32> %mask to <4 x i64>
  %values = tail call <8 x i32> @llvm.x86.avx2.gather.d.d.256(
      <8 x i32> zeroinitializer, ptr %src, <8 x i32> %indices,
      <8 x i32> %mask, i8 4)
  store <8 x i32> %values, ptr %dst
+ %all_valid.next = and <4 x i64> %all_valid, %mask64

+ ; Reduced once after the SIMD loop.
+ %mask_bytes = bitcast <4 x i64> %all_valid.lcssa to <32 x i8>
+ %nonnegative = icmp sgt <32 x i8> %mask_bytes, splat (i8 -1)
+ %invalid_bits = bitcast <32 x i1> %nonnegative to i32
+ %all_valid = icmp eq i32 %invalid_bits, 0
+ br i1 %all_valid, label %remainder, label %panic

x86-64 assembly diff (u32/u32 hot loop, abridged; registers normalized)

  vmovdqu    ymm_indices, [indices]
- vpcmpgtd   ymm_mask, ymm_indices, ymm_len
- vpxor      ymm_mask, ymm_mask, ymm_all_ones
+ vpxor      ymm_biased_indices, ymm_indices, ymm_sign_bit
+ vpcmpgtd   ymm_mask, ymm_biased_len, ymm_biased_indices
+ vpand      ymm_all_valid, ymm_all_valid, ymm_mask
  vpgatherdd ymm_values, [values + 4*ymm_indices], ymm_mask
  vmovdqu    [output], ymm_values

  jne        .Lloop
+ vpmovmskb  ecx, ymm_all_valid
+ cmp        ecx, -1
+ jne        .Lpanic

The loop is unrolled twice, so the vpxor/vpand pair appears twice per loop body. The gather sequence remains the dominant work.

Codegen command: cargo rustc -p vortex-array --lib --profile release -- --emit=llvm-ir,asm.

[codspeed-hq]

Merging this PR will improve performance by 26.41%

⚡ 1 improved benchmark
✅ 1840 untouched benchmarks
🆕 1 new benchmark
⏩ 55 skipped benchmarks¹

Performance Changes

	Mode	Benchmark	BASE	HEAD	Efficiency
⚡	Simulation	dict_canonicalize_zipfian[16, 1000]	24.7 µs	19.5 µs	+26.41%
🆕	Simulation	primitive_take_u32	N/A	4.5 ms	N/A

Tip

Curious why this is faster? Comment @codspeedbot explain why this is faster on this PR, or directly use the CodSpeed MCP with your agent.

---

_{Comparing ct/avx-signed-take-bug (38430b8) with develop (81c36da)}

Footnotes

1. 55 benchmarks were skipped, so the baseline results were used instead. If they were deleted from the codebase, click here and archive them to remove them from the performance reports. ↩

[connortsui20]

this was somewhat questionable in the first place, I don't think this behavior is what we want?

[robert3005]

I'm no expert on UB but this look correct now