1use std::collections::BTreeMap;
2use std::sync::atomic::{AtomicU64, Ordering};
3use std::sync::Arc;
4
5use super::super::{
6 BatchWorkShape, ClaimedSubmissionWaveBacking, ContractVersion, DeviceDescriptor,
7 DeviceReusableAddressScope, DeviceReusableExecutionTopologyFingerprint, DeviceRuntime,
8 EncodedDeviceOperation, EncodedReusableExecutionBindings, ExecutablePlanView,
9 LogicalBackingSliceAuthority, MemoryPlan, NodeId, OperationId, PlanHash, PlanId, ProviderId,
10 ProviderWorkspaceRequirement, SemanticValue, VNextError,
11};
12use super::foundation::{canonical_sha256, invalid_operation};
13use super::invocation::PreparedOperationDispatchBinding;
14use super::resolved_value::resource_uses_packed_batch_coordinates;
15use super::{
16 AttributeId, BatchedOperationInvocation, CapabilityCatalog, EngineProviderDescriptor,
17 OperationContract, OperationDescriptor, OperationFailure, OperationProviderDescriptor,
18 ResolvedValueBinding, ResolvedValueRole,
19};
20
21pub struct OperationResourceEstimateRequest<'a> {
27 node_id: &'a NodeId,
28 operation: &'a OperationDescriptor,
29 values: &'a [ResolvedValueBinding],
30 attributes: &'a BTreeMap<AttributeId, SemanticValue>,
31 input_fingerprint: &'a str,
32}
33
34pub struct ReusableExecutionTopologyRequest<'a> {
42 node_id: &'a NodeId,
43 operation_id: &'a OperationId,
44 attributes: &'a BTreeMap<AttributeId, SemanticValue>,
45 bindings: &'a [ResolvedValueBinding],
46 scratch_resource: Option<&'a super::super::ResourceId>,
47 binding_resource: Option<&'a super::super::ResourceId>,
48 persistent_resource: Option<&'a super::super::ResourceId>,
49 memory: &'a MemoryPlan,
50 work_shape: &'a BatchWorkShape,
51 claimed_backing: &'a ClaimedSubmissionWaveBacking,
52 step_backing: &'a [LogicalBackingSliceAuthority],
53}
54
55#[derive(Debug, Clone, Copy, PartialEq, Eq)]
60pub enum ReusableExecutionValueAddress {
61 Captured {
62 role: ResolvedValueRole,
63 ordinal: u32,
64 },
65 ProgramBinding {
66 role: ResolvedValueRole,
67 ordinal: u32,
68 },
69}
70
71impl ReusableExecutionValueAddress {
72 pub const fn captured(role: ResolvedValueRole, ordinal: u32) -> Self {
73 Self::Captured { role, ordinal }
74 }
75
76 pub const fn program_binding(role: ResolvedValueRole, ordinal: u32) -> Self {
77 Self::ProgramBinding { role, ordinal }
78 }
79
80 const fn identity(self) -> (ResolvedValueRole, u32) {
81 match self {
82 Self::Captured { role, ordinal } | Self::ProgramBinding { role, ordinal } => {
83 (role, ordinal)
84 }
85 }
86 }
87}
88
89#[derive(Debug, Clone, Copy, PartialEq, Eq)]
91pub enum ReusableExecutionWorkspaceAddress {
92 Scratch,
93 Binding,
94 Persistent,
95}
96
97impl<'a> ReusableExecutionTopologyRequest<'a> {
98 pub(super) fn new(
99 node_id: &'a NodeId,
100 operation_id: &'a OperationId,
101 attributes: &'a BTreeMap<AttributeId, SemanticValue>,
102 bindings: &'a [ResolvedValueBinding],
103 scratch_resource: Option<&'a super::super::ResourceId>,
104 binding_resource: Option<&'a super::super::ResourceId>,
105 persistent_resource: Option<&'a super::super::ResourceId>,
106 memory: &'a MemoryPlan,
107 work_shape: &'a BatchWorkShape,
108 claimed_backing: &'a ClaimedSubmissionWaveBacking,
109 step_backing: &'a [LogicalBackingSliceAuthority],
110 ) -> Result<Self, VNextError> {
111 if work_shape.participants().is_empty() {
112 return Err(invalid_operation(
113 "reusable execution topology request has no participants",
114 ));
115 }
116 Ok(Self {
117 node_id,
118 operation_id,
119 attributes,
120 bindings,
121 scratch_resource,
122 binding_resource,
123 persistent_resource,
124 memory,
125 work_shape,
126 claimed_backing,
127 step_backing,
128 })
129 }
130
131 pub fn node_id(&self) -> &NodeId {
132 self.node_id
133 }
134
135 pub fn operation_id(&self) -> &OperationId {
136 self.operation_id
137 }
138
139 pub fn attributes(&self) -> &BTreeMap<AttributeId, SemanticValue> {
140 self.attributes
141 }
142
143 pub fn bindings(&self) -> &[ResolvedValueBinding] {
144 self.bindings
145 }
146
147 pub fn work_shape(&self) -> &BatchWorkShape {
148 self.work_shape
149 }
150
151 pub fn binding_uses_packed_batch_coordinates(
156 &self,
157 role: ResolvedValueRole,
158 ordinal: u32,
159 ) -> Result<bool, VNextError> {
160 let binding = self
161 .bindings
162 .iter()
163 .find(|binding| binding.role() == role && binding.ordinal() == ordinal)
164 .ok_or_else(|| {
165 invalid_operation("reusable topology requested an unknown value binding")
166 })?;
167 let [component] = binding.storage().components() else {
168 return Err(invalid_operation(
169 "reusable topology coordinate ownership requires one resource component",
170 ));
171 };
172 resource_uses_packed_batch_coordinates(self.memory, component.resource_id())
173 }
174
175 pub fn reusable_address_scope(
180 &self,
181 values: &[ReusableExecutionValueAddress],
182 workspaces: &[ReusableExecutionWorkspaceAddress],
183 ) -> Result<Option<DeviceReusableAddressScope>, VNextError> {
184 if values.len() != self.bindings.len()
185 || values.iter().enumerate().any(|(index, value)| {
186 values[..index]
187 .iter()
188 .any(|prior| prior.identity() == value.identity())
189 })
190 || self.bindings.iter().any(|binding| {
191 values
192 .iter()
193 .filter(|value| value.identity() == (binding.role(), binding.ordinal()))
194 .count()
195 != 1
196 })
197 || workspaces.iter().enumerate().any(|(index, workspace)| {
198 workspaces[..index].iter().any(|prior| prior == workspace)
199 })
200 {
201 return Err(invalid_operation(
202 "reusable topology address contract does not cover every value exactly once",
203 ));
204 }
205
206 let has_program_bound_values = values
207 .iter()
208 .any(|value| matches!(value, ReusableExecutionValueAddress::ProgramBinding { .. }));
209 if has_program_bound_values
210 && !workspaces.contains(&ReusableExecutionWorkspaceAddress::Binding)
211 {
212 return Err(invalid_operation(
213 "program-bound reusable values require a captured binding workspace",
214 ));
215 }
216
217 let mut aggregate = DeviceReusableAddressScope::Plan;
218 for value in values {
219 let ReusableExecutionValueAddress::Captured { role, ordinal } = value else {
220 continue;
221 };
222 let Some(scope) = self.binding_reusable_address_scope(*role, *ordinal)? else {
223 return Ok(None);
224 };
225 aggregate = merge_reusable_address_scope(aggregate, scope)?;
226 }
227 for workspace in workspaces {
228 let scope = match workspace {
229 ReusableExecutionWorkspaceAddress::Scratch => {
230 self.scratch_reusable_address_scope()?
231 }
232 ReusableExecutionWorkspaceAddress::Binding => {
233 self.binding_workspace_reusable_address_scope()?
234 }
235 ReusableExecutionWorkspaceAddress::Persistent => {
236 self.persistent_workspace_reusable_address_scope()?
237 }
238 };
239 let Some(scope) = scope else {
240 return Ok(None);
241 };
242 aggregate = merge_reusable_address_scope(aggregate, scope)?;
243 }
244 Ok(Some(aggregate))
245 }
246
247 pub fn binding_reusable_address_scope(
252 &self,
253 role: ResolvedValueRole,
254 ordinal: u32,
255 ) -> Result<Option<DeviceReusableAddressScope>, VNextError> {
256 let binding = self
257 .bindings
258 .iter()
259 .find(|binding| binding.role() == role && binding.ordinal() == ordinal)
260 .ok_or_else(|| {
261 invalid_operation("reusable topology requested an unknown value binding")
262 })?;
263 let mut aggregate = DeviceReusableAddressScope::Plan;
264 for component in binding.storage().components() {
265 let Some(component_scope) =
266 self.resource_reusable_address_scope(component.resource_id())?
267 else {
268 return Ok(None);
269 };
270 aggregate = merge_reusable_address_scope(aggregate, component_scope)?;
271 }
272 Ok(Some(aggregate))
273 }
274
275 pub fn scratch_reusable_address_scope(
278 &self,
279 ) -> Result<Option<DeviceReusableAddressScope>, VNextError> {
280 self.workspace_reusable_address_scope(self.scratch_resource, "scratch")
281 }
282
283 pub fn binding_workspace_reusable_address_scope(
286 &self,
287 ) -> Result<Option<DeviceReusableAddressScope>, VNextError> {
288 self.workspace_reusable_address_scope(self.binding_resource, "binding")
289 }
290
291 pub fn persistent_workspace_reusable_address_scope(
294 &self,
295 ) -> Result<Option<DeviceReusableAddressScope>, VNextError> {
296 self.workspace_reusable_address_scope(self.persistent_resource, "persistent")
297 }
298
299 fn workspace_reusable_address_scope(
300 &self,
301 resource_id: Option<&super::super::ResourceId>,
302 workspace: &str,
303 ) -> Result<Option<DeviceReusableAddressScope>, VNextError> {
304 let resource_id = resource_id.ok_or_else(|| {
305 invalid_operation(format!(
306 "reusable topology requested absent provider {workspace} workspace"
307 ))
308 })?;
309 self.resource_reusable_address_scope(resource_id)
310 }
311
312 fn resource_reusable_address_scope(
313 &self,
314 resource_id: &super::super::ResourceId,
315 ) -> Result<Option<DeviceReusableAddressScope>, VNextError> {
316 if self
317 .memory
318 .static_allocations()
319 .binary_search_by(|allocation| allocation.resource_id().cmp(resource_id))
320 .is_ok()
321 {
322 return Ok(Some(DeviceReusableAddressScope::Plan));
323 }
324
325 let mut resource_scope = None;
326 for backing_slices in [self.claimed_backing.backing_slices(), self.step_backing] {
327 let authority_start =
328 backing_slices.partition_point(|authority| authority.resource_id() < resource_id);
329 let authority_end = authority_start
330 + backing_slices[authority_start..]
331 .partition_point(|authority| authority.resource_id() == resource_id);
332 for authority in &backing_slices[authority_start..authority_end] {
333 let Some(authority_scope) = authority.reusable_address_scope() else {
334 return Ok(None);
335 };
336 resource_scope = Some(merge_reusable_address_scope(
337 resource_scope.unwrap_or(DeviceReusableAddressScope::Plan),
338 authority_scope,
339 )?);
340 }
341 }
342 if resource_scope.is_some() {
343 return Ok(resource_scope);
344 }
345 if self
346 .memory
347 .dynamic_descriptors()
348 .binary_search_by(|descriptor| descriptor.base_resource_id().cmp(resource_id))
349 .is_ok()
350 {
351 return Ok(None);
352 }
353 Err(invalid_operation(
354 "reusable topology references an unknown memory resource",
355 ))
356 }
357}
358
359fn merge_reusable_address_scope(
360 left: DeviceReusableAddressScope,
361 right: DeviceReusableAddressScope,
362) -> Result<DeviceReusableAddressScope, VNextError> {
363 match (left, right) {
364 (DeviceReusableAddressScope::Plan, scope) | (scope, DeviceReusableAddressScope::Plan) => {
365 Ok(scope)
366 }
367 (
368 DeviceReusableAddressScope::ExecutionLane(left),
369 DeviceReusableAddressScope::ExecutionLane(right),
370 ) if left == right => Ok(DeviceReusableAddressScope::ExecutionLane(left)),
371 _ => Err(invalid_operation(
372 "reusable topology value spans different execution lanes",
373 )),
374 }
375}
376
377impl<'a> OperationResourceEstimateRequest<'a> {
378 pub(crate) fn new(
379 node_id: &'a NodeId,
380 operation: &'a OperationDescriptor,
381 values: &'a [ResolvedValueBinding],
382 attributes: &'a BTreeMap<AttributeId, SemanticValue>,
383 input_fingerprint: &'a str,
384 ) -> Result<Self, VNextError> {
385 operation.validate()?;
386 operation.validate_attributes(attributes)?;
387 operation.validate_resolved_bindings(values)?;
388 if !canonical_sha256(input_fingerprint) {
389 return Err(invalid_operation(
390 "resource estimator request has invalid input fingerprint",
391 ));
392 }
393 Ok(Self {
394 node_id,
395 operation,
396 values,
397 attributes,
398 input_fingerprint,
399 })
400 }
401
402 pub fn node_id(&self) -> &NodeId {
403 self.node_id
404 }
405
406 pub fn operation(&self) -> &OperationDescriptor {
407 self.operation
408 }
409
410 pub fn values(&self) -> &[ResolvedValueBinding] {
411 self.values
412 }
413
414 pub fn attributes(&self) -> &BTreeMap<AttributeId, SemanticValue> {
415 self.attributes
416 }
417
418 pub fn input_fingerprint(&self) -> &str {
419 self.input_fingerprint
420 }
421}
422
423#[derive(Debug, Clone, PartialEq, Eq)]
427pub struct OperationResourceEstimate {
428 estimator_id: String,
429 estimator_version: ContractVersion,
430 estimator_implementation_fingerprint: String,
431 claimed_input_fingerprint: String,
432 value_alignment_bytes: u64,
433 scratch: Option<ProviderWorkspaceRequirement>,
434 binding: Option<ProviderWorkspaceRequirement>,
435 persistent: Option<ProviderWorkspaceRequirement>,
436}
437
438impl OperationResourceEstimate {
439 #[allow(clippy::too_many_arguments)]
440 pub fn new(
441 estimator_id: impl Into<String>,
442 estimator_version: ContractVersion,
443 estimator_implementation_fingerprint: impl Into<String>,
444 claimed_input_fingerprint: impl Into<String>,
445 value_alignment_bytes: u64,
446 scratch: Option<ProviderWorkspaceRequirement>,
447 persistent: Option<ProviderWorkspaceRequirement>,
448 ) -> Self {
449 Self {
450 estimator_id: estimator_id.into(),
451 estimator_version,
452 estimator_implementation_fingerprint: estimator_implementation_fingerprint.into(),
453 claimed_input_fingerprint: claimed_input_fingerprint.into(),
454 value_alignment_bytes,
455 scratch,
456 binding: None,
457 persistent,
458 }
459 }
460
461 pub fn with_binding(mut self, binding: ProviderWorkspaceRequirement) -> Self {
462 self.binding = Some(binding);
463 self
464 }
465
466 pub fn estimator_id(&self) -> &str {
467 &self.estimator_id
468 }
469
470 pub const fn estimator_version(&self) -> ContractVersion {
471 self.estimator_version
472 }
473
474 pub fn estimator_implementation_fingerprint(&self) -> &str {
475 &self.estimator_implementation_fingerprint
476 }
477
478 pub fn claimed_input_fingerprint(&self) -> &str {
479 &self.claimed_input_fingerprint
480 }
481
482 pub const fn value_alignment_bytes(&self) -> u64 {
483 self.value_alignment_bytes
484 }
485
486 pub fn scratch(&self) -> Option<&ProviderWorkspaceRequirement> {
487 self.scratch.as_ref()
488 }
489
490 pub fn binding(&self) -> Option<&ProviderWorkspaceRequirement> {
491 self.binding.as_ref()
492 }
493
494 pub fn persistent(&self) -> Option<&ProviderWorkspaceRequirement> {
495 self.persistent.as_ref()
496 }
497}
498
499pub trait OperationResourceEstimator: Send + Sync {
503 fn descriptor(&self) -> &OperationProviderDescriptor;
504
505 fn estimate_resources(
506 &self,
507 request: OperationResourceEstimateRequest<'_>,
508 ) -> Result<OperationResourceEstimate, VNextError>;
509}
510
511pub trait OperationPlanningRegistry: Send + Sync {
515 fn contracts_for(&self, operation_id: &OperationId) -> Vec<&dyn OperationContract>;
516
517 fn estimators_for(&self, provider_id: &ProviderId) -> Vec<&dyn OperationResourceEstimator>;
518}
519
520#[derive(Debug, Clone, PartialEq, Eq)]
524pub(crate) struct OperationRegistryAuthority(u64);
525
526impl OperationRegistryAuthority {
527 fn mint() -> Result<Self, VNextError> {
528 static NEXT_AUTHORITY: AtomicU64 = AtomicU64::new(1);
529 let id = NEXT_AUTHORITY
530 .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |current| {
531 current.checked_add(1)
532 })
533 .map_err(|_| invalid_operation("operation registry authority space exhausted"))?;
534 Ok(Self(id))
535 }
536}
537
538pub struct OperationPlanningHandle<'registry> {
542 registry: &'registry dyn OperationPlanningRegistry,
543 authority: OperationRegistryAuthority,
544}
545
546impl OperationPlanningHandle<'_> {
547 pub(crate) fn authority(&self) -> &OperationRegistryAuthority {
548 &self.authority
549 }
550}
551
552impl OperationPlanningRegistry for OperationPlanningHandle<'_> {
553 fn contracts_for(&self, operation_id: &OperationId) -> Vec<&dyn OperationContract> {
554 self.registry.contracts_for(operation_id)
555 }
556
557 fn estimators_for(&self, provider_id: &ProviderId) -> Vec<&dyn OperationResourceEstimator> {
558 self.registry.estimators_for(provider_id)
559 }
560}
561
562#[derive(Debug, Clone, Copy, PartialEq, Eq)]
571pub enum ReusableExecutionTopology {
572 Static,
573 Dynamic(DeviceReusableExecutionTopologyFingerprint),
574 EagerBoundary,
575}
576
577pub trait OperationProvider<R: DeviceRuntime>: OperationResourceEstimator {
580 fn reusable_execution_topology(
592 &self,
593 request: ReusableExecutionTopologyRequest<'_>,
594 ) -> Result<ReusableExecutionTopology, VNextError>;
595
596 fn encode_selected(
597 &self,
598 invocation: BatchedOperationInvocation<'_, R::Buffer>,
599 ) -> Result<EncodedDeviceOperation<R::Command>, OperationFailure>;
600
601 fn encode_reusable_execution_bindings(
608 &self,
609 invocation: BatchedOperationInvocation<'_, R::Buffer>,
610 ) -> Result<EncodedReusableExecutionBindings<R::Command>, OperationFailure> {
611 self.encode_selected(invocation)
612 .map(EncodedReusableExecutionBindings::from_operation)
613 }
614}
615
616pub struct OperationRuntimeRegistry<R>
620where
621 R: DeviceRuntime,
622{
623 authority: OperationRegistryAuthority,
624 contracts: BTreeMap<OperationId, Box<dyn OperationContract>>,
625 providers: BTreeMap<ProviderId, Arc<dyn OperationProvider<R>>>,
626}
627
628impl<R> OperationRuntimeRegistry<R>
629where
630 R: DeviceRuntime,
631{
632 pub fn new(
633 contracts: Vec<Box<dyn OperationContract>>,
634 providers: Vec<Box<dyn OperationProvider<R>>>,
635 ) -> Result<Self, VNextError> {
636 if contracts.is_empty() || providers.is_empty() {
637 return Err(invalid_operation(
638 "operation runtime registry requires contracts and providers",
639 ));
640 }
641 let mut contract_map = BTreeMap::new();
642 for contract in contracts {
643 let descriptor = contract.descriptor();
644 descriptor.validate()?;
645 let operation_id = descriptor.id.clone();
646 if contract_map
647 .insert(operation_id.clone(), contract)
648 .is_some()
649 {
650 return Err(invalid_operation(format!(
651 "operation runtime registry has duplicate contract `{operation_id}`"
652 )));
653 }
654 }
655 let mut provider_map: BTreeMap<ProviderId, Arc<dyn OperationProvider<R>>> = BTreeMap::new();
656 for provider in providers {
657 let descriptor = provider.descriptor();
658 let contract = contract_map.get(descriptor.operation_id()).ok_or_else(|| {
659 invalid_operation(format!(
660 "runtime provider `{}` has no registered operation contract",
661 descriptor.provider_id()
662 ))
663 })?;
664 if descriptor.operation_fingerprint() != contract.descriptor().fingerprint()? {
665 return Err(invalid_operation(format!(
666 "runtime provider `{}` differs from its registered operation contract",
667 descriptor.provider_id()
668 )));
669 }
670 let provider_id = descriptor.provider_id().clone();
671 if provider_map
672 .insert(provider_id.clone(), Arc::from(provider))
673 .is_some()
674 {
675 return Err(invalid_operation(format!(
676 "operation runtime registry has duplicate or byte-identical provider `{provider_id}`"
677 )));
678 }
679 }
680 Ok(Self {
681 authority: OperationRegistryAuthority::mint()?,
682 contracts: contract_map,
683 providers: provider_map,
684 })
685 }
686
687 pub fn capability_catalog(
691 &self,
692 device: DeviceDescriptor,
693 engine_providers: Vec<EngineProviderDescriptor>,
694 ) -> Result<CapabilityCatalog, VNextError> {
695 let operations = self
696 .contracts
697 .values()
698 .map(|contract| contract.descriptor().clone())
699 .collect::<Vec<_>>();
700 let mut providers = self
701 .contracts
702 .keys()
703 .cloned()
704 .map(|operation_id| (operation_id, Vec::new()))
705 .collect::<BTreeMap<_, _>>();
706 for provider in self.providers.values() {
707 providers
708 .get_mut(provider.descriptor().operation_id())
709 .ok_or_else(|| {
710 invalid_operation(
711 "runtime provider operation is absent while deriving its catalog",
712 )
713 })?
714 .push(provider.descriptor().clone());
715 }
716 CapabilityCatalog::new(device, operations, providers, engine_providers)
717 }
718
719 pub fn planning(&self) -> OperationPlanningHandle<'_> {
720 OperationPlanningHandle {
721 registry: self,
722 authority: self.authority.clone(),
723 }
724 }
725
726 pub fn bind<'registry>(
727 &'registry self,
728 resolved: &dyn ExecutablePlanView,
729 node_id: &NodeId,
730 ) -> Result<BoundOperationProvider<'registry, R>, VNextError> {
731 let provider = self.selected_provider(resolved, node_id)?;
732 let plan = resolved.execution_plan();
733 let dispatch =
734 PreparedOperationDispatchBinding::prepare(resolved, provider.descriptor(), node_id)?;
735 Ok(BoundOperationProvider {
736 provider: BoundOperationProviderSource::Borrowed(provider.as_ref()),
737 plan_id: plan.payload().plan_id().clone(),
738 plan_hash: plan.plan_hash().clone(),
739 node_id: node_id.clone(),
740 dispatch,
741 })
742 }
743
744 pub fn bind_plan(
750 &self,
751 resolved: &dyn ExecutablePlanView,
752 ) -> Result<BoundOperationProviderSet<R>, VNextError> {
753 let providers = resolved
754 .execution_plan()
755 .payload()
756 .nodes()
757 .iter()
758 .map(|node| {
759 let provider = self.selected_provider(resolved, node.id())?;
760 let plan = resolved.execution_plan();
761 let dispatch = PreparedOperationDispatchBinding::prepare(
762 resolved,
763 provider.descriptor(),
764 node.id(),
765 )?;
766 Ok(BoundOperationProvider {
767 provider: BoundOperationProviderSource::Owned(Arc::clone(provider)),
768 plan_id: plan.payload().plan_id().clone(),
769 plan_hash: plan.plan_hash().clone(),
770 node_id: node.id().clone(),
771 dispatch,
772 })
773 })
774 .collect::<Result<Vec<BoundOperationProvider<'static, R>>, _>>()?;
775 if providers.is_empty() {
776 return Err(invalid_operation(
777 "executable plan cannot bind an empty provider set",
778 ));
779 }
780 Ok(BoundOperationProviderSet { providers })
781 }
782
783 fn selected_provider(
784 &self,
785 resolved: &dyn ExecutablePlanView,
786 node_id: &NodeId,
787 ) -> Result<&Arc<dyn OperationProvider<R>>, VNextError> {
788 let plan = resolved.execution_plan();
789 if plan.operation_registry_authority() != &self.authority {
790 return Err(invalid_operation(
791 "resolved plan belongs to a different operation runtime registry",
792 ));
793 }
794 let node = plan
795 .payload()
796 .nodes()
797 .iter()
798 .find(|node| node.id() == node_id)
799 .ok_or_else(|| invalid_operation(format!("plan has no node `{node_id}`")))?;
800 let provider = self
801 .providers
802 .get(node.selection().selected_provider())
803 .ok_or_else(|| {
804 invalid_operation(format!(
805 "runtime registry has no selected provider `{}`",
806 node.selection().selected_provider()
807 ))
808 })?;
809 let catalog_provider = resolved
810 .capabilities()
811 .providers_for(node.operation_id())?
812 .iter()
813 .find(|candidate| candidate.provider_id() == provider.descriptor().provider_id())
814 .ok_or_else(|| invalid_operation("runtime provider is absent from resolved catalog"))?;
815 if provider.descriptor() != catalog_provider
816 || provider.descriptor().provider_id() != node.selection().selected_provider()
817 || provider.descriptor().provider_implementation_fingerprint()
818 != node.provider_implementation_fingerprint()
819 {
820 return Err(invalid_operation(
821 "runtime provider is not the exact registry object selected by the resolved plan",
822 ));
823 }
824 Ok(provider)
825 }
826}
827
828impl<R> OperationPlanningRegistry for OperationRuntimeRegistry<R>
829where
830 R: DeviceRuntime,
831{
832 fn contracts_for(&self, operation_id: &OperationId) -> Vec<&dyn OperationContract> {
833 self.contracts
834 .get(operation_id)
835 .map(|contract| vec![contract.as_ref()])
836 .unwrap_or_default()
837 }
838
839 fn estimators_for(&self, provider_id: &ProviderId) -> Vec<&dyn OperationResourceEstimator> {
840 self.providers
841 .get(provider_id)
842 .map(|provider| vec![provider.as_ref() as &dyn OperationResourceEstimator])
843 .unwrap_or_default()
844 }
845}
846
847enum BoundOperationProviderSource<'registry, R>
848where
849 R: DeviceRuntime,
850{
851 Borrowed(&'registry dyn OperationProvider<R>),
852 Owned(Arc<dyn OperationProvider<R>>),
853}
854
855impl<R> BoundOperationProviderSource<'_, R>
856where
857 R: DeviceRuntime,
858{
859 fn provider(&self) -> &dyn OperationProvider<R> {
860 match self {
861 Self::Borrowed(provider) => *provider,
862 Self::Owned(provider) => provider.as_ref(),
863 }
864 }
865}
866
867pub struct BoundOperationProvider<'registry, R>
871where
872 R: DeviceRuntime,
873{
874 provider: BoundOperationProviderSource<'registry, R>,
875 plan_id: PlanId,
876 plan_hash: PlanHash,
877 node_id: NodeId,
878 dispatch: PreparedOperationDispatchBinding,
879}
880
881impl<R> BoundOperationProvider<'_, R>
882where
883 R: DeviceRuntime,
884{
885 pub(super) fn provider(&self) -> &dyn OperationProvider<R> {
886 self.provider.provider()
887 }
888
889 pub(super) fn validate_binding(
890 &self,
891 resolved: &dyn ExecutablePlanView,
892 node_id: &NodeId,
893 ) -> Result<(), VNextError> {
894 let plan = resolved.execution_plan();
895 if self.plan_id != *plan.payload().plan_id()
896 || self.plan_hash != *plan.plan_hash()
897 || &self.node_id != node_id
898 {
899 return Err(invalid_operation(
900 "bound operation provider belongs to a different plan or node",
901 ));
902 }
903 self.dispatch.node(resolved, node_id)?;
904 Ok(())
905 }
906
907 pub(super) fn matches_plan_node(
908 &self,
909 plan_id: &PlanId,
910 plan_hash: &PlanHash,
911 node_id: &NodeId,
912 ) -> bool {
913 &self.plan_id == plan_id && &self.plan_hash == plan_hash && &self.node_id == node_id
914 }
915
916 pub(super) fn dispatch(&self) -> &PreparedOperationDispatchBinding {
917 &self.dispatch
918 }
919
920 pub fn descriptor(&self) -> &OperationProviderDescriptor {
921 self.provider().descriptor()
922 }
923}
924
925pub struct BoundOperationProviderSet<R>
931where
932 R: DeviceRuntime,
933{
934 providers: Vec<BoundOperationProvider<'static, R>>,
935}
936
937impl<R> BoundOperationProviderSet<R>
938where
939 R: DeviceRuntime,
940{
941 pub fn providers(&self) -> &[BoundOperationProvider<'static, R>] {
942 &self.providers
943 }
944
945 pub fn len(&self) -> usize {
946 self.providers.len()
947 }
948
949 pub fn is_empty(&self) -> bool {
950 self.providers.is_empty()
951 }
952}