tenferro_cpu/
provider.rs

1//! Object-safe CPU contraction provider contracts.
2//!
3//! Providers synchronously write into engine-owned outputs. Request
4//! constructors are crate-private because only the CPU engine may attest that
5//! tensor metadata and reachable ranges have already been validated.
6
7use core::fmt;
8use std::num::NonZeroUsize;
9use std::sync::Arc;
10
11use tenferro_tensor::backend::GroupedGemmJob;
12use tenferro_tensor::{
13    DType, DotGeneralAccumulation, Tensor, TensorRead, TensorView, TensorViewMut, TensorWrite,
14};
15
16use crate::arbiter::{with_execution_owner, ResourcePermit};
17use crate::backend::CpuBackendKind;
18use crate::buffer_pool::BufferPool;
19use crate::domain_executor::{indexed_jobs, install_scoped};
20#[cfg(feature = "cpu-blas")]
21use crate::provider_capability::builtin_blas_execution_capabilities;
22#[cfg(not(feature = "cpu-blas"))]
23use crate::provider_capability::serial_capabilities;
24use crate::provider_capability::{engine_worker_capabilities, CpuProviderExecutionCapabilities};
25use crate::resource_domain::CpuResourceDomain;
26use crate::{
27    CpuDomainExecutorError, CpuDomainId, CpuInnerParallelism, CpuPlacementGuarantee, CpuSet,
28};
29
30/// Operand named by a provider capability reason.
31///
32/// # Examples
33///
34/// ```
35/// use tenferro_cpu::provider::CpuOperand;
36/// assert_eq!(CpuOperand::Lhs, CpuOperand::Lhs);
37/// ```
38#[derive(Clone, Copy, Debug, PartialEq, Eq)]
39pub enum CpuOperand {
40    /// Left input operand.
41    Lhs,
42    /// Right input operand.
43    Rhs,
44    /// Writable output operand.
45    Output,
46}
47
48/// Allocation-free reason that a provider cannot execute a validated request.
49///
50/// An unsupported outcome must be reported before the provider mutates the
51/// output.
52///
53/// # Examples
54///
55/// ```
56/// use tenferro_cpu::provider::CpuProviderUnsupported;
57/// assert_eq!(
58///     CpuProviderUnsupported::RuntimeUnavailable,
59///     CpuProviderUnsupported::RuntimeUnavailable,
60/// );
61/// ```
62#[derive(Clone, Copy, Debug, PartialEq, Eq)]
63#[non_exhaustive]
64pub enum CpuProviderUnsupported {
65    /// The provider does not implement the scalar dtype.
66    DType(DType),
67    /// The provider does not implement the input ranks.
68    Rank {
69        /// Left input rank.
70        lhs: usize,
71        /// Right input rank.
72        rhs: usize,
73    },
74    /// The provider does not implement an operand layout.
75    Layout(CpuOperand),
76    /// The provider cannot implement the requested conjugation.
77    Conjugation,
78    /// The provider cannot implement the requested alpha/beta update.
79    Accumulation,
80    /// The provider cannot implement a strided batch.
81    StridedBatch,
82    /// The provider cannot implement grouped GEMM.
83    Grouped,
84    /// The optional provider runtime is not available in this process.
85    RuntimeUnavailable,
86}
87
88/// Result of attempting a validated request through one provider slot.
89///
90/// # Examples
91///
92/// ```
93/// use tenferro_cpu::provider::{CpuProviderOutcome, CpuProviderUnsupported};
94/// let outcome = CpuProviderOutcome::Unsupported(
95///     CpuProviderUnsupported::RuntimeUnavailable,
96/// );
97/// assert!(matches!(outcome, CpuProviderOutcome::Unsupported(_)));
98/// ```
99#[derive(Clone, Copy, Debug, PartialEq, Eq)]
100#[must_use]
101pub enum CpuProviderOutcome {
102    /// The provider fully executed the request.
103    Executed,
104    /// The provider did not mutate the output and resolution may continue.
105    Unsupported(CpuProviderUnsupported),
106}
107
108/// Parallel scheduling mode selected for one CPU operation.
109///
110/// # Examples
111///
112/// ```
113/// use tenferro_cpu::provider::ParallelMode;
114/// assert_ne!(ParallelMode::Sequential, ParallelMode::Outer);
115/// assert_ne!(ParallelMode::Outer, ParallelMode::Inner);
116/// ```
117#[derive(Clone, Copy, Debug, PartialEq, Eq)]
118pub enum ParallelMode {
119    /// Neither the engine nor the provider may fan out this operation.
120    Sequential,
121    /// The engine owns outer fan-out and delegates Sequential child contexts.
122    /// Providers do not receive an Outer context.
123    Outer,
124    /// One provider kernel may use the selected executor's inner region.
125    Inner,
126}
127
128/// Borrowed execution policy for an already-entered CPU operation.
129///
130/// The context exposes immutable domain facts while keeping the resource lease,
131/// executor object, and the checked executor-entry boundary private. Providers
132/// cannot install or submit work through this value.
133///
134/// # Examples
135///
136/// Providers inspect this value inside a trait method:
137///
138/// ```
139/// use tenferro_cpu::provider::CpuExecutionContext;
140/// # fn inspect(context: &CpuExecutionContext<'_>) {
141/// assert!(context.thread_budget().get() >= 1);
142/// # }
143/// ```
144#[derive(Clone, Copy)]
145pub struct CpuExecutionContext<'a> {
146    domain: &'a CpuResourceDomain,
147    parallel_mode: ParallelMode,
148}
149
150impl fmt::Debug for CpuExecutionContext<'_> {
151    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
152        formatter
153            .debug_struct("CpuExecutionContext")
154            .field("domain_id", &self.domain_id())
155            .field("cpus", &self.cpus())
156            .field("thread_budget", &self.thread_budget())
157            .field("placement_guarantee", &self.placement_guarantee())
158            .field("parallel_mode", &self.parallel_mode())
159            .finish_non_exhaustive()
160    }
161}
162
163impl<'a> CpuExecutionContext<'a> {
164    fn entered(domain: &'a CpuResourceDomain, parallel_mode: ParallelMode) -> Self {
165        Self {
166            domain,
167            parallel_mode,
168        }
169    }
170
171    /// Return the stable identity of the selected CPU resource domain.
172    ///
173    /// # Examples
174    ///
175    /// ```
176    /// use tenferro_cpu::CpuExecutionContext;
177    /// # fn inspect(context: &CpuExecutionContext<'_>) {
178    /// let _domain_id = context.domain_id();
179    /// # }
180    /// ```
181    pub fn domain_id(&self) -> CpuDomainId {
182        self.domain.id()
183    }
184
185    /// Return the selected domain's declared logical CPU set.
186    ///
187    /// # Examples
188    ///
189    /// ```
190    /// use tenferro_cpu::CpuExecutionContext;
191    /// # fn inspect(context: &CpuExecutionContext<'_>) {
192    /// assert!(!context.cpus().is_empty());
193    /// # }
194    /// ```
195    pub fn cpus(&self) -> &CpuSet {
196        self.domain.cpus()
197    }
198
199    /// Return the non-zero maximum participating-thread budget.
200    ///
201    /// # Examples
202    ///
203    /// ```
204    /// use tenferro_cpu::CpuExecutionContext;
205    /// # fn inspect(context: &CpuExecutionContext<'_>) {
206    /// assert!(context.thread_budget().get() >= 1);
207    /// # }
208    /// ```
209    pub fn thread_budget(&self) -> NonZeroUsize {
210        self.domain.thread_budget()
211    }
212
213    /// Return the strength of the selected domain's placement guarantee.
214    ///
215    /// # Examples
216    ///
217    /// ```
218    /// use tenferro_cpu::CpuExecutionContext;
219    /// # fn inspect(context: &CpuExecutionContext<'_>) {
220    /// let _guarantee = context.placement_guarantee();
221    /// # }
222    /// ```
223    pub fn placement_guarantee(&self) -> CpuPlacementGuarantee {
224        self.domain.placement_guarantee()
225    }
226
227    /// Return the engine-selected scheduling mode for this entered provider call.
228    ///
229    /// Provider calls observe Sequential or Inner. Outer scheduling creates a
230    /// separate Sequential context inside every submitted child.
231    ///
232    /// # Examples
233    ///
234    /// ```
235    /// use tenferro_cpu::{CpuExecutionContext, ParallelMode};
236    /// # fn inspect(context: &CpuExecutionContext<'_>) {
237    /// assert!(matches!(
238    ///     context.parallel_mode(),
239    ///     ParallelMode::Sequential | ParallelMode::Inner
240    /// ));
241    /// # }
242    /// ```
243    pub fn parallel_mode(&self) -> ParallelMode {
244        self.parallel_mode
245    }
246
247    /// Materialize a borrowed tensor view for one scoped operation and reclaim
248    /// its temporary host buffer before returning.
249    ///
250    /// Owned tensor inputs are borrowed directly. View inputs are materialized
251    /// from `buffers`, passed to `operation`, and returned to the same pool on
252    /// both success and ordinary error. The receiver is an unforgeable proof
253    /// that the caller is already inside the selected CPU execution domain.
254    ///
255    /// # Errors
256    ///
257    /// Returns [`tenferro_tensor::Error::RuntimeState`] when the view is not
258    /// accessible from CPU host memory, propagates typed view-materialization
259    /// errors, and otherwise returns the error produced by `operation`.
260    ///
261    /// # Examples
262    ///
263    /// ```
264    /// use tenferro_cpu::CpuBackend;
265    /// use tenferro_tensor::{StridedSliceSpec, TensorRead, TensorView, TypedTensor};
266    ///
267    /// let mut backend = CpuBackend::with_threads(1)?;
268    /// let input = TypedTensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0])?;
269    /// backend.with_linalg_pool(|context, buffers| {
270    ///     let view = input
271    ///         .as_view()
272    ///         .try_slice(&[StridedSliceSpec::reverse()])?;
273    ///     context.with_materialized_tensor_read(
274    ///         buffers,
275    ///         "example",
276    ///         TensorRead::from_view(TensorView::F64(view)),
277    ///         |materialized, _| {
278    ///             assert_eq!(materialized.as_slice::<f64>().unwrap(), &[2.0, 1.0]);
279    ///             Ok(())
280    ///         },
281    ///     )
282    /// })?;
283    /// # Ok::<(), Box<dyn std::error::Error>>(())
284    /// ```
285    #[doc(hidden)]
286    pub fn with_materialized_tensor_read<R>(
287        &self,
288        buffers: &mut BufferPool,
289        op: &'static str,
290        input: TensorRead<'_>,
291        operation: impl FnOnce(&Tensor, &mut BufferPool) -> tenferro_tensor::Result<R>,
292    ) -> tenferro_tensor::Result<R> {
293        match input {
294            TensorRead::Tensor(tensor) => operation(tensor, buffers),
295            TensorRead::View(view) => {
296                let materialized = self.with_native_parallelism(|| {
297                    crate::materialize_tensor_read(buffers, op, TensorRead::View(view))
298                })?;
299                let result = operation(&materialized, buffers);
300                reclaim_tensor(buffers, materialized);
301                result
302            }
303        }
304    }
305
306    /// Reshape a compact tensor while retaining the current execution proof.
307    ///
308    /// This metadata-only helper does not enter an executor or borrow another
309    /// scratch pool.
310    ///
311    /// # Errors
312    ///
313    /// Returns [`tenferro_tensor::Error::Validation`] when the input and output
314    /// shapes have different element counts or the requested layout is invalid.
315    ///
316    /// # Examples
317    ///
318    /// ```
319    /// use tenferro_cpu::CpuBackend;
320    /// use tenferro_tensor::Tensor;
321    ///
322    /// let mut backend = CpuBackend::with_threads(1)?;
323    /// let input = Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0])?;
324    /// backend.with_linalg_pool(|context, _| {
325    ///     let output = context.reshape_tensor(&input, &[2, 1])?;
326    ///     assert_eq!(output.shape(), &[2, 1]);
327    ///     Ok(())
328    /// })?;
329    /// # Ok::<(), Box<dyn std::error::Error>>(())
330    /// ```
331    #[doc(hidden)]
332    pub fn reshape_tensor(
333        &self,
334        input: &Tensor,
335        shape: &[usize],
336    ) -> tenferro_tensor::Result<Tensor> {
337        crate::structural::reshape(input, shape)
338    }
339
340    /// Return the faer policy selected by this operation context.
341    ///
342    /// This hidden public method is the owner-scoped extension contract used by
343    /// operation-family crates such as `tenferro-linalg`. Keeping the mapping
344    /// here prevents sibling crates from deriving a second CPU threading
345    /// policy.
346    ///
347    /// # Examples
348    ///
349    /// ```
350    /// use tenferro_cpu::CpuExecutionContext;
351    /// # fn inspect(context: &CpuExecutionContext<'_>) {
352    /// let _policy = context.faer_parallelism();
353    /// # }
354    /// ```
355    #[cfg(feature = "cpu-faer")]
356    #[doc(hidden)]
357    pub fn faer_parallelism(self) -> faer::Par {
358        match (
359            self.parallel_mode,
360            self.domain.executor_capabilities().inner_parallelism,
361        ) {
362            (ParallelMode::Inner, CpuInnerParallelism::Rayon) if self.thread_budget().get() > 1 => {
363                faer::Par::rayon(self.thread_budget().get())
364            }
365            _ => faer::Par::Seq,
366        }
367    }
368
369    pub(crate) fn strided_exec_context(&self) -> strided_kernel::ExecContext {
370        match (
371            self.parallel_mode,
372            self.domain.executor_capabilities().inner_parallelism,
373        ) {
374            (ParallelMode::Inner, CpuInnerParallelism::Rayon) if self.thread_budget().get() > 1 => {
375                // This is only an operation-local thread limit for strided's
376                // replay policy. The Rayon pool itself is the already-entered
377                // CpuContext pool installed by `with_native_parallelism`.
378                match strided_kernel::ExecContext::max_threads(self.thread_budget().get()) {
379                    Ok(context) => context,
380                    // INVARIANT: CpuExecutionContext stores a NonZeroUsize
381                    // thread budget, and this branch passes that positive value.
382                    Err(_) => unreachable!("CpuExecutionContext has a non-zero thread budget"),
383                }
384            }
385            _ => strided_kernel::ExecContext::serial(),
386        }
387    }
388
389    pub(crate) fn with_native_parallelism<R>(&self, operation: impl FnOnce() -> R) -> R {
390        let policy = match (
391            self.parallel_mode,
392            self.domain.executor_capabilities().inner_parallelism,
393        ) {
394            (ParallelMode::Inner, CpuInnerParallelism::Rayon) if self.thread_budget().get() > 1 => {
395                strided_kernel::ExecutionPolicy::Rayon {
396                    max_threads: self.thread_budget(),
397                }
398            }
399            _ => strided_kernel::ExecutionPolicy::Sequential,
400        };
401        strided_kernel::with_execution_policy(policy, operation)
402    }
403}
404
405fn reclaim_tensor(buffers: &mut BufferPool, tensor: Tensor) {
406    match tensor {
407        Tensor::F32(tensor) => crate::backend::reclaim_typed(buffers, tensor),
408        Tensor::F64(tensor) => crate::backend::reclaim_typed(buffers, tensor),
409        Tensor::I32(tensor) => crate::backend::reclaim_typed(buffers, tensor),
410        Tensor::I64(tensor) => crate::backend::reclaim_typed(buffers, tensor),
411        Tensor::Bool(tensor) => crate::backend::reclaim_typed(buffers, tensor),
412        Tensor::C32(tensor) => crate::backend::reclaim_typed(buffers, tensor),
413        Tensor::C64(tensor) => crate::backend::reclaim_typed(buffers, tensor),
414    }
415}
416
417/// Crate-private unentered capability for one CPU operation.
418///
419/// This is the only type that owns the resource permit and may cross the
420/// selected domain executor boundary. A [`CpuExecutionContext`] is constructed
421/// only inside an installed job or an outer child job.
422#[derive(Clone, Copy)]
423pub(crate) struct CpuOperationEntry<'a> {
424    domain: &'a CpuResourceDomain,
425    permit: &'a ResourcePermit,
426}
427
428impl<'a> CpuOperationEntry<'a> {
429    pub(crate) fn new(domain: &'a CpuResourceDomain, permit: &'a ResourcePermit) -> Self {
430        Self { domain, permit }
431    }
432
433    pub(crate) fn domain_id(self) -> CpuDomainId {
434        self.domain.id()
435    }
436
437    pub(crate) fn enter<R: Send>(
438        self,
439        parallel_mode: ParallelMode,
440        operation: impl FnOnce(&CpuExecutionContext<'_>) -> R + Send,
441    ) -> Result<R, CpuDomainExecutorError> {
442        if parallel_mode == ParallelMode::Outer {
443            return Err(CpuDomainExecutorError::Scheduling {
444                message: "CPU executor install requires Sequential or Inner mode, got Outer"
445                    .to_owned(),
446            });
447        }
448        let owner = self.permit.owner();
449        with_execution_owner(owner, || {
450            install_scoped(self.domain.executor().as_ref(), || {
451                with_execution_owner(owner, || {
452                    let context = CpuExecutionContext::entered(self.domain, parallel_mode);
453                    operation(&context)
454                })
455            })
456        })
457    }
458
459    pub(crate) fn enter_or_reuse<R: Send>(
460        self,
461        entered: Option<&CpuExecutionContext<'_>>,
462        parallel_mode: ParallelMode,
463        operation: impl FnOnce(&CpuExecutionContext<'_>) -> R + Send,
464    ) -> Result<R, CpuDomainExecutorError> {
465        let Some(entered) = entered else {
466            return self.enter(parallel_mode, operation);
467        };
468        if parallel_mode == ParallelMode::Outer {
469            return Err(CpuDomainExecutorError::Scheduling {
470                message: "entered CPU session requires Sequential or Inner mode, got Outer"
471                    .to_owned(),
472            });
473        }
474        if entered.domain_id() != self.domain.id() {
475            return Err(CpuDomainExecutorError::Scheduling {
476                message: format!(
477                    "entered CPU session domain {:?} does not match operation domain {:?}",
478                    entered.domain_id(),
479                    self.domain.id()
480                ),
481            });
482        }
483        let owner = self.permit.owner();
484        Ok(with_execution_owner(owner, || {
485            let context = CpuExecutionContext::entered(self.domain, parallel_mode);
486            operation(&context)
487        }))
488    }
489
490    pub(crate) fn supports_infallible_session_entry(self) -> bool {
491        self.domain.ownership() == crate::CpuDomainOwnership::Managed
492    }
493
494    pub(crate) fn enter_managed_session<R: Send>(
495        self,
496        operation: impl FnOnce(CpuExecutionContext<'a>) -> R + Send,
497    ) -> R {
498        assert!(
499            self.supports_infallible_session_entry(),
500            "managed session entry requires a Tenferro-managed CPU domain"
501        );
502        let mode = self.preferred_engine_mode();
503        self.enter(mode, |_| {
504            operation(CpuExecutionContext::entered(self.domain, mode))
505        })
506        .unwrap_or_else(|error| {
507            panic!("Tenferro-managed CPU executor violated synchronous install contract: {error}")
508        })
509    }
510
511    pub(crate) fn submit_outer(
512        self,
513        len: usize,
514        operation: impl Fn(usize, &CpuExecutionContext<'_>) -> Result<(), CpuDomainExecutorError> + Sync,
515    ) -> Result<(), CpuDomainExecutorError> {
516        if !self.supports_outer() {
517            return Err(CpuDomainExecutorError::Scheduling {
518                message: format!(
519                    "CPU domain {:?} does not support Outer mode",
520                    self.domain.id()
521                ),
522            });
523        }
524        let owner = self.permit.owner();
525        let lane_count = len.min(self.domain.thread_budget().get());
526        let jobs = indexed_jobs(lane_count, |lane| {
527            // INVARIANT: valid lanes partition `0..len` by residue modulo the
528            // nonzero `lane_count`, so every logical job runs exactly once
529            // while the executor can schedule at most the domain budget.
530            let mut index = lane;
531            while index < len {
532                with_execution_owner(owner, || {
533                    let context =
534                        CpuExecutionContext::entered(self.domain, ParallelMode::Sequential);
535                    operation(index, &context)
536                })?;
537                let Some(next) = index.checked_add(lane_count) else {
538                    break;
539                };
540                index = next;
541            }
542            Ok(())
543        });
544        with_execution_owner(owner, || self.domain.executor().submit(&jobs))?;
545        if let Some(index) = jobs.invalid_index_attempt() {
546            return Err(CpuDomainExecutorError::Scheduling {
547                message: format!(
548                    "executor requested scoped CPU lane index {index}, but the submission has {lane_count} lanes for {len} logical jobs"
549                ),
550            });
551        }
552        Ok(())
553    }
554
555    pub(crate) fn preferred_engine_mode(self) -> ParallelMode {
556        if self.domain.thread_budget().get() > 1
557            && self.domain.executor_capabilities().inner_parallelism == CpuInnerParallelism::Rayon
558        {
559            ParallelMode::Inner
560        } else {
561            ParallelMode::Sequential
562        }
563    }
564
565    pub(crate) fn preferred_provider_mode(
566        self,
567        accepts: impl Fn(ParallelMode) -> bool,
568    ) -> Result<ParallelMode, crate::CpuProviderDomainError> {
569        if self.domain.thread_budget().get() == 1 {
570            return if accepts(ParallelMode::Sequential) {
571                Ok(ParallelMode::Sequential)
572            } else {
573                Err(crate::CpuProviderDomainError::ParallelModeNotSupported {
574                    mode: ParallelMode::Sequential,
575                })
576            };
577        }
578        if accepts(ParallelMode::Inner) {
579            return Ok(ParallelMode::Inner);
580        }
581        if accepts(ParallelMode::Sequential) {
582            return Ok(ParallelMode::Sequential);
583        }
584        Err(crate::CpuProviderDomainError::ParallelModeNotSupported {
585            mode: ParallelMode::Inner,
586        })
587    }
588
589    pub(crate) fn preferred_linalg_mode(self, kind: CpuBackendKind) -> ParallelMode {
590        if self.domain.thread_budget().get() == 1 {
591            return ParallelMode::Sequential;
592        }
593        match kind {
594            CpuBackendKind::Faer => self.preferred_engine_mode(),
595            // The linalg operation-family provider has not yet moved onto the
596            // provider capability traits. Preserve its existing ownership
597            // policy until that trait boundary is introduced.
598            CpuBackendKind::Blas => ParallelMode::Inner,
599        }
600    }
601
602    pub(crate) fn provider_default_compatibility_mode(self) -> ParallelMode {
603        if self.domain.thread_budget().get() == 1 {
604            ParallelMode::Sequential
605        } else {
606            ParallelMode::Inner
607        }
608    }
609
610    pub(crate) fn thread_budget(self) -> NonZeroUsize {
611        self.domain.thread_budget()
612    }
613
614    pub(crate) fn supports_outer(self) -> bool {
615        self.domain.thread_budget().get() > 1
616            && self.domain.executor_capabilities().outer_parallelism
617    }
618}
619
620/// Checked element offset and strides for a batched matrix operand.
621///
622/// # Examples
623///
624/// Providers receive this descriptor from a validated request:
625///
626/// ```
627/// use tenferro_cpu::provider::CpuGemmRequest;
628/// # fn inspect(request: &CpuGemmRequest<'_, '_, '_>) {
629/// assert!(request.lhs_layout().row_stride() != 0 || request.rows() <= 1);
630/// # }
631/// ```
632#[derive(Clone, Copy, Debug, PartialEq, Eq)]
633pub struct CpuBatchedMatrixLayout {
634    offset: isize,
635    row_stride: isize,
636    column_stride: isize,
637    batch_stride: isize,
638}
639
640impl CpuBatchedMatrixLayout {
641    #[allow(dead_code)]
642    pub(crate) fn new(
643        offset: isize,
644        row_stride: isize,
645        column_stride: isize,
646        batch_stride: isize,
647    ) -> Self {
648        Self {
649            offset,
650            row_stride,
651            column_stride,
652            batch_stride,
653        }
654    }
655
656    /// Return the checked base element offset.
657    pub fn offset(self) -> isize {
658        self.offset
659    }
660
661    /// Return the row element stride.
662    pub fn row_stride(self) -> isize {
663        self.row_stride
664    }
665
666    /// Return the column element stride.
667    pub fn column_stride(self) -> isize {
668        self.column_stride
669    }
670
671    /// Return the batch element stride.
672    pub fn batch_stride(self) -> isize {
673        self.batch_stride
674    }
675}
676
677/// Validated borrowed GEMM request.
678///
679/// A batch count of one is a single GEMM. A larger batch count is a strided
680/// batched GEMM.
681///
682/// # Examples
683///
684/// ```
685/// use tenferro_cpu::provider::CpuGemmRequest;
686/// # fn inspect(request: &CpuGemmRequest<'_, '_, '_>) {
687/// assert!(request.batch_count() >= 1);
688/// # }
689/// ```
690#[derive(Debug)]
691pub struct CpuGemmRequest<'request, 'input, 'output> {
692    lhs: &'request TensorRead<'input>,
693    rhs: &'request TensorRead<'input>,
694    output: &'request mut TensorWrite<'output>,
695    rows: usize,
696    columns: usize,
697    contracted: usize,
698    batch_count: usize,
699    lhs_layout: CpuBatchedMatrixLayout,
700    rhs_layout: CpuBatchedMatrixLayout,
701    output_layout: CpuBatchedMatrixLayout,
702    accumulation: DotGeneralAccumulation,
703}
704
705pub(crate) struct CpuGemmRequestParts<'request, 'input, 'output> {
706    pub(crate) lhs: &'request TensorRead<'input>,
707    pub(crate) rhs: &'request TensorRead<'input>,
708    pub(crate) output: &'request mut TensorWrite<'output>,
709    pub(crate) rows: usize,
710    pub(crate) columns: usize,
711    pub(crate) contracted: usize,
712    pub(crate) batch_count: usize,
713    pub(crate) lhs_layout: CpuBatchedMatrixLayout,
714    pub(crate) rhs_layout: CpuBatchedMatrixLayout,
715    pub(crate) output_layout: CpuBatchedMatrixLayout,
716    pub(crate) accumulation: DotGeneralAccumulation,
717}
718
719impl<'request, 'input, 'output> CpuGemmRequest<'request, 'input, 'output> {
720    #[allow(clippy::too_many_arguments, dead_code)]
721    pub(crate) fn new(
722        lhs: &'request TensorRead<'input>,
723        rhs: &'request TensorRead<'input>,
724        output: &'request mut TensorWrite<'output>,
725        rows: usize,
726        columns: usize,
727        contracted: usize,
728        batch_count: usize,
729        lhs_layout: CpuBatchedMatrixLayout,
730        rhs_layout: CpuBatchedMatrixLayout,
731        output_layout: CpuBatchedMatrixLayout,
732        accumulation: DotGeneralAccumulation,
733    ) -> Self {
734        Self {
735            lhs,
736            rhs,
737            output,
738            rows,
739            columns,
740            contracted,
741            batch_count,
742            lhs_layout,
743            rhs_layout,
744            output_layout,
745            accumulation,
746        }
747    }
748
749    /// Return the borrowed left input.
750    pub fn lhs(&self) -> &TensorRead<'input> {
751        self.lhs
752    }
753
754    /// Return the borrowed right input.
755    pub fn rhs(&self) -> &TensorRead<'input> {
756        self.rhs
757    }
758
759    /// Reborrow the writable output for the duration of the current call.
760    pub fn output(&mut self) -> &mut TensorWrite<'output> {
761        self.output
762    }
763
764    /// Return the number of output rows.
765    pub fn rows(&self) -> usize {
766        self.rows
767    }
768
769    /// Return the number of output columns.
770    pub fn columns(&self) -> usize {
771        self.columns
772    }
773
774    /// Return the contracted dimension.
775    pub fn contracted(&self) -> usize {
776        self.contracted
777    }
778
779    /// Return the number of matrices in this strided batch.
780    pub fn batch_count(&self) -> usize {
781        self.batch_count
782    }
783
784    /// Return the left input matrix layout.
785    pub fn lhs_layout(&self) -> CpuBatchedMatrixLayout {
786        self.lhs_layout
787    }
788
789    /// Return the right input matrix layout.
790    pub fn rhs_layout(&self) -> CpuBatchedMatrixLayout {
791        self.rhs_layout
792    }
793
794    /// Return the output matrix layout.
795    pub fn output_layout(&self) -> CpuBatchedMatrixLayout {
796        self.output_layout
797    }
798
799    /// Return conjugation and alpha/beta update semantics.
800    pub fn accumulation(&self) -> DotGeneralAccumulation {
801        self.accumulation
802    }
803
804    pub(crate) fn into_parts(self) -> CpuGemmRequestParts<'request, 'input, 'output> {
805        CpuGemmRequestParts {
806            lhs: self.lhs,
807            rhs: self.rhs,
808            output: self.output,
809            rows: self.rows,
810            columns: self.columns,
811            contracted: self.contracted,
812            batch_count: self.batch_count,
813            lhs_layout: self.lhs_layout,
814            rhs_layout: self.rhs_layout,
815            output_layout: self.output_layout,
816            accumulation: self.accumulation,
817        }
818    }
819}
820
821/// Validated borrowed grouped-GEMM request.
822///
823/// # Examples
824///
825/// ```
826/// use tenferro_cpu::provider::CpuGroupedGemmRequest;
827/// # fn inspect(request: &CpuGroupedGemmRequest<'_, '_, '_>) {
828/// assert_eq!(request.jobs().len(), request.jobs().iter().count());
829/// # }
830/// ```
831#[derive(Debug)]
832pub struct CpuGroupedGemmRequest<'request, 'input, 'output> {
833    lhs: &'request TensorRead<'input>,
834    rhs: &'request TensorRead<'input>,
835    output: &'request mut TensorWrite<'output>,
836    jobs: &'request [GroupedGemmJob],
837    accumulation: DotGeneralAccumulation,
838}
839
840impl<'request, 'input, 'output> CpuGroupedGemmRequest<'request, 'input, 'output> {
841    #[allow(dead_code)]
842    pub(crate) fn new(
843        lhs: &'request TensorRead<'input>,
844        rhs: &'request TensorRead<'input>,
845        output: &'request mut TensorWrite<'output>,
846        jobs: &'request [GroupedGemmJob],
847        accumulation: DotGeneralAccumulation,
848    ) -> Self {
849        Self {
850            lhs,
851            rhs,
852            output,
853            jobs,
854            accumulation,
855        }
856    }
857
858    /// Return the borrowed left input.
859    pub fn lhs(&self) -> &TensorRead<'input> {
860        self.lhs
861    }
862
863    /// Return the borrowed right input.
864    pub fn rhs(&self) -> &TensorRead<'input> {
865        self.rhs
866    }
867
868    /// Reborrow the writable output.
869    pub fn output(&mut self) -> &mut TensorWrite<'output> {
870        self.output
871    }
872
873    /// Return ordered, pairwise-disjoint validated jobs.
874    pub fn jobs(&self) -> &[GroupedGemmJob] {
875        self.jobs
876    }
877
878    /// Return shared conjugation and alpha/beta semantics.
879    pub fn accumulation(&self) -> DotGeneralAccumulation {
880        self.accumulation
881    }
882
883    pub(crate) fn into_parts(
884        self,
885    ) -> (
886        &'request TensorRead<'input>,
887        &'request TensorRead<'input>,
888        &'request mut TensorWrite<'output>,
889        &'request [GroupedGemmJob],
890        DotGeneralAccumulation,
891    ) {
892        (
893            self.lhs,
894            self.rhs,
895            self.output,
896            self.jobs,
897            self.accumulation,
898        )
899    }
900}
901
902/// Engine-requested layout materialization.
903///
904/// # Examples
905///
906/// ```
907/// use tenferro_cpu::provider::CpuLayoutTransformIntent;
908/// assert_eq!(
909///     CpuLayoutTransformIntent::CanonicalColumnMajor,
910///     CpuLayoutTransformIntent::CanonicalColumnMajor,
911/// );
912/// ```
913#[derive(Clone, Copy, Debug, PartialEq, Eq)]
914pub enum CpuLayoutTransformIntent {
915    /// Materialize a compact canonical column-major tensor.
916    CanonicalColumnMajor,
917}
918
919/// Validated borrowed layout-transform request.
920///
921/// # Examples
922///
923/// ```
924/// use tenferro_cpu::provider::{CpuLayoutTransformIntent, CpuLayoutTransformRequest};
925/// # fn inspect(request: &CpuLayoutTransformRequest<'_, '_, '_>) {
926/// assert_eq!(request.intent(), CpuLayoutTransformIntent::CanonicalColumnMajor);
927/// # }
928/// ```
929#[derive(Debug)]
930pub struct CpuLayoutTransformRequest<'request, 'input, 'output> {
931    input: &'request TensorRead<'input>,
932    output: &'request mut TensorWrite<'output>,
933    intent: CpuLayoutTransformIntent,
934    conjugate: bool,
935}
936
937impl<'request, 'input, 'output> CpuLayoutTransformRequest<'request, 'input, 'output> {
938    #[allow(dead_code)]
939    pub(crate) fn new(
940        input: &'request TensorRead<'input>,
941        output: &'request mut TensorWrite<'output>,
942        intent: CpuLayoutTransformIntent,
943        conjugate: bool,
944    ) -> Self {
945        Self {
946            input,
947            output,
948            intent,
949            conjugate,
950        }
951    }
952
953    /// Return the borrowed input.
954    pub fn input(&self) -> &TensorRead<'input> {
955        self.input
956    }
957
958    /// Reborrow the writable output.
959    pub fn output(&mut self) -> &mut TensorWrite<'output> {
960        self.output
961    }
962
963    /// Return the requested materialization intent.
964    pub fn intent(&self) -> CpuLayoutTransformIntent {
965        self.intent
966    }
967
968    /// Return whether materialization must conjugate each input element.
969    ///
970    /// # Examples
971    ///
972    /// ```
973    /// use tenferro_cpu::provider::CpuLayoutTransformRequest;
974    /// # fn inspect(request: &CpuLayoutTransformRequest<'_, '_, '_>) {
975    /// let _must_conjugate = request.conjugate();
976    /// # }
977    /// ```
978    pub fn conjugate(&self) -> bool {
979        self.conjugate
980    }
981
982    pub(crate) fn into_parts(
983        self,
984    ) -> (
985        &'request TensorRead<'input>,
986        &'request mut TensorWrite<'output>,
987        CpuLayoutTransformIntent,
988        bool,
989    ) {
990        (self.input, self.output, self.intent, self.conjugate)
991    }
992}
993
994/// Validated ordered contraction-role groups.
995///
996/// # Examples
997///
998/// ```
999/// use tenferro_cpu::provider::CpuDotGeneralRequest;
1000/// # fn inspect(request: &CpuDotGeneralRequest<'_, '_, '_>) {
1001/// let _pairs = request.axes().contracting_pairs().count();
1002/// # }
1003/// ```
1004#[derive(Clone, Copy, Debug)]
1005pub struct CpuContractionAxes<'a> {
1006    lhs_rank: usize,
1007    rhs_rank: usize,
1008    lhs_contracting: &'a [usize],
1009    rhs_contracting: &'a [usize],
1010    lhs_batch: &'a [usize],
1011    rhs_batch: &'a [usize],
1012    lhs_role_mask: Option<u64>,
1013    rhs_role_mask: Option<u64>,
1014}
1015
1016impl<'a> CpuContractionAxes<'a> {
1017    #[allow(clippy::too_many_arguments, dead_code)]
1018    pub(crate) fn new(
1019        lhs_rank: usize,
1020        rhs_rank: usize,
1021        lhs_contracting: &'a [usize],
1022        rhs_contracting: &'a [usize],
1023        lhs_batch: &'a [usize],
1024        rhs_batch: &'a [usize],
1025        lhs_role_mask: Option<u64>,
1026        rhs_role_mask: Option<u64>,
1027    ) -> Self {
1028        Self {
1029            lhs_rank,
1030            rhs_rank,
1031            lhs_contracting,
1032            rhs_contracting,
1033            lhs_batch,
1034            rhs_batch,
1035            lhs_role_mask,
1036            rhs_role_mask,
1037        }
1038    }
1039
1040    /// Return ordered contracting-axis pairs.
1041    pub fn contracting_pairs(&self) -> impl ExactSizeIterator<Item = (usize, usize)> + '_ {
1042        self.lhs_contracting
1043            .iter()
1044            .copied()
1045            .zip(self.rhs_contracting.iter().copied())
1046    }
1047
1048    /// Return ordered batch-axis pairs.
1049    pub fn batch_pairs(&self) -> impl ExactSizeIterator<Item = (usize, usize)> + '_ {
1050        self.lhs_batch
1051            .iter()
1052            .copied()
1053            .zip(self.rhs_batch.iter().copied())
1054    }
1055
1056    /// Return left axes that are neither contracting nor batch axes.
1057    pub fn lhs_free_axes(&self) -> impl Iterator<Item = usize> + '_ {
1058        (0..self.lhs_rank).filter(move |&axis| !self.lhs_axis_has_role(axis))
1059    }
1060
1061    /// Return right axes that are neither contracting nor batch axes.
1062    pub fn rhs_free_axes(&self) -> impl Iterator<Item = usize> + '_ {
1063        (0..self.rhs_rank).filter(move |&axis| !self.rhs_axis_has_role(axis))
1064    }
1065
1066    fn lhs_axis_has_role(&self, axis: usize) -> bool {
1067        self.lhs_role_mask.map_or_else(
1068            || self.lhs_contracting.contains(&axis) || self.lhs_batch.contains(&axis),
1069            |mask| mask & (1_u64 << axis) != 0,
1070        )
1071    }
1072
1073    fn rhs_axis_has_role(&self, axis: usize) -> bool {
1074        self.rhs_role_mask.map_or_else(
1075            || self.rhs_contracting.contains(&axis) || self.rhs_batch.contains(&axis),
1076            |mask| mask & (1_u64 << axis) != 0,
1077        )
1078    }
1079}
1080
1081/// Validated borrowed semantic binary `dot_general` request.
1082///
1083/// # Examples
1084///
1085/// ```
1086/// use tenferro_cpu::provider::CpuDotGeneralRequest;
1087/// # fn inspect(request: &CpuDotGeneralRequest<'_, '_, '_>) {
1088/// let _ = request.accumulation();
1089/// # }
1090/// ```
1091#[derive(Debug)]
1092pub struct CpuDotGeneralRequest<'request, 'input, 'output> {
1093    lhs: &'request TensorRead<'input>,
1094    rhs: &'request TensorRead<'input>,
1095    output: &'request mut TensorWrite<'output>,
1096    axes: CpuContractionAxes<'request>,
1097    accumulation: DotGeneralAccumulation,
1098}
1099
1100impl<'request, 'input, 'output> CpuDotGeneralRequest<'request, 'input, 'output> {
1101    #[allow(dead_code)]
1102    pub(crate) fn new(
1103        lhs: &'request TensorRead<'input>,
1104        rhs: &'request TensorRead<'input>,
1105        output: &'request mut TensorWrite<'output>,
1106        axes: CpuContractionAxes<'request>,
1107        accumulation: DotGeneralAccumulation,
1108    ) -> Self {
1109        Self {
1110            lhs,
1111            rhs,
1112            output,
1113            axes,
1114            accumulation,
1115        }
1116    }
1117
1118    /// Return the borrowed left input.
1119    pub fn lhs(&self) -> &TensorRead<'input> {
1120        self.lhs
1121    }
1122
1123    /// Return the borrowed right input.
1124    pub fn rhs(&self) -> &TensorRead<'input> {
1125        self.rhs
1126    }
1127
1128    /// Reborrow the writable output.
1129    pub fn output(&mut self) -> &mut TensorWrite<'output> {
1130        self.output
1131    }
1132
1133    /// Return the validated ordered axis groups.
1134    pub fn axes(&self) -> &CpuContractionAxes<'request> {
1135        &self.axes
1136    }
1137
1138    /// Return conjugation and alpha/beta update semantics.
1139    pub fn accumulation(&self) -> DotGeneralAccumulation {
1140        self.accumulation
1141    }
1142
1143    /// Consume the request and return the validated operand borrows.
1144    ///
1145    /// External general-contraction providers use this when they need
1146    /// simultaneous immutable access to both inputs and mutable access to the
1147    /// output.
1148    ///
1149    /// # Examples
1150    ///
1151    /// ```
1152    /// use tenferro_cpu::provider::CpuDotGeneralRequest;
1153    ///
1154    /// fn consume_request(request: CpuDotGeneralRequest<'_, '_, '_>) {
1155    ///     let (_lhs, _rhs, _output, _axes, _accumulation) = request.into_parts();
1156    /// }
1157    /// ```
1158    pub fn into_parts(
1159        self,
1160    ) -> (
1161        &'request TensorRead<'input>,
1162        &'request TensorRead<'input>,
1163        &'request mut TensorWrite<'output>,
1164        CpuContractionAxes<'request>,
1165        DotGeneralAccumulation,
1166    ) {
1167        (
1168            self.lhs,
1169            self.rhs,
1170            self.output,
1171            self.axes,
1172            self.accumulation,
1173        )
1174    }
1175}
1176
1177/// Provider for validated GEMM-family requests.
1178///
1179/// # Examples
1180///
1181/// Trait objects are supported directly:
1182///
1183/// ```
1184/// use tenferro_cpu::provider::CpuGemmProvider;
1185/// # fn accepts_provider(_: &dyn CpuGemmProvider) {}
1186/// ```
1187pub trait CpuGemmProvider: fmt::Debug + Send + Sync + 'static {
1188    /// Return immutable count, placement, and fan-out capabilities.
1189    ///
1190    /// This declaration must describe controls actually applied and restored
1191    /// by the provider adapter around each call. Merely discovering a runtime
1192    /// symbol is insufficient. A provider bundle samples this method exactly
1193    /// once during construction and keeps that snapshot for its lifetime; the
1194    /// returned contract must therefore remain valid for the provider object.
1195    fn execution_capabilities(&self) -> CpuProviderExecutionCapabilities;
1196
1197    /// Execute one validated GEMM.
1198    ///
1199    /// # Errors
1200    ///
1201    /// Returns [`tenferro_tensor::Error::BackendSource`] or
1202    /// [`tenferro_tensor::Error::BackendFailure`] when the provider runtime
1203    /// fails. A detected inconsistency in engine-attested request metadata is
1204    /// returned as [`tenferro_tensor::Error::Validation`]. Unsupported
1205    /// capabilities use [`CpuProviderOutcome::Unsupported`] instead.
1206    fn gemm(
1207        &self,
1208        context: &CpuExecutionContext<'_>,
1209        request: CpuGemmRequest<'_, '_, '_>,
1210    ) -> tenferro_tensor::Result<CpuProviderOutcome>;
1211
1212    /// Execute one validated strided-batched GEMM.
1213    ///
1214    /// # Errors
1215    ///
1216    /// Returns [`tenferro_tensor::Error::BackendSource`] or
1217    /// [`tenferro_tensor::Error::BackendFailure`] when the provider runtime
1218    /// fails. A detected inconsistency in engine-attested request metadata is
1219    /// returned as [`tenferro_tensor::Error::Validation`]. Unsupported
1220    /// capabilities use [`CpuProviderOutcome::Unsupported`] instead.
1221    fn strided_batched_gemm(
1222        &self,
1223        context: &CpuExecutionContext<'_>,
1224        request: CpuGemmRequest<'_, '_, '_>,
1225    ) -> tenferro_tensor::Result<CpuProviderOutcome>;
1226
1227    /// Execute one validated grouped GEMM.
1228    ///
1229    /// # Errors
1230    ///
1231    /// Returns [`tenferro_tensor::Error::BackendSource`] or
1232    /// [`tenferro_tensor::Error::BackendFailure`] when the provider runtime
1233    /// fails. A detected inconsistency in engine-attested request metadata is
1234    /// returned as [`tenferro_tensor::Error::Validation`]. Unsupported
1235    /// capabilities use [`CpuProviderOutcome::Unsupported`] instead.
1236    fn grouped_gemm(
1237        &self,
1238        context: &CpuExecutionContext<'_>,
1239        request: CpuGroupedGemmRequest<'_, '_, '_>,
1240    ) -> tenferro_tensor::Result<CpuProviderOutcome>;
1241}
1242
1243/// Provider for engine-owned tensor materialization.
1244///
1245/// # Examples
1246///
1247/// ```
1248/// use tenferro_cpu::provider::CpuLayoutTransformProvider;
1249/// # fn accepts_provider(_: &dyn CpuLayoutTransformProvider) {}
1250/// ```
1251pub trait CpuLayoutTransformProvider: fmt::Debug + Send + Sync + 'static {
1252    /// Return immutable count, placement, and fan-out capabilities.
1253    ///
1254    /// A provider bundle samples this method exactly once during construction
1255    /// and uses the stored descriptor for all validation and dispatch.
1256    fn execution_capabilities(&self) -> CpuProviderExecutionCapabilities;
1257
1258    /// Materialize one validated input into a preallocated output.
1259    ///
1260    /// # Errors
1261    ///
1262    /// Returns [`tenferro_tensor::Error::BackendSource`] or
1263    /// [`tenferro_tensor::Error::BackendFailure`] when execution fails. A
1264    /// detected inconsistency in engine-attested layout or range metadata is
1265    /// returned as [`tenferro_tensor::Error::Validation`]. Unsupported layouts
1266    /// use [`CpuProviderOutcome::Unsupported`] instead.
1267    fn materialize(
1268        &self,
1269        context: &CpuExecutionContext<'_>,
1270        request: CpuLayoutTransformRequest<'_, '_, '_>,
1271    ) -> tenferro_tensor::Result<CpuProviderOutcome>;
1272}
1273
1274/// Provider for complete validated binary `dot_general` requests.
1275///
1276/// # Examples
1277///
1278/// ```
1279/// use tenferro_cpu::provider::CpuGeneralContractionProvider;
1280/// # fn accepts_provider(_: &dyn CpuGeneralContractionProvider) {}
1281/// ```
1282pub trait CpuGeneralContractionProvider: fmt::Debug + Send + Sync + 'static {
1283    /// Return immutable count, placement, and fan-out capabilities.
1284    ///
1285    /// A provider bundle samples this method exactly once during construction
1286    /// and uses the stored descriptor for all validation and dispatch.
1287    fn execution_capabilities(&self) -> CpuProviderExecutionCapabilities;
1288
1289    /// Execute one complete semantic contraction.
1290    ///
1291    /// # Errors
1292    ///
1293    /// Returns [`tenferro_tensor::Error::BackendSource`] or
1294    /// [`tenferro_tensor::Error::BackendFailure`] when the provider runtime
1295    /// fails. A detected inconsistency in engine-attested axes or range
1296    /// metadata is returned as [`tenferro_tensor::Error::Validation`].
1297    /// Unsupported contractions use [`CpuProviderOutcome::Unsupported`]
1298    /// instead.
1299    fn dot_general(
1300        &self,
1301        context: &CpuExecutionContext<'_>,
1302        request: CpuDotGeneralRequest<'_, '_, '_>,
1303    ) -> tenferro_tensor::Result<CpuProviderOutcome>;
1304}
1305
1306/// Built-in faer GEMM provider.
1307///
1308/// # Examples
1309///
1310/// ```
1311/// use tenferro_cpu::provider::{CpuGemmProvider, FaerGemmProvider};
1312/// let provider: &dyn CpuGemmProvider = &FaerGemmProvider;
1313/// let _ = provider;
1314/// ```
1315#[derive(Clone, Copy, Debug, Default)]
1316pub struct FaerGemmProvider;
1317
1318impl CpuGemmProvider for FaerGemmProvider {
1319    fn execution_capabilities(&self) -> CpuProviderExecutionCapabilities {
1320        engine_worker_capabilities()
1321    }
1322
1323    fn gemm(
1324        &self,
1325        context: &CpuExecutionContext<'_>,
1326        request: CpuGemmRequest<'_, '_, '_>,
1327    ) -> tenferro_tensor::Result<CpuProviderOutcome> {
1328        #[cfg(feature = "cpu-faer")]
1329        {
1330            crate::gemm::execute_faer_gemm_request(context, request)
1331        }
1332        #[cfg(not(feature = "cpu-faer"))]
1333        {
1334            let _ = (context, request);
1335            Ok(CpuProviderOutcome::Unsupported(
1336                CpuProviderUnsupported::RuntimeUnavailable,
1337            ))
1338        }
1339    }
1340
1341    fn strided_batched_gemm(
1342        &self,
1343        context: &CpuExecutionContext<'_>,
1344        request: CpuGemmRequest<'_, '_, '_>,
1345    ) -> tenferro_tensor::Result<CpuProviderOutcome> {
1346        self.gemm(context, request)
1347    }
1348
1349    fn grouped_gemm(
1350        &self,
1351        context: &CpuExecutionContext<'_>,
1352        request: CpuGroupedGemmRequest<'_, '_, '_>,
1353    ) -> tenferro_tensor::Result<CpuProviderOutcome> {
1354        #[cfg(feature = "cpu-faer")]
1355        {
1356            crate::gemm::execute_faer_grouped_request(context, request)
1357        }
1358        #[cfg(not(feature = "cpu-faer"))]
1359        {
1360            let _ = (context, request);
1361            Ok(CpuProviderOutcome::Unsupported(
1362                CpuProviderUnsupported::RuntimeUnavailable,
1363            ))
1364        }
1365    }
1366}
1367
1368/// Built-in BLAS GEMM provider.
1369///
1370/// # Examples
1371///
1372/// ```
1373/// use tenferro_cpu::provider::{BlasGemmProvider, CpuGemmProvider};
1374/// let provider: &dyn CpuGemmProvider = &BlasGemmProvider;
1375/// let _ = provider;
1376/// ```
1377#[derive(Clone, Copy, Debug, Default)]
1378pub struct BlasGemmProvider;
1379
1380impl CpuGemmProvider for BlasGemmProvider {
1381    fn execution_capabilities(&self) -> CpuProviderExecutionCapabilities {
1382        #[cfg(feature = "cpu-blas")]
1383        {
1384            builtin_blas_execution_capabilities()
1385        }
1386        #[cfg(not(feature = "cpu-blas"))]
1387        {
1388            // This build returns RuntimeUnavailable without invoking BLAS.
1389            serial_capabilities()
1390        }
1391    }
1392
1393    fn gemm(
1394        &self,
1395        context: &CpuExecutionContext<'_>,
1396        request: CpuGemmRequest<'_, '_, '_>,
1397    ) -> tenferro_tensor::Result<CpuProviderOutcome> {
1398        #[cfg(feature = "cpu-blas")]
1399        {
1400            crate::gemm::execute_blas_gemm_request(context, request)
1401        }
1402        #[cfg(not(feature = "cpu-blas"))]
1403        {
1404            let _ = (context, request);
1405            Ok(CpuProviderOutcome::Unsupported(
1406                CpuProviderUnsupported::RuntimeUnavailable,
1407            ))
1408        }
1409    }
1410
1411    fn strided_batched_gemm(
1412        &self,
1413        context: &CpuExecutionContext<'_>,
1414        request: CpuGemmRequest<'_, '_, '_>,
1415    ) -> tenferro_tensor::Result<CpuProviderOutcome> {
1416        self.gemm(context, request)
1417    }
1418
1419    fn grouped_gemm(
1420        &self,
1421        context: &CpuExecutionContext<'_>,
1422        request: CpuGroupedGemmRequest<'_, '_, '_>,
1423    ) -> tenferro_tensor::Result<CpuProviderOutcome> {
1424        #[cfg(feature = "cpu-blas")]
1425        {
1426            crate::gemm::execute_blas_grouped_request(context, request)
1427        }
1428        #[cfg(not(feature = "cpu-blas"))]
1429        {
1430            let _ = (context, request);
1431            Ok(CpuProviderOutcome::Unsupported(
1432                CpuProviderUnsupported::RuntimeUnavailable,
1433            ))
1434        }
1435    }
1436}
1437
1438/// Built-in strided layout materialization provider.
1439///
1440/// # Examples
1441///
1442/// ```
1443/// use tenferro_cpu::provider::{CpuLayoutTransformProvider, StridedLayoutTransformProvider};
1444/// let provider: &dyn CpuLayoutTransformProvider = &StridedLayoutTransformProvider;
1445/// let _ = provider;
1446/// ```
1447#[derive(Clone, Copy, Debug, Default)]
1448pub struct StridedLayoutTransformProvider;
1449
1450impl CpuLayoutTransformProvider for StridedLayoutTransformProvider {
1451    fn execution_capabilities(&self) -> CpuProviderExecutionCapabilities {
1452        engine_worker_capabilities()
1453    }
1454
1455    fn materialize(
1456        &self,
1457        context: &CpuExecutionContext<'_>,
1458        request: CpuLayoutTransformRequest<'_, '_, '_>,
1459    ) -> tenferro_tensor::Result<CpuProviderOutcome> {
1460        context.with_native_parallelism(|| materialize_strided_layout(request))
1461    }
1462}
1463
1464fn materialize_strided_layout(
1465    request: CpuLayoutTransformRequest<'_, '_, '_>,
1466) -> tenferro_tensor::Result<CpuProviderOutcome> {
1467    let (input, output, _intent, conjugate) = request.into_parts();
1468    if conjugate {
1469        macro_rules! dispatch_conjugated {
1470            ($owned:ident, $view:ident) => {
1471                match (input, &mut *output) {
1472                    (
1473                        TensorRead::Tensor(Tensor::$owned(input)),
1474                        TensorWrite::Tensor(Tensor::$owned(output)),
1475                    ) => {
1476                        let input = input.as_view();
1477                        let mut output = output.as_view_mut();
1478                        crate::structural::typed_conjugate_view_into(
1479                            &input,
1480                            &mut output,
1481                            "cpu layout materialization",
1482                        )?;
1483                        return Ok(CpuProviderOutcome::Executed);
1484                    }
1485                    (
1486                        TensorRead::View(TensorView::$view(input)),
1487                        TensorWrite::Tensor(Tensor::$owned(output)),
1488                    ) => {
1489                        let mut output = output.as_view_mut();
1490                        crate::structural::typed_conjugate_view_into(
1491                            input,
1492                            &mut output,
1493                            "cpu layout materialization",
1494                        )?;
1495                        return Ok(CpuProviderOutcome::Executed);
1496                    }
1497                    (
1498                        TensorRead::Tensor(Tensor::$owned(input)),
1499                        TensorWrite::View(TensorViewMut::$view(output)),
1500                    ) => {
1501                        let input = input.as_view();
1502                        crate::structural::typed_conjugate_view_into(
1503                            &input,
1504                            output,
1505                            "cpu layout materialization",
1506                        )?;
1507                        return Ok(CpuProviderOutcome::Executed);
1508                    }
1509                    (
1510                        TensorRead::View(TensorView::$view(input)),
1511                        TensorWrite::View(TensorViewMut::$view(output)),
1512                    ) => {
1513                        crate::structural::typed_conjugate_view_into(
1514                            input,
1515                            output,
1516                            "cpu layout materialization",
1517                        )?;
1518                        return Ok(CpuProviderOutcome::Executed);
1519                    }
1520                    _ => {}
1521                }
1522            };
1523        }
1524        dispatch_conjugated!(F32, F32);
1525        dispatch_conjugated!(F64, F64);
1526        dispatch_conjugated!(C32, C32);
1527        dispatch_conjugated!(C64, C64);
1528        return Ok(CpuProviderOutcome::Unsupported(
1529            CpuProviderUnsupported::DType(input.dtype()),
1530        ));
1531    }
1532    macro_rules! dispatch {
1533        ($owned:ident, $view:ident) => {
1534            match (input, &mut *output) {
1535                (
1536                    TensorRead::Tensor(Tensor::$owned(input)),
1537                    TensorWrite::Tensor(Tensor::$owned(output)),
1538                ) => {
1539                    let input = input.as_view();
1540                    let mut output = output.as_view_mut();
1541                    crate::structural::typed_copy_view_into(
1542                        &input,
1543                        &mut output,
1544                        "cpu layout materialization",
1545                    )?;
1546                    return Ok(CpuProviderOutcome::Executed);
1547                }
1548                (
1549                    TensorRead::View(TensorView::$view(input)),
1550                    TensorWrite::Tensor(Tensor::$owned(output)),
1551                ) => {
1552                    let mut output = output.as_view_mut();
1553                    crate::structural::typed_copy_view_into(
1554                        input,
1555                        &mut output,
1556                        "cpu layout materialization",
1557                    )?;
1558                    return Ok(CpuProviderOutcome::Executed);
1559                }
1560                (
1561                    TensorRead::Tensor(Tensor::$owned(input)),
1562                    TensorWrite::View(TensorViewMut::$view(output)),
1563                ) => {
1564                    let input = input.as_view();
1565                    crate::structural::typed_copy_view_into(
1566                        &input,
1567                        output,
1568                        "cpu layout materialization",
1569                    )?;
1570                    return Ok(CpuProviderOutcome::Executed);
1571                }
1572                (
1573                    TensorRead::View(TensorView::$view(input)),
1574                    TensorWrite::View(TensorViewMut::$view(output)),
1575                ) => {
1576                    crate::structural::typed_copy_view_into(
1577                        input,
1578                        output,
1579                        "cpu layout materialization",
1580                    )?;
1581                    return Ok(CpuProviderOutcome::Executed);
1582                }
1583                _ => {}
1584            }
1585        };
1586    }
1587    dispatch!(F32, F32);
1588    dispatch!(F64, F64);
1589    dispatch!(I32, I32);
1590    dispatch!(I64, I64);
1591    dispatch!(Bool, Bool);
1592    dispatch!(C32, C32);
1593    dispatch!(C64, C64);
1594    Ok(CpuProviderOutcome::Unsupported(
1595        CpuProviderUnsupported::DType(input.dtype()),
1596    ))
1597}
1598
1599pub(crate) fn builtin_gemm_provider(kind: CpuBackendKind) -> Arc<dyn CpuGemmProvider> {
1600    match kind {
1601        CpuBackendKind::Faer => Arc::new(FaerGemmProvider),
1602        CpuBackendKind::Blas => Arc::new(BlasGemmProvider),
1603    }
1604}
1605
1606pub(crate) fn builtin_layout_provider() -> Arc<dyn CpuLayoutTransformProvider> {
1607    Arc::new(StridedLayoutTransformProvider)
1608}
1609
1610#[cfg(test)]
1611pub(crate) mod tests;
tenferro_cpu/provider.rs

tenferro_cpu/
provider.rs
