tenferro_cpu_fused/
lib.rs

1#![doc(hidden)]
2
3//! Fused CPU elementwise execution.
4
5/// Internal result alias for the fused CPU adapter.
6///
7/// # Examples
8///
9/// ```rust
10/// use tenferro_cpu_fused::Result;
11/// let result: Result<()> = Ok(());
12/// assert!(result.is_ok());
13/// ```
14pub type Result<T> = tenferro_tensor::Result<T>;
15pub use tenferro_tensor::{DType, Error, Tensor, TypedTensor};
16
17use std::mem::size_of_val;
18use strided_basic::{ErasedRawStridedPtr, ErasedRawStridedRef, Exec, KernelDType};
19use strided_fused::{ErasedFusedPlan, FusedInst, FusedOp, FusedPlan};
20use tenferro_cpu_basic::{
21    erased_raw_strided_ref, erased_raw_strided_uninit_mut, typed_host_data, BufferPool, PoolScalar,
22    PooledUninitOutput,
23};
24use tenferro_tensor::backend::{
25    ElementwiseFusionInputView, ElementwiseFusionOp, ElementwiseFusionPlan,
26};
27use tenferro_tensor::col_major_strides;
28use tenferro_tensor::{ShapeVec, StrideVec};
29
30const ELEMENTWISE_FUSION_OP: &str = "execute_elementwise_fusion";
31
32const ELEMENTWISE_FUSION_MIN_ELEMENTS: usize = 16 * 1024;
33
34/// Largest input count `strided_fused::ErasedFusedPlan` accepts.
35///
36/// [`elementwise_fusion_with_pool`] declines larger plans, and the CPU runtime
37/// engine reports this limit so the prepared planner does not plan such
38/// regions at all.
39///
40/// # Examples
41///
42/// ```rust
43/// assert_eq!(tenferro_cpu_fused::ERASED_FUSION_MAX_INPUTS, 4);
44/// ```
45pub const ERASED_FUSION_MAX_INPUTS: usize = 4;
46
47fn validate_elementwise_fusion_inputs(
48    inputs: &[&Tensor],
49    plan: &ElementwiseFusionPlan,
50) -> crate::Result<bool> {
51    if inputs.len() != plan.input_count() {
52        return Err(crate::Error::invalid_argument(
53            ELEMENTWISE_FUSION_OP,
54            "inputs",
55            format!(
56                "plan expects {} inputs but backend received {}",
57                plan.input_count(),
58                inputs.len()
59            ),
60        ));
61    }
62    if plan.outputs().is_empty() {
63        return Ok(false);
64    }
65    for input in inputs {
66        if input.dtype() != plan.dtype() {
67            return Err(crate::Error::dtype_mismatch(
68                ELEMENTWISE_FUSION_OP,
69                input.dtype(),
70                plan.dtype(),
71            ));
72        }
73    }
74    Ok(true)
75}
76
77fn strided_fused_op(op: ElementwiseFusionOp) -> FusedOp {
78    match op {
79        ElementwiseFusionOp::Add => FusedOp::Add,
80        ElementwiseFusionOp::Multiply => FusedOp::Multiply,
81        ElementwiseFusionOp::Negate => FusedOp::Negate,
82        ElementwiseFusionOp::Conj => FusedOp::Conj,
83        ElementwiseFusionOp::Divide => FusedOp::Divide,
84        ElementwiseFusionOp::Abs => FusedOp::Abs,
85        ElementwiseFusionOp::Exp => FusedOp::Exp,
86        ElementwiseFusionOp::Log => FusedOp::Log,
87        ElementwiseFusionOp::Sin => FusedOp::Sin,
88        ElementwiseFusionOp::Cos => FusedOp::Cos,
89        ElementwiseFusionOp::Tanh => FusedOp::Tanh,
90        ElementwiseFusionOp::Sqrt => FusedOp::Sqrt,
91        ElementwiseFusionOp::Rsqrt => FusedOp::Rsqrt,
92        ElementwiseFusionOp::Pow => FusedOp::Pow,
93        ElementwiseFusionOp::Expm1 => FusedOp::Expm1,
94        ElementwiseFusionOp::Log1p => FusedOp::Log1p,
95        // `strided-fused 0.6.0` propagates NaN through the ordered selects; the
96        // unfused kernels use the same `strided_basic` functions.
97        ElementwiseFusionOp::Maximum => FusedOp::Maximum,
98        ElementwiseFusionOp::Minimum => FusedOp::Minimum,
99        ElementwiseFusionOp::Clamp => FusedOp::Clamp,
100        // INVARIANT: `plan_uses_unfused_op` declines remainder/`erf` (which
101        // have no `strided_fused` instruction) before a strided plan is built.
102        ElementwiseFusionOp::Remainder | ElementwiseFusionOp::Erf => {
103            unreachable!(
104                "ops without a selected strided_fused instruction are declined before CPU fusion"
105            )
106        }
107    }
108}
109
110/// Whether the plan contains an op that `strided_fused` cannot replay.
111///
112/// `erf` has no `FusedOp` instruction yet, so a region containing it is
113/// declined and its ops run unfused.
114fn plan_uses_unfused_op(plan: &ElementwiseFusionPlan) -> bool {
115    plan.ops().iter().any(|inst| {
116        matches!(
117            inst.op(),
118            ElementwiseFusionOp::Remainder | ElementwiseFusionOp::Erf
119        )
120    })
121}
122
123fn should_defer_to_broadcast_multiply_special_case(plan: &ElementwiseFusionPlan) -> bool {
124    !plan.input_views().iter().all(|view| view.is_identity())
125        && plan.ops().len() == 1
126        && plan.outputs() == [plan.input_count()]
127        && plan.ops()[0].op() == ElementwiseFusionOp::Multiply
128}
129
130fn strided_fused_plan(plan: &ElementwiseFusionPlan) -> FusedPlan {
131    FusedPlan::new(
132        plan.input_count(),
133        plan.outputs().iter().copied(),
134        plan.ops()
135            .iter()
136            .map(|inst| FusedInst::new(strided_fused_op(inst.op()), inst.inputs().iter().copied())),
137    )
138}
139
140/// The two dtypes the CPU fused replay covers.
141#[derive(Clone, Copy)]
142enum FusedDType {
143    F32,
144    F64,
145}
146
147impl FusedDType {
148    /// The erased kernel tag for this dtype.
149    fn kernel(self) -> KernelDType {
150        match self {
151            FusedDType::F32 => KernelDType::F32,
152            FusedDType::F64 => KernelDType::F64,
153        }
154    }
155}
156
157/// The fused dtype a plan selects, or `None` when the region stays unfused.
158///
159/// Only `F32`/`F64` are selected: `strided-fused 0.6.0`'s vectorized chunk
160/// replay covers only those dtypes, and every other preset dtype would run the
161/// unmeasured per-element interpreter. Ordered ops (`Maximum`/`Minimum`/
162/// `Clamp`) are selected: `strided-fused 0.6.0` propagates NaN through them
163/// with the same `strided_basic` select the unfused kernels use, so the fused
164/// route matches the per-instruction route bit for bit (`R-numerics-01`,
165/// strided-rs#315).
166fn select_fused_dtype(plan: &ElementwiseFusionPlan) -> Option<FusedDType> {
167    match plan.dtype() {
168        DType::F32 => Some(FusedDType::F32),
169        DType::F64 => Some(FusedDType::F64),
170        _ => None,
171    }
172}
173
174fn typed_bytes<T>(data: &[T]) -> &[u8] {
175    // SAFETY: `data` is an aligned typed slice. The returned byte slice has
176    // the same lifetime and exact byte length, and is read-only.
177    unsafe { std::slice::from_raw_parts(data.as_ptr().cast::<u8>(), size_of_val(data)) }
178}
179
180struct ErasedFusionInput<'a> {
181    data: &'a [u8],
182    dims: ShapeVec,
183    strides: StrideVec,
184}
185
186fn tensor_host_bytes(dtype: FusedDType, input: &Tensor) -> crate::Result<&[u8]> {
187    let op = ELEMENTWISE_FUSION_OP;
188    macro_rules! bytes {
189        ($tensor:expr) => {
190            typed_host_data(op, $tensor).map(typed_bytes)
191        };
192    }
193
194    match dtype {
195        FusedDType::F32 => bytes!(fused_host::<f32>(op, input)?),
196        FusedDType::F64 => bytes!(fused_host::<f64>(op, input)?),
197    }
198}
199/// The typed tensor behind `input`, or the refusal this path produces for one.
200///
201/// Callers select `T` from the plan's fused dtype, so `None` means the tensor's
202/// runtime dtype disagrees with that selection rather than a caller mistake.
203fn fused_host<'a, T: tenferro_tensor::TensorScalar>(
204    op: &'static str,
205    input: &'a Tensor,
206) -> crate::Result<&'a TypedTensor<T>> {
207    input.as_typed::<T>().ok_or_else(|| {
208        crate::Error::unsupported_dtype(
209            op,
210            input.dtype(),
211            "an externally defined payload is not a fused input",
212        )
213    })
214}
215
216fn erased_fusion_input<'a>(
217    dtype: FusedDType,
218    input: &'a Tensor,
219    view: &ElementwiseFusionInputView,
220) -> crate::Result<ErasedFusionInput<'a>> {
221    let data = tensor_host_bytes(dtype, input)?;
222    let base_shape = input.shape();
223    let base_strides = col_major_strides(base_shape)?;
224    let ElementwiseFusionInputView::BroadcastInDim { shape, dims } = view else {
225        return Ok(ErasedFusionInput {
226            data,
227            dims: ShapeVec::from(base_shape),
228            strides: base_strides,
229        });
230    };
231
232    if dims.len() != base_shape.len() {
233        return Err(crate::Error::invalid_argument(
234            ELEMENTWISE_FUSION_OP,
235            "configuration",
236            format!(
237                "broadcast dims length {} does not match input rank {}",
238                dims.len(),
239                base_shape.len()
240            ),
241        ));
242    }
243
244    let mut strides: StrideVec = std::iter::repeat_n(0isize, shape.len()).collect();
245    let mut seen = vec![false; shape.len()];
246    for (source_axis, &target_axis) in dims.iter().enumerate() {
247        if target_axis >= shape.len() {
248            return Err(crate::Error::axis_out_of_bounds(
249                ELEMENTWISE_FUSION_OP,
250                target_axis,
251                shape.len(),
252            ));
253        }
254        if seen[target_axis] {
255            return Err(crate::Error::duplicate_axis(
256                ELEMENTWISE_FUSION_OP,
257                target_axis,
258                "broadcast dims",
259            ));
260        }
261        seen[target_axis] = true;
262        let source_dim = base_shape[source_axis];
263        let target_dim = shape[target_axis];
264        if source_dim != target_dim && source_dim != 1 {
265            return Err(crate::Error::shape_mismatch(
266                ELEMENTWISE_FUSION_OP,
267                shape.to_vec(),
268                base_shape.to_vec(),
269            ));
270        }
271        if source_dim == target_dim {
272            strides[target_axis] = base_strides[source_axis];
273        }
274    }
275
276    Ok(ErasedFusionInput {
277        data,
278        dims: ShapeVec::from(shape.as_slice()),
279        strides,
280    })
281}
282
283#[doc(hidden)]
284pub fn elementwise_fusion_with_pool(
285    buffers: &mut BufferPool,
286    inputs: &[&Tensor],
287    plan: &ElementwiseFusionPlan,
288    exec: Exec<'_>,
289) -> crate::Result<Option<Vec<Tensor>>> {
290    if !validate_elementwise_fusion_inputs(inputs, plan)? {
291        return Ok(None);
292    }
293    if inputs.is_empty() || inputs.len() > ERASED_FUSION_MAX_INPUTS {
294        return Ok(None);
295    }
296    if plan_uses_unfused_op(plan) {
297        return Ok(None);
298    }
299    if should_defer_to_broadcast_multiply_special_case(plan) {
300        return Ok(None);
301    }
302    let Some(dtype) = select_fused_dtype(plan) else {
303        return Ok(None);
304    };
305
306    let input_layouts = inputs
307        .iter()
308        .zip(plan.input_views())
309        .map(|(input, view)| erased_fusion_input(dtype, input, view))
310        .collect::<crate::Result<Vec<_>>>()?;
311    let shape = input_layouts[0].dims.clone();
312    if input_layouts
313        .iter()
314        .skip(1)
315        .any(|input| input.dims != shape)
316    {
317        return Ok(None);
318    }
319    let element_count =
320        tenferro_tensor::validate::checked_shape_product(ELEMENTWISE_FUSION_OP, "shape", &shape)?;
321    if element_count < ELEMENTWISE_FUSION_MIN_ELEMENTS {
322        return Ok(None);
323    }
324
325    let input_refs = input_layouts
326        .iter()
327        .map(|input| {
328            // SAFETY: fusion inputs are initialized typed storage with matching
329            // dtype and alignment; validated layouts bound every reachable read
330            // for the retained input borrow.
331            unsafe {
332                erased_raw_strided_ref(dtype.kernel(), input.data, &input.dims, &input.strides, 0)
333            }
334            .map_err(|err| crate::Error::backend_source(ELEMENTWISE_FUSION_OP, err))
335        })
336        .collect::<crate::Result<Vec<_>>>()?;
337
338    execute_erased_fused_outputs(buffers, dtype, &input_refs, &shape, plan, exec).map(Some)
339}
340
341fn execute_erased_fused_outputs(
342    buffers: &mut BufferPool,
343    dtype: FusedDType,
344    input_refs: &[ErasedRawStridedRef<'_>],
345    shape: &[usize],
346    plan: &ElementwiseFusionPlan,
347    exec: Exec<'_>,
348) -> crate::Result<Vec<Tensor>> {
349    let input_ptrs: Vec<_> = input_refs
350        .iter()
351        .map(ErasedRawStridedPtr::from_ref)
352        .collect();
353    let kernel = dtype.kernel();
354    match dtype {
355        FusedDType::F32 => {
356            execute_erased_fused_plan::<f32>(buffers, kernel, &input_ptrs, shape, plan, exec)
357        }
358        FusedDType::F64 => {
359            execute_erased_fused_plan::<f64>(buffers, kernel, &input_ptrs, shape, plan, exec)
360        }
361    }
362}
363
364/// Compile one strided plan for every live output and replay it once.
365///
366/// The plan's `outputs` are executed in order by
367/// [`ErasedFusedPlan::execute_outputs_uninit`], so a region with several live
368/// outputs replays its instruction list a single time instead of once per
369/// output.
370fn execute_erased_fused_plan<T>(
371    buffers: &mut BufferPool,
372    dtype: KernelDType,
373    input_ptrs: &[ErasedRawStridedPtr<'_>],
374    shape: &[usize],
375    plan: &ElementwiseFusionPlan,
376    exec: Exec<'_>,
377) -> crate::Result<Vec<Tensor>>
378where
379    T: Clone + PoolScalar,
380{
381    let erased_plan = ErasedFusedPlan::compile(dtype, strided_fused_plan(plan))
382        .map_err(|err| crate::Error::backend_source(ELEMENTWISE_FUSION_OP, err))?;
383    let output_strides = col_major_strides(shape)?;
384    let mut outputs = plan
385        .outputs()
386        .iter()
387        .map(|_| PooledUninitOutput::<T>::new(buffers, shape))
388        .collect::<crate::Result<Vec<_>>>()?;
389    // SAFETY: every pooled output exclusively owns its allocation with matching
390    // dtype/alignment, the destination descriptors borrow each output mutably
391    // for the exclusive lifetime of `dests`, and the fused plan overwrites
392    // every reachable element before `assume_init_recycled` exposes typed
393    // storage.
394    let mut dests = outputs
395        .iter_mut()
396        .map(|output| {
397            // SAFETY: as above.
398            unsafe {
399                erased_raw_strided_uninit_mut(
400                    dtype,
401                    output.as_uninit_bytes_mut(),
402                    shape,
403                    &output_strides,
404                    0,
405                )
406            }
407            .map_err(|err| crate::Error::backend_source(ELEMENTWISE_FUSION_OP, err))
408        })
409        .collect::<crate::Result<Vec<_>>>()?;
410    erased_plan
411        .execute_outputs_uninit(&mut dests, input_ptrs, exec)
412        .map_err(|err| crate::Error::backend_source(ELEMENTWISE_FUSION_OP, err))?;
413    drop(dests);
414    outputs
415        .drain(..)
416        .map(|output| {
417            // SAFETY: the fused replay wrote every logical destination element
418            // and retains no destination view.
419            let tensor = unsafe { output.assume_init_recycled()? };
420            Ok(Tensor::from_typed::<T>(tensor))
421        })
422        .collect()
423}
424
425#[cfg(test)]
426mod tests;
tenferro_cpu_fused/lib.rs

tenferro_cpu_fused/
lib.rs
