Rust · Matrix Kernels

Matrix Kernels

use numkong::Dots;use numkong::DotsPackedMatrix;use numkong::DotsPackedOps;use numkong::SymmetricDotsOps;use numkong::Angulars;use numkong::AngularsPackedOps;use numkong::Euclideans;use numkong::EuclideansPackedOps;use numkong::SymmetricAngularsOps;use numkong::SymmetricEuclideansOps;use numkong::Hammings;use numkong::HammingsPackedOps;use numkong::Jaccards;use numkong::JaccardsPackedOps;use numkong::SymmetricHammingsOps;use numkong::SymmetricJaccardsOps;use numkong::Attention;use numkong::AttentionPackedMatrix;use numkong::AttentionRope; 8.0.0-rc.1 · b1f866ab

Packed kernels are the main “matrix throughput” path in the crate. They are GEMM-like in workload shape. They are not a thin BLAS clone.

use numkong::{DotsPackedMatrix, Tensor};

let a = Tensor::<f32>::full(&[1024, 512], 1.0).unwrap();
let b = Tensor::<f32>::full(&[256, 512], 1.0).unwrap();

let b_packed = DotsPackedMatrix::new(&b).unwrap();
let c = a.dots_packed(&b_packed).unwrap();

assert_eq!(c.shape(), &[1024, 256]);

The useful economics are:

  • pack B once
  • reuse it across many A batches
  • convert or pad once during packing instead of on every multiply
  • reuse precomputed norms for angulars_packed and euclideans_packed

The crate checks row contiguity because these kernels assume contiguous rows. Caller-side source alignment is not required. The owned packed buffer handles its own aligned allocation internally.

§ Symmetric Kernels for SYRK-Like Workloads

Symmetric kernels are for self-similarity and self-distance. They are SYRK-like in shape. They avoid duplicate (i, j) and (j, i) work.

use numkong::Tensor;

let vectors = Tensor::<f32>::full(&[100, 768], 1.0).unwrap();
let gram = vectors.view().dots_symmetric().unwrap();

assert_eq!(gram.shape(), &[100, 100]);

This family is also where row-window partitioning becomes the natural parallel model. That is structurally different from packed GEMM-style work against a shared packed RHS.

API

TraitDotsLow-level trait for batched dot product computation using pre-packed matrices.6
dots.rs:823 ↗
pub trait Dots: StorageElement + Sealed

Given A ∈ ℝᵐˣᵏ and packed B ∈ ℝⁿˣᵏ, computes C ∈ ℝᵐˣⁿ where: Cᵢⱼ = aᵢ · bⱼ

B is pre-packed into a backend-specific layout for optimal access; strides are all in bytes.

When to use

Reach for this trait, or the matching Tensor::dots_packed* wrappers, whenever you multiply many query rows against the same matrix of database rows: pre-packing B once amortises layout-conversion cost across every subsequent query. The accumulator type is intentionally widened to avoid precision loss — f32 × f32 → f64, f16 × f16 → f32, i8 × i8 → i32, u8 × u8 → u32, and so on; see each impl's Dots::Accumulator. On u1x8 the multiply degenerates to a bitwise AND and the accumulator counts set bits into a u32.

Errors

Error::KernelFailed when the kernel refuses to run, like on a buffer packed under other Capabilities than the current ones.

Implementors

AccumulatorAccumulator type for the multiplication.
dots.rs:825 ↗
type Accumulator: StorageElement;
dots_pack_sizeReturns the size in bytes needed for the packed B matrix buffer.
dots.rs:828 ↗
fn dots_pack_size(columns: usize, depth: usize) -> Result<usize, Error>
dots_packed_shapeReads the packed B matrix's columns and depth from its header. # Safety packed must point to a buffer produced by dots_pack.
dots.rs:833 ↗
unsafe fn dots_packed_shape(packed: *const u8) -> Result<(usize, usize), Error>
dots_packPacks the B matrix into an optimized backend-specific layout.
dots.rs:841 ↗
unsafe fn dots_pack(
    matrix: *const Self,
    columns: usize,
    depth: usize,
    matrix_stride: usize,
    packed: *mut u8,
    columns_begin: usize,
    columns_end: usize,
) -> Result<(), Error>

Packing fuses no transposition, so a transposed B is transposed in a separate pass first.

Safety

  • b must point to valid memory for columns * depth elements
  • packed must point to a buffer of at least dots_pack_size(columns, depth) bytes
dots_packedComputes C = A × Bᵀ using packed B.
dots.rs:857 ↗
unsafe fn dots_packed(
    queries: *const Self,
    packed: *const u8,
    output: *mut Self::Accumulator,
    rows: usize,
    columns: usize,
    depth: usize,
    query_stride: usize,
    output_stride: usize,
) -> Result<(), Error>

Safety

  • a must point to valid memory for rows * depth elements with given stride
  • packed must be a buffer filled by dots_pack
  • c must point to valid memory for rows × columns elements with given stride
dots_symmetricComputes C = A × Aᵀ where C is symmetric.
dots.rs:878 ↗
unsafe fn dots_symmetric(
    vectors: *const Self,
    vector_count: usize,
    depth: usize,
    stride: usize,
    result: *mut Self::Accumulator,
    result_stride: usize,
    rows_begin: usize,
    rows_end: usize,
) -> Result<(), Error>

Given input matrix A of shape [n,k], computes the symmetric matrix of all pairwise dot products. Only the upper triangle is computed, then mirrored to the lower triangle.

Safety

  • vectors must point to valid memory for vector_count × depth elements with given stride
  • result must point to valid memory for vector_count×vector_count elements, given stride
  • Strides are in bytes, not elements
  • rows_begin <= rows_end; rows_end is clamped to vector_count
StructDotsPackedMatrixPre-packed B matrix for efficient repeated GEMM operations.16
dots.rs:2219 ↗
pub struct DotsPackedMatrix<Scalar: Dots, Alloc: Allocator = Global>

Uses raw memory allocation, not std::Vec, for maximum control.

When multiplying A × Bᵀ multiple times with the same B matrix, packing B once and reusing it is much faster than packing each time.

Usage

For C = A × Bᵀ where B is [n,k]:

// Requires linking against libnumkong C library
let b_packed = DotsPackedMatrix::new(&b_array).unwrap();
let c = a_array.dots_packed(&b_packed).unwrap();

For C = A × B where B is [k,n] in standard GEMM layout:

// Requires linking against libnumkong C library
let b_packed = DotsPackedMatrix::new_transposed(&b_array).unwrap();
let c = a_array.dots_packed(&b_packed).unwrap();

Implements

Debug, Send, Sync
cloneClone this packed matrix, returning an error on allocation failure.
dots.rs:2235 ↗
pub fn clone(&self) -> Result<Self, Error>
empty_inAn empty packed matrix owning no allocation, using a custom allocator.
dots.rs:2250 ↗
pub fn empty_in(alloc: Alloc) -> Self

Fill it with pack_into.

new_inPacks the matrix into a buffer allocated by alloc.
dots.rs:2260 ↗
pub fn new_in<Matrix, const MAX_RANK: usize>(matrix: &Matrix, alloc: Alloc) -> Result<Self, Error>
where
    Matrix: TensorRef<Scalar, MAX_RANK> + ?Sized,
pack_intoRepack b into this matrix's existing buffer, reusing the allocation when the packed size fits capacity and reallocating through the stored allocator only when it must grow.
dots.rs:2273 ↗
pub fn pack_into<Matrix, const MAX_RANK: usize>(&mut self, matrix: &Matrix) -> Result<(), Error>
where
    Matrix: TensorRef<Scalar, MAX_RANK> + ?Sized,

A steady-state loop over same-shaped operands then allocates at most once. b must be 2D with contiguous rows.

reservePre-grow the buffer to hold a columns-by-depth B matrix, so a later pack_into that fits stays allocation-free with a stable pointer — hoist this out of a repeated-pack loop.
dots.rs:2319 ↗
pub fn reserve(&mut self, columns: usize, depth: usize) -> Result<(), Error>
new_transposed_inPack Bᵀ where B is [k,n] row-major, the standard GEMM layout, using a custom allocator.
dots.rs:2405 ↗
pub fn new_transposed_in<Matrix, const MAX_RANK: usize>(
    matrix: &Matrix,
    alloc: Alloc,
) -> Result<Self, Error>
where
    Matrix: TensorRef<Scalar, MAX_RANK> + ?Sized,

Materializes the transpose into a contiguous buffer, then packs normally. Result computes: C = A × B

Returns Err if:

  • b is not 2D
  • b is a sub-byte type — transpose unsupported
  • allocation fails
*× 2Returns a reference to the allocator.
dots.rs:2436 ↗
pub fn allocator(&self) -> &Alloc
shapeReturns the shape [columns,depth] of the original B matrix.
dots.rs:2439 ↗
pub fn shape(&self) -> (usize, usize)
pack_sizeBytes a packed buffer occupies for a columns-by-depth B matrix of this scalar type under the active backend's layout.
dots.rs:2444 ↗
pub fn pack_size(columns: usize, depth: usize) -> Result<usize, Error>

Mirrors crate::attention::AttentionPackedMatrix::pack_size and lets a caller pre-size an external buffer without packing.

from_packed_bytes_inAdopt an externally-produced packed buffer by copying bytes into a container-owned allocation tagged with the given columns and depth.
dots.rs:2455 ↗
pub unsafe fn from_packed_bytes_in(
    bytes: &[u8],
    columns: usize,
    depth: usize,
    alloc: Alloc,
) -> Result<Self, Error>

The dots packed layout is not self-describing — it carries no dtype tag and its AMX variant drops the exact depth — so the caller must supply the shape the bytes were packed for.

Safety

bytes must be a valid packing of a columns-by-depth matrix of Scalar, produced by this build's packer for a compatible backend; anything else makes a later dots_packed read out of bounds or return garbage.

capacityBytes currently allocated (>= the live packed size).
dots.rs:2469 ↗
pub fn capacity(&self) -> usize
*× 2Reset to logically empty, keeping the allocation so the next pack_into reuses it.
dots.rs:2472 ↗
pub fn clear(&mut self)
as_*_bytes× 2Returns the packed data buffer.
dots.rs:2475 ↗
pub fn as__bytes(&self) -> &[u8]
as_*_ptr× 2Returns a pointer to the packed data.
dots.rs:2478 ↗
pub fn as__ptr(&self) -> *const u8
newPack B matrix where B is [n,k] row-major using the global allocator.
dots.rs:2486 ↗
pub fn new<Matrix, const MAX_RANK: usize>(matrix: &Matrix) -> Result<Self, Error>
where
    Matrix: TensorRef<Scalar, MAX_RANK> + ?Sized,

Result computes: C = A × Bᵀ

new_transposedPack Bᵀ where B is [k,n] row-major, standard GEMM layout, using the global allocator.
dots.rs:2496 ↗
pub fn new_transposed<Matrix, const MAX_RANK: usize>(matrix: &Matrix) -> Result<Self, Error>
where
    Matrix: TensorRef<Scalar, MAX_RANK> + ?Sized,

Result computes: C = A × B

TraitDotsPackedOpsExtension trait: packed GEMM (C = A × Bᵀ) for any immutable tensor reference — owned Tensor, borrowed TensorView, or TensorSpan.2
dots.rs:2682 ↗
pub trait DotsPackedOps<Scalar: Dots, const MAX_RANK: usize>: TensorRef<Scalar, MAX_RANK>

Blanket-implemented for every TensorRef, so an A operand backed by an mmap'd view can multiply against a pre-packed DotsPackedMatrix without first materializing an owned copy. The allocating entry point allocates its result through a clone of self's allocator.

Implementors

dots_packedDot-product multiply: C = self × packed_rightᵀ
dots.rs:2694 ↗
fn dots_packed<PackedAlloc: Allocator>(
    &self,
    packed_right: &DotsPackedMatrix<Scalar, PackedAlloc>,
) -> Result<Tensor<Scalar::Accumulator, Self::Alloc, MAX_RANK>, Error>
where
    Self::Alloc: Clone,

self must be 2D [m,k] with contiguous rows. packed_right contains B [n,k] packed. Returns C [m,n], allocated like self.

Returns Err if:

  • self is not 2D
  • self has non-contiguous rows
  • inner dimensions don't match
  • output allocation fails
  • the kernel refuses packed_right, like one packed under other capabilities
dots_packed_intoDot-product multiply into an existing output, avoiding allocation.
dots.rs:2727 ↗
fn dots_packed_into<PackedAlloc, OutputTensor, const OUTPUT_MAX_RANK: usize>(
    &self,
    packed_right: &DotsPackedMatrix<Scalar, PackedAlloc>,
    output: &mut OutputTensor,
) -> Result<(), Error>
where
    PackedAlloc: Allocator,
    OutputTensor: TensorMut<Scalar::Accumulator, OUTPUT_MAX_RANK>,

The output may be a &mut Tensor<...> or &mut TensorSpan<...>; any writable tensor container that implements TensorMut works. The kernel overwrites c entirely, so it need not arrive pre-initialized.

TraitSymmetricDotsOpsExtension trait: symmetric dot-product matrix for any TensorRef implementor.2
dots.rs:3102 ↗
pub trait SymmetricDotsOps<Scalar: Dots, const MAX_RANK: usize>: TensorRef<Scalar, MAX_RANK>
where
    Scalar::Accumulator: 'static,

Blanket-implemented for every TensorRef<Scalar, R>, so calling vectors.dots_symmetric() works on both owned Tensor and borrowed TensorView / TensorSpan. Only the upper triangle, including the diagonal, of the output is written by the kernel; the lower triangle is queries untouched — callers that need a full dense matrix must mirror it.

Prefer this extension trait when you have a generic TensorRef; use the inherent TensorView::dots_symmetric form instead when you already hold a view and want to skip the extra trait import.

Implementors

  • impl<Scalar: Dots, const R: usize, OutputTensor: TensorRef<Scalar, R>> SymmetricDotsOps<Scalar, R> for OutputTensor where Scalar::Accumulator: 'static
dots_symmetricComputes the upper triangle of the symmetric dot-product matrix of the rows with themselves.
dots.rs:3107 ↗
fn dots_symmetric(&self) -> Result<Tensor<Scalar::Accumulator, Self::Alloc, MAX_RANK>, Error>
where
    Self::Alloc: Clone,
dots_symmetric_intoWrites the symmetric dot-product matrix into pre-allocated output, touching only the upper triangle of it.
dots.rs:3116 ↗
fn dots_symmetric_into<Out, const OUTPUT_MAX_RANK: usize>(
    &self,
    output: &mut Out,
) -> Result<(), Error>
where
    Out: TensorMut<Scalar::Accumulator, OUTPUT_MAX_RANK>,
TraitAngularsLow-level trait for batched angular distance operations.3
spatials.rs:789 ↗
pub trait Angulars: Dots

Given A ∈ ℝᵐˣᵏ and packed B ∈ ℝⁿˣᵏ, computes C ∈ ℝᵐˣⁿ where: Cᵢⱼ = 1 − cos(θᵢⱼ) = 1 − (aᵢ · bⱼ) / (‖aᵢ‖ × ‖bⱼ‖)

Packing reuses Dots::dots_pack for optimal memory layout.

When to use

Angular distance, also called cosine, is the standard similarity metric for embeddings from neural language / vision models. The accumulator is widened during the dot-product step — f32 inputs accumulate in f64, f16 / bf16 in f32, i8 / u8 in i32 / u32 — then normalised back down to Self::SpatialResult. Pre-packing the corpus is the standard pattern for repeated-query retrieval.

Errors

Error::KernelFailed when the kernel refuses to run, like on a buffer packed under other Capabilities than the current ones.

Implementors

SpatialResultResult type for angular distances.
spatials.rs:791 ↗
type SpatialResult: StorageElement;
angulars_packedComputes angular distances between A rows and packed B columns.
spatials.rs:799 ↗
unsafe fn angulars_packed(
    queries: *const Self,
    packed: *const u8,
    output: *mut Self::SpatialResult,
    rows: usize,
    columns: usize,
    depth: usize,
    query_stride: usize,
    output_stride: usize,
) -> Result<(), Error>

Safety

  • a must point to valid memory for rows * depth elements with given stride
  • packed must be a buffer filled by Dots::dots_pack
  • c must point to valid memory for rows × columns result elements with given stride
angulars_symmetricComputes symmetric angular distance matrix.
spatials.rs:815 ↗
unsafe fn angulars_symmetric(
    vectors: *const Self,
    vector_count: usize,
    depth: usize,
    stride: usize,
    result: *mut Self::SpatialResult,
    result_stride: usize,
    rows_begin: usize,
    rows_end: usize,
) -> Result<(), Error>

Safety

  • vectors must point to valid memory for vector_count * depth elements
  • result must point to valid memory for vector_count * vector_count result elements
TraitAngularsPackedOpsExtension trait: packed angular distances (C = angular(A, Bᵀ)) for any immutable tensor reference — owned Tensor, borrowed TensorView, or TensorSpan.2
spatials.rs:2574 ↗
pub trait AngularsPackedOps<Scalar: Angulars, const MAX_RANK: usize>: TensorRef<Scalar, MAX_RANK>

Blanket-implemented for every TensorRef, so an A operand backed by an mmap'd view can score against a pre-packed DotsPackedMatrix without first materializing an owned copy. The allocating entry point allocates its result through a clone of self's allocator.

Implementors

angulars_packedAngular distances: C = angular(self, packed_rightᵀ)
spatials.rs:2586 ↗
fn angulars_packed<PackedAlloc: Allocator>(
    &self,
    packed_right: &DotsPackedMatrix<Scalar, PackedAlloc>,
) -> Result<Tensor<Scalar::SpatialResult, Self::Alloc, MAX_RANK>, Error>
where
    Self::Alloc: Clone,

self must be 2D [m,k] with contiguous rows. packed_right contains B [n,k] packed. Returns C [m,n], allocated like self.

Returns Err if:

  • self is not 2D
  • self has non-contiguous rows
  • inner dimensions don't match
  • output allocation fails
  • the kernel refuses packed_right, like one packed under other capabilities
angulars_packed_intoAngular distances into an existing output, avoiding allocation.
spatials.rs:2619 ↗
fn angulars_packed_into<PackedAlloc, OutputTensor, const OUTPUT_MAX_RANK: usize>(
    &self,
    packed_right: &DotsPackedMatrix<Scalar, PackedAlloc>,
    output: &mut OutputTensor,
) -> Result<(), Error>
where
    PackedAlloc: Allocator,
    OutputTensor: TensorMut<Scalar::SpatialResult, OUTPUT_MAX_RANK>,

The output may be a &mut Tensor<...> or &mut TensorSpan<...>; any writable tensor container that implements TensorMut works. The kernel overwrites c entirely, so it need not arrive pre-initialized.

TraitEuclideansLow-level trait for batched euclidean distance operations.3
spatials.rs:845 ↗
pub trait Euclideans: Dots

Given A ∈ ℝᵐˣᵏ and packed B ∈ ℝⁿˣᵏ, computes C ∈ ℝᵐˣⁿ where: Cᵢⱼ = √(max(0, ‖aᵢ‖² + ‖bⱼ‖² − 2 · aᵢ · bⱼ))

Packing reuses Dots::dots_pack for optimal memory layout.

When to use

Use for ℓ₂ distance in k-NN / clustering / metric learning. The kernel computes ‖a‖² + ‖b‖² − 2 a·b in the widened accumulator domain (f32 × f32 → f64, f16 × f16 → f32, i8 × i8 → i32, …) before the square-root, so numerical cancellation is bounded. Pre-pack the reference matrix once when queries are repeated against the same corpus.

Errors

Error::KernelFailed when the kernel refuses to run, like on a buffer packed under other Capabilities than the current ones.

Implementors

SpatialResultResult type for euclidean distances.
spatials.rs:847 ↗
type SpatialResult: StorageElement;
euclideans_packedComputes euclidean distances between A rows and packed B columns.
spatials.rs:855 ↗
unsafe fn euclideans_packed(
    queries: *const Self,
    packed: *const u8,
    output: *mut Self::SpatialResult,
    rows: usize,
    columns: usize,
    depth: usize,
    query_stride: usize,
    output_stride: usize,
) -> Result<(), Error>

Safety

  • a must point to valid memory for rows * depth elements with given stride
  • packed must be a buffer filled by Dots::dots_pack
  • c must point to valid memory for rows × columns result elements with given stride
euclideans_symmetricComputes symmetric euclidean distance matrix.
spatials.rs:871 ↗
unsafe fn euclideans_symmetric(
    vectors: *const Self,
    vector_count: usize,
    depth: usize,
    stride: usize,
    result: *mut Self::SpatialResult,
    result_stride: usize,
    rows_begin: usize,
    rows_end: usize,
) -> Result<(), Error>

Safety

  • vectors must point to valid memory for vector_count * depth elements
  • result must point to valid memory for vector_count * vector_count result elements
TraitEuclideansPackedOpsExtension trait: packed euclidean distances (C = euclidean(A, Bᵀ)) for any immutable tensor reference — owned Tensor, borrowed TensorView, or TensorSpan.2
spatials.rs:2780 ↗
pub trait EuclideansPackedOps<Scalar: Euclideans, const MAX_RANK: usize>: TensorRef<Scalar, MAX_RANK>

Blanket-implemented for every TensorRef, so an A operand backed by an mmap'd view can score against a pre-packed DotsPackedMatrix without first materializing an owned copy. The allocating entry point allocates its result through a clone of self's allocator.

Implementors

euclideans_packedEuclidean distances: C = euclidean(self, packed_rightᵀ)
spatials.rs:2792 ↗
fn euclideans_packed<PackedAlloc: Allocator>(
    &self,
    packed_right: &DotsPackedMatrix<Scalar, PackedAlloc>,
) -> Result<Tensor<Scalar::SpatialResult, Self::Alloc, MAX_RANK>, Error>
where
    Self::Alloc: Clone,

self must be 2D [m,k] with contiguous rows. packed_right contains B [n,k] packed. Returns C [m,n], allocated like self.

Returns Err if:

  • self is not 2D
  • self has non-contiguous rows
  • inner dimensions don't match
  • output allocation fails
  • the kernel refuses packed_right, like one packed under other capabilities
euclideans_packed_intoEuclidean distances into an existing output, avoiding allocation.
spatials.rs:2825 ↗
fn euclideans_packed_into<PackedAlloc, OutputTensor, const OUTPUT_MAX_RANK: usize>(
    &self,
    packed_right: &DotsPackedMatrix<Scalar, PackedAlloc>,
    output: &mut OutputTensor,
) -> Result<(), Error>
where
    PackedAlloc: Allocator,
    OutputTensor: TensorMut<Scalar::SpatialResult, OUTPUT_MAX_RANK>,

The output may be a &mut Tensor<...> or &mut TensorSpan<...>; any writable tensor container that implements TensorMut works. The kernel overwrites c entirely, so it need not arrive pre-initialized.

TraitSymmetricAngularsOpsExtension trait: symmetric angular-distance matrix for any TensorRef implementor.2
spatials.rs:3073 ↗
pub trait SymmetricAngularsOps<Scalar: Angulars, const MAX_RANK: usize>: TensorRef<Scalar, MAX_RANK>

Implementors

angulars_symmetricComputes the upper triangle of the symmetric angular-distance matrix of the rows with themselves.
spatials.rs:3076 ↗
fn angulars_symmetric(&self) -> Result<Tensor<Scalar::SpatialResult, Self::Alloc, MAX_RANK>, Error>
where
    Self::Alloc: Clone,
angulars_symmetric_intoWrites the symmetric angular-distance matrix into pre-allocated output, touching only the upper triangle of it.
spatials.rs:3085 ↗
fn angulars_symmetric_into<Out, const OUTPUT_MAX_RANK: usize>(
    &self,
    output: &mut Out,
) -> Result<(), Error>
where
    Out: TensorMut<Scalar::SpatialResult, OUTPUT_MAX_RANK>,
TraitSymmetricEuclideansOpsExtension trait: symmetric euclidean distance matrix for any TensorRef implementor.2
spatials.rs:3107 ↗
pub trait SymmetricEuclideansOps<Scalar: Euclideans, const MAX_RANK: usize>: TensorRef<Scalar, MAX_RANK>

Blanket-implemented for every TensorRef<Scalar, R>, which means vectors.euclideans_symmetric() compiles whether vectors is an owned Tensor or a borrowed view. The kernel only writes the upper triangle, including the diagonal — the lower triangle is queries alone and callers should mirror it themselves if required.

Prefer this trait when working through a generic TensorRef; reach for the inherent TensorView::euclideans_symmetric method when you already hold a view.

Implementors

euclideans_symmetricComputes the upper triangle of the symmetric Euclidean-distance matrix of the rows with themselves.
spatials.rs:3110 ↗
fn euclideans_symmetric(
    &self,
) -> Result<Tensor<Scalar::SpatialResult, Self::Alloc, MAX_RANK>, Error>
where
    Self::Alloc: Clone,
euclideans_symmetric_intoWrites the symmetric euclidean-distance matrix into pre-allocated output, touching only the upper triangle of it.
spatials.rs:3119 ↗
fn euclideans_symmetric_into<Out, const OUTPUT_MAX_RANK: usize>(
    &self,
    output: &mut Out,
) -> Result<(), Error>
where
    Out: TensorMut<Scalar::SpatialResult, OUTPUT_MAX_RANK>,
TraitHammingsLow-level trait for batched Hamming distance operations.2
sets.rs:114 ↗
pub trait Hammings: Dots

Given A ∈ {0,1}ᵐˣᵏ and packed B ∈ {0,1}ⁿˣᵏ, computes C ∈ ℕᵐˣⁿ where: Cᵢⱼ = popcount(aᵢ ⊕ bⱼ)

Packing is inherited from the Dots supertrait.

When to use

Binary feature vectors, bit-packed into u1x8, are common in approximate nearest-neighbour search and bloom-filter-style retrieval. Packing the query set once and running Hamming distance against many candidate rows is the hot path. Results accumulate in u32, wide enough for any practical binary vector length.

Errors

Error::KernelFailed when the kernel refuses to run, like on a buffer packed under other Capabilities than the current ones.

Implementors

hammings_packedComputes Hamming distances between values matrix rows and packed query rows.
sets.rs:121 ↗
unsafe fn hammings_packed(
    queries: *const Self,
    packed: *const u8,
    result: *mut u32,
    rows: usize,
    columns: usize,
    depth: usize,
    v_stride: usize,
    r_stride: usize,
) -> Result<(), Error>

Safety

  • a must point to valid memory for the values matrix
  • packed must be a buffer filled by Dots::dots_pack
  • result must point to valid memory for rows × columns u32 elements
hammings_symmetricComputes symmetric Gram matrix of Hamming distances: C = A × Aᵀ.
sets.rs:137 ↗
unsafe fn hammings_symmetric(
    vectors: *const Self,
    vector_count: usize,
    depth: usize,
    stride: usize,
    result: *mut u32,
    result_stride: usize,
    rows_begin: usize,
    rows_end: usize,
) -> Result<(), Error>

Safety

  • vectors must point to valid memory for the input matrix
  • result must point to valid memory for vector_count * vector_count u32 elements
TraitHammingsPackedOpsExtension trait: packed Hamming distances (C = hamming(A, Bᵀ)) for any immutable tensor reference — owned Tensor, borrowed TensorView, or TensorSpan.2
sets.rs:356 ↗
pub trait HammingsPackedOps<Scalar: Hammings, const MAX_RANK: usize>: TensorRef<Scalar, MAX_RANK>

Blanket-implemented for every TensorRef, so an A operand backed by an mmap'd view can score against a pre-packed DotsPackedMatrix without first materializing an owned copy. The allocating entry point allocates its result through a clone of self's allocator.

Implementors

hammings_packedHamming distances: C = self × packed_rightᵀ
sets.rs:368 ↗
fn hammings_packed<PackedAlloc: Allocator>(
    &self,
    packed_right: &DotsPackedMatrix<Scalar, PackedAlloc>,
) -> Result<Tensor<u32, Self::Alloc, MAX_RANK>, Error>
where
    Self::Alloc: Clone,

self must be 2D [m,k] with contiguous rows. packed_right contains B [n,k] packed. Returns C [m,n], allocated like self.

Returns Err if:

  • self is not 2D
  • self has non-contiguous rows
  • inner dimensions don't match
  • output allocation fails
  • the kernel refuses packed_right, like one packed under other capabilities
hammings_packed_intoHamming distances into an existing output, avoiding allocation.
sets.rs:397 ↗
fn hammings_packed_into<PackedAlloc, OutputTensor, const OUTPUT_MAX_RANK: usize>(
    &self,
    packed_right: &DotsPackedMatrix<Scalar, PackedAlloc>,
    output: &mut OutputTensor,
) -> Result<(), Error>
where
    PackedAlloc: Allocator,
    OutputTensor: TensorMut<u32, OUTPUT_MAX_RANK>,

The output may be a &mut Tensor<...> or &mut TensorSpan<...>; any writable tensor container that implements TensorMut works. The kernel overwrites c entirely, so it need not arrive pre-initialized.

TraitJaccardsLow-level trait for batched Jaccard distance operations.3
sets.rs:226 ↗
pub trait Jaccards: Dots

Given A ∈ {0,1}ᵐˣᵏ and packed B ∈ {0,1}ⁿˣᵏ, computes C ∈ ℝᵐˣⁿ where: Cᵢⱼ = 1 − popcount(aᵢ ∧ bⱼ) / popcount(aᵢ ∨ bⱼ)

Packing is inherited from the Dots supertrait.

When to use

Jaccard distance measures set dissimilarity and is the natural metric for binary feature presence/absence. Pack the query set once and query many candidates in a single batched kernel call. The result type is f32 because the ratio is inherently fractional.

Errors

Error::KernelFailed when the kernel refuses to run, like on a buffer packed under other Capabilities than the current ones.

Implementors

JaccardResultResult type for Jaccard distances.
sets.rs:228 ↗
type JaccardResult: StorageElement;
jaccards_packedComputes Jaccard distances between values matrix rows and packed query rows.
sets.rs:236 ↗
unsafe fn jaccards_packed(
    queries: *const Self,
    packed: *const u8,
    result: *mut Self::JaccardResult,
    rows: usize,
    columns: usize,
    depth: usize,
    v_stride: usize,
    r_stride: usize,
) -> Result<(), Error>

Safety

  • a must point to valid memory for the values matrix
  • packed must be a buffer filled by Dots::dots_pack
  • result must point to valid memory for rows × columns elements
jaccards_symmetricComputes symmetric Gram matrix of Jaccard distances.
sets.rs:252 ↗
unsafe fn jaccards_symmetric(
    vectors: *const Self,
    vector_count: usize,
    depth: usize,
    stride: usize,
    result: *mut Self::JaccardResult,
    result_stride: usize,
    rows_begin: usize,
    rows_end: usize,
) -> Result<(), Error>

Safety

  • vectors must point to valid memory for the input matrix
  • result must point to valid memory for vector_count * vector_count elements
TraitJaccardsPackedOpsExtension trait: packed Jaccard distances (C = jaccard(A, Bᵀ)) for any immutable tensor reference — owned Tensor, borrowed TensorView, or TensorSpan.2
sets.rs:459 ↗
pub trait JaccardsPackedOps<Scalar: Jaccards, const MAX_RANK: usize>: TensorRef<Scalar, MAX_RANK>

Blanket-implemented for every TensorRef, so an A operand backed by an mmap'd view can score against a pre-packed DotsPackedMatrix without first materializing an owned copy. The allocating entry point allocates its result through a clone of self's allocator.

Implementors

jaccards_packedJaccard distances: C = self × packed_rightᵀ
sets.rs:471 ↗
fn jaccards_packed<PackedAlloc: Allocator>(
    &self,
    packed_right: &DotsPackedMatrix<Scalar, PackedAlloc>,
) -> Result<Tensor<Scalar::JaccardResult, Self::Alloc, MAX_RANK>, Error>
where
    Self::Alloc: Clone,

self must be 2D [m,k] with contiguous rows. packed_right contains B [n,k] packed. Returns C [m,n], allocated like self.

Returns Err if:

  • self is not 2D
  • self has non-contiguous rows
  • inner dimensions don't match
  • output allocation fails
  • the kernel refuses packed_right, like one packed under other capabilities
jaccards_packed_intoJaccard distances into an existing output, avoiding allocation.
sets.rs:504 ↗
fn jaccards_packed_into<PackedAlloc, OutputTensor, const OUTPUT_MAX_RANK: usize>(
    &self,
    packed_right: &DotsPackedMatrix<Scalar, PackedAlloc>,
    output: &mut OutputTensor,
) -> Result<(), Error>
where
    PackedAlloc: Allocator,
    OutputTensor: TensorMut<Scalar::JaccardResult, OUTPUT_MAX_RANK>,

The output may be a &mut Tensor<...> or &mut TensorSpan<...>; any writable tensor container that implements TensorMut works. The kernel overwrites c entirely, so it need not arrive pre-initialized.

TraitSymmetricHammingsOpsExtension trait: symmetric Hamming distance matrix for any TensorRef implementor.2
sets.rs:950 ↗
pub trait SymmetricHammingsOps<Scalar: Hammings, const MAX_RANK: usize>: TensorRef<Scalar, MAX_RANK>

Blanket-implemented for every TensorRef<Scalar, R>, exposing hammings_symmetric on owned Tensor as well as borrowed views. The kernel writes only the upper triangle, including the diagonal — the lower triangle is not touched, so mirror it if you need a fully-populated matrix.

Prefer this trait when writing generic code over TensorRef; use the inherent TensorView::hammings_symmetric when you already hold a view.

Implementors

hammings_symmetricComputes the upper triangle of the symmetric Hamming-distance matrix of the rows with themselves.
sets.rs:953 ↗
fn hammings_symmetric(&self) -> Result<Tensor<u32, Self::Alloc, MAX_RANK>, Error>
where
    Self::Alloc: Clone,
hammings_symmetric_intoWrites the symmetric Hamming-distance matrix into pre-allocated output, touching only the upper triangle of it.
sets.rs:962 ↗
fn hammings_symmetric_into<Out, const OUTPUT_MAX_RANK: usize>(
    &self,
    output: &mut Out,
) -> Result<(), Error>
where
    Out: TensorMut<u32, OUTPUT_MAX_RANK>,
TraitSymmetricJaccardsOpsExtension trait: symmetric Jaccard distance matrix for any TensorRef implementor.2
sets.rs:984 ↗
pub trait SymmetricJaccardsOps<Scalar: Jaccards, const MAX_RANK: usize>: TensorRef<Scalar, MAX_RANK>

Blanket-implemented for every TensorRef<Scalar, R>, so vectors.jaccards_symmetric() is available on both owned Tensor and borrowed views. The kernel writes only the upper triangle, including the diagonal; mirror to the lower triangle yourself if a dense symmetric result is required.

Prefer this trait when writing generic code over TensorRef; use the inherent TensorView::jaccards_symmetric when you already hold a view.

Implementors

jaccards_symmetricComputes the upper triangle of the symmetric Jaccard-distance matrix of the rows with themselves.
sets.rs:987 ↗
fn jaccards_symmetric(&self) -> Result<Tensor<Scalar::JaccardResult, Self::Alloc, MAX_RANK>, Error>
where
    Self::Alloc: Clone,
jaccards_symmetric_intoWrites the symmetric Jaccard-distance matrix into pre-allocated output, touching only the upper triangle of it.
sets.rs:996 ↗
fn jaccards_symmetric_into<Out, const OUTPUT_MAX_RANK: usize>(
    &self,
    output: &mut Out,
) -> Result<(), Error>
where
    Out: TensorMut<Scalar::JaccardResult, OUTPUT_MAX_RANK>,
TraitAttentionTrait abstracting ragged-attention pack/compute operations per scalar type.4
attention.rs:365 ↗
pub trait Attention: StorageElement + Clone

Errors

Error::KernelFailed when the kernel refuses to run, like on a buffer packed under other Capabilities than the current ones.

Implementors

attention_pack_sizeReturns bytes enough to pack token_count keys in segment_count segments, however split.
attention.rs:367 ↗
fn attention_pack_size(
    heads: usize,
    depth: usize,
    token_count: usize,
    segment_count: usize,
) -> Result<usize, Error>
attention_packed_shapeReads the KV-cache geometry — heads, depth, segments — from the packed header. # Safety packed must point to a buffer produced by attention_pack.
attention.rs:377 ↗
unsafe fn attention_packed_shape(packed: *const u8) -> Result<(usize, usize, usize), Error>
attention_packPack a window of the (segment, kv_head) task grid into the KV-cache blob. # Safety - k / v must point to token matrices with key_stride / value_stride byte rows covering every token addressed by key_offsets + key_lengths - key_offsets holds segment_count + 1 non-decreasing slot boundaries, and key_lengths is null or holds one count per segment, at most its slot's width - key_value_packed must have at least attention_pack_size(..) bytes
attention.rs:390 ↗
unsafe fn attention_pack(
    heads: usize,
    depth: usize,
    key_offsets: *const u32,
    key_lengths: *const u32,
    segment_count: usize,
    keys: *const Self,
    key_stride: usize,
    values: *const Self,
    value_stride: usize,
    key_value_packed: *mut u8,
    tasks_begin: usize,
    tasks_end: usize,
) -> Result<(), Error>

Windows run in any order: the one starting at task 0 also writes the header and directory, which no window reads.

attention_packedCompute attention over the half-open task window tasks_begin..tasks_end of the (query token, head) grid, task token * head_count + head, the end clipped to the grid.
attention.rs:420 ↗
unsafe fn attention_packed(
    head_count: usize,
    heads: usize,
    depth: usize,
    query_offsets: *const u32,
    query_token_count: usize,
    scale: f32,
    keys_before: usize,
    keys_after: usize,
    queries: *const Self,
    query_stride: usize,
    key_value_packed: *const u8,
    output: *mut f32,
    output_stride: usize,
    log_sum_exp: *mut f32,
    tasks_begin: usize,
    tasks_end: usize,
) -> Result<(), Error>

Each segment's queries align to the end of its keys: row r of a segment with q queries and k keys sits at position p = r + k - q and sees key j when p - keys_before <= j <= p + keys_after. usize::MAX leaves a side unbounded, so causal is (usize::MAX, 0), a sliding window of w keys is (w - 1, 0) and bidirectional is (usize::MAX, usize::MAX). Rows that see no key are written as zeros.

Safety

  • key_value_packed must have been produced by attention_pack with matching geometry
  • queries rows addressed by query_offsets must be valid, output writable with output_stride byte rows
  • query_token_count equals query_offsets[segments]
  • log_sum_exp is null or writable for one value per query token and head
StructAttentionPackedMatrixPre-packed ragged KV-cache for scaled-dot-product attention.17
attention.rs:991 ↗
pub struct AttentionPackedMatrix<Scalar: Attention, Alloc: Allocator = Global>

Owns the backend-opaque blob plus the geometry needed to validate query batches: segment count, KV head count, head width, and total token count.

Implements

Debug, Send, Sync
cloneClone this packed KV-cache, returning an error on allocation failure.
attention.rs:1007 ↗
pub fn clone(&self) -> Result<Self, Error>
empty_inAn empty cache that owns no allocation, holding only the given allocator.
attention.rs:1154 ↗
pub fn empty_in(alloc: Alloc) -> Self

Nothing is packed yet: every geometry field reads zero and as_bytes is empty until the first pack_into, which allocates on demand and can then be re-run each decode step to reuse the buffer.

new_inPack keys/values into a freshly allocated cache using the given allocator.
attention.rs:1166 ↗
pub fn new_in<KeysTensor, ValuesTensor, const MAX_RANK: usize>(
    keys: &KeysTensor,
    values: &ValuesTensor,
    depth: usize,
    key_offsets: &[u32],
    key_lengths: Option<&[u32]>,
    alloc: Alloc,
) -> Result<Self, Error>
where
    KeysTensor: TensorRef<Scalar, MAX_RANK> + ?Sized,
    ValuesTensor: TensorRef<Scalar, MAX_RANK> + ?Sized,
pack_intoRepack keys/values into this cache's existing blob, reusing the allocation when the packed size fits capacity and reallocating through the stored allocator only when it must grow.
attention.rs:1192 ↗
pub fn pack_into<KeysTensor, ValuesTensor, const MAX_RANK: usize>(
    &mut self,
    keys: &KeysTensor,
    values: &ValuesTensor,
    depth: usize,
    key_offsets: &[u32],
    key_lengths: Option<&[u32]>,
) -> Result<(), Error>
where
    KeysTensor: TensorRef<Scalar, MAX_RANK> + ?Sized,
    ValuesTensor: TensorRef<Scalar, MAX_RANK> + ?Sized,

The decode-loop path: allocate the cache once at layer-init, then refresh it every forward step with no further allocation. Packing overwrites, so a grow discards the old contents rather than copying them.

key_offsets holds the segment_count + 1 boundaries of the segments' key slots, and key_lengths, when given, how many keys each slot holds, at most its width; None fills every slot.

reservePre-grow the cache to hold the given ragged geometry, so a later pack_into that fits stays allocation-free with a stable pointer — allocate once at layer-init for the maximum sequence length, then refresh every decode step with no further allocation.
attention.rs:1239 ↗
pub fn reserve(&mut self, heads: usize, depth: usize, key_lengths: &[u32]) -> Result<(), Error>

key_lengths carries one token count per segment.

attention_intoRagged attention into a caller-provided f32 output tensor of the same logical shape as q — [tokens,heads * depth], contiguous rows.
attention.rs:1257 ↗
pub fn attention_into<QueriesTensor, OutTensor, const MAX_RANK: usize, const OUT_MAX_RANK: usize>(
    &self,
    queries: &QueriesTensor,
    query_offsets: &[u32],
    scale: Option<f32>,
    keys_before: usize,
    keys_after: usize,
    output: &mut OutTensor,
) -> Result<(), Error>
where
    QueriesTensor: TensorRef<Scalar, MAX_RANK> + ?Sized,
    OutTensor: TensorMut<f32, OUT_MAX_RANK> + ?Sized,

Each segment's queries align to the end of its keys, and query row r of a segment with q queries and k keys sees key j when r + k - q - keys_before <= j <= r + k - q + keys_after. usize::MAX leaves a side unbounded: causal is (usize::MAX, 0), a sliding window of w keys is (w - 1, 0) and bidirectional is (usize::MAX, usize::MAX). Rows that see no key are written as zeros.

pack_sizeBytes a packed cache needs for token_count keys in segment_count segments however they split, the size query mirroring Dots::dots_pack_size.
attention.rs:1302 ↗
pub fn pack_size(
    heads: usize,
    depth: usize,
    token_count: usize,
    segment_count: usize,
) -> Result<usize, Error>

Useful for pre-sizing an external buffer before packing, without constructing a cache.

peek_shapeRead the geometry of an externally-produced packed KV blob straight from its self-describing header, returning (heads, depth, segments), or Error::InvalidShape if the slice is too short to hold one.
attention.rs:1312 ↗
pub fn peek_shape(packed: &[u8]) -> Result<(usize, usize, usize), Error>

Unlike the dots packed matrix, the attention blob records its own shape, so no caller-supplied dimensions are needed.

Reads via the C nk_attention_packed_shape_serial_<dtype> accessor for this scalar type.

from_packed_bytes_inAdopt an externally-produced packed KV blob by copying bytes into a container-owned allocation.
attention.rs:1332 ↗
pub unsafe fn from_packed_bytes_in(bytes: &[u8], alloc: Alloc) -> Result<Self, Error>

The attention blob is self-describing, so (heads, depth, segments) are read back from its header; tokens reports 0 for an adopted blob because the total token count is not part of the queryable header.

Safety

bytes must be a valid attention packing for Scalar, produced by this build's packer; anything else makes a later attention read out of bounds.

shapeReturns the shape (heads, depth, segments) of the packed cache.
attention.rs:1343 ↗
pub fn shape(&self) -> (usize, usize, usize)
segmentsReturns the number of ragged segments in the packed cache.
attention.rs:1346 ↗
pub fn segments(&self) -> usize
headsReturns the number of KV heads per token.
attention.rs:1349 ↗
pub fn heads(&self) -> usize
depthReturns the number of channels per head.
attention.rs:1352 ↗
pub fn depth(&self) -> usize
tokensReturns the total KV token count across segments.
attention.rs:1355 ↗
pub fn tokens(&self) -> usize
capacityBytes currently allocated for the packed blob (>= the live packed size).
attention.rs:1358 ↗
pub fn capacity(&self) -> usize
newPack ragged K/V token matrices using the global allocator.
attention.rs:1380 ↗
pub fn new<KeysTensor, ValuesTensor, const MAX_RANK: usize>(
    keys: &KeysTensor,
    values: &ValuesTensor,
    depth: usize,
    key_offsets: &[u32],
    key_lengths: Option<&[u32]>,
) -> Result<Self, Error>
where
    KeysTensor: TensorRef<Scalar, MAX_RANK> + ?Sized,
    ValuesTensor: TensorRef<Scalar, MAX_RANK> + ?Sized,
attentionRagged attention allocating a fresh f32 output tensor through a clone of this cache's allocator.
attention.rs:1398 ↗
pub fn attention<QueriesTensor, const MAX_RANK: usize>(
    &self,
    queries: &QueriesTensor,
    query_offsets: &[u32],
    scale: Option<f32>,
    keys_before: usize,
    keys_after: usize,
) -> Result<Tensor<f32, Alloc>, Error>
where
    QueriesTensor: TensorRef<Scalar, MAX_RANK> + ?Sized,
    Alloc: Clone,

The allocating twin of attention_into.

TraitAttentionRopeIn-place NeoX split-half RoPE over a row-major [rows,head_count × depth] tensor.1
attention.rs:1786 ↗
pub trait AttentionRope: Sized + StorageElement

Implementors

attention_rope_intoRotates a 2D [rows,head_count × depth] tensor in place using the [rows,depth/2] cos/sin angle grids, shared across heads, for an even depth of channels per head.
attention.rs:1792 ↗
fn attention_rope_into<XMut, const RX: usize>(
    x: &mut XMut,
    cos: &[f32],
    sin: &[f32],
    head_count: usize,
    depth: usize,
) -> Result<(), Error>
where
    XMut: TensorMut<Self, RX> + ?Sized,

The row stride is read from the tensor, so x may be a non-contiguous sub-span, for example the Q or K column-section of a fused QKV buffer. Returns Err on a shape mismatch.