Skip to content

PiPNN kernels: follow-ups deferred from the #1287 review #1439

Description

@SeliMeli

Follow-ups from the review of #1287 (PiPNN numerical kernels) that did not land in that PR.

  • Choose the SIMD group width per architecture. LANES in diskann/src/graph/pipnn/simd.rs is one build-time constant for every architecture. Moving the choice into each architecture needs slice-based grouping (as_chunks::<{A::Vector::LANES}>() needs the unstable generic_const_exprs), a new safety argument for load_group, and an AVX2 comparison of 8 and 16 lanes. Related: Validate PiPNN SIMD width on AArch64/NEON #1423. (review)
  • Iterate mask lanes efficiently on aarch64. The top-k kernels visit the set lanes of a comparison through an integer bitmask. aarch64 has no instruction that builds this bitmask directly. A generic mask-lane iterator in diskann-wide could use the best method on each architecture. (review)
  • Remove the runtime-k allocation in select_top_k_symmetric. For k > 3, each call collects a Vec<&mut [Candidate]> with one slice per point. A small abstraction over the flat output and k could avoid it. Low priority: the benchmarked configurations use leaf_k 2 or 3. (review)
  • Keep large finite L2 distances finite. Leaf L2 writes ‖x‖² + ‖y‖², and GEMM adds -2·xᵀy. Partition L2 does the same with the leader norms. When the norm terms overflow, the result is +inf or NaN, so top-k drops the pair even when the true distance is finite. No test covers large inputs. (review)
  • Define empty-input behavior. The kernels document an empty leader set and an empty leaf as unsupported, and partitioning always samples at least one leader. The behavior differs by metric: some metrics panic in row_iter_mut on a matrix without columns. Decide whether these inputs return an error or produce no output.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions