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, ExecContext, 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;
28
29const ELEMENTWISE_FUSION_OP: &str = "execute_elementwise_fusion";
30
31const ELEMENTWISE_FUSION_MIN_ELEMENTS: usize = 16 * 1024;
32
33fn validate_elementwise_fusion_inputs(
34    inputs: &[&Tensor],
35    plan: &ElementwiseFusionPlan,
36) -> crate::Result<bool> {
37    if inputs.len() != plan.input_count() {
38        return Err(crate::Error::invalid_argument(
39            ELEMENTWISE_FUSION_OP,
40            "inputs",
41            format!(
42                "plan expects {} inputs but backend received {}",
43                plan.input_count(),
44                inputs.len()
45            ),
46        ));
47    }
48    if plan.input_views().len() != plan.input_count() {
49        return Err(crate::Error::invalid_argument(
50            ELEMENTWISE_FUSION_OP,
51            "input_views",
52            format!(
53                "plan has {} input views for {} inputs",
54                plan.input_views().len(),
55                plan.input_count()
56            ),
57        ));
58    }
59    if plan.outputs().is_empty() {
60        return Ok(false);
61    }
62    for input in inputs {
63        if input.dtype() != plan.dtype() {
64            return Err(crate::Error::dtype_mismatch(
65                ELEMENTWISE_FUSION_OP,
66                input.dtype(),
67                plan.dtype(),
68            ));
69        }
70    }
71    Ok(true)
72}
73
74fn strided_fused_op(op: ElementwiseFusionOp) -> FusedOp {
75    match op {
76        ElementwiseFusionOp::Add => FusedOp::Add,
77        ElementwiseFusionOp::Multiply => FusedOp::Multiply,
78        ElementwiseFusionOp::Negate => FusedOp::Negate,
79        ElementwiseFusionOp::Conj => FusedOp::Conj,
80        ElementwiseFusionOp::Divide => FusedOp::Divide,
81        ElementwiseFusionOp::Abs => FusedOp::Abs,
82        ElementwiseFusionOp::Maximum => FusedOp::Maximum,
83        ElementwiseFusionOp::Minimum => FusedOp::Minimum,
84        ElementwiseFusionOp::Clamp => FusedOp::Clamp,
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        ElementwiseFusionOp::Remainder => {
96            unreachable!("remainder must be filtered before CPU elementwise fusion")
97        }
98    }
99}
100
101fn plan_uses_unfused_op(plan: &ElementwiseFusionPlan) -> bool {
102    plan.ops()
103        .iter()
104        .any(|inst| inst.op() == ElementwiseFusionOp::Remainder)
105}
106
107fn plan_uses_ordered_op(plan: &ElementwiseFusionPlan) -> bool {
108    plan.ops().iter().any(|inst| {
109        matches!(
110            inst.op(),
111            ElementwiseFusionOp::Maximum
112                | ElementwiseFusionOp::Minimum
113                | ElementwiseFusionOp::Clamp
114        )
115    })
116}
117
118fn should_defer_to_broadcast_multiply_special_case(plan: &ElementwiseFusionPlan) -> bool {
119    !plan.input_views().iter().all(|view| view.is_identity())
120        && plan.ops().len() == 1
121        && plan.outputs() == [plan.input_count()]
122        && plan.ops()[0].op() == ElementwiseFusionOp::Multiply
123}
124
125fn single_output_strided_fused_plan(plan: &ElementwiseFusionPlan, output: usize) -> FusedPlan {
126    FusedPlan {
127        input_count: plan.input_count(),
128        outputs: vec![output],
129        ops: plan
130            .ops()
131            .iter()
132            .map(|inst| FusedInst {
133                op: strided_fused_op(inst.op()),
134                inputs: inst.inputs().to_vec(),
135            })
136            .collect(),
137    }
138}
139
140fn kernel_dtype(dtype: DType) -> KernelDType {
141    match dtype {
142        DType::F32 => KernelDType::F32,
143        DType::F64 => KernelDType::F64,
144        DType::I32 => KernelDType::I32,
145        DType::I64 => KernelDType::I64,
146        DType::Bool => KernelDType::Bool,
147        DType::C32 => KernelDType::C32,
148        DType::C64 => KernelDType::C64,
149    }
150}
151
152fn typed_bytes<T>(data: &[T]) -> &[u8] {
153    // SAFETY: `data` is an aligned typed slice. The returned byte slice has
154    // the same lifetime and exact byte length, and is read-only.
155    unsafe { std::slice::from_raw_parts(data.as_ptr().cast::<u8>(), size_of_val(data)) }
156}
157
158struct ErasedFusionInput<'a> {
159    data: &'a [u8],
160    dims: Vec<usize>,
161    strides: Vec<isize>,
162}
163
164fn tensor_host_bytes<'a>(op: &'static str, input: &'a Tensor) -> crate::Result<&'a [u8]> {
165    macro_rules! bytes {
166        ($tensor:expr) => {
167            typed_host_data(op, $tensor).map(typed_bytes)
168        };
169    }
170
171    match input {
172        Tensor::F32(tensor) => bytes!(tensor),
173        Tensor::F64(tensor) => bytes!(tensor),
174        Tensor::I32(tensor) => bytes!(tensor),
175        Tensor::I64(tensor) => bytes!(tensor),
176        Tensor::Bool(tensor) => bytes!(tensor),
177        Tensor::C32(tensor) => bytes!(tensor),
178        Tensor::C64(tensor) => bytes!(tensor),
179    }
180}
181
182fn erased_fusion_input<'a>(
183    input: &'a Tensor,
184    view: &ElementwiseFusionInputView,
185) -> crate::Result<ErasedFusionInput<'a>> {
186    let data = tensor_host_bytes(ELEMENTWISE_FUSION_OP, input)?;
187    let base_shape = input.shape();
188    let base_strides = col_major_strides(base_shape)?;
189    let ElementwiseFusionInputView::BroadcastInDim { shape, dims } = view else {
190        return Ok(ErasedFusionInput {
191            data,
192            dims: base_shape.to_vec(),
193            strides: base_strides,
194        });
195    };
196
197    if dims.len() != base_shape.len() {
198        return Err(crate::Error::invalid_argument(
199            ELEMENTWISE_FUSION_OP,
200            "configuration",
201            format!(
202                "broadcast dims length {} does not match input rank {}",
203                dims.len(),
204                base_shape.len()
205            ),
206        ));
207    }
208
209    let mut strides = vec![0; shape.len()];
210    let mut seen = vec![false; shape.len()];
211    for (source_axis, &target_axis) in dims.iter().enumerate() {
212        if target_axis >= shape.len() {
213            return Err(crate::Error::axis_out_of_bounds(
214                ELEMENTWISE_FUSION_OP,
215                target_axis,
216                shape.len(),
217            ));
218        }
219        if seen[target_axis] {
220            return Err(crate::Error::duplicate_axis(
221                ELEMENTWISE_FUSION_OP,
222                target_axis,
223                "broadcast dims",
224            ));
225        }
226        seen[target_axis] = true;
227        let source_dim = base_shape[source_axis];
228        let target_dim = shape[target_axis];
229        if source_dim != target_dim && source_dim != 1 {
230            return Err(crate::Error::shape_mismatch(
231                ELEMENTWISE_FUSION_OP,
232                shape.to_vec(),
233                base_shape.to_vec(),
234            ));
235        }
236        if source_dim == target_dim {
237            strides[target_axis] = base_strides[source_axis];
238        }
239    }
240
241    Ok(ErasedFusionInput {
242        data,
243        dims: shape.to_vec(),
244        strides,
245    })
246}
247
248#[doc(hidden)]
249pub fn elementwise_fusion_with_pool(
250    buffers: &mut BufferPool,
251    exec_context: &ExecContext,
252    inputs: &[&Tensor],
253    plan: &ElementwiseFusionPlan,
254) -> crate::Result<Option<Vec<Tensor>>> {
255    if !validate_elementwise_fusion_inputs(inputs, plan)? {
256        return Ok(None);
257    }
258    if inputs.is_empty() {
259        return Ok(None);
260    }
261    if plan_uses_unfused_op(plan) {
262        return Ok(None);
263    }
264    if should_defer_to_broadcast_multiply_special_case(plan) {
265        return Ok(None);
266    }
267    if !dtype_supports_erased_fusion(plan.dtype(), plan) {
268        return Ok(None);
269    }
270
271    let input_layouts = inputs
272        .iter()
273        .zip(plan.input_views())
274        .map(|(input, view)| erased_fusion_input(input, view))
275        .collect::<crate::Result<Vec<_>>>()?;
276    let shape = input_layouts[0].dims.clone();
277    if input_layouts
278        .iter()
279        .skip(1)
280        .any(|input| input.dims != shape)
281    {
282        return Ok(None);
283    }
284    let element_count =
285        tenferro_tensor::validate::checked_shape_product(ELEMENTWISE_FUSION_OP, "shape", &shape)?;
286    if element_count < ELEMENTWISE_FUSION_MIN_ELEMENTS {
287        return Ok(None);
288    }
289
290    let dtype = kernel_dtype(plan.dtype());
291    let input_refs = input_layouts
292        .iter()
293        .map(|input| {
294            // SAFETY: fusion inputs are initialized typed storage with matching
295            // dtype and alignment; validated layouts bound every reachable read
296            // for the retained input borrow.
297            unsafe { erased_raw_strided_ref(dtype, input.data, &input.dims, &input.strides, 0) }
298                .map_err(|err| crate::Error::backend_source(ELEMENTWISE_FUSION_OP, err))
299        })
300        .collect::<crate::Result<Vec<_>>>()?;
301
302    execute_erased_fused_outputs(buffers, exec_context, dtype, &input_refs, &shape, plan).map(Some)
303}
304
305fn dtype_supports_erased_fusion(dtype: DType, plan: &ElementwiseFusionPlan) -> bool {
306    match dtype {
307        DType::F32 | DType::F64 => true,
308        DType::C32 | DType::C64 => !plan_uses_ordered_op(plan),
309        DType::I32 | DType::I64 => plan.ops().iter().all(|inst| {
310            matches!(
311                inst.op(),
312                ElementwiseFusionOp::Add
313                    | ElementwiseFusionOp::Multiply
314                    | ElementwiseFusionOp::Negate
315                    | ElementwiseFusionOp::Conj
316                    | ElementwiseFusionOp::Abs
317                    | ElementwiseFusionOp::Maximum
318                    | ElementwiseFusionOp::Minimum
319                    | ElementwiseFusionOp::Clamp
320            )
321        }),
322        DType::Bool => plan
323            .ops()
324            .iter()
325            .all(|inst| inst.op() == ElementwiseFusionOp::Conj),
326    }
327}
328
329fn execute_erased_fused_outputs(
330    buffers: &mut BufferPool,
331    exec_context: &ExecContext,
332    dtype: KernelDType,
333    input_refs: &[ErasedRawStridedRef<'_>],
334    shape: &[usize],
335    plan: &ElementwiseFusionPlan,
336) -> crate::Result<Vec<Tensor>> {
337    let input_ptrs: Vec<_> = input_refs
338        .iter()
339        .map(ErasedRawStridedPtr::from_ref)
340        .collect();
341    match dtype {
342        KernelDType::F32 => plan
343            .outputs()
344            .iter()
345            .map(|&output| {
346                execute_erased_fused_output::<f32>(
347                    buffers,
348                    exec_context,
349                    dtype,
350                    &input_ptrs,
351                    shape,
352                    plan,
353                    output,
354                    Tensor::F32,
355                )
356            })
357            .collect(),
358        KernelDType::F64 => plan
359            .outputs()
360            .iter()
361            .map(|&output| {
362                execute_erased_fused_output::<f64>(
363                    buffers,
364                    exec_context,
365                    dtype,
366                    &input_ptrs,
367                    shape,
368                    plan,
369                    output,
370                    Tensor::F64,
371                )
372            })
373            .collect(),
374        KernelDType::I32 => plan
375            .outputs()
376            .iter()
377            .map(|&output| {
378                execute_erased_fused_output::<i32>(
379                    buffers,
380                    exec_context,
381                    dtype,
382                    &input_ptrs,
383                    shape,
384                    plan,
385                    output,
386                    Tensor::I32,
387                )
388            })
389            .collect(),
390        KernelDType::I64 => plan
391            .outputs()
392            .iter()
393            .map(|&output| {
394                execute_erased_fused_output::<i64>(
395                    buffers,
396                    exec_context,
397                    dtype,
398                    &input_ptrs,
399                    shape,
400                    plan,
401                    output,
402                    Tensor::I64,
403                )
404            })
405            .collect(),
406        KernelDType::Bool => plan
407            .outputs()
408            .iter()
409            .map(|&output| {
410                execute_erased_fused_output::<bool>(
411                    buffers,
412                    exec_context,
413                    dtype,
414                    &input_ptrs,
415                    shape,
416                    plan,
417                    output,
418                    Tensor::Bool,
419                )
420            })
421            .collect(),
422        KernelDType::C32 => plan
423            .outputs()
424            .iter()
425            .map(|&output| {
426                execute_erased_fused_output::<num_complex::Complex32>(
427                    buffers,
428                    exec_context,
429                    dtype,
430                    &input_ptrs,
431                    shape,
432                    plan,
433                    output,
434                    Tensor::C32,
435                )
436            })
437            .collect(),
438        KernelDType::C64 => plan
439            .outputs()
440            .iter()
441            .map(|&output| {
442                execute_erased_fused_output::<num_complex::Complex64>(
443                    buffers,
444                    exec_context,
445                    dtype,
446                    &input_ptrs,
447                    shape,
448                    plan,
449                    output,
450                    Tensor::C64,
451                )
452            })
453            .collect(),
454        _ => Err(crate::Error::unsupported(
455            ELEMENTWISE_FUSION_OP,
456            format!(
457                "unsupported dtype {}; supported dtypes: F32/F64/I32/I64/Bool/C32/C64",
458                dtype.label()
459            ),
460        )),
461    }
462}
463
464#[allow(clippy::too_many_arguments)]
465fn execute_erased_fused_output<T>(
466    buffers: &mut BufferPool,
467    exec_context: &ExecContext,
468    dtype: KernelDType,
469    input_ptrs: &[ErasedRawStridedPtr<'_>],
470    shape: &[usize],
471    plan: &ElementwiseFusionPlan,
472    output: usize,
473    wrap: fn(TypedTensor<T>) -> Tensor,
474) -> crate::Result<Tensor>
475where
476    T: Clone + PoolScalar,
477{
478    let fused_plan = single_output_strided_fused_plan(plan, output);
479    let erased_plan = ErasedFusedPlan::compile(dtype, fused_plan)
480        .map_err(|err| crate::Error::backend_source(ELEMENTWISE_FUSION_OP, err))?;
481    let mut out = PooledUninitOutput::<T>::new(buffers, shape.to_vec())?;
482    let output_strides = col_major_strides(shape)?;
483    // SAFETY: `out` exclusively owns the output allocation with matching
484    // dtype/alignment and the fused plan overwrites every reachable element
485    // before `assume_init` exposes typed storage.
486    let mut dest = unsafe {
487        erased_raw_strided_uninit_mut(dtype, out.as_uninit_bytes_mut(), shape, &output_strides, 0)
488    }
489    .map_err(|err| crate::Error::backend_source(ELEMENTWISE_FUSION_OP, err))?;
490    erased_plan
491        .execute_uninit(exec_context, &mut dest, input_ptrs)
492        .map_err(|err| crate::Error::backend_source(ELEMENTWISE_FUSION_OP, err))?;
493    // SAFETY: the fused replay writes every logical destination element and retains no destination view.
494    Ok(wrap(unsafe { out.assume_init()? }))
495}
496
497#[cfg(test)]
498mod tests;
tenferro_cpu_fused/lib.rs

tenferro_cpu_fused/
lib.rs
