nvalchemiops.segment_ops: Segment Operations

The segment-ops module provides the framework-agnostic, Warp-level implementation of GPU-accelerated segmented reductions and per-segment algebra. All operations act on warp.array objects and expect the segment indices idx to be sorted in non-decreasing order.

Tip

This is the low-level Warp interface that operates on warp.array objects. For tensor bindings see nvalchemiops.torch.segment_ops: Segment Operations and nvalchemiops.jax.segment_ops: Segment Operations. For usage guidance and CUDA graph capture patterns, see Segment Operations.

Segmented operations for sorted segment indices.

Provides segmented reductions (sum, component_sum, dot, max_norm) and broadcasts (mul, add, matvec) over arrays grouped by sorted segment indices idx.

Reduction kernels use a run-length approach: each thread processes a contiguous chunk of elements, accumulates within runs of identical segment ids, and emits one atomic per segment boundary. The chunk size is auto-tuned based on N and the GPU’s SM count.

Reductions

Reduce the elements of each segment to a single per-segment result.

nvalchemiops.segment_ops.segmented_sum(x, idx, out)

Compute per-segment sum using run-length encoded reduction.

out[s] = sum(x[i] for i where idx[i] == s)

Requires sorted idx in non-decreasing order.

Parameters:

• x (wp.array, shape (N,)) – Input values. Supported dtypes: float32, float64, vec3f, vec3d.

• idx (wp.array(dtype=int32), shape (N,)) – Sorted segment indices in [0, M).

• out (wp.array, shape (M,), dtype matches x) – Per-segment sums. Zeroed internally before each use.

Return type:

None

nvalchemiops.segment_ops.segmented_component_sum(x, idx, out)

Sum all vector components to scalar per segment using RLE reduction.

out[s] = sum(x[i][0] + x[i][1] + x[i][2] for i where idx[i] == s)

Requires sorted idx in non-decreasing order.

Parameters:

• x (wp.array, shape (N,), dtype vec3f / vec3d) – Input 3D vectors.

• idx (wp.array(dtype=int32), shape (N,)) – Sorted segment indices in [0, M).

• out (wp.array, shape (M,), dtype float32 / float64) – Per-segment scalar sums. Zeroed internally before each use.

Return type:

None

See also

segmented_sum

Element-wise sum (preserves vector dtype)

segmented_dot

Dot product per segment

nvalchemiops.segment_ops.segmented_dot(x, y, idx, out)

Compute per-segment dot product reduction using RLE.

• Scalar types: out[s] = sum(x[i] * y[i] for i where idx[i] == s)

• Vector types: out[s] = sum(dot(x[i], y[i]) for i where idx[i] == s)

Requires sorted idx in non-decreasing order.

Parameters:

• x (wp.array, shape (N,)) – Input arrays. Must have same dtype. Supported: float32, float64, vec3f, vec3d.

• y (wp.array, shape (N,)) – Input arrays. Must have same dtype. Supported: float32, float64, vec3f, vec3d.

• idx (wp.array(dtype=int32), shape (N,)) – Sorted segment indices in [0, M).

• out (wp.array, shape (M,), dtype float32 / float64) – Per-segment dot products. Zeroed internally before each use.

Return type:

None

See also

segmented_inner_products

Compute x*y, x*x, and y*y in one pass

segmented_sum

Element-wise sum per segment

nvalchemiops.segment_ops.segmented_inner_products(x, y, idx, out_xy, out_xx, out_yy)

Compute three inner products per segment in one fused pass using RLE.

Computes out_xy[s], out_xx[s], out_yy[s] (the x*y, x*x, y*y inner products) in a single kernel launch, roughly 3x faster than three separate segmented_dot calls.

Requires sorted idx in non-decreasing order.

Parameters:

• x (wp.array, shape (N,)) – Input arrays. Must have same dtype. Supported: float32, float64, vec3f, vec3d.

• y (wp.array, shape (N,)) – Input arrays. Must have same dtype. Supported: float32, float64, vec3f, vec3d.

• idx (wp.array(dtype=int32), shape (N,)) – Sorted segment indices in [0, M).

• out_xy (wp.array, shape (M,), dtype float32 / float64) – Per-segment inner products. Zeroed internally before each use.

• out_xx (wp.array, shape (M,), dtype float32 / float64) – Per-segment inner products. Zeroed internally before each use.

• out_yy (wp.array, shape (M,), dtype float32 / float64) – Per-segment inner products. Zeroed internally before each use.

Return type:

None

See also

segmented_dot

Single dot product per segment

nvalchemiops.segment_ops.segmented_max_norm(x, idx, out)

Compute maximum vector norm per segment using RLE reduction.

out[s] = max(length(x[i]) for i where idx[i] == s)

Requires sorted idx in non-decreasing order. For single-segment cases (M=1) with N >= 8192, an optimized fast path reduces atomic contention.

Parameters:

• x (wp.array, shape (N,), dtype vec3f / vec3d) – Input 3D vectors.

• idx (wp.array(dtype=int32), shape (N,)) – Sorted segment indices in [0, M).

• out (wp.array, shape (M,), dtype float32 / float64) – Maximum norms per segment. Zeroed internally before each use.

Return type:

None

See also

segmented_dot

Squared norms (v*v) per segment

nvalchemiops.segment_ops.segmented_max(x, idx, out)

Compute maximum scalar value per segment using RLE reduction.

out[s] = max(x[i] for i where idx[i] == s)

Requires sorted idx in non-decreasing order. Caller must initialize out to -inf before calling.

Parameters:

• x (wp.array, shape (N,), dtype float32 / float64) – Input scalar values.

• idx (wp.array(dtype=int32), shape (N,)) – Sorted segment indices in [0, M).

• out (wp.array, shape (M,), dtype matches x) – Maximum values per segment. Must be initialized to -inf by caller.

Return type:

None

Notes

float16 is not supported (Warp atomic_max requires float32/float64).

nvalchemiops.segment_ops.segmented_min(x, idx, out)

Compute minimum scalar value per segment using RLE reduction.

out[s] = min(x[i] for i where idx[i] == s)

Requires sorted idx in non-decreasing order. Caller must initialize out to +inf before calling.

Parameters:

• x (wp.array, shape (N,), dtype float32 / float64) – Input scalar values.

• idx (wp.array(dtype=int32), shape (N,)) – Sorted segment indices in [0, M).

• out (wp.array, shape (M,), dtype matches x) – Minimum values per segment. Must be initialized to +inf by caller.

Return type:

None

Notes

float16 is not supported (Warp atomic_min requires float32/float64).

nvalchemiops.segment_ops.segmented_mean(x, idx, sums, counts, out)

Compute per-segment mean by composing sum + count + divide.

out[s] = sum(x[i] for i where idx[i] == s) / count(s)

For vector types (vec3f/vec3d), computes the mean vector per segment. Requires sorted idx in non-decreasing order.

All scratch arrays must be pre-allocated. This avoids hidden allocations and allows reuse across calls.

Parameters:

• x (wp.array, shape (N,), dtype float32 / float64 / vec3f / vec3d) – Input values to average per segment.

• idx (wp.array(dtype=int32), shape (N,)) – Sorted segment indices in [0, M).

• sums (wp.array, shape (M,), dtype matches x) – Scratch array for per-segment sums. Zeroed internally before each use.

• counts (wp.array(dtype=int32), shape (M,)) – Scratch array for per-segment element counts. Zeroed internally before each use.

• out (wp.array, shape (M,), dtype matches x) – Output mean values per segment.

Return type:

None

Notes

• float16/vec3h not supported (reduction requires atomics).

• Composes segmented_sum + segmented_count + divide internally.

• sums and counts are written to as scratch space.

nvalchemiops.segment_ops.segmented_rms_norm(x, idx, sum_sq, counts, out)

Compute RMS (root mean square) vector norm per segment.

out[s] = sqrt(mean(dot(x[i], x[i]) for i where idx[i] == s))

This is sqrt(sum_of_squared_norms / count) for each segment. Requires sorted idx in non-decreasing order.

All scratch arrays must be pre-allocated.

Parameters:

• x (wp.array, shape (N,), dtype vec3f / vec3d) – Input 3D vectors.

• idx (wp.array(dtype=int32), shape (N,)) – Sorted segment indices in [0, M).

• sum_sq (wp.array, shape (M,), dtype float32 / float64) – Scratch array for per-segment sum of squared norms. Zeroed internally before each use.

• counts (wp.array(dtype=int32), shape (M,)) – Scratch array for per-segment element counts. Zeroed internally before each use.

• out (wp.array, shape (M,), dtype float32 / float64) – Output RMS norms per segment. Precision matches x.

Return type:

None

Notes

• float16/vec3h not supported (reduction requires atomics).

• Composes segmented_dot + segmented_count + finalize internally.

• sum_sq and counts are written to as scratch space.

nvalchemiops.segment_ops.segmented_count(idx, out)

Count elements per segment using run-length encoding.

out[s] = count(i where idx[i] == s)

Requires sorted idx in non-decreasing order.

Parameters:

• idx (wp.array(dtype=int32), shape (N,)) – Sorted segment indices in [0, M).

• out (wp.array(dtype=int32), shape (M,)) – Per-segment element counts. Zeroed internally before each use.

Return type:

None

Broadcasts and Element-wise Algebra

Apply a per-segment value across every element of its segment.

nvalchemiops.segment_ops.segmented_mul(x, y, idx, out)

Broadcast multiply: out[i] = x[i] * y[idx[i]].

Supports mixed types (vec3 * scalar) as well as same-type (scalar * scalar). idx need not be sorted.

Parameters:

• x (wp.array, shape (N,)) – Per-atom input values. Supported: vec3h, vec3f, vec3d, float16, float32, float64.

• y (wp.array, shape (M,)) – Per-segment scalar multipliers (precision matches x).

• idx (wp.array(dtype=int32), shape (N,)) – Segment indices in [0, M).

• out (wp.array, shape (N,), dtype matches x) – Scaled output.

Return type:

None

See also

segmented_add

Broadcast addition

segmented_axpy

Fused multiply-add (y += a*x)

nvalchemiops.segment_ops.segmented_add(x, y, idx, out)

Broadcast add: out[i] = x[i] + y[idx[i]].

Supports mixed types (vec3 + scalar, scalar + vec3) as well as same-type. idx need not be sorted.

Parameters:

• x (wp.array, shape (N,)) – Per-atom input values. Supported: vec3h, vec3f, vec3d, float16, float32, float64.

• y (wp.array, shape (M,)) – Per-segment values to broadcast and add.

• idx (wp.array(dtype=int32), shape (N,)) – Segment indices in [0, M).

• out (wp.array, shape (N,)) – Output array.

Return type:

None

See also

segmented_mul

Broadcast multiplication

segmented_axpy

Fused multiply-add

nvalchemiops.segment_ops.segmented_matvec(v, m, idx, out)

Per-segment matrix-vector multiply: out[i] = M[idx[i]]^T @ v[i].

Computes the transpose matrix-vector product for each element using its segment’s matrix. idx need not be sorted.

Dispatches through _segmented_mul_overloads since Warp’s native mul(vec, mat) computes v^T @ M = M^T @ v.

Parameters:

• v (wp.array, shape (N,), dtype vec3h / vec3f / vec3d) – Per-atom input vectors.

• m (wp.array, shape (M,), dtype mat33h / mat33f / mat33d) – Per-segment 3x3 matrices (precision must match v).

• idx (wp.array(dtype=int32), shape (N,)) – Segment indices in [0, M).

• out (wp.array, shape (N,), dtype matches v) – Output transformed vectors.

Return type:

None

Notes

Uses transpose convention: M^T @ v, not M @ v.

See also

segmented_mul

Broadcast scalar multiplication

segmented_add

Broadcast vector addition

nvalchemiops.segment_ops.segmented_broadcast(values, idx, out)

Broadcast per-segment values to per-element array: out[i] = values[idx[i]].

Pure gather operation that copies a per-segment value to every element belonging to that segment.

Parameters:

• values (wp.array, shape (M,)) – Per-segment values to broadcast. Supported dtypes: float16, float32, float64, vec3h, vec3f, vec3d.

• idx (wp.array(dtype=int32), shape (N,)) – Segment indices in [0, M). Need not be sorted.

• out (wp.array, shape (N,), dtype matches values) – Output array with broadcast values.

Return type:

None

nvalchemiops.segment_ops.segment_div(numerator, denominator, result)

Element-wise division with zero-denominator guard.

Computes result[i] = numerator[i] / denominator[i] where zero denominators produce zero results instead of NaN/inf.

Parameters:

• numerator (wp.array, shape (N,), dtype float16 / float32 / float64) – Numerator values.

• denominator (wp.array, shape (N,), dtype int32) – Denominator values (e.g., segment element counts).

• result (wp.array, shape (N,), dtype matches numerator) – Output division results.

Return type:

None

Fused Updates

In-place broadcast linear-combination updates (AXPY / AXPBY).

nvalchemiops.segment_ops.segmented_axpy(y, x, a, idx)

In-place segmented broadcast FMA: y[i] += x[i] * a[idx[i]].

Modifies y in-place. idx need not be sorted.

Parameters:

• y (wp.array, shape (N,)) – Accumulator, modified in-place. Supported dtypes: vec3f, vec3d, float32, float64 (also vec3h, float16 for half-precision).

• x (wp.array, shape (N,), dtype matches y) – Input array to scale and add.

• a (wp.array, shape (M,)) – Per-segment scalar multipliers (precision matches y).

• idx (wp.array(dtype=int32), shape (N,)) – Segment indices in [0, M).

Return type:

None

See also

segmented_axpby

Generalized linear combination (a*x + b*y)

segmented_mul

Broadcast multiply (out = x * y[idx])

nvalchemiops.segment_ops.segmented_axpby(out, a, x, b, y, idx)

Broadcast linear combination: out[i] = a[idx[i]] * x[i] + b[idx[i]] * y[i].

idx need not be sorted.

Parameters:

• out (wp.array, shape (N,)) – Output array. Supported dtypes: vec3f, vec3d, float32, float64 (also vec3h, float16).

• a (wp.array, shape (M,)) – Per-segment scalar multipliers for x.

• x (wp.array, shape (N,), dtype matches out) – First input array.

• b (wp.array, shape (M,)) – Per-segment scalar multipliers for y.

• y (wp.array, shape (N,), dtype matches out) – Second input array.

• idx (wp.array(dtype=int32), shape (N,)) – Segment indices in [0, M).

Return type:

None

See also

segmented_axpy

Simpler one-term version (y += a*x)

segmented_mul

Broadcast multiply only

segmented_add

Broadcast add only