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