tenferro_cpu/
runtime_adapter.rs

1use std::fmt;
2use std::mem::{size_of, size_of_val};
3use std::sync::Arc;
4
5use tenferro_runtime::program::{CoreSemanticOp, SemanticOpRef, SemanticOperationView};
6use tenferro_runtime::{
7    assemble_executable_engine_registration, CoreCapabilityBundle, CoreCapabilityKind,
8    CorePrepareContext, DotGeneralPreparation, DotGeneralPrepareRequest, ElementwisePrepareRequest,
9    ElementwiseRuntime, EngineId, EngineRegistration, EngineRegistrationMetadata,
10    ExecutableEngineRegistrationConfig, ExecutionContextIdentity, HardwareClassId,
11    IndexingPrepareRequest, IndexingRuntime, InputIngressContract, InputPlacementContract,
12    InputSignature, InputSignatureContract, InputSpecializationProjection,
13    InputSpecializationRequirements, LayoutPrepareRequest, LayoutProjection, LayoutRuntime,
14    LayoutSpecialization, MemoryKind, PrepareCapability, PrepareError, PreparedOperation,
15    PreparedOperationBinding, PreparedOperationPlan, ProviderContractError, ProviderDeviceIdentity,
16    ProviderId, ReductionPrepareRequest, ReductionRuntime, ResidentOutputContract,
17    RuntimeConfigError, RuntimeInputContract, SpecializationError, SpecializationProjection,
18    SpecializationRequirements, StorageClass, TensorRead, UnsupportedReason,
19};
20
21use crate::CpuEngine;
22
23/// Stateless CPU preparation witness: CPU preparation reads no engine state.
24#[derive(Debug, Clone, Copy, Default)]
25struct CpuPreparer;
26
27const CPU_ENGINE_ID: &str = "tenferro-cpu.default.v1";
28const CPU_HARDWARE_CLASS_ID: &str = "tenferro-cpu.host.v1";
29const CPU_STORAGE_CLASS_ID: &str = "tenferro-cpu.host.v1";
30const UNKNOWN_CORE_OPERATION: &str = "unknown-core-operation";
31
32/// Return the canonical CPU runtime engine identifier.
33///
34/// # Errors
35///
36/// Returns [`RuntimeConfigError`] if the built-in CPU engine identifier violates
37/// runtime identifier validation.
38pub fn runtime_engine_id() -> Result<EngineId, RuntimeConfigError> {
39    EngineId::new(CPU_ENGINE_ID).map_err(RuntimeConfigError::from)
40}
41
42/// The built-in CPU engine id, which is a valid identifier by construction.
43pub(crate) fn default_runtime_engine_id() -> EngineId {
44    match EngineId::new(CPU_ENGINE_ID) {
45        Ok(id) => id,
46        Err(_) => unreachable!("the built-in CPU engine id is a valid identifier"),
47    }
48}
49
50/// Return the canonical CPU runtime hardware class.
51///
52/// # Errors
53///
54/// Returns [`RuntimeConfigError`] if the built-in CPU hardware class violates
55/// runtime identifier validation.
56pub fn runtime_hardware_class() -> Result<HardwareClassId, RuntimeConfigError> {
57    HardwareClassId::new(CPU_HARDWARE_CLASS_ID).map_err(RuntimeConfigError::from)
58}
59
60/// Build a runtime engine registration for a [`CpuEngine`].
61///
62/// The registration exposes the stateless CPU core preparation capabilities
63/// and the engine's input contract.
64///
65/// # Errors
66///
67/// Returns [`RuntimeConfigError`] if one of the built-in CPU runtime identifiers
68/// violates runtime validation or if the registration is internally invalid.
69pub fn runtime_engine_registration(
70    engine: &CpuEngine,
71) -> Result<EngineRegistration, RuntimeConfigError> {
72    runtime_engine_registration_with_id(engine, runtime_engine_id()?)
73}
74
75/// Build a runtime engine registration for a [`CpuEngine`] under `engine_id`.
76///
77/// # Errors
78///
79/// Returns [`RuntimeConfigError`] if `engine_id` or a built-in CPU runtime
80/// identifier violates runtime validation, or if the registration is invalid.
81pub fn runtime_engine_registration_with_id(
82    engine: &CpuEngine,
83    engine_id: EngineId,
84) -> Result<EngineRegistration, RuntimeConfigError> {
85    let preparer = Arc::new(CpuPreparer);
86    let elementwise: Arc<dyn ElementwiseRuntime> = preparer.clone();
87    let reduction: Arc<dyn ReductionRuntime> = preparer.clone();
88    let indexing: Arc<dyn IndexingRuntime> = preparer.clone();
89    let dot_general: Arc<dyn DotGeneralPreparation> = preparer.clone();
90    let layout: Arc<dyn LayoutRuntime> = preparer;
91
92    let mut capabilities = CoreCapabilityBundle::builder();
93    capabilities
94        .elementwise(elementwise)
95        .reduction(reduction)
96        .indexing(indexing)
97        .dot_general(dot_general)
98        .layout(layout);
99
100    let storage = runtime_storage_class()?;
101    let default_storage = storage.clone();
102    let placement_storage = storage.clone();
103    let signature_storage = storage.clone();
104    let runtime_storage = storage.clone();
105    let resident_storage = storage.clone();
106    let allocation_domain = engine.allocation_domain();
107    let provider_id = crate::cpu_provider_id();
108    let provider_device_identity =
109        ProviderDeviceIdentity::new(ProviderId::new(provider_id)?, "host".to_owned())?;
110    let ingress = InputIngressContract::new(
111        InputPlacementContract::new(move |placement, candidate| {
112            candidate == &placement_storage
113                && (cpu_input_placement(placement)
114                    || (placement.memory_kind == MemoryKind::Managed
115                        && allocation_domain.is_some()))
116        }),
117        InputSignatureContract::new(move |placement, family, domain, candidate| {
118            candidate == &signature_storage
119                && cpu_input_signature(placement, family, domain, allocation_domain)
120        }),
121        RuntimeInputContract::new(move |input: &TensorRead<'_>, candidate| {
122            candidate == &runtime_storage && cpu_runtime_input(input, allocation_domain)
123        }),
124        ResidentOutputContract::new(move |input: &TensorRead<'_>, candidate| {
125            candidate == &resident_storage && cpu_runtime_input(input, allocation_domain)
126        }),
127    );
128    let metadata = EngineRegistrationMetadata::new(
129        engine_id,
130        provider_device_identity,
131        runtime_hardware_class()?,
132        Arc::from(vec![storage]),
133        default_storage,
134        capabilities.build(),
135    );
136    assemble_executable_engine_registration(ExecutableEngineRegistrationConfig::new(
137        metadata,
138        ExecutionContextIdentity::of::<CpuEngine>(),
139        ingress,
140    ))
141}
142
143fn cpu_input_signature(
144    placement: &tenferro_tensor::Placement,
145    backend_family: Option<&'static str>,
146    input_domain: Option<tenferro_tensor::AllocationDomainId>,
147    allocation_domain: Option<tenferro_tensor::AllocationDomainId>,
148) -> bool {
149    if placement.memory_kind == MemoryKind::Managed {
150        return backend_family.is_some()
151            && allocation_domain.is_some()
152            && input_domain == allocation_domain;
153    }
154    cpu_input_placement(placement)
155        && match backend_family {
156            None => input_domain.is_none(),
157            Some(_) => allocation_domain.is_some() && input_domain == allocation_domain,
158        }
159}
160
161fn cpu_input_placement(placement: &tenferro_tensor::Placement) -> bool {
162    matches!(
163        placement.memory_kind,
164        MemoryKind::PinnedHost | MemoryKind::UnpinnedHost
165    )
166}
167
168fn cpu_runtime_input(
169    input: &TensorRead<'_>,
170    allocation_domain: Option<tenferro_tensor::AllocationDomainId>,
171) -> bool {
172    cpu_input_signature(
173        input.placement(),
174        input.backend_family(),
175        input.allocation_domain(),
176        allocation_domain,
177    )
178}
179
180fn runtime_storage_class() -> Result<StorageClass, RuntimeConfigError> {
181    StorageClass::new(CPU_STORAGE_CLASS_ID).map_err(RuntimeConfigError::from)
182}
183
184#[derive(Clone, Copy, Debug, Eq, PartialEq)]
185enum CpuPreparedKind {
186    Elementwise,
187    Reduction,
188    Indexing,
189    DotGeneral,
190    Layout,
191}
192
193#[derive(Debug)]
194struct CpuPreparedOperation {
195    binding: PreparedOperationBinding,
196    specialization: SpecializationProjection,
197    #[allow(dead_code, reason = "bounded Debug records the selected CPU family")]
198    kind: CpuPreparedKind,
199}
200
201impl PreparedOperation for CpuPreparedOperation {
202    fn binding(&self) -> &PreparedOperationBinding {
203        &self.binding
204    }
205
206    fn specialization(&self) -> &SpecializationProjection {
207        &self.specialization
208    }
209
210    fn retained_bytes(&self) -> usize {
211        checked_specialization_heap_retained_bytes(&self.specialization).unwrap_or(usize::MAX)
212    }
213}
214
215impl ElementwiseRuntime for CpuPreparer {
216    fn prepare(
217        &self,
218        request: ElementwisePrepareRequest<'_>,
219    ) -> Result<PrepareCapability, PrepareError> {
220        prepare_cpu(
221            request.operation(),
222            request.context(),
223            CpuPreparedKind::Elementwise,
224        )
225    }
226
227    fn max_fused_region_inputs(&self) -> Option<usize> {
228        Some(tenferro_cpu_fused::ERASED_FUSION_MAX_INPUTS)
229    }
230}
231
232impl ReductionRuntime for CpuPreparer {
233    fn prepare(
234        &self,
235        request: ReductionPrepareRequest<'_>,
236    ) -> Result<PrepareCapability, PrepareError> {
237        prepare_cpu(
238            request.operation(),
239            request.context(),
240            CpuPreparedKind::Reduction,
241        )
242    }
243}
244
245impl IndexingRuntime for CpuPreparer {
246    fn prepare(
247        &self,
248        request: IndexingPrepareRequest<'_>,
249    ) -> Result<PrepareCapability, PrepareError> {
250        prepare_cpu(
251            request.operation(),
252            request.context(),
253            CpuPreparedKind::Indexing,
254        )
255    }
256}
257
258impl DotGeneralPreparation for CpuPreparer {
259    fn prepare(
260        &self,
261        request: DotGeneralPrepareRequest<'_>,
262    ) -> Result<PrepareCapability, PrepareError> {
263        prepare_cpu(
264            request.operation(),
265            request.context(),
266            CpuPreparedKind::DotGeneral,
267        )
268    }
269}
270
271impl LayoutRuntime for CpuPreparer {
272    fn prepare(
273        &self,
274        request: LayoutPrepareRequest<'_>,
275    ) -> Result<PrepareCapability, PrepareError> {
276        prepare_cpu(
277            request.operation(),
278            request.context(),
279            CpuPreparedKind::Layout,
280        )
281    }
282
283    fn folds_metadata_views(&self) -> bool {
284        // Every CPU `*_read` entry reads a strided view in place (#2221).
285        true
286    }
287}
288
289impl tenferro_runtime::ExecutionEngine for CpuEngine {
290    fn target(&self) -> tenferro_runtime::EngineTarget {
291        tenferro_runtime::EngineTarget::new(
292            self.runtime_engine_id.clone(),
293            ExecutionContextIdentity::of::<CpuEngine>(),
294            self.allocation_domain(),
295        )
296    }
297
298    fn take_extension_caches(&mut self) -> tenferro_runtime::ExtensionCacheStore {
299        std::mem::take(&mut self.extension_caches)
300    }
301
302    fn restore_extension_caches(&mut self, caches: tenferro_runtime::ExtensionCacheStore) {
303        self.extension_caches = caches;
304    }
305
306    fn mark_broken(&mut self) {
307        self.broken = true;
308    }
309
310    fn as_any_mut(&mut self) -> &mut dyn std::any::Any {
311        self
312    }
313}
314
315fn prepare_cpu(
316    operation: SemanticOperationView<'_>,
317    context: &CorePrepareContext<'_>,
318    expected_kind: CpuPreparedKind,
319) -> Result<PrepareCapability, PrepareError> {
320    validate_cpu_runtime_context(context)?;
321    let SemanticOpRef::Core(op) = operation.op() else {
322        return Err(wrong_family_error(expected_kind, "extension"));
323    };
324    let Some(actual_kind) = cpu_operation_kind(op) else {
325        return Ok(PrepareCapability::Unsupported(
326            UnsupportedReason::Operation {
327                operation: UNKNOWN_CORE_OPERATION,
328            },
329        ));
330    };
331    if actual_kind != expected_kind {
332        return Err(wrong_family_error(expected_kind, core_operation_name(op)));
333    }
334
335    let minimum = minimum_specialization_requirements(actual_kind, context.inputs())?;
336    let merged =
337        merge_specialization_requirements(context.specialization().requirements(), &minimum);
338    if &merged != context.specialization().requirements() {
339        return Ok(PrepareCapability::NeedsSpecialization(merged));
340    }
341
342    Ok(PrepareCapability::Prepared(
343        PreparedOperationPlan::metadata(Arc::new(CpuPreparedOperation {
344            binding: context.binding().clone(),
345            specialization: context.specialization().clone(),
346            kind: actual_kind,
347        })),
348    ))
349}
350
351fn validate_cpu_runtime_context(context: &CorePrepareContext<'_>) -> Result<(), PrepareError> {
352    let expected_context = ExecutionContextIdentity::of::<CpuEngine>();
353    if context.binding().context_identity() != expected_context {
354        return Err(PrepareError::ProviderContract {
355            source: ProviderContractError::WrongOperationFamily {
356                expected: CoreCapabilityKind::Elementwise,
357                operation: "cpu-context-mismatch",
358            },
359        });
360    }
361    if context.binding().hardware_class().as_str() != CPU_HARDWARE_CLASS_ID {
362        return Err(PrepareError::ProviderContract {
363            source: ProviderContractError::WrongOperationFamily {
364                expected: CoreCapabilityKind::Elementwise,
365                operation: "cpu-hardware-mismatch",
366            },
367        });
368    }
369    if context.resolved_placement().storage_class().as_str() != CPU_STORAGE_CLASS_ID {
370        return Err(PrepareError::Unsupported {
371            reason: UnsupportedReason::StorageClass {
372                storage_class: context.resolved_placement().storage_class().clone(),
373            },
374        });
375    }
376    Ok(())
377}
378
379fn cpu_operation_kind(op: &CoreSemanticOp) -> Option<CpuPreparedKind> {
380    Some(match op {
381        CoreSemanticOp::Add
382        | CoreSemanticOp::Sub
383        | CoreSemanticOp::Mul
384        | CoreSemanticOp::Neg
385        | CoreSemanticOp::Conj
386        | CoreSemanticOp::Div
387        | CoreSemanticOp::Rem
388        | CoreSemanticOp::Abs
389        | CoreSemanticOp::Sign
390        | CoreSemanticOp::Maximum
391        | CoreSemanticOp::Minimum
392        | CoreSemanticOp::Compare(_)
393        | CoreSemanticOp::Select
394        | CoreSemanticOp::Clamp
395        | CoreSemanticOp::Exp
396        | CoreSemanticOp::Log
397        | CoreSemanticOp::Sin
398        | CoreSemanticOp::Cos
399        | CoreSemanticOp::Tanh
400        | CoreSemanticOp::Sqrt
401        | CoreSemanticOp::Rsqrt
402        | CoreSemanticOp::Pow
403        | CoreSemanticOp::Expm1
404        | CoreSemanticOp::Log1p
405        | CoreSemanticOp::Erf => CpuPreparedKind::Elementwise,
406        CoreSemanticOp::ReduceSum { .. }
407        | CoreSemanticOp::ReduceSumSquares { .. }
408        | CoreSemanticOp::ReduceProd { .. }
409        | CoreSemanticOp::ReduceMax { .. }
410        | CoreSemanticOp::ReduceMin { .. } => CpuPreparedKind::Reduction,
411        CoreSemanticOp::Gather(_)
412        | CoreSemanticOp::GatherDynamicSliceSizes { .. }
413        | CoreSemanticOp::Scatter(_)
414        | CoreSemanticOp::Slice(_)
415        | CoreSemanticOp::DynamicSlice { .. }
416        | CoreSemanticOp::DynamicUpdateSlice
417        | CoreSemanticOp::Pad(_)
418        | CoreSemanticOp::Concatenate { .. }
419        | CoreSemanticOp::Reverse { .. }
420        | CoreSemanticOp::ShapeOf { .. }
421        | CoreSemanticOp::DynamicTruncate { .. }
422        | CoreSemanticOp::PadToMatch { .. } => CpuPreparedKind::Indexing,
423        CoreSemanticOp::DotGeneral { .. } => CpuPreparedKind::DotGeneral,
424        CoreSemanticOp::Transpose { .. }
425        | CoreSemanticOp::Reshape { .. }
426        | CoreSemanticOp::BroadcastInDim { .. }
427        | CoreSemanticOp::Convert { .. }
428        | CoreSemanticOp::Constant { .. }
429        | CoreSemanticOp::ExtractDiag { .. }
430        | CoreSemanticOp::EmbedDiag { .. }
431        | CoreSemanticOp::Tril { .. }
432        | CoreSemanticOp::Triu { .. } => CpuPreparedKind::Layout,
433        _ => return None,
434    })
435}
436
437fn minimum_specialization_requirements(
438    kind: CpuPreparedKind,
439    inputs: &InputSignature,
440) -> Result<SpecializationRequirements, PrepareError> {
441    let mut requirements = Vec::with_capacity(inputs.entries().len());
442    for (input, entry) in inputs.entries().iter().enumerate() {
443        let mut builder = InputSpecializationRequirements::builder();
444        builder.dtype(true).rank(true);
445        match kind {
446            CpuPreparedKind::Indexing => {
447                builder.concrete_dimensions(concrete_axes_for_rank(input, entry.shape().len())?);
448            }
449            CpuPreparedKind::DotGeneral => {
450                builder
451                    .concrete_dimensions(concrete_axes_for_rank(input, entry.shape().len())?)
452                    .layout(LayoutSpecialization::Class);
453            }
454            CpuPreparedKind::Elementwise | CpuPreparedKind::Reduction | CpuPreparedKind::Layout => {
455            }
456        }
457        requirements.push(
458            builder
459                .build()
460                .expect("CPU minimum specialization requirements are internally valid"),
461        );
462    }
463    Ok(SpecializationRequirements::new(requirements))
464}
465
466fn concrete_axes_for_rank(input: usize, rank: usize) -> Result<Vec<u32>, PrepareError> {
467    if u32::try_from(rank).is_err() {
468        return Err(PrepareError::Specialization {
469            source: SpecializationError::ProjectionOverflow { input, rank },
470        });
471    }
472    Ok((0..rank)
473        .map(|axis| u32::try_from(axis).expect("rank precheck keeps axes encodable"))
474        .collect())
475}
476
477fn merge_specialization_requirements(
478    current: &SpecializationRequirements,
479    minimum: &SpecializationRequirements,
480) -> SpecializationRequirements {
481    debug_assert_eq!(current.inputs().len(), minimum.inputs().len());
482    let inputs = current
483        .inputs()
484        .iter()
485        .zip(minimum.inputs())
486        .map(|(current, minimum)| merge_input_requirements(current, minimum))
487        .collect::<Vec<_>>();
488    SpecializationRequirements::new(inputs)
489}
490
491fn merge_input_requirements(
492    current: &InputSpecializationRequirements,
493    minimum: &InputSpecializationRequirements,
494) -> InputSpecializationRequirements {
495    let mut axes = current.concrete_dimensions().to_vec();
496    for axis in minimum.concrete_dimensions() {
497        if !axes.contains(axis) {
498            axes.push(*axis);
499        }
500    }
501    let layout = current.layout().max(minimum.layout());
502    let rank = current.specializes_rank()
503        || minimum.specializes_rank()
504        || !axes.is_empty()
505        || layout == LayoutSpecialization::ExactStrides;
506    let alignment = match (current.alignment_log2(), minimum.alignment_log2()) {
507        (Some(left), Some(right)) => Some(left.max(right)),
508        (Some(value), None) | (None, Some(value)) => Some(value),
509        (None, None) => None,
510    };
511    let mut builder = InputSpecializationRequirements::builder();
512    builder
513        .dtype(current.specializes_dtype() || minimum.specializes_dtype())
514        .rank(rank)
515        .concrete_dimensions(axes)
516        .placement(current.placement().max(minimum.placement()))
517        .layout(layout)
518        .alignment_log2(alignment);
519    builder
520        .build()
521        .expect("merged CPU specialization requirements preserve builder invariants")
522}
523
524fn wrong_family_error(expected_kind: CpuPreparedKind, operation: &'static str) -> PrepareError {
525    PrepareError::ProviderContract {
526        source: ProviderContractError::WrongOperationFamily {
527            expected: expected_kind.core_capability(),
528            operation,
529        },
530    }
531}
532
533impl CpuPreparedKind {
534    fn core_capability(self) -> CoreCapabilityKind {
535        match self {
536            Self::Elementwise => CoreCapabilityKind::Elementwise,
537            Self::Reduction => CoreCapabilityKind::Reduction,
538            Self::Indexing => CoreCapabilityKind::Indexing,
539            Self::DotGeneral => CoreCapabilityKind::DotGeneral,
540            Self::Layout => CoreCapabilityKind::Layout,
541        }
542    }
543}
544
545fn core_operation_name(op: &CoreSemanticOp) -> &'static str {
546    match op {
547        CoreSemanticOp::Add => "add",
548        CoreSemanticOp::Sub => "sub",
549        CoreSemanticOp::Mul => "mul",
550        CoreSemanticOp::Neg => "neg",
551        CoreSemanticOp::Conj => "conj",
552        CoreSemanticOp::DotGeneral { .. } => "dot_general",
553        CoreSemanticOp::Transpose { .. } => "transpose",
554        CoreSemanticOp::Reshape { .. } => "reshape",
555        CoreSemanticOp::BroadcastInDim { .. } => "broadcast_in_dim",
556        CoreSemanticOp::Convert { .. } => "convert",
557        CoreSemanticOp::Constant { .. } => "constant",
558        CoreSemanticOp::ReduceSum { .. } => "reduce_sum",
559        CoreSemanticOp::ReduceSumSquares { .. } => "reduce_sum_squares",
560        CoreSemanticOp::Div => "div",
561        CoreSemanticOp::Rem => "rem",
562        CoreSemanticOp::Abs => "abs",
563        CoreSemanticOp::Sign => "sign",
564        CoreSemanticOp::Maximum => "maximum",
565        CoreSemanticOp::Minimum => "minimum",
566        CoreSemanticOp::Compare(_) => "compare",
567        CoreSemanticOp::Select => "select",
568        CoreSemanticOp::Clamp => "clamp",
569        CoreSemanticOp::Exp => "exp",
570        CoreSemanticOp::Log => "log",
571        CoreSemanticOp::Sin => "sin",
572        CoreSemanticOp::Cos => "cos",
573        CoreSemanticOp::Tanh => "tanh",
574        CoreSemanticOp::Sqrt => "sqrt",
575        CoreSemanticOp::Rsqrt => "rsqrt",
576        CoreSemanticOp::Pow => "pow",
577        CoreSemanticOp::Expm1 => "expm1",
578        CoreSemanticOp::Log1p => "log1p",
579        CoreSemanticOp::Erf => "erf",
580        CoreSemanticOp::ExtractDiag { .. } => "extract_diag",
581        CoreSemanticOp::EmbedDiag { .. } => "embed_diag",
582        CoreSemanticOp::Tril { .. } => "tril",
583        CoreSemanticOp::Triu { .. } => "triu",
584        CoreSemanticOp::Gather(_) => "gather",
585        CoreSemanticOp::GatherDynamicSliceSizes { .. } => "gather_dynamic_slice_sizes",
586        CoreSemanticOp::Scatter(_) => "scatter",
587        CoreSemanticOp::Slice(_) => "slice",
588        CoreSemanticOp::DynamicSlice { .. } => "dynamic_slice",
589        CoreSemanticOp::DynamicUpdateSlice => "dynamic_update_slice",
590        CoreSemanticOp::Pad(_) => "pad",
591        CoreSemanticOp::Concatenate { .. } => "concatenate",
592        CoreSemanticOp::Reverse { .. } => "reverse",
593        CoreSemanticOp::ShapeOf { .. } => "shape_of",
594        CoreSemanticOp::DynamicTruncate { .. } => "dynamic_truncate",
595        CoreSemanticOp::PadToMatch { .. } => "pad_to_match",
596        CoreSemanticOp::ReduceProd { .. } => "reduce_prod",
597        CoreSemanticOp::ReduceMax { .. } => "reduce_max",
598        CoreSemanticOp::ReduceMin { .. } => "reduce_min",
599        _ => UNKNOWN_CORE_OPERATION,
600    }
601}
602
603fn checked_specialization_heap_retained_bytes(
604    specialization: &SpecializationProjection,
605) -> Option<usize> {
606    let requirements = specialization.requirements();
607    checked_sum([
608        requirements
609            .inputs()
610            .len()
611            .checked_mul(size_of::<InputSpecializationRequirements>())?,
612        checked_sum(
613            requirements
614                .inputs()
615                .iter()
616                .map(|input| size_of_val(input.concrete_dimensions())),
617        )?,
618        specialization
619            .inputs()
620            .len()
621            .checked_mul(size_of::<InputSpecializationProjection>())?,
622        checked_sum_options(
623            specialization
624                .inputs()
625                .iter()
626                .map(input_projection_retained_bytes),
627        )?,
628    ])
629}
630
631fn input_projection_retained_bytes(projection: &InputSpecializationProjection) -> Option<usize> {
632    size_of_val(projection.concrete_dimensions()).checked_add(match projection.layout() {
633        Some(LayoutProjection::ExactStrides(strides)) if strides.is_heap() => {
634            size_of_val(strides.as_slice())
635        }
636        _ => 0,
637    })
638}
639
640fn checked_sum(values: impl IntoIterator<Item = usize>) -> Option<usize> {
641    values
642        .into_iter()
643        .try_fold(0usize, |sum, value| sum.checked_add(value))
644}
645
646fn checked_sum_options(values: impl IntoIterator<Item = Option<usize>>) -> Option<usize> {
647    values
648        .into_iter()
649        .try_fold(0usize, |sum, value| sum.checked_add(value?))
650}
651
652impl fmt::Display for CpuPreparedKind {
653    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
654        formatter.write_str(match self {
655            Self::Elementwise => "elementwise",
656            Self::Reduction => "reduction",
657            Self::Indexing => "indexing",
658            Self::DotGeneral => "dot_general",
659            Self::Layout => "layout",
660        })
661    }
662}
663
664#[cfg(test)]
665mod tests;
tenferro_cpu/runtime_adapter.rs

tenferro_cpu/
runtime_adapter.rs
