pub struct ErasedReducePlan { /* private fields */ }Expand description
Dtype-erased reduction wrapper.
This is the erased replay boundary for full-tensor scalar reductions and axis reductions with a fixed output layout. It supports only operations with an unambiguous identity value in the selected dtype.
§Evaluation order
The op is matched once per execution; every loop runs a kernel
specialized for one (dtype, op) pair. Integer results are exact
(wrapping) and ReduceOp::Max / ReduceOp::Min do not depend on
order, so the rules below only affect the rounding of floating point and
complex ReduceOp::Sum, ReduceOp::Product, ReduceOp::StableProduct
and ReduceOp::SumSquares. Two details of float max and min follow the
visit order: a NaN input makes the result NaN, but which NaN (payload and
sign) is returned is unspecified, and when the extremum is a zero its sign
is unspecified if both +0.0 and -0.0 occur.
ReduceOp::Product is the reassociated product, matching PyTorch: it
splits each run into lanes, uses the SIMD product where available, threads a
full reduction, and combines the partials in a fixed tree. It is fast and
never recomputes an output, but reassociating can overflow a lane to inf
and underflow another to 0, whose inf * 0 combine is NaN; an
alternating 2.0/0.5 input at 2^600/2^-600 scale can therefore give
NaN where the source-order product is representable (tenferro-rs#2150).
ReduceOp::StableProduct avoids that. For f32/f64 it tracks a signed
mantissa and an exact integer exponent per lane and across runs and workers,
so a partial cannot overflow or underflow; the final value is scaled with a
correctly rounded ldexp. An input that is zero, subnormal, infinite or NaN
cannot be represented this way, so an output that saw one is recomputed in
source order, preserving IEEE propagation. It is reassociated and threaded
like ReduceOp::Product otherwise. For integers it is the reassociated
wrapping product and for complex dtypes the source-order product.
- Full reductions (from
Self::compile, orSelf::compile_axeswith every axis reduced) visit the source in a layout derived order: extent one axes are dropped, the rest are ordered by increasing absolute stride and contiguous axes are fused into runs. Each run is reduced by the SIMD kernel (unit stride,simdfeature), a sixteen lane kernel (other unit stride runs) or an eight lane kernel (strided runs), and runs are combined left to right. A dense column major source therefore reduces exactly like one contiguous slice. - With a parallel context and more than
MINTHREADLENGTHelements, a full reduction splits the traversal into one deterministic range per worker thread. Each range runs the same run kernel as the serial path and the partials are combined in a fixed tree. The result is bitwise repeatable for a fixed layout and thread count, but may differ between thread counts and from the serial result. - Axis reductions with kept axes compute every output independently, so the result never depends on the context or thread count. Each output folds its reduced elements sequentially in caller reduced axis order, except that a unit stride leading reduced run (after fusing adjacent reduced axes) of length at least 16 is reduced by the contiguous run kernel, with runs folded in order. When instead the leading kept axis has unit stride, adjacent outputs are accumulated as a block from one contiguous source run per reduced element; each output still folds in caller reduced axis order.
Implementations§
Source§impl ErasedReducePlan
impl ErasedReducePlan
Sourcepub fn compile(
dtype: KernelDType,
op: ReduceOp,
dims: &[usize],
src_strides: &[isize],
) -> Result<Self>
pub fn compile( dtype: KernelDType, op: ReduceOp, dims: &[usize], src_strides: &[isize], ) -> Result<Self>
Validate and store a full-reduction plan for one dtype and source layout.
Sourcepub fn compile_axes(
dtype: KernelDType,
op: ReduceOp,
src_dims: &[usize],
src_strides: &[isize],
dest_dims: &[usize],
dest_strides: &[isize],
axes: &[usize],
) -> Result<Self>
pub fn compile_axes( dtype: KernelDType, op: ReduceOp, src_dims: &[usize], src_strides: &[isize], dest_dims: &[usize], dest_strides: &[isize], axes: &[usize], ) -> Result<Self>
Validate and store an axis-reduction plan for one dtype and fixed source/output layouts.
axes names the source axes reduced away. Output dimensions must be the
remaining source dimensions in source-axis order. When all axes are
reduced, any output layout with exactly one reachable element is accepted.
pub fn dtype(&self) -> KernelDType
pub fn op(&self) -> ReduceOp
Sourcepub fn execute(
&self,
ctx: &ExecContext,
dest: &mut ErasedRawStridedMut<'_>,
src: &ErasedRawStridedRef<'_>,
) -> Result<()>
pub fn execute( &self, ctx: &ExecContext, dest: &mut ErasedRawStridedMut<'_>, src: &ErasedRawStridedRef<'_>, ) -> Result<()>
Execute the reduction into an erased output descriptor.
Sourcepub fn execute_uninit(
&self,
ctx: &ExecContext,
dest: &mut ErasedRawStridedUninitMut<'_>,
src: &ErasedRawStridedPtr<'_>,
) -> Result<()>
pub fn execute_uninit( &self, ctx: &ExecContext, dest: &mut ErasedRawStridedUninitMut<'_>, src: &ErasedRawStridedPtr<'_>, ) -> Result<()>
On success, every reachable destination slot is fully overwritten;
unreachable holes are neither read nor initialized. Validation errors
are returned before any destination write. A panic during execution
may leave a partially initialized MaybeUninit destination, which is
still safely droppable; no readable value is promised for unwritten
reachable slots.