tenferro_gpu/cuda/
runtime_adapter.rs

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

tenferro_gpu/cuda/
runtime_adapter.rs
