strided_fused/
erased.rs

1use crate::*;
2use core::mem::MaybeUninit;
3use num_complex::{Complex32, Complex64};
4use strided_basic::execution::{
5    check_dtype, ensure_same_shape, erased_view, validate_destination_layout_without_alloc,
6    validate_uninit_no_overlap, ValidatedDestinationLayout,
7};
8const ERASED_FUSED_INPUT_LIMIT: usize = 4;
9
10/// Dtype-erased wrapper around [`FusedPlan`].
11///
12/// This is the erased replay boundary for unary map and zip-map elementwise
13/// families. It supports the same runtime op-code vocabulary as [`FusedPlan`],
14/// but only for the scalar dtypes currently implementing [`FusedScalar`]. A
15/// plan may name one or more outputs; [`ErasedFusedPlan::execute_outputs`] and
16/// [`ErasedFusedPlan::execute_outputs_uninit`] write them all in one replay.
17#[derive(Clone, Debug)]
18pub struct ErasedFusedPlan {
19    dtype: KernelDType,
20    plan: FusedPlan,
21}
22
23impl ErasedFusedPlan {
24    /// Validate and store a fused elementwise plan for one dtype.
25    ///
26    /// The plan must name at least one output; repeated output value ids are
27    /// allowed and write the same value to each named destination.
28    ///
29    /// # Errors
30    ///
31    /// Returns an unsupported dtype, an input arity outside `1..=4`, a
32    /// zero-output plan, a malformed SSA DAG, or an operation the dtype does
33    /// not support.
34    pub fn compile(dtype: KernelDType, plan: FusedPlan) -> Result<Self> {
35        check_fused_dtype(dtype)?;
36        if plan.input_count == 0 || plan.input_count > ERASED_FUSED_INPUT_LIMIT {
37            return Err(StridedError::UnsupportedArity {
38                arity: plan.input_count,
39                max: ERASED_FUSED_INPUT_LIMIT,
40            });
41        }
42        validate_fused_plan_for_dtype(dtype, &plan)?;
43        Ok(Self { dtype, plan })
44    }
45
46    #[inline]
47    pub fn dtype(&self) -> KernelDType {
48        self.dtype
49    }
50
51    #[inline]
52    pub fn plan(&self) -> &FusedPlan {
53        &self.plan
54    }
55
56    /// Execute a single-output fused elementwise plan through erased descriptors.
57    ///
58    /// Forwarding front end for [`Self::execute_outputs`]; the plan must name
59    /// exactly one output.
60    ///
61    /// # Errors
62    ///
63    /// Returns a typed dtype, input-count, shape, destination injectivity, or
64    /// unsupported-operation error before any destination is written.
65    pub fn execute(
66        &self,
67        dest: &mut ErasedRawStridedMut<'_>,
68        inputs: &[ErasedRawStridedRef<'_>],
69        exec: Exec<'_>,
70    ) -> Result<()> {
71        self.execute_outputs(std::slice::from_mut(dest), inputs, exec)
72    }
73
74    /// Execute a plan into one or more destinations in a single replay.
75    ///
76    /// `plan.outputs` are written to `dests` in order. Every destination is
77    /// validated before any write: dtype, shape against the inputs, and output
78    /// injectivity. Overlapping initialized input/output memory is not
79    /// supported; the mutable destination slices are exclusive, so only
80    /// destination-versus-input overlap is a caller concern.
81    ///
82    /// # Errors
83    ///
84    /// Returns a typed dtype, input-count, destination-count, shape,
85    /// destination injectivity, or unsupported-operation error before any
86    /// destination is written.
87    ///
88    /// # Examples
89    ///
90    /// ```
91    /// use strided_basic::{
92    ///     ErasedRawStridedMut, ErasedRawStridedRef, Exec, KernelDType,
93    /// };
94    /// use strided_fused::{ErasedFusedPlan, FusedInst, FusedOp, FusedPlan};
95    ///
96    /// let plan = ErasedFusedPlan::compile(
97    ///     KernelDType::F64,
98    ///     FusedPlan::new(
99    ///         1,
100    ///         vec![0, 1],
101    ///         vec![FusedInst::new(FusedOp::Negate, vec![0])],
102    ///     ),
103    /// )
104    /// .unwrap();
105    /// let input = [1.0_f64, -2.0];
106    /// let dims = [2];
107    /// let strides = [1];
108    /// let inputs = [ErasedRawStridedRef::from_slice(&input, &dims, &strides, 0).unwrap()];
109    /// let mut pass_through = [0.0_f64; 2];
110    /// let mut negated = [0.0_f64; 2];
111    /// {
112    ///     let mut dests = [
113    ///         ErasedRawStridedMut::from_slice_mut(&mut pass_through, &dims, &strides, 0).unwrap(),
114    ///         ErasedRawStridedMut::from_slice_mut(&mut negated, &dims, &strides, 0).unwrap(),
115    ///     ];
116    ///     plan.execute_outputs(&mut dests, &inputs, Exec::serial())
117    ///         .unwrap();
118    /// }
119    /// assert_eq!(pass_through, [1.0, -2.0]);
120    /// assert_eq!(negated, [-1.0, 2.0]);
121    /// ```
122    pub fn execute_outputs(
123        &self,
124        dests: &mut [ErasedRawStridedMut<'_>],
125        inputs: &[ErasedRawStridedRef<'_>],
126        exec: Exec<'_>,
127    ) -> Result<()> {
128        if inputs.len() != self.plan.input_count {
129            return Err(StridedError::RankMismatch(
130                inputs.len(),
131                self.plan.input_count,
132            ));
133        }
134        if dests.len() != self.plan.outputs.len() {
135            return Err(StridedError::RankMismatch(
136                dests.len(),
137                self.plan.outputs.len(),
138            ));
139        }
140        for dest in dests.iter() {
141            check_dtype(self.dtype, dest.dtype())?;
142        }
143        for input in inputs {
144            check_dtype(self.dtype, input.dtype())?;
145        }
146        let dims = dests[0].dims();
147        for dest in &dests[1..] {
148            ensure_same_shape(dims, dest.dims())?;
149        }
150        for input in inputs {
151            ensure_same_shape(dims, input.dims())?;
152        }
153
154        match self.dtype {
155            KernelDType::F32 => execute_fused_multi::<f32>(&self.plan, dests, inputs, exec),
156            KernelDType::F64 => execute_fused_multi::<f64>(&self.plan, dests, inputs, exec),
157            KernelDType::I32 => execute_fused_multi::<i32>(&self.plan, dests, inputs, exec),
158            KernelDType::I64 => execute_fused_multi::<i64>(&self.plan, dests, inputs, exec),
159            KernelDType::Bool => execute_fused_multi::<bool>(&self.plan, dests, inputs, exec),
160            KernelDType::C32 => execute_fused_multi::<Complex32>(&self.plan, dests, inputs, exec),
161            KernelDType::C64 => execute_fused_multi::<Complex64>(&self.plan, dests, inputs, exec),
162            _ => Err(StridedError::UnsupportedDType {
163                dtype: self.dtype.label(),
164            }),
165        }
166    }
167
168    /// Execute a single-output fused plan into fully overwritten uninitialized
169    /// storage.
170    ///
171    /// Forwarding front end for [`Self::execute_outputs_uninit`]; the plan must
172    /// name exactly one output.
173    ///
174    /// # Errors
175    ///
176    /// Returns a typed dtype, input-count, shape, bounds, destination
177    /// injectivity, unsupported-operation, or input/output-overlap error before
178    /// any destination is written.
179    pub fn execute_uninit(
180        &self,
181        dest: &mut ErasedRawStridedUninitMut<'_>,
182        inputs: &[ErasedRawStridedPtr<'_>],
183        exec: Exec<'_>,
184    ) -> Result<()> {
185        self.execute_outputs_uninit(std::slice::from_mut(dest), inputs, exec)
186    }
187
188    /// Execute a plan into one or more fully overwritten uninitialized
189    /// destinations in a single replay.
190    ///
191    /// Dtype, shape, destination injectivity, bounds, and input/output overlap
192    /// are validated for every destination before any shared typed input
193    /// descriptor is formed or any destination byte is written. On `Ok(())`,
194    /// every logical destination element is initialized. An error leaves every
195    /// destination untouched; a panic during execution may leave partial
196    /// initialization, but the backing `MaybeUninit` storage remains safe to
197    /// drop.
198    ///
199    /// # Errors
200    ///
201    /// Returns a typed dtype, input-count, destination-count, shape, bounds,
202    /// destination injectivity, unsupported-operation, or input/output-overlap
203    /// error. All error-producing validation completes before execution starts.
204    ///
205    /// # Examples
206    ///
207    /// ```
208    /// use core::mem::MaybeUninit;
209    /// use strided_basic::{
210    ///     ErasedRawStridedPtr, ErasedRawStridedRef, ErasedRawStridedUninitMut,
211    ///     Exec, KernelDType,
212    /// };
213    /// use strided_fused::{ErasedFusedPlan, FusedInst, FusedOp, FusedPlan};
214    ///
215    /// let plan = ErasedFusedPlan::compile(
216    ///     KernelDType::F32,
217    ///     FusedPlan::new(
218    ///         1,
219    ///         vec![0, 1],
220    ///         vec![FusedInst::new(FusedOp::Negate, vec![0])],
221    ///     ),
222    /// )
223    /// .unwrap();
224    /// let input = [1.0_f32, -2.0];
225    /// let dims = [2];
226    /// let strides = [1];
227    /// let input = ErasedRawStridedRef::from_slice(&input, &dims, &strides, 0).unwrap();
228    /// let pointers = [ErasedRawStridedPtr::from_ref(&input)];
229    /// let mut pass_through = [MaybeUninit::<f32>::uninit(); 2];
230    /// let mut negated = [MaybeUninit::<f32>::uninit(); 2];
231    /// {
232    ///     let mut dests = [
233    ///         ErasedRawStridedUninitMut::from_uninit_slice(
234    ///             &mut pass_through,
235    ///             &dims,
236    ///             &strides,
237    ///             0,
238    ///         )
239    ///         .unwrap(),
240    ///         ErasedRawStridedUninitMut::from_uninit_slice(&mut negated, &dims, &strides, 0).unwrap(),
241    ///     ];
242    ///     plan.execute_outputs_uninit(&mut dests, &pointers, Exec::serial())
243    ///         .unwrap();
244    /// }
245    /// // SAFETY: `execute_outputs_uninit` overwrites every reachable element.
246    /// assert_eq!(unsafe { pass_through[1].assume_init() }, -2.0);
247    /// // SAFETY: as above.
248    /// assert_eq!(unsafe { negated[1].assume_init() }, 2.0);
249    /// ```
250    pub fn execute_outputs_uninit(
251        &self,
252        dests: &mut [ErasedRawStridedUninitMut<'_>],
253        inputs: &[ErasedRawStridedPtr<'_>],
254        exec: Exec<'_>,
255    ) -> Result<()> {
256        if inputs.len() != self.plan.input_count {
257            return Err(StridedError::RankMismatch(
258                inputs.len(),
259                self.plan.input_count,
260            ));
261        }
262        if dests.len() != self.plan.outputs.len() {
263            return Err(StridedError::RankMismatch(
264                dests.len(),
265                self.plan.outputs.len(),
266            ));
267        }
268        for dest in dests.iter() {
269            check_dtype(self.dtype, dest.dtype())?;
270        }
271        for input in inputs {
272            check_dtype(self.dtype, input.dtype())?;
273        }
274        let dims = dests[0].dims();
275        for dest in &dests[1..] {
276            ensure_same_shape(dims, dest.dims())?;
277        }
278        for input in inputs {
279            ensure_same_shape(dims, input.dims())?;
280        }
281        for (index, input) in inputs.iter().enumerate() {
282            for dest in dests.iter() {
283                validate_uninit_no_overlap(dest, input, index)?;
284            }
285        }
286        // Store the validation tokens inline: the destination-count bound is
287        // the five-stream operand bound, so two outputs plus one input allocate
288        // nothing. Above the bound the vector spills by contract.
289        let mut validated: strided_basic::OperandVec<ValidatedDestinationLayout> =
290            strided_basic::OperandVec::with_capacity(dests.len());
291        for dest in dests.iter() {
292            validated.push(validate_destination_layout_without_alloc(
293                dest.dims(),
294                dest.strides(),
295            )?);
296        }
297        self.dispatch_uninit(dests, inputs, validated.as_slice(), exec)
298    }
299
300    fn dispatch_uninit(
301        &self,
302        dests: &mut [ErasedRawStridedUninitMut<'_>],
303        inputs: &[ErasedRawStridedPtr<'_>],
304        validated: &[ValidatedDestinationLayout],
305        exec: Exec<'_>,
306    ) -> Result<()> {
307        match self.dtype {
308            KernelDType::F32 => {
309                execute_fused_uninit_multi::<f32>(&self.plan, dests, inputs, validated, exec)
310            }
311            KernelDType::F64 => {
312                execute_fused_uninit_multi::<f64>(&self.plan, dests, inputs, validated, exec)
313            }
314            KernelDType::I32 => {
315                execute_fused_uninit_multi::<i32>(&self.plan, dests, inputs, validated, exec)
316            }
317            KernelDType::I64 => {
318                execute_fused_uninit_multi::<i64>(&self.plan, dests, inputs, validated, exec)
319            }
320            KernelDType::Bool => {
321                execute_fused_uninit_multi::<bool>(&self.plan, dests, inputs, validated, exec)
322            }
323            KernelDType::C32 => {
324                execute_fused_uninit_multi::<Complex32>(&self.plan, dests, inputs, validated, exec)
325            }
326            KernelDType::C64 => {
327                execute_fused_uninit_multi::<Complex64>(&self.plan, dests, inputs, validated, exec)
328            }
329            _ => Err(StridedError::UnsupportedDType {
330                dtype: self.dtype.label(),
331            }),
332        }
333    }
334}
335
336fn check_fused_dtype(dtype: KernelDType) -> Result<()> {
337    match dtype {
338        KernelDType::F32
339        | KernelDType::F64
340        | KernelDType::I32
341        | KernelDType::I64
342        | KernelDType::Bool
343        | KernelDType::C32
344        | KernelDType::C64 => Ok(()),
345        _ => Err(StridedError::UnsupportedDType {
346            dtype: dtype.label(),
347        }),
348    }
349}
350
351fn validate_fused_plan_for_dtype(dtype: KernelDType, plan: &FusedPlan) -> Result<()> {
352    let output_count = plan.outputs.len();
353    match dtype {
354        KernelDType::F32 => {
355            crate::fused::validate_plan_for_scalar::<f32>(plan, plan.input_count, output_count)
356        }
357        KernelDType::F64 => {
358            crate::fused::validate_plan_for_scalar::<f64>(plan, plan.input_count, output_count)
359        }
360        KernelDType::I32 => {
361            crate::fused::validate_plan_for_scalar::<i32>(plan, plan.input_count, output_count)
362        }
363        KernelDType::I64 => {
364            crate::fused::validate_plan_for_scalar::<i64>(plan, plan.input_count, output_count)
365        }
366        KernelDType::Bool => {
367            crate::fused::validate_plan_for_scalar::<bool>(plan, plan.input_count, output_count)
368        }
369        KernelDType::C32 => crate::fused::validate_plan_for_scalar::<Complex32>(
370            plan,
371            plan.input_count,
372            output_count,
373        ),
374        KernelDType::C64 => crate::fused::validate_plan_for_scalar::<Complex64>(
375            plan,
376            plan.input_count,
377            output_count,
378        ),
379        _ => Err(StridedError::UnsupportedDType {
380            dtype: dtype.label(),
381        }),
382    }
383}
384
385fn execute_fused_multi<T>(
386    plan: &FusedPlan,
387    dests: &mut [ErasedRawStridedMut<'_>],
388    inputs: &[ErasedRawStridedRef<'_>],
389    exec: Exec<'_>,
390) -> Result<()>
391where
392    T: FusedScalar + KernelStorageElement,
393{
394    // The single-output case uses stack storage so the hot path allocates
395    // nothing; several destinations stay inline up to the operand bound and
396    // spill above it.
397    if dests.len() == 1 {
398        let dims = dests[0].dims();
399        let strides = dests[0].strides();
400        let offset = dests[0].offset();
401        let data = dests[0].data_as_mut::<T>()?;
402        let mut dest_views =
403            [unsafe { StridedViewMut::new_unchecked(data, dims, strides, offset) }];
404        return execute_fused_inputs(&mut dest_views, inputs, plan, exec);
405    }
406
407    if dests.len() <= strided_basic::OPERAND_CAP {
408        let mut storage =
409            [const { MaybeUninit::<StridedViewMut<'_, T>>::uninit() }; strided_basic::OPERAND_CAP];
410        // SAFETY: `storage` has `OPERAND_CAP` writable slots and outlives the
411        // guard (declared after it); at most `OPERAND_CAP` views are pushed.
412        let mut views = unsafe { crate::fused::InitializedViewsMut::new(storage.as_mut_ptr()) };
413        for dest in dests.iter_mut() {
414            let dims = dest.dims();
415            let strides = dest.strides();
416            let offset = dest.offset();
417            let data = dest.data_as_mut::<T>()?;
418            views.push(unsafe { StridedViewMut::new_unchecked(data, dims, strides, offset) });
419        }
420        return execute_fused_inputs(views.as_mut_slice(), inputs, plan, exec);
421    }
422
423    let mut dest_views = Vec::with_capacity(dests.len());
424    for dest in dests.iter_mut() {
425        let dims = dest.dims();
426        let strides = dest.strides();
427        let offset = dest.offset();
428        let data = dest.data_as_mut::<T>()?;
429        dest_views.push(unsafe { StridedViewMut::new_unchecked(data, dims, strides, offset) });
430    }
431    execute_fused_inputs(&mut dest_views, inputs, plan, exec)
432}
433
434fn execute_fused_inputs<T>(
435    dests: &mut [StridedViewMut<'_, T>],
436    inputs: &[ErasedRawStridedRef<'_>],
437    plan: &FusedPlan,
438    exec: Exec<'_>,
439) -> Result<()>
440where
441    T: FusedScalar + KernelStorageElement,
442{
443    match inputs {
444        [a] => {
445            let views = [erased_view::<T>(a)?];
446            execute_fused_views(dests, &views, plan, exec)
447        }
448        [a, b] => {
449            let views = [erased_view::<T>(a)?, erased_view::<T>(b)?];
450            execute_fused_views(dests, &views, plan, exec)
451        }
452        [a, b, c] => {
453            let views = [
454                erased_view::<T>(a)?,
455                erased_view::<T>(b)?,
456                erased_view::<T>(c)?,
457            ];
458            execute_fused_views(dests, &views, plan, exec)
459        }
460        [a, b, c, d] => {
461            let views = [
462                erased_view::<T>(a)?,
463                erased_view::<T>(b)?,
464                erased_view::<T>(c)?,
465                erased_view::<T>(d)?,
466            ];
467            execute_fused_views(dests, &views, plan, exec)
468        }
469        _ => Err(StridedError::UnsupportedArity {
470            arity: inputs.len(),
471            max: ERASED_FUSED_INPUT_LIMIT,
472        }),
473    }
474}
475
476fn execute_fused_uninit_multi<T>(
477    plan: &FusedPlan,
478    dests: &mut [ErasedRawStridedUninitMut<'_>],
479    inputs: &[ErasedRawStridedPtr<'_>],
480    validated: &[ValidatedDestinationLayout],
481    exec: Exec<'_>,
482) -> Result<()>
483where
484    T: FusedScalar + KernelStorageElement,
485{
486    if dests.len() == 1 {
487        let dims = dests[0].dims();
488        let strides = dests[0].strides();
489        let offset = dests[0].offset();
490        let data = dests[0].data_as_uninit_mut::<T>()?;
491        let mut dest_views =
492            [unsafe { StridedViewMut::new_unchecked(data, dims, strides, offset) }];
493        return execute_fused_uninit_inputs(plan, &mut dest_views, inputs, validated, exec);
494    }
495
496    if dests.len() <= strided_basic::OPERAND_CAP {
497        let mut storage = [const { MaybeUninit::<StridedViewMut<'_, MaybeUninit<T>>>::uninit() };
498            strided_basic::OPERAND_CAP];
499        // SAFETY: `storage` has `OPERAND_CAP` writable slots and outlives the
500        // guard (declared after it); at most `OPERAND_CAP` views are pushed.
501        let mut views = unsafe { crate::fused::InitializedViewsMut::new(storage.as_mut_ptr()) };
502        for dest in dests.iter_mut() {
503            let dims = dest.dims();
504            let strides = dest.strides();
505            let offset = dest.offset();
506            let data = dest.data_as_uninit_mut::<T>()?;
507            views.push(unsafe { StridedViewMut::new_unchecked(data, dims, strides, offset) });
508        }
509        return execute_fused_uninit_inputs(plan, views.as_mut_slice(), inputs, validated, exec);
510    }
511
512    let mut dest_views = Vec::with_capacity(dests.len());
513    for dest in dests.iter_mut() {
514        let dims = dest.dims();
515        let strides = dest.strides();
516        let offset = dest.offset();
517        let data = dest.data_as_uninit_mut::<T>()?;
518        dest_views.push(unsafe { StridedViewMut::new_unchecked(data, dims, strides, offset) });
519    }
520    execute_fused_uninit_inputs(plan, &mut dest_views, inputs, validated, exec)
521}
522
523fn execute_fused_uninit_inputs<T>(
524    plan: &FusedPlan,
525    dests: &mut [StridedViewMut<'_, MaybeUninit<T>>],
526    inputs: &[ErasedRawStridedPtr<'_>],
527    validated: &[ValidatedDestinationLayout],
528    exec: Exec<'_>,
529) -> Result<()>
530where
531    T: FusedScalar + KernelStorageElement,
532{
533    match inputs {
534        // SAFETY: `ErasedFusedPlan::execute_outputs_uninit` checked every input
535        // against every destination before this helper.
536        [a] => {
537            let refs = [unsafe { a.try_as_ref_after_no_overlap()? }];
538            let views = [erased_view::<T>(&refs[0])?];
539            crate::fused::fused_elementwise_into_uninit_multi(dests, &views, plan, validated, exec)
540        }
541        [a, b] => {
542            let refs = [
543                // SAFETY: as above.
544                unsafe { a.try_as_ref_after_no_overlap()? },
545                // SAFETY: as above.
546                unsafe { b.try_as_ref_after_no_overlap()? },
547            ];
548            let views = [erased_view::<T>(&refs[0])?, erased_view::<T>(&refs[1])?];
549            crate::fused::fused_elementwise_into_uninit_multi(dests, &views, plan, validated, exec)
550        }
551        [a, b, c] => {
552            let refs = [
553                // SAFETY: as above.
554                unsafe { a.try_as_ref_after_no_overlap()? },
555                // SAFETY: as above.
556                unsafe { b.try_as_ref_after_no_overlap()? },
557                // SAFETY: as above.
558                unsafe { c.try_as_ref_after_no_overlap()? },
559            ];
560            let views = [
561                erased_view::<T>(&refs[0])?,
562                erased_view::<T>(&refs[1])?,
563                erased_view::<T>(&refs[2])?,
564            ];
565            crate::fused::fused_elementwise_into_uninit_multi(dests, &views, plan, validated, exec)
566        }
567        [a, b, c, d] => {
568            let refs = [
569                // SAFETY: as above.
570                unsafe { a.try_as_ref_after_no_overlap()? },
571                // SAFETY: as above.
572                unsafe { b.try_as_ref_after_no_overlap()? },
573                // SAFETY: as above.
574                unsafe { c.try_as_ref_after_no_overlap()? },
575                // SAFETY: as above.
576                unsafe { d.try_as_ref_after_no_overlap()? },
577            ];
578            let views = [
579                erased_view::<T>(&refs[0])?,
580                erased_view::<T>(&refs[1])?,
581                erased_view::<T>(&refs[2])?,
582                erased_view::<T>(&refs[3])?,
583            ];
584            crate::fused::fused_elementwise_into_uninit_multi(dests, &views, plan, validated, exec)
585        }
586        _ => Err(StridedError::UnsupportedArity {
587            arity: inputs.len(),
588            max: ERASED_FUSED_INPUT_LIMIT,
589        }),
590    }
591}
592
593fn execute_fused_views<T>(
594    dests: &mut [StridedViewMut<'_, T>],
595    inputs: &[StridedView<'_, T>],
596    plan: &FusedPlan,
597    exec: Exec<'_>,
598) -> Result<()>
599where
600    T: FusedScalar + KernelStorageElement,
601{
602    fused_elementwise_into(dests, inputs, plan, exec)
603}
strided_fused/erased.rs

strided_fused/
erased.rs
