1use super::{
2 canonical_fingerprint, canonical_runtime_policy_fingerprint, invalid_plan, is_canonical_sha256,
3 joint_candidate_components, joint_partial_precedes, node_weight_requirements,
4 provider_resource_estimator_input_fingerprint, static_contiguous_storage_profile,
5 storage_incompatible_resource_ids, tensor_storage_layout_fingerprint,
6 validate_active_sequence_ceiling, validate_program_bindings, validate_scheduled_token_ceiling,
7 validate_semantic_binding, workspace_base_id, workspace_storage_layout_fingerprint,
8 AliasPolicy, AllocationKind, AllocationLifetime, BTreeMap, BTreeSet, BufferUsage,
9 CanonicalValueBinding, CapabilityCatalog, CapabilityId, DimensionConstraint,
10 DynamicResourceDemand, DynamicResourceDescriptor, DynamicStorageContract,
11 DynamicStorageProfile, DynamicStorageRequirement, ElementType, ExecutionPlanPayload,
12 ExecutionWeightPlan, GlobalValueRange, JointComponentSolution, JointPartialSelection,
13 JointProviderCandidate, JointProviderStorageSelection, JointSelectionObjective, MemoryPlan,
14 NodeId, NodeTokenBindingProjection, NodeWorkContract, OperationDescriptor,
15 OperationRegistryAuthority, PlanBuildRequest, PlanExactAlias, PlanExactAliasKind, PlanHash,
16 PlanHashMaterial, PlanId, PlanNode, PlanNodeResolution, PlanProviderRejectReason,
17 PlanStateEffect, PreparedModelFamily, ProgramNode, ProgramNodeWorkSpec, ProgramValueId,
18 ProviderCompatibilityRequest, ProviderId, ProviderResourcePlan, ProviderSelection,
19 ProviderSelectionReason, ProviderWorkspaceScope, QuantizationFormatId, RejectedProvider,
20 ResolvedValueBinding, ResolvedValueRole, ResourceAllocation, ResourceId,
21 ReusableExecutionMemoryPlan, ReusableExecutionPolicy, RuntimePolicy, Serialize,
22 StateCapacityDemand, StateDependencyTracker, StateInitialization, StateLifetime,
23 StaticWeightTransformPlan, TensorAccess, TrustedExecutionWeightPlan, VNextError,
24 ValueAllocationAccumulator, ValueResourceDemand, WeightFormatId, WeightSchema,
25 EXECUTION_PLAN_SCHEMA, STATIC_WEIGHT_TRANSFORM_SCRATCH_ALIGNMENT_BYTES,
26};
27use super::{resolve_retained_completion_values, CompletionRetentionSpec, RetainedCompletionValue};
28use crate::vnext::{
29 CompletionReadbackRequest, ExecutionDeterminismRequirement, HostTransferLayout,
30 ResourceWorkShape, WeightComponentPayload, WeightComponentSource, WeightComponentSpec,
31};
32
33#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
34pub struct ExecutionPlan {
35 pub(super) payload: ExecutionPlanPayload,
36 pub(super) plan_hash: PlanHash,
37 #[serde(skip)]
38 pub(super) operation_registry_authority: OperationRegistryAuthority,
39 #[serde(skip)]
40 pub(super) trusted_execution_weights: TrustedExecutionWeightPlan,
41 #[serde(skip)]
42 pub(super) checkpoint_unsupported_reasons: Vec<super::SequenceCheckpointUnsupportedReason>,
43}
44
45impl ExecutionPlan {
46 pub fn build<P: RuntimePolicy>(request: PlanBuildRequest<'_, P>) -> Result<Self, VNextError> {
47 request.policy.validate()?;
48 let maximum_active_sequences = request.policy.maximum_active_sequences();
49 validate_active_sequence_ceiling(maximum_active_sequences)?;
50 let maximum_scheduled_tokens = request.policy.maximum_scheduled_tokens();
51 validate_scheduled_token_ceiling(maximum_scheduled_tokens)?;
52 let operation_registry_authority = request
53 .node_resolutions
54 .first()
55 .ok_or_else(|| invalid_plan("plan build request has no node resolutions"))?
56 .operation_registry_authority
57 .clone();
58 if request.node_resolutions.iter().any(|resolution| {
59 resolution.operation_registry_authority != operation_registry_authority
60 }) {
61 return Err(invalid_plan(
62 "node resolutions belong to different operation runtime registries",
63 ));
64 }
65 let policy_capacity = request.policy.memory_capacity_bytes();
66 let memory_reserve = request.policy.memory_reserve_bytes();
67 let device_capacity = request.capabilities.device().total_memory_bytes;
68 if policy_capacity == 0
69 || policy_capacity > device_capacity
70 || memory_reserve >= policy_capacity
71 {
72 return Err(invalid_plan(
73 "runtime policy raw capacity, reserve, or typed admission concurrency is invalid for the device descriptor",
74 ));
75 }
76 let family = request.family;
77 let program = family.program();
78 request
79 .execution_weights
80 .validate_against_catalog(family, request.capabilities)?;
81 let execution_weight_schema = request.execution_weights.plan().schema();
82 let prepared_family_fingerprint = family.fingerprint()?;
83 let program_fingerprint = program.fingerprint()?;
84 let capability_catalog_fingerprint = request.capabilities.fingerprint()?;
85 let device_runtime_implementation_fingerprint = request
86 .capabilities
87 .device()
88 .runtime_implementation_fingerprint
89 .clone();
90 let policy_fingerprint = canonical_runtime_policy_fingerprint(request.policy)?;
91 let weight_format = execution_weight_schema.format_id.clone();
92 let quantization_formats = execution_weight_schema.quantization_formats();
93
94 let mut resolutions = BTreeMap::new();
95 for resolution in request.node_resolutions {
96 let node_id = resolution.node_id.clone();
97 if resolutions.insert(node_id.clone(), resolution).is_some() {
98 return Err(invalid_plan(format!(
99 "node `{node_id}` has duplicate physical resolutions"
100 )));
101 }
102 }
103
104 let program_nodes = program
105 .blocks()
106 .iter()
107 .flat_map(|block| &block.nodes)
108 .collect::<Vec<_>>();
109 let joint_storage = Self::select_joint_provider_storage(
110 &program_nodes,
111 &resolutions,
112 request.capabilities,
113 request.policy,
114 )?;
115 let mut selected_node_resources = joint_storage.node_resources;
116 let selected_resource_profiles = joint_storage.resource_profiles;
117 let mut storage_rejections = joint_storage.storage_rejections;
118 let producers = program_nodes
119 .iter()
120 .flat_map(|node| {
121 node.outputs
122 .iter()
123 .map(move |output| (output.clone(), node.id.clone()))
124 })
125 .collect::<BTreeMap<_, _>>();
126 let mut last_consumers = BTreeMap::<ProgramValueId, usize>::new();
127 for (node_index, node) in program_nodes.iter().enumerate() {
128 for input in &node.inputs {
129 last_consumers.insert(input.clone(), node_index);
130 }
131 }
132 let program_outputs = program.outputs().iter().cloned().collect::<BTreeSet<_>>();
133 let mut canonical_values = BTreeMap::new();
134 let mut bound_values = BTreeSet::new();
135 let mut nodes = Vec::new();
136 let mut state_dependencies = StateDependencyTracker::default();
137 for (node_index, program_node) in program_nodes.into_iter().enumerate() {
138 let resolution = resolutions.remove(&program_node.id).ok_or_else(|| {
139 invalid_plan(format!(
140 "node `{}` has no physical resolution",
141 program_node.id
142 ))
143 })?;
144 let provider_resources = selected_node_resources
145 .remove(&program_node.id)
146 .ok_or_else(|| invalid_plan("joint solver omitted one program node"))?;
147 let storage_rejection = storage_rejections.remove(&program_node.id);
148 let node = Self::build_node(
149 family,
150 execution_weight_schema,
151 &prepared_family_fingerprint,
152 program_node,
153 resolution,
154 provider_resources,
155 storage_rejection,
156 request.capabilities,
157 request.policy.execution_determinism_requirement(),
158 &producers,
159 node_index,
160 &last_consumers,
161 &program_outputs,
162 &mut state_dependencies,
163 &mut canonical_values,
164 &mut bound_values,
165 )?;
166 nodes.push(node);
167 }
168 if !resolutions.is_empty() {
169 return Err(invalid_plan(format!(
170 "physical resolutions contain unknown nodes: {:?}",
171 resolutions.keys().collect::<Vec<_>>()
172 )));
173 }
174 if !selected_node_resources.is_empty() {
175 return Err(invalid_plan(
176 "joint solver returned resources for unknown program nodes",
177 ));
178 }
179 if !storage_rejections.is_empty() {
180 return Err(invalid_plan(
181 "joint solver returned storage rejections for unknown program nodes",
182 ));
183 }
184 Self::validate_semantic_coverage(family, &bound_values)?;
185 Self::validate_global_storage_aliasing(&canonical_values, &nodes)?;
186 let retained_completion_values =
187 resolve_retained_completion_values(&nodes, &request.completion_retention)?;
188 let terminal_output_resources = nodes
189 .iter()
190 .flat_map(|node| node.values())
191 .filter(|binding| program_outputs.contains(binding.value_id()))
192 .flat_map(|binding| binding.storage().components())
193 .map(|component| component.resource_id().clone())
194 .collect::<BTreeSet<_>>()
195 .into_iter()
196 .collect::<Vec<_>>();
197 let retained_completion_resources = retained_completion_values
198 .iter()
199 .map(|value| value.resource_id().clone())
200 .chain(terminal_output_resources.iter().cloned())
201 .collect::<BTreeSet<_>>();
202 let memory = Self::build_memory_plan(
203 family,
204 request.execution_weights.plan(),
205 device_capacity,
206 policy_capacity,
207 memory_reserve,
208 maximum_active_sequences,
209 maximum_scheduled_tokens,
210 &nodes,
211 &selected_resource_profiles,
212 request.policy.reusable_execution_policy(),
213 &retained_completion_resources,
214 )?;
215
216 let (sequence_checkpoint_layout, checkpoint_unsupported_reasons) =
217 super::sequence_checkpoint::derive_sequence_checkpoint(
218 program,
219 &nodes,
220 memory.dynamic_descriptors(),
221 request.capabilities,
222 )?;
223 let memory = memory.with_checkpoint_capacity(
224 request.policy.checkpoint_capacity_policy().copied(),
225 sequence_checkpoint_layout.as_ref(),
226 &nodes,
227 &retained_completion_resources,
228 )?;
229 let mut payload = ExecutionPlanPayload {
230 schema: EXECUTION_PLAN_SCHEMA,
231 plan_id: PlanId::new("plan/unset")?,
232 family_id: family.family_id().clone(),
233 device_id: request.capabilities.device().id.clone(),
234 device_runtime_implementation_fingerprint,
235 prepared_family_fingerprint,
236 program_fingerprint,
237 capability_catalog_fingerprint,
238 policy_version: request.policy.version(),
239 policy_fingerprint,
240 maximum_scheduled_tokens,
241 execution_weights: request.execution_weights.plan().clone(),
242 weight_format,
243 quantization_formats,
244 retained_completion_values,
245 terminal_output_resources,
246 nodes,
247 memory,
248 sequence_checkpoint_layout,
249 };
250 let plan_hash = PlanHash::new(canonical_fingerprint(
251 &PlanHashMaterial::from(&payload),
252 "fingerprint execution plan",
253 )?)?;
254 payload.plan_id = Self::plan_id_for_hash(&plan_hash)?;
255 let plan = Self {
256 payload,
257 plan_hash,
258 operation_registry_authority,
259 trusted_execution_weights: request.execution_weights,
260 checkpoint_unsupported_reasons,
261 };
262 plan.validate_internal()?;
263 Ok(plan)
264 }
265
266 #[allow(clippy::too_many_arguments)]
267 pub(super) fn build_node(
268 family: &PreparedModelFamily,
269 execution_weight_schema: &WeightSchema,
270 prepared_family_fingerprint: &str,
271 program_node: &ProgramNode,
272 resolution: PlanNodeResolution,
273 provider_resources: ProviderResourcePlan,
274 storage_rejection: Option<RejectedProvider>,
275 catalog: &CapabilityCatalog,
276 execution_determinism: ExecutionDeterminismRequirement,
277 producers: &BTreeMap<ProgramValueId, NodeId>,
278 node_index: usize,
279 last_consumers: &BTreeMap<ProgramValueId, usize>,
280 program_outputs: &BTreeSet<ProgramValueId>,
281 state_dependencies: &mut StateDependencyTracker,
282 canonical_values: &mut BTreeMap<ProgramValueId, CanonicalValueBinding>,
283 bound_values: &mut BTreeSet<ProgramValueId>,
284 ) -> Result<PlanNode, VNextError> {
285 let operation = catalog.operation_for_node(&program_node.id, &program_node.operation_id)?;
286 if !operation.version.satisfies(program_node.required_version) {
287 return Err(VNextError::IncompatibleOperationVersion {
288 node_id: Some(program_node.id.to_string()),
289 operation_id: program_node.operation_id.to_string(),
290 required_major: program_node.required_version.major,
291 required_minor: program_node.required_version.minor,
292 available_major: operation.version.major,
293 available_minor: operation.version.minor,
294 });
295 }
296 operation.validate_attributes(&program_node.attributes)?;
297 operation.validate_resolved_bindings(&resolution.values)?;
298 validate_program_bindings(program_node, &resolution.values)?;
299 let exact_aliases = Self::extract_exact_aliases(operation, &resolution.values)?;
300 let work = Self::derive_node_work_contract(program_node, operation, &resolution.values)?;
301 Self::validate_alias_liveness(
302 family,
303 program_node,
304 node_index,
305 last_consumers,
306 program_outputs,
307 &exact_aliases,
308 &resolution.values,
309 )?;
310 let state_effects = Self::derive_state_effects(family, &resolution.values)?;
311 for binding in &resolution.values {
312 validate_semantic_binding(family, execution_weight_schema, binding)?;
313 Self::validate_cross_node_value(binding, canonical_values)?;
314 bound_values.insert(binding.value_id().clone());
315 }
316
317 let (required_weight_formats, required_quantization_formats) =
318 node_weight_requirements(family, &resolution.values)?;
319 let selection = Self::select_provider(
320 program_node,
321 operation,
322 catalog,
323 &resolution.required_capabilities,
324 resolution.preferred_provider.as_ref(),
325 &provider_resources.provider_id,
326 storage_rejection,
327 &required_weight_formats,
328 &required_quantization_formats,
329 execution_determinism,
330 )?;
331 provider_resources.validate_shape()?;
332 if provider_resources.provider_id != selection.selected_provider {
333 return Err(invalid_plan(format!(
334 "node `{}` resource estimate belongs to provider `{}` instead of selected provider `{}`",
335 program_node.id,
336 provider_resources.provider_id,
337 selection.selected_provider
338 )));
339 }
340 let selected_provider = catalog
341 .providers_for_node(&program_node.id, &program_node.operation_id)?
342 .iter()
343 .find(|provider| provider.provider_id() == &selection.selected_provider)
344 .ok_or_else(|| {
345 invalid_plan(format!(
346 "node `{}` selected provider is absent from the catalog",
347 program_node.id
348 ))
349 })?;
350 if provider_resources.estimator_id != selected_provider.resource_estimator_id()
351 || provider_resources.estimator_version
352 != selected_provider.resource_estimator_version()
353 || provider_resources.estimator_implementation_fingerprint
354 != selected_provider.resource_estimator_implementation_fingerprint()
355 {
356 return Err(invalid_plan(format!(
357 "node `{}` provider resource estimate is not issued by the selected catalog provider's estimator",
358 program_node.id
359 )));
360 }
361 let minimum_value_alignment = operation.resources.minimum_value_alignment_bytes;
362 if provider_resources.value_alignment_bytes < minimum_value_alignment
363 || provider_resources.value_alignment_bytes % minimum_value_alignment != 0
364 || !operation
365 .resources
366 .scratch
367 .accepts(provider_resources.scratch.is_some())
368 || !operation
369 .resources
370 .binding
371 .accepts(provider_resources.binding.is_some())
372 || !operation
373 .resources
374 .persistent
375 .accepts(provider_resources.persistent.is_some())
376 {
377 return Err(invalid_plan(format!(
378 "node `{}` provider resource estimate violates the operation's alignment or workspace-presence contract",
379 program_node.id
380 )));
381 }
382 let expected_estimator_input = provider_resource_estimator_input_fingerprint(
383 family,
384 prepared_family_fingerprint,
385 operation,
386 program_node,
387 &selection.selected_provider,
388 &resolution.values,
389 &resolution.required_capabilities,
390 )?;
391 if provider_resources.estimator_input_fingerprint != expected_estimator_input {
392 return Err(invalid_plan(format!(
393 "node `{}` provider resource estimate is not bound to its selected provider, shape, attributes, and bindings",
394 program_node.id
395 )));
396 }
397 let mut dependencies = program_node
398 .inputs
399 .iter()
400 .filter_map(|input| producers.get(input))
401 .cloned()
402 .collect::<BTreeSet<_>>();
403 Self::add_state_dependencies(
404 &program_node.id,
405 &state_effects,
406 state_dependencies,
407 &mut dependencies,
408 );
409 let dependencies = dependencies.into_iter().collect::<Vec<_>>();
410 let scratch_resource = provider_resources
411 .scratch
412 .as_ref()
413 .map(|_| {
414 workspace_base_id(
415 &program_node.id,
416 "scratch",
417 &provider_resources.estimate_fingerprint,
418 )
419 })
420 .transpose()?;
421 let binding_resource = provider_resources
422 .binding
423 .as_ref()
424 .map(|_| {
425 workspace_base_id(
426 &program_node.id,
427 "binding",
428 &provider_resources.estimate_fingerprint,
429 )
430 })
431 .transpose()?;
432 let persistent_resource = provider_resources
433 .persistent
434 .as_ref()
435 .map(|_| {
436 workspace_base_id(
437 &program_node.id,
438 "persistent",
439 &provider_resources.estimate_fingerprint,
440 )
441 })
442 .transpose()?;
443 let resources = resolution
444 .values
445 .iter()
446 .flat_map(|binding| binding.storage().components())
447 .map(|component| component.resource_id().clone())
448 .chain(scratch_resource.iter().cloned())
449 .chain(binding_resource.iter().cloned())
450 .chain(persistent_resource.iter().cloned())
451 .collect::<BTreeSet<_>>()
452 .into_iter()
453 .collect();
454 Ok(PlanNode {
455 id: program_node.id.clone(),
456 dependencies,
457 operation_id: program_node.operation_id.clone(),
458 operation_version: program_node.required_version,
459 operation_fingerprint: operation.fingerprint()?,
460 provider_implementation_fingerprint: selected_provider
461 .provider_implementation_fingerprint()
462 .to_owned(),
463 provider_execution_semantics: selected_provider.execution_semantics(),
464 required_capabilities: resolution.required_capabilities,
465 attributes: program_node.attributes.clone(),
466 work,
467 selection,
468 provider_resources,
469 values: resolution.values,
470 exact_aliases,
471 state_effects,
472 scratch_resource,
473 binding_resource,
474 persistent_resource,
475 resources,
476 })
477 }
478
479 pub(super) fn derive_node_work_contract(
480 node: &ProgramNode,
481 operation: &OperationDescriptor,
482 bindings: &[ResolvedValueBinding],
483 ) -> Result<NodeWorkContract, VNextError> {
484 let ProgramNodeWorkSpec::Tokens {
485 value_id,
486 axis: source_axis,
487 } = &node.work
488 else {
489 return Ok(NodeWorkContract::Fixed);
490 };
491 let source_binding = bindings
492 .iter()
493 .find(|binding| binding.value_id() == value_id)
494 .ok_or_else(|| invalid_plan("node token work source has no resolved binding"))?;
495 let source_contract = match source_binding.role() {
496 ResolvedValueRole::Input => operation.inputs.get(source_binding.ordinal() as usize),
497 ResolvedValueRole::Output => operation.outputs.get(source_binding.ordinal() as usize),
498 }
499 .ok_or_else(|| invalid_plan("node token work source ordinal is outside its operation"))?;
500 let source_axis_index = usize::try_from(*source_axis)
501 .map_err(|_| invalid_plan("node token work axis exceeds usize"))?;
502 let source_symbol = match source_contract.dimensions().get(source_axis_index) {
503 Some(DimensionConstraint::Symbol(symbol)) => symbol,
504 _ => {
505 return Err(invalid_plan(
506 "node token work source axis is not one symbolic operation dimension",
507 ))
508 }
509 };
510 if source_binding.usage() != BufferUsage::Activations
511 || source_binding
512 .tensor()
513 .dimensions()
514 .get(source_axis_index)
515 .is_none()
516 {
517 return Err(invalid_plan(
518 "node token work source is not an in-bounds activation axis",
519 ));
520 }
521
522 let mut projections = Vec::new();
523 for binding in bindings {
524 let contract = match binding.role() {
525 ResolvedValueRole::Input => operation.inputs.get(binding.ordinal() as usize),
526 ResolvedValueRole::Output => operation.outputs.get(binding.ordinal() as usize),
527 }
528 .ok_or_else(|| {
529 invalid_plan("resolved work binding ordinal is outside its operation")
530 })?;
531 let matching_axes = contract
532 .dimensions()
533 .iter()
534 .enumerate()
535 .filter(|(_, dimension)| {
536 matches!(dimension, DimensionConstraint::Symbol(symbol) if symbol == source_symbol)
537 })
538 .map(|(axis, _)| axis)
539 .collect::<Vec<_>>();
540 if matching_axes.len() > 1 {
541 return Err(invalid_plan(
542 "one resolved binding repeats the node token work dimension",
543 ));
544 }
545 let Some(axis) = matching_axes.first().copied() else {
546 continue;
547 };
548 let dimensions = binding.tensor().dimensions();
549 if binding.usage() != BufferUsage::Activations || dimensions.get(axis).is_none() {
550 return Err(invalid_plan(
551 "node token work projection is not an in-bounds activation axis",
552 ));
553 }
554 projections.push(NodeTokenBindingProjection {
555 value_id: binding.value_id().clone(),
556 role: binding.role(),
557 ordinal: binding.ordinal(),
558 axis: u32::try_from(axis)
559 .map_err(|_| invalid_plan("node token projection axis exceeds u32"))?,
560 rank: u32::try_from(dimensions.len())
561 .map_err(|_| invalid_plan("node token projection rank exceeds u32"))?,
562 canonical_extent: dimensions[axis],
563 });
564 }
565 projections.sort();
566 if projections.is_empty()
567 || projections
568 .windows(2)
569 .any(|pair| pair[0].role == pair[1].role && pair[0].ordinal == pair[1].ordinal)
570 {
571 return Err(invalid_plan(
572 "node token work projections are empty or non-canonical",
573 ));
574 }
575 let source = projections
576 .iter()
577 .find(|projection| projection.value_id == *value_id && projection.axis == *source_axis)
578 .cloned()
579 .ok_or_else(|| invalid_plan("node token work source did not resolve exactly"))?;
580 if projections
581 .iter()
582 .any(|projection| projection.canonical_extent != source.canonical_extent)
583 {
584 return Err(invalid_plan(
585 "node token work projections disagree on canonical extent",
586 ));
587 }
588 Ok(NodeWorkContract::Tokens {
589 source,
590 projections,
591 })
592 }
593
594 fn validate_node_work_contract(node: &PlanNode) -> Result<(), VNextError> {
595 let NodeWorkContract::Tokens {
596 source,
597 projections,
598 } = &node.work
599 else {
600 return Ok(());
601 };
602 if projections.is_empty()
603 || projections.windows(2).any(|pair| pair[0] >= pair[1])
604 || !projections.contains(source)
605 || projections
606 .iter()
607 .any(|projection| projection.canonical_extent != source.canonical_extent)
608 {
609 return Err(invalid_plan(format!(
610 "node `{}` token work contract is empty or non-canonical",
611 node.id
612 )));
613 }
614 for projection in projections {
615 let binding = node
616 .values
617 .iter()
618 .find(|binding| {
619 binding.role() == projection.role
620 && binding.ordinal() == projection.ordinal
621 && binding.value_id() == &projection.value_id
622 })
623 .ok_or_else(|| {
624 invalid_plan(format!(
625 "node `{}` token projection has no exact value binding",
626 node.id
627 ))
628 })?;
629 let axis = usize::try_from(projection.axis)
630 .map_err(|_| invalid_plan("node token projection axis exceeds usize"))?;
631 if binding.usage() != BufferUsage::Activations
632 || usize::try_from(projection.rank).ok()
633 != Some(binding.tensor().dimensions().len())
634 || binding.tensor().dimensions().get(axis) != Some(&projection.canonical_extent)
635 {
636 return Err(invalid_plan(format!(
637 "node `{}` token projection differs from its resolved tensor",
638 node.id
639 )));
640 }
641 }
642 Ok(())
643 }
644
645 pub(super) fn extract_exact_aliases(
646 operation: &OperationDescriptor,
647 bindings: &[ResolvedValueBinding],
648 ) -> Result<Vec<PlanExactAlias>, VNextError> {
649 let inputs = &bindings[..operation.inputs.len()];
650 let outputs = &bindings[operation.inputs.len()..];
651 let mut aliases = Vec::new();
652 for (output_ordinal, output) in outputs.iter().enumerate() {
653 let (input_ordinal, kind) = match output.alias() {
654 AliasPolicy::NoAlias => continue,
655 AliasPolicy::MayAlias { tensor_index } => {
656 (*tensor_index, PlanExactAliasKind::MayAlias)
657 }
658 AliasPolicy::MustAlias { tensor_index } => {
659 (*tensor_index, PlanExactAliasKind::MustAlias)
660 }
661 };
662 let input = inputs.get(input_ordinal as usize).ok_or_else(|| {
663 invalid_plan(format!(
664 "operation `{}` alias input ordinal is out of range after validation",
665 operation.id
666 ))
667 })?;
668 if output.storage() == input.storage() {
669 aliases.push(PlanExactAlias {
670 output_value_id: output.value_id().clone(),
671 output_ordinal: output_ordinal as u32,
672 input_value_id: input.value_id().clone(),
673 input_ordinal,
674 kind,
675 });
676 } else if kind == PlanExactAliasKind::MustAlias {
677 return Err(invalid_plan(format!(
678 "operation `{}` lost its mandatory exact alias proof",
679 operation.id
680 )));
681 }
682 }
683 Ok(aliases)
684 }
685
686 #[allow(clippy::too_many_arguments)]
687 pub(super) fn validate_alias_liveness(
688 family: &PreparedModelFamily,
689 node: &ProgramNode,
690 node_index: usize,
691 last_consumers: &BTreeMap<ProgramValueId, usize>,
692 program_outputs: &BTreeSet<ProgramValueId>,
693 aliases: &[PlanExactAlias],
694 bindings: &[ResolvedValueBinding],
695 ) -> Result<(), VNextError> {
696 for alias in aliases {
697 let input = bindings
698 .iter()
699 .find(|binding| {
700 binding.role() == ResolvedValueRole::Input
701 && binding.ordinal() == alias.input_ordinal
702 && binding.value_id() == &alias.input_value_id
703 })
704 .ok_or_else(|| invalid_plan("exact alias input proof has no matching binding"))?;
705 if family
706 .program()
707 .states()
708 .iter()
709 .any(|state| state.value_id == alias.input_value_id)
710 {
711 return Err(invalid_plan(format!(
712 "node `{}` output aliases state `{}` without a typed state transition contract",
713 node.id, alias.input_value_id
714 )));
715 }
716 if input.usage() != BufferUsage::Activations
717 || last_consumers.get(&alias.input_value_id) != Some(&node_index)
718 || program_outputs.contains(&alias.input_value_id)
719 {
720 return Err(invalid_plan(format!(
721 "node `{}` aliases activation `{}` before its final legal consumer",
722 node.id, alias.input_value_id
723 )));
724 }
725 }
726 Ok(())
727 }
728
729 pub(super) fn derive_state_effects(
730 family: &PreparedModelFamily,
731 bindings: &[ResolvedValueBinding],
732 ) -> Result<Vec<PlanStateEffect>, VNextError> {
733 let mut effects = Vec::new();
734 for state in family.program().states() {
735 let mut reads = false;
736 let mut writes = false;
737 let state_bindings = bindings
738 .iter()
739 .filter(|binding| binding.value_id() == &state.value_id)
740 .collect::<Vec<_>>();
741 for binding in &state_bindings {
742 match binding.access() {
743 TensorAccess::Read => reads = true,
744 TensorAccess::Write => writes = true,
745 TensorAccess::ReadWrite => {
746 reads = true;
747 writes = true;
748 }
749 }
750 }
751 let access = match (reads, writes) {
752 (false, false) => continue,
753 (true, false) => TensorAccess::Read,
754 (false, true) => TensorAccess::Write,
755 (true, true) => TensorAccess::ReadWrite,
756 };
757 let lifetime = match state.lifetime {
758 StateLifetime::Request => AllocationLifetime::Request,
759 StateLifetime::Sequence => AllocationLifetime::Sequence,
760 StateLifetime::Step => AllocationLifetime::Step,
761 };
762 let resource_ids = state_bindings
763 .iter()
764 .flat_map(|binding| binding.storage().components())
765 .map(|component| component.resource_id().clone())
766 .collect::<BTreeSet<_>>()
767 .into_iter()
768 .collect::<Vec<_>>();
769 if resource_ids.is_empty() {
770 return Err(invalid_plan(format!(
771 "state `{}` effect has no physical resource closure",
772 state.id
773 )));
774 }
775 effects.push(PlanStateEffect {
776 state_id: state.id.clone(),
777 state_value_id: state.value_id.clone(),
778 lifetime,
779 access,
780 resource_ids,
781 });
782 }
783 if effects
784 .windows(2)
785 .any(|pair| pair[0].state_id >= pair[1].state_id)
786 {
787 return Err(invalid_plan("state effects are not canonical"));
788 }
789 Ok(effects)
790 }
791
792 pub(super) fn add_state_dependencies(
793 node_id: &NodeId,
794 effects: &[PlanStateEffect],
795 tracker: &mut StateDependencyTracker,
796 dependencies: &mut BTreeSet<NodeId>,
797 ) {
798 for effect in effects {
799 let state_id = &effect.state_id;
800 match effect.access {
801 TensorAccess::Read => {
802 if let Some(writer) = tracker.last_writer.get(state_id) {
803 dependencies.insert(writer.clone());
804 }
805 tracker
806 .readers_since_write
807 .entry(state_id.clone())
808 .or_default()
809 .insert(node_id.clone());
810 }
811 TensorAccess::Write | TensorAccess::ReadWrite => {
812 if let Some(writer) = tracker.last_writer.get(state_id) {
813 dependencies.insert(writer.clone());
814 }
815 if let Some(readers) = tracker.readers_since_write.remove(state_id) {
816 dependencies.extend(readers);
817 }
818 tracker
819 .last_writer
820 .insert(state_id.clone(), node_id.clone());
821 }
822 }
823 }
824 }
825
826 pub(super) fn validate_cross_node_value(
827 binding: &ResolvedValueBinding,
828 values: &mut BTreeMap<ProgramValueId, CanonicalValueBinding>,
829 ) -> Result<(), VNextError> {
830 let canonical = CanonicalValueBinding {
831 tensor: binding.tensor().clone(),
832 usage: binding.usage(),
833 storage: binding.storage().clone(),
834 readonly_weight: (binding.usage() == BufferUsage::Weights
835 && binding.access() == TensorAccess::Read)
836 .then(|| binding.weight().cloned())
837 .flatten(),
838 };
839 match values.get(binding.value_id()) {
840 Some(previous) if previous != &canonical => Err(invalid_plan(format!(
841 "value `{}` changes tensor or physical storage between nodes",
842 binding.value_id()
843 ))),
844 Some(_) => Ok(()),
845 None => {
846 values.insert(binding.value_id().clone(), canonical);
847 Ok(())
848 }
849 }
850 }
851
852 pub(super) fn validate_global_storage_aliasing(
853 values: &BTreeMap<ProgramValueId, CanonicalValueBinding>,
854 nodes: &[PlanNode],
855 ) -> Result<(), VNextError> {
856 let alias_classes = Self::alias_classes(nodes)?;
857 let mut by_resource = BTreeMap::<ResourceId, Vec<GlobalValueRange>>::new();
858 for (value_id, binding) in values {
859 for component in binding.storage.components() {
860 let end_bytes = component
861 .offset_bytes()
862 .checked_add(component.length_bytes())
863 .ok_or_else(|| invalid_plan("global value storage range overflows u64"))?;
864 let ranges = by_resource
865 .entry(component.resource_id().clone())
866 .or_default();
867 for previous in ranges.iter().filter(|previous| {
868 previous.value_id != *value_id
869 && previous.offset_bytes < end_bytes
870 && component.offset_bytes() < previous.end_bytes
871 }) {
872 let same_alias_class = alias_classes.get(&previous.value_id).is_some()
873 && alias_classes.get(&previous.value_id) == alias_classes.get(value_id);
874 let previous_binding = values.get(&previous.value_id).ok_or_else(|| {
875 invalid_plan("global alias range has no canonical value binding")
876 })?;
877 let shared_signs =
878 binding
879 .readonly_weight
880 .as_ref()
881 .zip(previous_binding.readonly_weight.as_ref())
882 .is_some_and(|(weight, previous_weight)| {
883 previous_binding.storage.components().iter().any(
884 |previous_component| {
885 previous_component.offset_bytes() == previous.offset_bytes
886 && previous_component
887 .offset_bytes()
888 .checked_add(previous_component.length_bytes())
889 == Some(previous.end_bytes)
890 && crate::vnext::same_shared_transform_sign_component(
891 weight,
892 component,
893 previous_weight,
894 previous_component,
895 )
896 },
897 )
898 });
899 if (!same_alias_class || previous_binding.storage != binding.storage)
900 && !shared_signs
901 {
902 return Err(invalid_plan(format!(
903 "values `{}` and `{value_id}` have undeclared, partial, or non-equivalent overlap in physical resource `{}`",
904 previous.value_id,
905 component.resource_id()
906 )));
907 }
908 }
909 ranges.push(GlobalValueRange {
910 value_id: value_id.clone(),
911 offset_bytes: component.offset_bytes(),
912 end_bytes,
913 });
914 }
915 }
916 Ok(())
917 }
918
919 pub(super) fn alias_classes(
920 nodes: &[PlanNode],
921 ) -> Result<BTreeMap<ProgramValueId, ProgramValueId>, VNextError> {
922 let mut graph = BTreeMap::<ProgramValueId, BTreeSet<ProgramValueId>>::new();
923 for node in nodes {
924 let mut previous_output_ordinal = None;
925 for alias in &node.exact_aliases {
926 if previous_output_ordinal.is_some_and(|ordinal| ordinal >= alias.output_ordinal) {
927 return Err(invalid_plan(format!(
928 "node `{}` exact aliases are not canonical",
929 node.id
930 )));
931 }
932 previous_output_ordinal = Some(alias.output_ordinal);
933 let input = node
934 .values
935 .iter()
936 .find(|binding| {
937 binding.role() == ResolvedValueRole::Input
938 && binding.ordinal() == alias.input_ordinal
939 && binding.value_id() == &alias.input_value_id
940 })
941 .ok_or_else(|| invalid_plan("plan exact alias input binding is missing"))?;
942 let output = node
943 .values
944 .iter()
945 .find(|binding| {
946 binding.role() == ResolvedValueRole::Output
947 && binding.ordinal() == alias.output_ordinal
948 && binding.value_id() == &alias.output_value_id
949 })
950 .ok_or_else(|| invalid_plan("plan exact alias output binding is missing"))?;
951 let policy_matches = matches!(
952 (output.alias(), alias.kind),
953 (
954 AliasPolicy::MayAlias { tensor_index },
955 PlanExactAliasKind::MayAlias
956 ) if *tensor_index == alias.input_ordinal
957 ) || matches!(
958 (output.alias(), alias.kind),
959 (
960 AliasPolicy::MustAlias { tensor_index },
961 PlanExactAliasKind::MustAlias
962 ) if *tensor_index == alias.input_ordinal
963 );
964 if !policy_matches
965 || input.storage() != output.storage()
966 || input.usage() != BufferUsage::Activations
967 || output.usage() != BufferUsage::Activations
968 {
969 return Err(invalid_plan(format!(
970 "node `{}` exact alias proof differs from its bindings",
971 node.id
972 )));
973 }
974 graph
975 .entry(alias.input_value_id.clone())
976 .or_default()
977 .insert(alias.output_value_id.clone());
978 graph
979 .entry(alias.output_value_id.clone())
980 .or_default()
981 .insert(alias.input_value_id.clone());
982 }
983 }
984
985 let mut classes = BTreeMap::new();
986 let mut visited = BTreeSet::new();
987 for start in graph.keys() {
988 if visited.contains(start) {
989 continue;
990 }
991 let mut pending = vec![start.clone()];
992 let mut members = BTreeSet::new();
993 while let Some(value) = pending.pop() {
994 if !visited.insert(value.clone()) {
995 continue;
996 }
997 members.insert(value.clone());
998 if let Some(neighbors) = graph.get(&value) {
999 pending.extend(neighbors.iter().cloned());
1000 }
1001 }
1002 let representative = members
1003 .first()
1004 .cloned()
1005 .ok_or_else(|| invalid_plan("empty alias equivalence class"))?;
1006 for member in members {
1007 classes.insert(member, representative.clone());
1008 }
1009 }
1010 Ok(classes)
1011 }
1012
1013 pub(super) fn validate_semantic_coverage(
1014 family: &PreparedModelFamily,
1015 bound: &BTreeSet<ProgramValueId>,
1016 ) -> Result<(), VNextError> {
1017 let required = family
1018 .program()
1019 .inputs()
1020 .iter()
1021 .cloned()
1022 .chain(
1023 family
1024 .program()
1025 .weights()
1026 .iter()
1027 .map(|weight| weight.value_id.clone()),
1028 )
1029 .chain(
1030 family
1031 .program()
1032 .states()
1033 .iter()
1034 .map(|state| state.value_id.clone()),
1035 )
1036 .chain(family.program().outputs().iter().cloned())
1037 .collect::<BTreeSet<_>>();
1038 if !required.is_subset(bound) {
1039 return Err(invalid_plan(format!(
1040 "semantic values lack physical bindings: {:?}",
1041 required.difference(bound).collect::<Vec<_>>()
1042 )));
1043 }
1044 Ok(())
1045 }
1046
1047 pub(super) fn available_storage_profiles<P: RuntimePolicy>(
1048 requirement: &DynamicStorageRequirement,
1049 catalog: &CapabilityCatalog,
1050 policy: &P,
1051 ) -> BTreeSet<DynamicStorageProfile> {
1052 policy
1053 .dynamic_storage_profile_order()
1054 .iter()
1055 .copied()
1056 .filter(|profile| {
1057 catalog.device().dynamic_storage_profiles.contains(profile)
1058 && requirement.accepts(*profile)
1059 })
1060 .collect()
1061 }
1062
1063 pub(super) fn merge_storage_constraint(
1064 constraints: &mut BTreeMap<ResourceId, BTreeSet<DynamicStorageProfile>>,
1065 resource_id: ResourceId,
1066 accepted: BTreeSet<DynamicStorageProfile>,
1067 ) -> bool {
1068 if accepted.is_empty() {
1069 return false;
1070 }
1071 match constraints.get_mut(&resource_id) {
1072 Some(existing) => {
1073 existing.retain(|profile| accepted.contains(profile));
1074 !existing.is_empty()
1075 }
1076 None => {
1077 constraints.insert(resource_id, accepted);
1078 true
1079 }
1080 }
1081 }
1082
1083 pub(super) fn select_joint_provider_storage<P: RuntimePolicy>(
1084 program_nodes: &[&ProgramNode],
1085 resolutions: &BTreeMap<NodeId, PlanNodeResolution>,
1086 catalog: &CapabilityCatalog,
1087 policy: &P,
1088 ) -> Result<JointProviderStorageSelection, VNextError> {
1089 let mut candidate_sets = Vec::with_capacity(program_nodes.len());
1090 for node in program_nodes {
1091 let resolution = resolutions.get(&node.id).ok_or_else(|| {
1092 invalid_plan(format!("node `{}` has no physical resolution", node.id))
1093 })?;
1094 let providers = catalog.providers_for_node(&node.id, &node.operation_id)?;
1095 let mut candidates = Vec::new();
1096 for resources in &resolution.provider_resource_candidates {
1097 let provider = providers
1098 .iter()
1099 .find(|provider| provider.provider_id() == resources.provider_id())
1100 .ok_or_else(|| {
1101 invalid_plan("provider resource candidate is absent from the catalog")
1102 })?;
1103 let mut constraints = BTreeMap::new();
1104 let mut compatible = true;
1105 for binding in resolution
1106 .values
1107 .iter()
1108 .filter(|binding| binding.usage() != BufferUsage::Weights)
1109 {
1110 let Some(requirement) =
1111 provider.dynamic_storage_for(binding.role(), binding.ordinal())
1112 else {
1113 compatible = false;
1114 break;
1115 };
1116 let accepted = Self::available_storage_profiles(requirement, catalog, policy);
1117 for component in binding.storage().components() {
1118 if !Self::merge_storage_constraint(
1119 &mut constraints,
1120 component.resource_id().clone(),
1121 accepted.clone(),
1122 ) {
1123 compatible = false;
1124 break;
1125 }
1126 }
1127 if !compatible {
1128 break;
1129 }
1130 }
1131 for (kind, workspace) in [
1132 ("scratch", resources.scratch()),
1133 ("binding", resources.binding()),
1134 ("persistent", resources.persistent()),
1135 ] {
1136 let Some(workspace) = workspace else {
1137 continue;
1138 };
1139 let resource_id =
1140 workspace_base_id(&node.id, kind, resources.estimate_fingerprint())?;
1141 if !Self::merge_storage_constraint(
1142 &mut constraints,
1143 resource_id,
1144 Self::available_storage_profiles(workspace.storage(), catalog, policy),
1145 ) {
1146 compatible = false;
1147 break;
1148 }
1149 }
1150 if compatible {
1151 let is_preferred = resolution
1152 .preferred_provider
1153 .as_ref()
1154 .is_some_and(|preferred| preferred == resources.provider_id());
1155 candidates.push(JointProviderCandidate {
1156 resources: resources.clone(),
1157 allowed_profiles: constraints,
1158 is_preferred,
1159 });
1160 }
1161 }
1162 candidates.sort_by(|left, right| {
1163 let left_preferred = resolution
1164 .preferred_provider
1165 .as_ref()
1166 .is_some_and(|preferred| preferred == left.resources.provider_id());
1167 let right_preferred = resolution
1168 .preferred_provider
1169 .as_ref()
1170 .is_some_and(|preferred| preferred == right.resources.provider_id());
1171 right_preferred.cmp(&left_preferred).then(
1172 left.resources
1173 .provider_id()
1174 .cmp(right.resources.provider_id()),
1175 )
1176 });
1177 if candidates.is_empty() {
1178 return Err(invalid_plan(format!(
1179 "node `{}` has no provider candidate with an available storage profile",
1180 node.id
1181 )));
1182 }
1183 candidate_sets.push(candidates);
1184 }
1185
1186 let (chosen, resource_profiles) = Self::solve_joint_provider_candidates(
1187 &candidate_sets,
1188 policy.dynamic_storage_profile_order(),
1189 )?;
1190
1191 let mut storage_rejections = BTreeMap::new();
1192 for (index, node) in program_nodes.iter().enumerate() {
1193 let resolution = resolutions
1194 .get(&node.id)
1195 .ok_or_else(|| invalid_plan("joint storage resolution disappeared"))?;
1196 let Some(preferred) = resolution.preferred_provider.as_ref() else {
1197 continue;
1198 };
1199 if chosen[index].provider_id() == preferred {
1200 continue;
1201 }
1202 let Some(preferred_candidate) = candidate_sets[index]
1203 .iter()
1204 .find(|candidate| candidate.resources.provider_id() == preferred)
1205 else {
1206 if let Some(reason) = resolution.provider_resolution_rejections.get(preferred) {
1207 storage_rejections.insert(
1208 node.id.clone(),
1209 RejectedProvider {
1210 provider_id: preferred.clone(),
1211 reasons: reason.clone(),
1212 },
1213 );
1214 }
1215 continue;
1216 };
1217 let resource_ids =
1218 storage_incompatible_resource_ids(preferred_candidate, &resource_profiles);
1219 if resource_ids.is_empty() {
1220 return Err(invalid_plan(format!(
1221 "preferred provider `{preferred}` was not selected for node `{}` without a storage conflict",
1222 node.id
1223 )));
1224 }
1225 storage_rejections.insert(
1226 node.id.clone(),
1227 RejectedProvider {
1228 provider_id: preferred.clone(),
1229 reasons: PlanProviderRejectReason::StorageIncompatible { resource_ids },
1230 },
1231 );
1232 }
1233
1234 let node_resources = program_nodes
1235 .iter()
1236 .zip(chosen)
1237 .map(|(node, resources)| (node.id.clone(), resources))
1238 .collect();
1239 Ok(JointProviderStorageSelection {
1240 node_resources,
1241 resource_profiles,
1242 storage_rejections,
1243 })
1244 }
1245
1246 pub(super) fn solve_joint_provider_candidates(
1247 candidate_sets: &[Vec<JointProviderCandidate>],
1248 profile_order: &[DynamicStorageProfile],
1249 ) -> Result<
1250 (
1251 Vec<ProviderResourcePlan>,
1252 BTreeMap<ResourceId, DynamicStorageProfile>,
1253 ),
1254 VNextError,
1255 > {
1256 if candidate_sets.is_empty() || candidate_sets.iter().any(Vec::is_empty) {
1257 return Err(invalid_plan(
1258 "joint provider/storage search has an empty candidate set",
1259 ));
1260 }
1261 if profile_order.is_empty() {
1262 return Err(invalid_plan(
1263 "joint provider/storage search has an empty profile order",
1264 ));
1265 }
1266
1267 let components = joint_candidate_components(candidate_sets);
1268 let mut chosen = vec![None; candidate_sets.len()];
1269 let mut resource_profiles = BTreeMap::new();
1270 for component in components {
1271 let solution =
1272 Self::solve_joint_provider_component(&component, candidate_sets, profile_order)?;
1273 for (node_index, resources) in component.iter().copied().zip(solution.chosen) {
1274 if chosen[node_index].replace(resources).is_some() {
1275 return Err(invalid_plan(
1276 "joint storage component assigned one node more than once",
1277 ));
1278 }
1279 }
1280 for (resource_id, profile) in solution.resource_profiles {
1281 if resource_profiles.insert(resource_id, profile).is_some() {
1282 return Err(invalid_plan(
1283 "joint storage components overlap one resource",
1284 ));
1285 }
1286 }
1287 }
1288 let chosen = chosen
1289 .into_iter()
1290 .collect::<Option<Vec<_>>>()
1291 .ok_or_else(|| invalid_plan("joint storage components omitted one node"))?;
1292 Ok((chosen, resource_profiles))
1293 }
1294
1295 pub(super) fn solve_joint_provider_component(
1296 component: &[usize],
1297 candidate_sets: &[Vec<JointProviderCandidate>],
1298 profile_order: &[DynamicStorageProfile],
1299 ) -> Result<JointComponentSolution, VNextError> {
1300 let mut frontier = BTreeMap::from([(
1301 BTreeMap::<ResourceId, BTreeSet<DynamicStorageProfile>>::new(),
1302 JointPartialSelection::default(),
1303 )]);
1304 for node_index in component {
1305 let mut next = BTreeMap::<
1306 BTreeMap<ResourceId, BTreeSet<DynamicStorageProfile>>,
1307 JointPartialSelection,
1308 >::new();
1309 for (constraints, partial) in frontier {
1310 for candidate in &candidate_sets[*node_index] {
1311 let mut next_constraints = constraints.clone();
1312 if candidate
1313 .allowed_profiles
1314 .iter()
1315 .any(|(resource_id, accepted)| {
1316 !Self::merge_storage_constraint(
1317 &mut next_constraints,
1318 resource_id.clone(),
1319 accepted.clone(),
1320 )
1321 })
1322 {
1323 continue;
1324 }
1325 let mut next_partial = partial.clone();
1326 next_partial.chosen.push(candidate.resources.clone());
1327 next_partial.preferred.push(candidate.is_preferred);
1328 match next.entry(next_constraints) {
1329 std::collections::btree_map::Entry::Vacant(entry) => {
1330 entry.insert(next_partial);
1331 }
1332 std::collections::btree_map::Entry::Occupied(mut entry) => {
1333 if joint_partial_precedes(&next_partial, entry.get()) {
1334 entry.insert(next_partial);
1335 }
1336 }
1337 }
1338 }
1339 }
1340 if next.is_empty() {
1341 return Err(invalid_plan(
1342 "no joint provider/storage assignment satisfies shared resource constraints",
1343 ));
1344 }
1345 frontier = next;
1346 }
1347
1348 let mut best: Option<(JointSelectionObjective, JointComponentSolution)> = None;
1349 for (constraints, partial) in frontier {
1350 let resource_profiles = constraints
1351 .into_iter()
1352 .map(|(resource_id, accepted)| {
1353 let (rank, profile) = profile_order
1354 .iter()
1355 .copied()
1356 .enumerate()
1357 .find(|(_, profile)| accepted.contains(profile))
1358 .ok_or_else(|| {
1359 invalid_plan("joint storage solution lost policy-ordered profile")
1360 })?;
1361 Ok((resource_id, (rank, profile)))
1362 })
1363 .collect::<Result<BTreeMap<_, _>, VNextError>>()?;
1364 let objective = JointSelectionObjective::new(
1365 &partial,
1366 resource_profiles.values().map(|(rank, _)| *rank),
1367 profile_order.len(),
1368 )?;
1369 let solution = JointComponentSolution {
1370 chosen: partial.chosen,
1371 resource_profiles: resource_profiles
1372 .into_iter()
1373 .map(|(resource_id, (_, profile))| (resource_id, profile))
1374 .collect(),
1375 };
1376 if best
1377 .as_ref()
1378 .is_none_or(|(current, _)| objective.precedes(current))
1379 {
1380 best = Some((objective, solution));
1381 }
1382 }
1383 best.map(|(_, solution)| solution).ok_or_else(|| {
1384 invalid_plan("no joint provider/storage assignment satisfies one component")
1385 })
1386 }
1387
1388 pub(super) fn select_provider(
1389 node: &ProgramNode,
1390 operation: &OperationDescriptor,
1391 catalog: &CapabilityCatalog,
1392 resolution_required_capabilities: &BTreeSet<CapabilityId>,
1393 preferred_provider: Option<&ProviderId>,
1394 storage_selected_provider: &ProviderId,
1395 storage_rejection: Option<RejectedProvider>,
1396 required_weight_formats: &BTreeSet<WeightFormatId>,
1397 required_quantization_formats: &BTreeSet<QuantizationFormatId>,
1398 execution_determinism: ExecutionDeterminismRequirement,
1399 ) -> Result<ProviderSelection, VNextError> {
1400 let required_capabilities = operation
1401 .provider
1402 .required_capabilities
1403 .union(resolution_required_capabilities)
1404 .cloned()
1405 .collect::<BTreeSet<_>>();
1406 let request = ProviderCompatibilityRequest::new(
1407 node.operation_id.clone(),
1408 node.required_version,
1409 required_capabilities,
1410 required_weight_formats.clone(),
1411 required_quantization_formats.clone(),
1412 execution_determinism,
1413 )?;
1414 let report = catalog.provider_compatibility(request)?;
1415 report.require_compatible_for_node(&catalog.device().id, &node.id)?;
1416 if !report
1417 .compatible_provider_ids()
1418 .contains(storage_selected_provider)
1419 {
1420 return Err(invalid_plan(format!(
1421 "joint storage solver selected incompatible provider `{storage_selected_provider}`"
1422 )));
1423 }
1424 let selected_provider = storage_selected_provider.clone();
1425 let selection_reason = match preferred_provider {
1426 Some(preferred) if preferred == &selected_provider => {
1427 ProviderSelectionReason::PreferredCompatible
1428 }
1429 Some(_) => ProviderSelectionReason::FallbackFromPreferred,
1430 None => ProviderSelectionReason::CanonicalCompatible,
1431 };
1432 let mut rejected_providers = report
1433 .rejected()
1434 .iter()
1435 .map(|rejection| RejectedProvider {
1436 provider_id: rejection.provider_id.clone(),
1437 reasons: PlanProviderRejectReason::Incompatible(rejection.reasons.clone()),
1438 })
1439 .collect::<Vec<_>>();
1440 if let Some(preferred) = preferred_provider {
1441 let registered = catalog
1442 .providers_for_node(&node.id, &node.operation_id)?
1443 .iter()
1444 .any(|provider| provider.provider_id() == preferred);
1445 if !registered {
1446 rejected_providers.push(RejectedProvider {
1447 provider_id: preferred.clone(),
1448 reasons: PlanProviderRejectReason::NotRegistered,
1449 });
1450 }
1451 }
1452 if let Some(rejection) = storage_rejection {
1453 if rejected_providers
1454 .iter()
1455 .any(|existing| existing.provider_id == rejection.provider_id)
1456 {
1457 return Err(invalid_plan(
1458 "provider has duplicate compatibility and storage rejection evidence",
1459 ));
1460 }
1461 rejected_providers.push(rejection);
1462 }
1463 if let Some(preferred) =
1464 preferred_provider.filter(|preferred| *preferred != &selected_provider)
1465 {
1466 if !rejected_providers
1467 .iter()
1468 .any(|rejection| &rejection.provider_id == preferred)
1469 {
1470 return Err(invalid_plan(format!(
1471 "preferred provider `{preferred}` fallback lacks typed rejection evidence"
1472 )));
1473 }
1474 }
1475 rejected_providers.sort_by(|left, right| left.provider_id.cmp(&right.provider_id));
1476 Ok(ProviderSelection {
1477 requested_provider: preferred_provider.cloned(),
1478 selected_provider,
1479 selection_reason,
1480 rejected_providers,
1481 })
1482 }
1483
1484 pub(super) fn validate_provider_selection_evidence(
1485 selection: &ProviderSelection,
1486 ) -> Result<(), VNextError> {
1487 if selection
1488 .rejected_providers
1489 .windows(2)
1490 .any(|pair| pair[0].provider_id >= pair[1].provider_id)
1491 || selection
1492 .rejected_providers
1493 .iter()
1494 .any(|rejection| rejection.provider_id == selection.selected_provider)
1495 || selection.rejected_providers.iter().any(|rejection| {
1496 matches!(
1497 &rejection.reasons,
1498 PlanProviderRejectReason::StorageIncompatible { resource_ids }
1499 if resource_ids.is_empty()
1500 || resource_ids.windows(2).any(|pair| pair[0] >= pair[1])
1501 )
1502 })
1503 {
1504 return Err(invalid_plan(
1505 "provider rejection evidence is duplicate, non-canonical, or rejects the selected provider",
1506 ));
1507 }
1508 match (
1509 selection.requested_provider.as_ref(),
1510 selection.selection_reason,
1511 ) {
1512 (None, ProviderSelectionReason::CanonicalCompatible) => {}
1513 (Some(requested), ProviderSelectionReason::PreferredCompatible)
1514 if requested == &selection.selected_provider => {}
1515 (Some(requested), ProviderSelectionReason::FallbackFromPreferred)
1516 if requested != &selection.selected_provider
1517 && selection
1518 .rejected_providers
1519 .iter()
1520 .any(|rejection| &rejection.provider_id == requested) => {}
1521 _ => return Err(invalid_plan(
1522 "provider selection reason is inconsistent with preference and rejection evidence",
1523 )),
1524 }
1525 Ok(())
1526 }
1527
1528 pub(super) fn build_memory_plan(
1529 family: &PreparedModelFamily,
1530 execution_weights: &ExecutionWeightPlan,
1531 device_capacity_bytes: u64,
1532 policy_capacity_bytes: u64,
1533 reserve_bytes: u64,
1534 maximum_active_sequences: u32,
1535 maximum_scheduled_tokens: u64,
1536 nodes: &[PlanNode],
1537 selected_resource_profiles: &BTreeMap<ResourceId, DynamicStorageProfile>,
1538 reusable_execution_policy: Option<&ReusableExecutionPolicy>,
1539 retained_completion_resources: &BTreeSet<ResourceId>,
1540 ) -> Result<MemoryPlan, VNextError> {
1541 validate_scheduled_token_ceiling(maximum_scheduled_tokens)?;
1542 let program_inputs = family
1543 .program()
1544 .inputs()
1545 .iter()
1546 .cloned()
1547 .collect::<BTreeSet<_>>();
1548 let program_outputs = family
1549 .program()
1550 .outputs()
1551 .iter()
1552 .cloned()
1553 .collect::<BTreeSet<_>>();
1554 let mut product_io_resources = nodes
1555 .iter()
1556 .flat_map(|node| node.values())
1557 .filter(|binding| {
1558 program_inputs.contains(binding.value_id())
1559 || program_outputs.contains(binding.value_id())
1560 })
1561 .flat_map(|binding| binding.storage().components())
1562 .map(|component| component.resource_id().clone())
1563 .collect::<BTreeSet<_>>();
1564 product_io_resources.extend(retained_completion_resources.iter().cloned());
1565 let state_initializations = family
1566 .program()
1567 .states()
1568 .iter()
1569 .map(|state| (state.value_id.clone(), state.initialization))
1570 .collect::<BTreeMap<_, _>>();
1571 let mut values = BTreeMap::<ResourceId, ValueAllocationAccumulator>::new();
1572 let mut static_allocations = Vec::new();
1573 let mut dynamic_descriptors = Vec::new();
1574 let workspace_layout_fingerprint = workspace_storage_layout_fingerprint()?;
1575 if let Some(resource_id) =
1576 execution_weights.static_weight_transform_scratch_resource_id()?
1577 {
1578 static_allocations.push(ResourceAllocation::new(
1579 resource_id,
1580 execution_weights.maximum_static_weight_transform_scratch_bytes()?,
1581 STATIC_WEIGHT_TRANSFORM_SCRATCH_ALIGNMENT_BYTES,
1582 BufferUsage::Scratch,
1583 ElementType::U8,
1584 AllocationKind::InitializationScratch,
1585 DynamicStorageContract::new(
1586 static_contiguous_storage_profile()?,
1587 workspace_layout_fingerprint.clone(),
1588 )?,
1589 )?);
1590 }
1591 for node in nodes {
1592 let value_alignment = node.provider_resources.value_alignment_bytes;
1593 for binding in &node.values {
1594 let logical_layout_fingerprint =
1595 tensor_storage_layout_fingerprint(binding.tensor().layout())?;
1596 for component in binding.storage().components() {
1597 if component.offset_bytes() % value_alignment != 0 {
1598 return Err(invalid_plan(format!(
1599 "resource `{}` offset is not aligned for provider `{}`",
1600 component.resource_id(),
1601 node.provider_resources.provider_id
1602 )));
1603 }
1604 let end = component
1605 .offset_bytes()
1606 .checked_add(component.length_bytes())
1607 .ok_or_else(|| invalid_plan("resource byte range overflows u64"))?;
1608 let token_projection = node
1609 .work
1610 .token_projection(binding.role(), binding.ordinal())
1611 .map(|projection| {
1612 if component.offset_bytes() != 0
1613 || component.length_bytes() % projection.canonical_extent() != 0
1614 {
1615 return Err(invalid_plan(format!(
1616 "token-scaled resource `{}` is not one exact canonical tensor range",
1617 component.resource_id()
1618 )));
1619 }
1620 Ok((
1621 component.length_bytes() / projection.canonical_extent(),
1622 projection.canonical_extent(),
1623 ))
1624 })
1625 .transpose()?;
1626 let demand = Self::value_resource_demand(
1627 family,
1628 binding.value_id(),
1629 binding.usage(),
1630 end,
1631 token_projection,
1632 maximum_active_sequences,
1633 maximum_scheduled_tokens,
1634 product_io_resources.contains(component.resource_id()),
1635 )?;
1636 let initialization = state_initializations
1637 .get(binding.value_id())
1638 .copied()
1639 .unwrap_or(StateInitialization::None);
1640 values
1641 .entry(component.resource_id().clone())
1642 .and_modify(|allocation| {
1643 allocation.merge_result =
1644 allocation.merge_result.take().and_then(|_| {
1645 allocation.merge(
1646 end,
1647 value_alignment,
1648 binding.usage(),
1649 component.element_type(),
1650 demand,
1651 initialization,
1652 logical_layout_fingerprint.clone(),
1653 )
1654 });
1655 })
1656 .or_insert_with(|| ValueAllocationAccumulator {
1657 end_bytes: end,
1658 alignment_bytes: value_alignment,
1659 usage: binding.usage(),
1660 element_type: component.element_type(),
1661 demand,
1662 initialization,
1663 logical_layout_fingerprints: BTreeSet::from([
1664 logical_layout_fingerprint.clone(),
1665 ]),
1666 merge_result: Some(()),
1667 });
1668 }
1669 }
1670 if let Some(workspace) = &node.provider_resources.scratch {
1671 let resource_id = node.scratch_resource.clone().ok_or_else(|| {
1672 invalid_plan(format!(
1673 "node `{}` scratch base identity is missing",
1674 node.id
1675 ))
1676 })?;
1677 if workspace.scope != ProviderWorkspaceScope::Invocation {
1678 return Err(invalid_plan(format!(
1679 "node `{}` scratch workspace is not invocation scoped",
1680 node.id
1681 )));
1682 }
1683 let storage = DynamicStorageContract::new(
1684 *selected_resource_profiles
1685 .get(&resource_id)
1686 .ok_or_else(|| {
1687 invalid_plan(format!(
1688 "scratch resource `{resource_id}` has no selected storage profile"
1689 ))
1690 })?,
1691 workspace_layout_fingerprint.clone(),
1692 )?;
1693 dynamic_descriptors.push(DynamicResourceDescriptor::new(
1694 resource_id,
1695 workspace
1696 .size_formula
1697 .bind_runtime_limits(maximum_active_sequences, maximum_scheduled_tokens)?,
1698 workspace.alignment_bytes,
1699 BufferUsage::Scratch,
1700 ElementType::U8,
1701 AllocationLifetime::Invocation,
1702 AllocationKind::Scratch {
1703 node_id: node.id.clone(),
1704 },
1705 storage,
1706 StateInitialization::None,
1707 maximum_active_sequences,
1708 )?);
1709 } else if node.scratch_resource.is_some() {
1710 return Err(invalid_plan(format!(
1711 "node `{}` has scratch resources without a provider estimate",
1712 node.id
1713 )));
1714 }
1715 if let Some(workspace) = &node.provider_resources.binding {
1716 let resource_id = node.binding_resource.clone().ok_or_else(|| {
1717 invalid_plan(format!(
1718 "node `{}` binding workspace base identity is missing",
1719 node.id
1720 ))
1721 })?;
1722 if workspace.scope != ProviderWorkspaceScope::Invocation {
1723 return Err(invalid_plan(format!(
1724 "node `{}` binding workspace is not invocation scoped",
1725 node.id
1726 )));
1727 }
1728 let storage = DynamicStorageContract::new(
1729 *selected_resource_profiles
1730 .get(&resource_id)
1731 .ok_or_else(|| {
1732 invalid_plan(format!(
1733 "binding resource `{resource_id}` has no selected storage profile"
1734 ))
1735 })?,
1736 workspace_layout_fingerprint.clone(),
1737 )?;
1738 dynamic_descriptors.push(DynamicResourceDescriptor::new(
1739 resource_id,
1740 workspace
1741 .size_formula
1742 .bind_runtime_limits(maximum_active_sequences, maximum_scheduled_tokens)?,
1743 workspace.alignment_bytes,
1744 BufferUsage::Binding,
1745 ElementType::U8,
1746 AllocationLifetime::Invocation,
1747 AllocationKind::Binding {
1748 node_id: node.id.clone(),
1749 },
1750 storage,
1751 StateInitialization::None,
1752 maximum_active_sequences,
1753 )?);
1754 } else if node.binding_resource.is_some() {
1755 return Err(invalid_plan(format!(
1756 "node `{}` has binding resources without a provider estimate",
1757 node.id
1758 )));
1759 }
1760 if let Some(workspace) = &node.provider_resources.persistent {
1761 let resource_id = node.persistent_resource.clone().ok_or_else(|| {
1762 invalid_plan(format!(
1763 "node `{}` persistent base identity is missing",
1764 node.id
1765 ))
1766 })?;
1767 match workspace.scope {
1768 ProviderWorkspaceScope::Plan => {
1769 let bytes = workspace.fixed_bytes().ok_or_else(|| {
1770 invalid_plan(format!(
1771 "node `{}` plan workspace does not have a fixed formula",
1772 node.id
1773 ))
1774 })?;
1775 let storage = DynamicStorageContract::new(
1776 *selected_resource_profiles
1777 .get(&resource_id)
1778 .ok_or_else(|| {
1779 invalid_plan(format!(
1780 "plan workspace `{resource_id}` has no selected storage profile"
1781 ))
1782 })?,
1783 workspace_layout_fingerprint.clone(),
1784 )?;
1785 static_allocations.push(ResourceAllocation::new(
1786 resource_id,
1787 bytes,
1788 workspace.alignment_bytes,
1789 BufferUsage::Persistent,
1790 ElementType::U8,
1791 AllocationKind::Persistent {
1792 node_id: node.id.clone(),
1793 },
1794 storage,
1795 )?);
1796 }
1797 scope @ (ProviderWorkspaceScope::Request
1798 | ProviderWorkspaceScope::Sequence
1799 | ProviderWorkspaceScope::Step) => {
1800 let lifetime = match scope {
1801 ProviderWorkspaceScope::Request => AllocationLifetime::Request,
1802 ProviderWorkspaceScope::Sequence => AllocationLifetime::Sequence,
1803 ProviderWorkspaceScope::Step => AllocationLifetime::Step,
1804 ProviderWorkspaceScope::Plan | ProviderWorkspaceScope::Invocation => {
1805 unreachable!()
1806 }
1807 };
1808 let storage = DynamicStorageContract::new(
1809 *selected_resource_profiles.get(&resource_id).ok_or_else(|| {
1810 invalid_plan(format!(
1811 "persistent resource `{resource_id}` has no selected storage profile"
1812 ))
1813 })?,
1814 workspace_layout_fingerprint.clone(),
1815 )?;
1816 dynamic_descriptors.push(DynamicResourceDescriptor::new(
1817 resource_id,
1818 workspace.size_formula.bind_runtime_limits(
1819 maximum_active_sequences,
1820 maximum_scheduled_tokens,
1821 )?,
1822 workspace.alignment_bytes,
1823 BufferUsage::Persistent,
1824 ElementType::U8,
1825 lifetime,
1826 AllocationKind::Persistent {
1827 node_id: node.id.clone(),
1828 },
1829 storage,
1830 StateInitialization::None,
1831 maximum_active_sequences,
1832 )?);
1833 }
1834 ProviderWorkspaceScope::Invocation => {
1835 return Err(invalid_plan(format!(
1836 "node `{}` persistent workspace cannot be invocation scoped",
1837 node.id
1838 )));
1839 }
1840 }
1841 } else if node.persistent_resource.is_some() {
1842 return Err(invalid_plan(format!(
1843 "node `{}` has persistent resources without a provider estimate",
1844 node.id
1845 )));
1846 }
1847 }
1848 for (resource_id, accumulator) in values {
1849 accumulator.merge_result.ok_or_else(|| {
1850 invalid_plan(format!(
1851 "resource `{resource_id}` has conflicting usage, dtype, lifetime, or demand"
1852 ))
1853 })?;
1854 let logical_layout_fingerprint = canonical_fingerprint(
1855 &accumulator.logical_layout_fingerprints,
1856 "fingerprint dynamic resource tensor layout classes",
1857 )?;
1858 match accumulator.demand {
1859 ValueResourceDemand::PlanStatic => {
1860 let storage = DynamicStorageContract::new(
1861 static_contiguous_storage_profile()?,
1862 logical_layout_fingerprint,
1863 )?;
1864 static_allocations.push(ResourceAllocation::new(
1865 resource_id,
1866 accumulator.end_bytes,
1867 accumulator.alignment_bytes,
1868 accumulator.usage,
1869 accumulator.element_type,
1870 AllocationKind::Value,
1871 storage,
1872 )?);
1873 }
1874 demand => {
1875 let storage = DynamicStorageContract::new(
1876 *selected_resource_profiles.get(&resource_id).ok_or_else(|| {
1877 invalid_plan(format!(
1878 "dynamic value resource `{resource_id}` has no selected storage profile"
1879 ))
1880 })?,
1881 logical_layout_fingerprint,
1882 )?;
1883 dynamic_descriptors.push(DynamicResourceDescriptor::new(
1884 resource_id,
1885 demand
1886 .dynamic_demand(accumulator.end_bytes, accumulator.alignment_bytes)?,
1887 accumulator.alignment_bytes,
1888 accumulator.usage,
1889 accumulator.element_type,
1890 demand.lifetime().ok_or_else(|| {
1891 invalid_plan("dynamic value demand lost its scoped lifetime")
1892 })?,
1893 AllocationKind::Value,
1894 storage,
1895 accumulator.initialization,
1896 maximum_active_sequences,
1897 )?);
1898 }
1899 }
1900 }
1901 MemoryPlan::from_core_with_completion_retention(
1902 device_capacity_bytes,
1903 policy_capacity_bytes,
1904 reserve_bytes,
1905 maximum_active_sequences,
1906 static_allocations,
1907 dynamic_descriptors,
1908 nodes,
1909 reusable_execution_policy,
1910 retained_completion_resources,
1911 )
1912 }
1913
1914 pub(super) fn value_resource_demand(
1915 family: &PreparedModelFamily,
1916 value_id: &ProgramValueId,
1917 usage: BufferUsage,
1918 minimum_bytes: u64,
1919 token_projection: Option<(u64, u64)>,
1920 maximum_active_sequences: u32,
1921 maximum_scheduled_tokens: u64,
1922 is_product_io_resource: bool,
1923 ) -> Result<ValueResourceDemand, VNextError> {
1924 if family
1925 .program()
1926 .weights()
1927 .iter()
1928 .any(|weight| &weight.value_id == value_id)
1929 {
1930 return Ok(ValueResourceDemand::PlanStatic);
1931 }
1932 let state = family
1933 .program()
1934 .states()
1935 .iter()
1936 .find(|state| &state.value_id == value_id);
1937 let Some(state) = state else {
1938 if usage != BufferUsage::Activations {
1939 return Err(invalid_plan(format!(
1940 "non-state value `{value_id}` is not backed by activation memory"
1941 )));
1942 }
1943 let lifetime = AllocationLifetime::Step;
1944 if let Some((bytes_per_token, canonical_tokens)) = token_projection {
1945 if bytes_per_token == 0 || canonical_tokens == 0 {
1946 return Err(invalid_plan(
1947 "token-scaled activation has zero bytes or canonical tokens",
1948 ));
1949 }
1950 let maximum_tokens = Self::activation_token_capacity(
1951 lifetime,
1952 canonical_tokens,
1953 maximum_scheduled_tokens,
1954 )?;
1955 return Ok(ValueResourceDemand::TokenScaled {
1956 lifetime,
1957 bytes_per_token,
1958 maximum_tokens,
1959 });
1960 }
1961 if is_product_io_resource {
1962 return Ok(ValueResourceDemand::ParticipantFixed {
1963 lifetime,
1964 maximum_participants: maximum_active_sequences,
1965 });
1966 }
1967 return Ok(ValueResourceDemand::Fixed { lifetime });
1968 };
1969 let lifetime = match state.lifetime {
1970 StateLifetime::Request => AllocationLifetime::Request,
1971 StateLifetime::Sequence => AllocationLifetime::Sequence,
1972 StateLifetime::Step => AllocationLifetime::Step,
1973 };
1974 state.capacity_demand.validate(state.tensor.byte_len()?)?;
1975 match state.capacity_demand {
1976 StateCapacityDemand::FixedPerScope => Ok(ValueResourceDemand::Fixed { lifetime }),
1977 StateCapacityDemand::TokenScaled {
1978 bytes_per_token,
1979 maximum_tokens,
1980 } => {
1981 if bytes_per_token < minimum_bytes {
1982 return Err(invalid_plan(
1983 "token-scaled state demand is smaller than its resolved resource range",
1984 ));
1985 }
1986 Ok(ValueResourceDemand::TokenScaled {
1987 lifetime,
1988 bytes_per_token,
1989 maximum_tokens,
1990 })
1991 }
1992 }
1993 }
1994
1995 pub(super) fn activation_token_capacity(
1996 lifetime: AllocationLifetime,
1997 canonical_tokens: u64,
1998 maximum_scheduled_tokens: u64,
1999 ) -> Result<u64, VNextError> {
2000 if canonical_tokens == 0 {
2001 return Err(invalid_plan(
2002 "token-scaled activation has zero canonical tokens",
2003 ));
2004 }
2005 if lifetime == AllocationLifetime::Request {
2006 Ok(canonical_tokens)
2007 } else {
2008 validate_scheduled_token_ceiling(maximum_scheduled_tokens)?;
2009 Ok(maximum_scheduled_tokens)
2010 }
2011 }
2012
2013 pub(super) fn plan_id_for_hash(hash: &PlanHash) -> Result<PlanId, VNextError> {
2014 PlanId::new(format!("plan/sha256/{}", hash.as_str()))
2015 }
2016
2017 pub(super) fn validate_internal(&self) -> Result<(), VNextError> {
2018 if self.payload.schema != EXECUTION_PLAN_SCHEMA {
2019 return Err(VNextError::UnsupportedPlanSchema {
2020 expected_major: EXECUTION_PLAN_SCHEMA.major,
2021 expected_minor: EXECUTION_PLAN_SCHEMA.minor,
2022 actual_major: self.payload.schema.major,
2023 actual_minor: self.payload.schema.minor,
2024 });
2025 }
2026 let computed = PlanHash::new(canonical_fingerprint(
2027 &PlanHashMaterial::from(&self.payload),
2028 "validate execution plan hash",
2029 )?)?;
2030 if computed != self.plan_hash {
2031 return Err(VNextError::PlanHashMismatch {
2032 expected: computed.to_string(),
2033 actual: self.plan_hash.to_string(),
2034 });
2035 }
2036 if self.payload.plan_id != Self::plan_id_for_hash(&computed)? {
2037 return Err(invalid_plan(
2038 "plan id is not derived from the semantic plan hash",
2039 ));
2040 }
2041 if self.payload.nodes.is_empty()
2042 || !is_canonical_sha256(&self.payload.prepared_family_fingerprint)
2043 || !is_canonical_sha256(&self.payload.program_fingerprint)
2044 || !is_canonical_sha256(&self.payload.capability_catalog_fingerprint)
2045 || !is_canonical_sha256(&self.payload.device_runtime_implementation_fingerprint)
2046 || !is_canonical_sha256(&self.payload.policy_fingerprint)
2047 || self.payload.maximum_scheduled_tokens == 0
2048 {
2049 return Err(invalid_plan("plan provenance or node set is invalid"));
2050 }
2051 self.payload
2052 .execution_weights
2053 .validate_structure(&self.payload.family_id)?;
2054 if &self.payload.execution_weights != self.trusted_execution_weights.plan()
2055 || self.payload.weight_format != self.payload.execution_weights.schema().format_id
2056 || self.payload.quantization_formats
2057 != self
2058 .payload
2059 .execution_weights
2060 .schema()
2061 .quantization_formats()
2062 {
2063 return Err(invalid_plan(
2064 "execution weight summary differs from the execution weight plan",
2065 ));
2066 }
2067 let retention_spec = CompletionRetentionSpec::new(
2068 self.payload
2069 .retained_completion_values
2070 .iter()
2071 .map(|value| value.value_id().clone())
2072 .collect(),
2073 );
2074 let expected_retained_completion_values =
2075 resolve_retained_completion_values(&self.payload.nodes, &retention_spec)?;
2076 if self.payload.retained_completion_values != expected_retained_completion_values {
2077 return Err(invalid_plan(
2078 "retained completion values are not derived from plan outputs",
2079 ));
2080 }
2081 if self.payload.terminal_output_resources.is_empty()
2082 || self
2083 .payload
2084 .terminal_output_resources
2085 .windows(2)
2086 .any(|pair| pair[0] >= pair[1])
2087 {
2088 return Err(invalid_plan(
2089 "terminal output resource evidence is empty or non-canonical",
2090 ));
2091 }
2092 let retained_completion_resources = self
2093 .payload
2094 .retained_completion_values
2095 .iter()
2096 .map(|value| value.resource_id().clone())
2097 .chain(self.payload.terminal_output_resources.iter().cloned())
2098 .collect::<BTreeSet<_>>();
2099 self.payload.memory.validate()?;
2100 let dynamic_capacity_bytes = self
2101 .payload
2102 .memory
2103 .usable_capacity_bytes
2104 .checked_sub(self.payload.memory.static_bytes)
2105 .ok_or_else(|| invalid_plan("static memory exceeds usable capacity"))?;
2106 let base_pools = MemoryPlan::derive_dynamic_pools_with_completion_retention(
2107 &self.payload.memory.dynamic_descriptors,
2108 &self.payload.nodes,
2109 dynamic_capacity_bytes,
2110 &retained_completion_resources,
2111 )?;
2112 let expected_reusable_execution = self
2113 .payload
2114 .memory
2115 .reusable_execution
2116 .as_ref()
2117 .map(|actual| {
2118 MemoryPlan::derive_reusable_execution(
2119 &actual.policy()?,
2120 self.payload.nodes.len(),
2121 &self.payload.memory.dynamic_descriptors,
2122 &base_pools,
2123 )
2124 })
2125 .transpose()?;
2126 if self.payload.memory.reusable_execution != expected_reusable_execution {
2127 return Err(invalid_plan(
2128 "reusable execution budgets are not derived from plan resources",
2129 ));
2130 }
2131 let reusable_workspace_ceilings = expected_reusable_execution
2132 .as_ref()
2133 .map(ReusableExecutionMemoryPlan::pool_workspace_ceilings)
2134 .transpose()?
2135 .unwrap_or_default();
2136 let checkpoint_growth_ceilings =
2137 super::checkpoint_capacity::derive_checkpoint_growth_ceilings(
2138 self.payload.memory.checkpoint_capacity.as_ref(),
2139 self.payload.sequence_checkpoint_layout.as_ref(),
2140 &self.payload.memory.dynamic_descriptors,
2141 )?;
2142 let expected_pools = MemoryPlan::derive_dynamic_pools_with_checkpoint(
2143 &self.payload.memory.dynamic_descriptors,
2144 &self.payload.nodes,
2145 dynamic_capacity_bytes,
2146 &reusable_workspace_ceilings,
2147 &retained_completion_resources,
2148 &checkpoint_growth_ceilings,
2149 )?;
2150 if self.payload.memory.dynamic_pools != expected_pools {
2151 return Err(invalid_plan(
2152 "memory pools or invocation reuse are not derived from plan dependencies",
2153 ));
2154 }
2155 let static_allocations = self
2156 .payload
2157 .memory
2158 .static_allocations
2159 .iter()
2160 .map(|allocation| (allocation.resource_id.clone(), allocation))
2161 .collect::<BTreeMap<_, _>>();
2162 let dynamic_descriptors = self
2163 .payload
2164 .memory
2165 .dynamic_descriptors
2166 .iter()
2167 .map(|descriptor| (descriptor.base_resource_id.clone(), descriptor))
2168 .collect::<BTreeMap<_, _>>();
2169 let initialization_scratch = static_allocations
2170 .values()
2171 .filter(|allocation| allocation.kind == AllocationKind::InitializationScratch)
2172 .collect::<Vec<_>>();
2173 match self
2174 .payload
2175 .execution_weights
2176 .static_weight_transform_scratch_resource_id()?
2177 {
2178 Some(expected_id) => {
2179 let [allocation] = initialization_scratch.as_slice() else {
2180 return Err(invalid_plan(
2181 "static weight transforms require exactly one initialization scratch allocation",
2182 ));
2183 };
2184 if allocation.resource_id != expected_id
2185 || allocation.per_instance_bytes
2186 != self
2187 .payload
2188 .execution_weights
2189 .maximum_static_weight_transform_scratch_bytes()?
2190 || allocation.alignment_bytes != STATIC_WEIGHT_TRANSFORM_SCRATCH_ALIGNMENT_BYTES
2191 || allocation.usage != BufferUsage::Scratch
2192 || allocation.element_type != ElementType::U8
2193 {
2194 return Err(invalid_plan(
2195 "static weight transform scratch differs from its execution weight plan",
2196 ));
2197 }
2198 }
2199 None if initialization_scratch.is_empty() => {}
2200 None => {
2201 return Err(invalid_plan(
2202 "initialization scratch exists without a static weight transform",
2203 ));
2204 }
2205 }
2206 let mut seen_nodes = BTreeSet::new();
2207 let mut canonical_values = BTreeMap::new();
2208 for node in &self.payload.nodes {
2209 node.provider_resources.validate_shape()?;
2210 Self::validate_provider_selection_evidence(&node.selection)?;
2211 Self::validate_node_work_contract(node)?;
2212 if !seen_nodes.insert(node.id.clone())
2213 || !is_canonical_sha256(&node.provider_implementation_fingerprint)
2214 || node
2215 .dependencies
2216 .iter()
2217 .any(|dependency| dependency == &node.id || !seen_nodes.contains(dependency))
2218 || node.dependencies.windows(2).any(|pair| pair[0] >= pair[1])
2219 || node
2220 .state_effects
2221 .windows(2)
2222 .any(|pair| pair[0].state_id >= pair[1].state_id)
2223 || node.resources.iter().collect::<BTreeSet<_>>().len() != node.resources.len()
2224 || node.resources.iter().any(|resource| {
2225 !static_allocations.contains_key(resource)
2226 && !dynamic_descriptors.contains_key(resource)
2227 })
2228 || node.provider_resources.provider_id != node.selection.selected_provider
2229 {
2230 return Err(invalid_plan(format!(
2231 "node `{}` identity, dependency, or resource closure is invalid",
2232 node.id
2233 )));
2234 }
2235 let expected_resources = node
2236 .values
2237 .iter()
2238 .flat_map(|binding| binding.storage().components())
2239 .map(|component| component.resource_id().clone())
2240 .chain(node.scratch_resource.iter().cloned())
2241 .chain(node.binding_resource.iter().cloned())
2242 .chain(node.persistent_resource.iter().cloned())
2243 .collect::<BTreeSet<_>>()
2244 .into_iter()
2245 .collect::<Vec<_>>();
2246 if node.resources != expected_resources {
2247 return Err(invalid_plan(format!(
2248 "node `{}` resource closure is not canonical",
2249 node.id
2250 )));
2251 }
2252 for effect in &node.state_effects {
2253 if !matches!(
2254 effect.lifetime,
2255 AllocationLifetime::Request
2256 | AllocationLifetime::Sequence
2257 | AllocationLifetime::Step
2258 ) || effect.resource_ids.is_empty()
2259 || effect
2260 .resource_ids
2261 .windows(2)
2262 .any(|pair| pair[0] >= pair[1])
2263 {
2264 return Err(invalid_plan(format!(
2265 "node `{}` state effect has an invalid lifetime or resource closure",
2266 node.id
2267 )));
2268 }
2269 let matching = node
2270 .values
2271 .iter()
2272 .filter(|binding| binding.value_id() == &effect.state_value_id)
2273 .collect::<Vec<_>>();
2274 let expected_effect_resources = matching
2275 .iter()
2276 .flat_map(|binding| binding.storage().components())
2277 .map(|component| component.resource_id().clone())
2278 .collect::<BTreeSet<_>>()
2279 .into_iter()
2280 .collect::<Vec<_>>();
2281 let reads = matching.iter().any(|binding| {
2282 matches!(
2283 binding.access(),
2284 TensorAccess::Read | TensorAccess::ReadWrite
2285 )
2286 });
2287 let writes = matching.iter().any(|binding| {
2288 matches!(
2289 binding.access(),
2290 TensorAccess::Write | TensorAccess::ReadWrite
2291 )
2292 });
2293 let expected_access = match (reads, writes) {
2294 (true, false) => Some(TensorAccess::Read),
2295 (false, true) => Some(TensorAccess::Write),
2296 (true, true) => Some(TensorAccess::ReadWrite),
2297 (false, false) => None,
2298 };
2299 if effect.resource_ids != expected_effect_resources
2300 || expected_access != Some(effect.access)
2301 || effect.resource_ids.iter().any(|resource_id| {
2302 dynamic_descriptors
2303 .get(resource_id)
2304 .is_none_or(|descriptor| descriptor.lifetime != effect.lifetime)
2305 })
2306 {
2307 return Err(invalid_plan(format!(
2308 "node `{}` state effect is not derived from its typed bindings",
2309 node.id
2310 )));
2311 }
2312 }
2313 if node.scratch_resource.is_some() != node.provider_resources.scratch.is_some()
2314 || node.binding_resource.is_some() != node.provider_resources.binding.is_some()
2315 || node.persistent_resource.is_some()
2316 != node.provider_resources.persistent.is_some()
2317 {
2318 return Err(invalid_plan(format!(
2319 "node `{}` workspace base identity presence differs from its provider estimate",
2320 node.id
2321 )));
2322 }
2323 if let Some(resource_id) = &node.scratch_resource {
2324 let descriptor = dynamic_descriptors.get(resource_id).ok_or_else(|| {
2325 invalid_plan(format!("node `{}` scratch descriptor is missing", node.id))
2326 })?;
2327 let workspace = node.provider_resources.scratch.as_ref().ok_or_else(|| {
2328 invalid_plan(format!("node `{}` scratch estimate is missing", node.id))
2329 })?;
2330 if descriptor.demand
2331 != workspace.size_formula.bind_runtime_limits(
2332 self.payload.memory.maximum_active_sequences,
2333 self.payload.maximum_scheduled_tokens,
2334 )?
2335 || descriptor.alignment_bytes != workspace.alignment_bytes
2336 || descriptor.usage != BufferUsage::Scratch
2337 || descriptor.lifetime != AllocationLifetime::Invocation
2338 || descriptor.theoretical_maximum_instances
2339 != self.payload.memory.maximum_active_sequences
2340 || descriptor.kind
2341 != (AllocationKind::Scratch {
2342 node_id: node.id.clone(),
2343 })
2344 {
2345 return Err(invalid_plan(format!(
2346 "node `{}` scratch descriptor differs from its provider estimate",
2347 node.id
2348 )));
2349 }
2350 }
2351 if let Some(resource_id) = &node.binding_resource {
2352 let descriptor = dynamic_descriptors.get(resource_id).ok_or_else(|| {
2353 invalid_plan(format!("node `{}` binding descriptor is missing", node.id))
2354 })?;
2355 let workspace = node.provider_resources.binding.as_ref().ok_or_else(|| {
2356 invalid_plan(format!("node `{}` binding estimate is missing", node.id))
2357 })?;
2358 if descriptor.demand
2359 != workspace.size_formula.bind_runtime_limits(
2360 self.payload.memory.maximum_active_sequences,
2361 self.payload.maximum_scheduled_tokens,
2362 )?
2363 || descriptor.alignment_bytes != workspace.alignment_bytes
2364 || descriptor.usage != BufferUsage::Binding
2365 || descriptor.lifetime != AllocationLifetime::Invocation
2366 || descriptor.theoretical_maximum_instances
2367 != self.payload.memory.maximum_active_sequences
2368 || descriptor.kind
2369 != (AllocationKind::Binding {
2370 node_id: node.id.clone(),
2371 })
2372 {
2373 return Err(invalid_plan(format!(
2374 "node `{}` binding descriptor differs from its provider estimate",
2375 node.id
2376 )));
2377 }
2378 }
2379 if let Some(resource_id) = &node.persistent_resource {
2380 let workspace = node.provider_resources.persistent.as_ref().ok_or_else(|| {
2381 invalid_plan(format!("node `{}` persistent estimate is missing", node.id))
2382 })?;
2383 match workspace.scope {
2384 ProviderWorkspaceScope::Plan => {
2385 let allocation = static_allocations.get(resource_id).ok_or_else(|| {
2386 invalid_plan(format!(
2387 "node `{}` plan-static persistent allocation is missing",
2388 node.id
2389 ))
2390 })?;
2391 if Some(allocation.per_instance_bytes) != workspace.fixed_bytes()
2392 || allocation.alignment_bytes != workspace.alignment_bytes
2393 || allocation.usage != BufferUsage::Persistent
2394 || !workspace.storage.accepts(allocation.storage.profile())
2395 || allocation.storage.logical_layout_fingerprint()
2396 != workspace_storage_layout_fingerprint()?
2397 || allocation.kind
2398 != (AllocationKind::Persistent {
2399 node_id: node.id.clone(),
2400 })
2401 {
2402 return Err(invalid_plan(format!(
2403 "node `{}` plan-static persistent allocation differs from its provider estimate",
2404 node.id
2405 )));
2406 }
2407 }
2408 scope @ (ProviderWorkspaceScope::Request
2409 | ProviderWorkspaceScope::Sequence
2410 | ProviderWorkspaceScope::Step) => {
2411 let expected_lifetime = match scope {
2412 ProviderWorkspaceScope::Request => AllocationLifetime::Request,
2413 ProviderWorkspaceScope::Sequence => AllocationLifetime::Sequence,
2414 ProviderWorkspaceScope::Step => AllocationLifetime::Step,
2415 ProviderWorkspaceScope::Plan | ProviderWorkspaceScope::Invocation => {
2416 unreachable!()
2417 }
2418 };
2419 let descriptor = dynamic_descriptors.get(resource_id).ok_or_else(|| {
2420 invalid_plan(format!(
2421 "node `{}` dynamic persistent descriptor is missing",
2422 node.id
2423 ))
2424 })?;
2425 if descriptor.demand
2426 != workspace.size_formula.bind_runtime_limits(
2427 self.payload.memory.maximum_active_sequences,
2428 self.payload.maximum_scheduled_tokens,
2429 )?
2430 || descriptor.alignment_bytes != workspace.alignment_bytes
2431 || descriptor.usage != BufferUsage::Persistent
2432 || descriptor.lifetime != expected_lifetime
2433 || descriptor.theoretical_maximum_instances
2434 != self.payload.memory.maximum_active_sequences
2435 || descriptor.kind
2436 != (AllocationKind::Persistent {
2437 node_id: node.id.clone(),
2438 })
2439 {
2440 return Err(invalid_plan(format!(
2441 "node `{}` dynamic persistent descriptor differs from its provider estimate",
2442 node.id
2443 )));
2444 }
2445 }
2446 ProviderWorkspaceScope::Invocation => {
2447 return Err(invalid_plan(format!(
2448 "node `{}` persistent workspace cannot be invocation scoped",
2449 node.id
2450 )));
2451 }
2452 }
2453 }
2454 for binding in &node.values {
2455 Self::validate_cross_node_value(binding, &mut canonical_values)?;
2456 }
2457 }
2458 Self::validate_global_storage_aliasing(&canonical_values, &self.payload.nodes)?;
2459 Ok(())
2460 }
2461
2462 pub fn payload(&self) -> &ExecutionPlanPayload {
2463 &self.payload
2464 }
2465
2466 pub fn completion_checkpoint(
2467 &self,
2468 value_id: &ProgramValueId,
2469 ) -> Result<&RetainedCompletionValue, VNextError> {
2470 self.payload
2471 .retained_completion_values
2472 .binary_search_by(|value| value.value_id().cmp(value_id))
2473 .map(|index| &self.payload.retained_completion_values[index])
2474 .map_err(|_| {
2475 invalid_plan(format!(
2476 "semantic value `{value_id}` is not retained for completion readback"
2477 ))
2478 })
2479 }
2480
2481 pub fn completion_checkpoint_readback_for_work(
2485 &self,
2486 value_id: &ProgramValueId,
2487 participant_index: u32,
2488 work: &ResourceWorkShape,
2489 ) -> Result<CompletionReadbackRequest, VNextError> {
2490 let checkpoint = self.completion_checkpoint(value_id)?;
2491 if checkpoint.logical_offset_bytes() != 0 {
2492 return Err(invalid_plan(
2493 "work-shaped completion readback requires a whole-resource activation",
2494 ));
2495 }
2496 let descriptor = self
2497 .payload
2498 .memory
2499 .dynamic_descriptors()
2500 .iter()
2501 .find(|descriptor| descriptor.base_resource_id() == checkpoint.resource_id())
2502 .ok_or_else(|| {
2503 invalid_plan(format!(
2504 "retained completion resource `{}` has no dynamic descriptor",
2505 checkpoint.resource_id()
2506 ))
2507 })?;
2508 if descriptor.element_type() != checkpoint.tensor().element_type() {
2509 return Err(invalid_plan(
2510 "retained completion resource element type differs from its tensor",
2511 ));
2512 }
2513 let byte_len = match descriptor.demand() {
2514 DynamicResourceDemand::ActualSequences { .. } => {
2517 checkpoint.tensor().minimum_storage_bytes()?
2518 }
2519 _ => descriptor.evaluate_logical_request_bytes(work)?,
2520 };
2521 let element_bytes = checkpoint.tensor().element_type().size_bytes();
2522 if byte_len % element_bytes != 0 {
2523 return Err(invalid_plan(
2524 "retained completion byte extent is not element aligned",
2525 ));
2526 }
2527 checkpoint.readback_request(
2528 participant_index,
2529 HostTransferLayout::new(checkpoint.tensor().element_type(), byte_len / element_bytes)?,
2530 )
2531 }
2532
2533 pub fn plan_hash(&self) -> &PlanHash {
2534 &self.plan_hash
2535 }
2536
2537 pub(crate) fn operation_registry_authority(&self) -> &OperationRegistryAuthority {
2538 &self.operation_registry_authority
2539 }
2540
2541 pub(crate) fn materialize_weight_components<'source>(
2542 &self,
2543 family: &PreparedModelFamily,
2544 source: &'source dyn WeightComponentSource,
2545 components: &[&WeightComponentSpec],
2546 ) -> Result<Vec<WeightComponentPayload<'source>>, VNextError> {
2547 self.trusted_execution_weights
2548 .materialize_components(family, source, components)
2549 }
2550
2551 pub(crate) fn static_weight_transform_for_components(
2552 &self,
2553 components: &[&WeightComponentSpec],
2554 ) -> Result<Option<&StaticWeightTransformPlan>, VNextError> {
2555 self.trusted_execution_weights
2556 .static_weight_transform_for_components(components)
2557 }
2558}