#![forbid(unsafe_code)]
mod artifact;
mod candidate;
pub mod cost;
mod envelope;
mod facts;
pub mod legality;
mod normalize;
mod search;
mod select;
pub mod target;
pub use envelope::{
ArtifactEnvelope, TargetEntryPoint, TargetPayload, TargetPayloadFormat, TargetProfile,
TargetResourceAccess, TargetResourceBinding, TargetResourceMemory,
ARTIFACT_ENVELOPE_SCHEMA_VERSION, TARGET_PAYLOAD_SCHEMA_VERSION,
};
pub use target::{
attach_target, compile_selected_modules, EmittedTargetModule, SelectedLowering,
TargetCompileError, TargetCompiler, TargetModuleBundle, TargetModuleImage,
TARGET_MODULE_BUNDLE_SCHEMA_VERSION,
};
use std::collections::{BTreeMap, BTreeSet};
use serde::{Deserialize, Serialize};
use thiserror::Error;
pub use vyre_foundation::diagnostics::Diagnostic;
use vyre_foundation::diagnostics::{DiagnosticStage, OpLocation, RetryClass};
use vyre_foundation::ir::{
BufferAccess, DataType, GraphValueId, ProgramGraph, ShapeDim, ValueLifetime,
};
use vyre_foundation::validate::{validate_with_options, BackendCapabilities, ValidationOptions};
pub const ARTIFACT_SCHEMA_VERSION: u16 = 4;
const ARTIFACT_MAGIC: &[u8; 4] = b"VMK0";
const ARTIFACT_HEADER_BYTES: usize = 10;
const ARTIFACT_DIGEST_BYTES: usize = 32;
const ARTIFACT_DIGEST_DOMAIN: &[u8] = b"vyre-megakernel-artifact-v4\0";
const SOURCE_DIGEST_DOMAIN: &[u8] = b"vyre-megakernel-source-v2\0";
const REQUEST_DIGEST_DOMAIN: &[u8] = b"vyre-megakernel-request-v2\0";
const COMPILER_IR_CAPABILITIES: BackendCapabilities = BackendCapabilities {
supports_subgroup_ops: true,
supports_indirect_dispatch: true,
supports_specialization_constants: true,
supports_distributed_collectives: true,
has_mul_high: true,
has_dual_issue_fp32_int32: true,
has_tensor_core_int: true,
has_native_f16: true,
has_warp_shuffle: true,
has_shared_memory: true,
has_transcendental_polynomial_emit: true,
max_native_int_width: u32::MAX,
};
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize)]
pub struct Digest(pub [u8; 32]);
impl Digest {
#[must_use]
pub const fn as_bytes(&self) -> &[u8; 32] {
&self.0
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize)]
pub struct ArtifactNodeId(pub u32);
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize)]
pub struct ArtifactValueId(pub u32);
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize)]
pub struct FusionGroupId(pub u32);
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum DependencyEndpoint {
Node(ArtifactNodeId),
Value(ArtifactValueId),
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum DependencyKind {
Data,
Retained,
Materialization,
}
#[derive(Clone, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize)]
pub struct DependencyEdge {
pub from: DependencyEndpoint,
pub to: DependencyEndpoint,
pub kind: DependencyKind,
pub value: Option<ArtifactValueId>,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct SearchBudget {
pub max_candidates: u32,
pub max_cpu_work: u64,
pub max_target_compilations: u32,
pub max_measurements: u32,
pub max_elapsed_ns: u64,
}
impl SearchBudget {
#[must_use]
pub const fn new(
max_candidates: u32,
max_cpu_work: u64,
max_target_compilations: u32,
max_measurements: u32,
max_elapsed_ns: u64,
) -> Self {
Self {
max_candidates,
max_cpu_work,
max_target_compilations,
max_measurements,
max_elapsed_ns,
}
}
}
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct SearchWork {
pub candidates_explored: u32,
pub cpu_work: u64,
pub target_compilations: u32,
pub measurements: u32,
pub elapsed_ns: u64,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct ExternalFacts {
pub configuration_digest: Digest,
pub symbolic_bindings: BTreeMap<String, u64>,
pub constant_identities: BTreeMap<GraphValueId, Digest>,
}
impl ExternalFacts {
#[must_use]
pub fn new(configuration_digest: Digest, symbolic_bindings: BTreeMap<String, u64>) -> Self {
Self {
configuration_digest,
symbolic_bindings,
constant_identities: BTreeMap::new(),
}
}
}
pub struct CompileRequest {
graph: ProgramGraph,
facts: ExternalFacts,
search_budget: SearchBudget,
max_artifact_bytes: u64,
}
impl CompileRequest {
#[must_use]
pub const fn new(
graph: ProgramGraph,
facts: ExternalFacts,
search_budget: SearchBudget,
max_artifact_bytes: u64,
) -> Self {
Self {
graph,
facts,
search_budget,
max_artifact_bytes,
}
}
pub fn validate(self) -> Result<ValidatedCompileRequest, CompileError> {
if self.max_artifact_bytes == 0 {
return Err(failure(
CompilerFailureKind::ArtifactLimit,
"request.max_artifact_bytes",
"artifact byte limit must be greater than zero",
"supply a positive bounded artifact byte limit",
));
}
if self.search_budget.max_candidates == 0
|| self.search_budget.max_cpu_work == 0
|| self.search_budget.max_elapsed_ns == 0
{
return Err(failure(
CompilerFailureKind::InvalidSearchBudget,
"request.search_budget",
"candidate, CPU-work, and elapsed-work bounds must be positive",
"supply explicit positive bounds for every mandatory search dimension",
));
}
self.graph.analyze().map_err(|error| {
failure(
CompilerFailureKind::InvalidProgram,
"request.graph",
error.to_string(),
"supply a structurally valid acyclic ProgramGraph",
)
})?;
for node in self.graph.nodes() {
let report = validate_with_options(
&node.program,
ValidationOptions::universal().with_backend_capabilities(COMPILER_IR_CAPABILITIES),
);
if let Some(issue) = report.errors.into_iter().next() {
let path = format!("request.graph.nodes[{}].program", node.id.0);
let mut diagnostic = issue.diagnostic();
if let Some(location) = diagnostic.location.as_mut() {
location.path = Some(path);
location.graph_node = Some(node.id.0);
}
return Err(CompileError { diagnostic });
}
}
validate_bindings(&self.graph, &self.facts.symbolic_bindings)?;
validate_constant_identities(&self.graph, &self.facts.constant_identities)?;
Ok(ValidatedCompileRequest {
graph: self.graph,
facts: self.facts,
search_budget: self.search_budget,
max_artifact_bytes: self.max_artifact_bytes,
})
}
}
pub struct ValidatedCompileRequest {
graph: ProgramGraph,
facts: ExternalFacts,
search_budget: SearchBudget,
max_artifact_bytes: u64,
}
impl ValidatedCompileRequest {
#[must_use]
pub const fn graph(&self) -> &ProgramGraph {
&self.graph
}
#[must_use]
pub const fn facts(&self) -> &ExternalFacts {
&self.facts
}
#[must_use]
pub const fn search_budget(&self) -> SearchBudget {
self.search_budget
}
#[must_use]
pub const fn max_artifact_bytes(&self) -> u64 {
self.max_artifact_bytes
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub(crate) enum CompilerFailureKind {
InvalidProgram,
MissingSymbol,
UnknownSymbol,
DependencyCycle,
ResourceOverflow,
UnsizedResource,
ArtifactLimit,
MalformedArtifact,
VersionSkew,
DigestMismatch,
MalformedTargetPayload,
TargetPayloadVersionSkew,
TargetPayloadDigestMismatch,
TargetPayloadAssociationMismatch,
IncompatibleTargetPayload,
InvalidSearchBudget,
MissingConstantIdentity,
UnknownConstantIdentity,
}
impl CompilerFailureKind {
#[must_use]
const fn as_str(self) -> &'static str {
match self {
Self::InvalidProgram => "MKC001_INVALID_PROGRAM",
Self::MissingSymbol => "MKC002_MISSING_SYMBOL",
Self::UnknownSymbol => "MKC003_UNKNOWN_SYMBOL",
Self::DependencyCycle => "MKC010_DEPENDENCY_CYCLE",
Self::ResourceOverflow => "MKC011_RESOURCE_OVERFLOW",
Self::UnsizedResource => "MKC012_UNSIZED_RESOURCE",
Self::ArtifactLimit => "MKC013_ARTIFACT_LIMIT",
Self::MalformedArtifact => "MKC014_MALFORMED_ARTIFACT",
Self::VersionSkew => "MKC015_VERSION_SKEW",
Self::DigestMismatch => "MKC016_DIGEST_MISMATCH",
Self::MalformedTargetPayload => "MKC017_MALFORMED_TARGET_PAYLOAD",
Self::TargetPayloadVersionSkew => "MKC018_TARGET_PAYLOAD_VERSION_SKEW",
Self::TargetPayloadDigestMismatch => "MKC019_TARGET_PAYLOAD_DIGEST_MISMATCH",
Self::TargetPayloadAssociationMismatch => "MKC020_TARGET_PAYLOAD_ASSOCIATION_MISMATCH",
Self::IncompatibleTargetPayload => "MKC021_INCOMPATIBLE_TARGET_PAYLOAD",
Self::InvalidSearchBudget => "MKC022_INVALID_SEARCH_BUDGET",
Self::MissingConstantIdentity => "MKC023_MISSING_CONSTANT_IDENTITY",
Self::UnknownConstantIdentity => "MKC024_UNKNOWN_CONSTANT_IDENTITY",
}
}
}
const fn diagnostic_stage(code: CompilerFailureKind) -> DiagnosticStage {
match code {
CompilerFailureKind::InvalidProgram
| CompilerFailureKind::MissingSymbol
| CompilerFailureKind::UnknownSymbol
| CompilerFailureKind::InvalidSearchBudget
| CompilerFailureKind::MissingConstantIdentity
| CompilerFailureKind::UnknownConstantIdentity => DiagnosticStage::Validate,
CompilerFailureKind::DependencyCycle => DiagnosticStage::Plan,
CompilerFailureKind::ResourceOverflow | CompilerFailureKind::UnsizedResource => {
DiagnosticStage::Lower
}
CompilerFailureKind::ArtifactLimit => DiagnosticStage::Emit,
CompilerFailureKind::MalformedArtifact
| CompilerFailureKind::VersionSkew
| CompilerFailureKind::DigestMismatch
| CompilerFailureKind::MalformedTargetPayload
| CompilerFailureKind::TargetPayloadVersionSkew
| CompilerFailureKind::TargetPayloadDigestMismatch
| CompilerFailureKind::TargetPayloadAssociationMismatch
| CompilerFailureKind::IncompatibleTargetPayload => DiagnosticStage::Admit,
}
}
const fn diagnostic_retry(code: CompilerFailureKind) -> RetryClass {
match code {
CompilerFailureKind::VersionSkew
| CompilerFailureKind::TargetPayloadVersionSkew
| CompilerFailureKind::IncompatibleTargetPayload => RetryClass::RecompileSource,
_ => RetryClass::Never,
}
}
#[derive(Clone, Debug, PartialEq, Eq, Error)]
#[error("{diagnostic}")]
pub struct CompileError {
pub diagnostic: Diagnostic,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct NodeRecord {
pub id: ArtifactNodeId,
pub name: String,
pub program: Vec<u8>,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct GeometryRecord {
pub node: ArtifactNodeId,
pub workgroup_size: [u32; 3],
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ResourceLifetime {
Constant,
Invocation,
Retained,
Output,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct ResourceRecord {
pub value: ArtifactValueId,
pub name: String,
pub element_count: u64,
pub byte_count: u64,
pub lifetime: ResourceLifetime,
pub first_stage: u32,
pub last_stage: u32,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct ResourceEnvelope {
pub total_bytes: u64,
pub peak_live_bytes: u64,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum AbiAccess {
ReadOnly,
WriteOnly,
ReadWrite,
Uniform,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct ResourceAbiRecord {
pub slot: u32,
pub value: ArtifactValueId,
pub dtype: DataType,
pub access: AbiAccess,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct EntryAbiRecord {
pub node: ArtifactNodeId,
pub inputs: Vec<ArtifactValueId>,
pub outputs: Vec<ArtifactValueId>,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct ArtifactAbi {
pub resources: Vec<ResourceAbiRecord>,
pub entries: Vec<EntryAbiRecord>,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct FusionRecord {
pub id: FusionGroupId,
pub members: Vec<ArtifactNodeId>,
pub stage: u32,
pub legality: Vec<Digest>,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct FusionRejection {
pub from: ArtifactNodeId,
pub to: ArtifactNodeId,
pub value: ArtifactValueId,
pub reason: legality::FusionRejectionReason,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct BarrierRecord {
pub before_stage: u32,
pub after_stage: u32,
pub dependencies: Vec<u32>,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum MaterializationReason {
CrossGroupUse,
Output,
Retained,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct MaterializationRecord {
pub value: ArtifactValueId,
pub producer: FusionGroupId,
pub stage: u32,
pub reason: MaterializationReason,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct SelectedPlan {
pub fusion: Vec<FusionRecord>,
pub barriers: Vec<BarrierRecord>,
pub materializations: Vec<MaterializationRecord>,
pub candidates_explored: u32,
pub search_budget: SearchBudget,
pub search_work: SearchWork,
pub selection_cost: cost::CostBreakdown,
pub pruned_fusions: Vec<FusionRejection>,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct Provenance {
pub source_graph: Digest,
pub request: Digest,
pub compiler_version: String,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
struct ArtifactPayload {
schema_version: u16,
nodes: Vec<NodeRecord>,
dependencies: Vec<DependencyEdge>,
selected_plan: SelectedPlan,
abi: ArtifactAbi,
resources: Vec<ResourceRecord>,
resource_envelope: ResourceEnvelope,
geometry: Vec<GeometryRecord>,
provenance: Provenance,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Artifact {
payload: ArtifactPayload,
digest: Digest,
}
impl Artifact {
#[must_use]
pub const fn schema_version(&self) -> u16 {
self.payload.schema_version
}
#[must_use]
pub fn nodes(&self) -> &[NodeRecord] {
&self.payload.nodes
}
#[must_use]
pub fn dependencies(&self) -> &[DependencyEdge] {
&self.payload.dependencies
}
#[must_use]
pub fn fusion(&self) -> &[FusionRecord] {
&self.payload.selected_plan.fusion
}
#[must_use]
pub fn barriers(&self) -> &[BarrierRecord] {
&self.payload.selected_plan.barriers
}
#[must_use]
pub fn resources(&self) -> &[ResourceRecord] {
&self.payload.resources
}
#[must_use]
pub const fn resource_envelope(&self) -> ResourceEnvelope {
self.payload.resource_envelope
}
#[must_use]
pub fn geometry(&self) -> &[GeometryRecord] {
&self.payload.geometry
}
#[must_use]
pub fn materializations(&self) -> &[MaterializationRecord] {
&self.payload.selected_plan.materializations
}
#[must_use]
pub const fn selected_plan(&self) -> &SelectedPlan {
&self.payload.selected_plan
}
#[must_use]
pub const fn abi(&self) -> &ArtifactAbi {
&self.payload.abi
}
#[must_use]
pub const fn provenance(&self) -> &Provenance {
&self.payload.provenance
}
#[must_use]
pub const fn digest(&self) -> Digest {
self.digest
}
pub fn to_bytes(&self) -> Result<Vec<u8>, CompileError> {
encode_payload(&self.payload)
}
pub fn from_bytes(bytes: &[u8]) -> Result<Self, CompileError> {
if bytes.len() < ARTIFACT_HEADER_BYTES + ARTIFACT_DIGEST_BYTES {
return Err(failure(
CompilerFailureKind::MalformedArtifact,
"artifact.header",
"artifact is shorter than its fixed framing",
"supply complete VMK0 bytes",
));
}
if &bytes[..4] != ARTIFACT_MAGIC {
return Err(failure(
CompilerFailureKind::MalformedArtifact,
"artifact.magic",
"artifact magic is not VMK0",
"supply canonical megakernel artifact bytes",
));
}
let version = u16::from_le_bytes([bytes[4], bytes[5]]);
if version != ARTIFACT_SCHEMA_VERSION {
return Err(failure(
CompilerFailureKind::VersionSkew,
"artifact.schema_version",
format!("schema {version} is unsupported; expected {ARTIFACT_SCHEMA_VERSION}"),
"recompile the source graph with this compiler version",
));
}
let body_len = u32::from_le_bytes(bytes[6..10].try_into().expect("fixed slice")) as usize;
let expected_len = ARTIFACT_HEADER_BYTES
.checked_add(body_len)
.and_then(|len| len.checked_add(ARTIFACT_DIGEST_BYTES))
.ok_or_else(|| {
failure(
CompilerFailureKind::MalformedArtifact,
"artifact.body_length",
"framed body length overflowed addressable memory",
"supply bounded canonical artifact bytes",
)
})?;
if bytes.len() != expected_len {
return Err(failure(
CompilerFailureKind::MalformedArtifact,
"artifact.body_length",
format!(
"framing declares {expected_len} bytes but received {}",
bytes.len()
),
"supply exactly one complete canonical artifact",
));
}
let body = &bytes[ARTIFACT_HEADER_BYTES..ARTIFACT_HEADER_BYTES + body_len];
let expected_digest = artifact_digest(version, body);
let encoded_digest: [u8; 32] = bytes[ARTIFACT_HEADER_BYTES + body_len..]
.try_into()
.expect("validated digest length");
if expected_digest.0 != encoded_digest {
return Err(failure(
CompilerFailureKind::DigestMismatch,
"artifact.digest",
"artifact body does not match its content identity",
"discard the corrupted artifact and recompile",
));
}
let payload: ArtifactPayload = serde_json::from_slice(body).map_err(|error| {
failure(
CompilerFailureKind::MalformedArtifact,
"artifact.body",
error.to_string(),
"supply a canonical body emitted by this crate",
)
})?;
if payload.schema_version != version {
return Err(failure(
CompilerFailureKind::VersionSkew,
"artifact.body.schema_version",
"body schema disagrees with framing schema",
"recompile instead of rewriting artifact framing",
));
}
let canonical = serde_json::to_vec(&payload).map_err(serialization_failure)?;
if canonical != body {
return Err(failure(
CompilerFailureKind::MalformedArtifact,
"artifact.body",
"artifact body is valid JSON but not canonical JSON",
"use the canonical bytes emitted by Artifact::to_bytes",
));
}
Ok(Self {
payload,
digest: expected_digest,
})
}
}
pub fn compile(request: &ValidatedCompileRequest) -> Result<Artifact, CompileError> {
let canonical_wire = request.graph.to_wire().map_err(|error| {
failure(
CompilerFailureKind::InvalidProgram,
"request.graph",
error.to_string(),
"supply a graph representable by the canonical foundation wire format",
)
})?;
let source_graph = domain_digest(SOURCE_DIGEST_DOMAIN, &canonical_wire);
let nodes = request
.graph
.nodes()
.iter()
.map(|node| {
let program = node.program.canonical_wire_bytes().map_err(|error| {
failure(
CompilerFailureKind::InvalidProgram,
format!("request.graph.nodes[{}].program", node.id.0),
error.to_string(),
"supply canonical-wire-compatible typed IR",
)
})?;
Ok(NodeRecord {
id: ArtifactNodeId(node.id.0),
name: node.name.clone(),
program,
})
})
.collect::<Result<Vec<_>, CompileError>>()?;
let geometry = request
.graph
.nodes()
.iter()
.map(|node| GeometryRecord {
node: ArtifactNodeId(node.id.0),
workgroup_size: node.program.workgroup_size,
})
.collect::<Vec<_>>();
let normalized = normalize::normalize(&request.graph)?;
let dependencies = normalized.dependencies;
let artifact::ArtifactPlan {
node_groups,
stages,
selected_plan,
} = artifact::plan(&request.graph, &dependencies, request.search_budget)?;
let (resources, resource_envelope) = build_resources(
&request.graph,
&request.facts.symbolic_bindings,
&node_groups,
&stages,
)?;
let abi = build_abi(&request.graph)?;
let request_bytes =
serde_json::to_vec(&RequestIdentity::from(request)).map_err(serialization_failure)?;
let provenance = Provenance {
source_graph,
request: domain_digest(REQUEST_DIGEST_DOMAIN, &request_bytes),
compiler_version: env!("CARGO_PKG_VERSION").to_string(),
};
let payload = ArtifactPayload {
schema_version: ARTIFACT_SCHEMA_VERSION,
nodes,
dependencies,
selected_plan,
abi,
resources,
resource_envelope,
geometry,
provenance,
};
let bytes = encode_payload(&payload)?;
let byte_len = u64::try_from(bytes.len())
.map_err(|_| overflow("artifact", "artifact length exceeds u64"))?;
if byte_len > request.max_artifact_bytes {
return Err(failure(
CompilerFailureKind::ArtifactLimit,
"artifact",
format!(
"canonical artifact is {byte_len} bytes; limit is {}",
request.max_artifact_bytes
),
"raise the explicit artifact bound or reduce the source graph",
));
}
let digest: [u8; 32] = bytes[bytes.len() - ARTIFACT_DIGEST_BYTES..]
.try_into()
.expect("encoded digest length");
Ok(Artifact {
payload,
digest: Digest(digest),
})
}
#[derive(Serialize)]
struct RequestIdentity<'a> {
configuration_digest: Digest,
symbolic_bindings: &'a BTreeMap<String, u64>,
constant_identities: Vec<(u32, Digest)>,
search_budget: SearchBudget,
}
impl<'a> From<&'a ValidatedCompileRequest> for RequestIdentity<'a> {
fn from(request: &'a ValidatedCompileRequest) -> Self {
Self {
configuration_digest: request.facts.configuration_digest,
symbolic_bindings: &request.facts.symbolic_bindings,
constant_identities: request
.facts
.constant_identities
.iter()
.map(|(id, digest)| (id.0, *digest))
.collect(),
search_budget: request.search_budget,
}
}
}
fn validate_bindings(
graph: &ProgramGraph,
bindings: &BTreeMap<String, u64>,
) -> Result<(), CompileError> {
let symbols: BTreeSet<&str> = graph
.values()
.iter()
.flat_map(|value| &value.contract.shape)
.filter_map(|dim| match dim {
ShapeDim::Known(_) => None,
ShapeDim::Symbol(symbol) => Some(symbol.as_str()),
})
.collect();
if let Some(symbol) = symbols
.iter()
.find(|symbol| !bindings.contains_key(**symbol))
{
return Err(failure(
CompilerFailureKind::MissingSymbol,
format!("request.facts.symbolic_bindings.{symbol}"),
"graph symbol has no exact extent",
"bind every symbolic graph dimension before compilation",
));
}
if let Some(symbol) = bindings
.keys()
.find(|symbol| !symbols.contains(symbol.as_str()))
{
return Err(failure(
CompilerFailureKind::UnknownSymbol,
format!("request.facts.symbolic_bindings.{symbol}"),
"binding does not occur in the graph",
"remove stale bindings or use the graph's exact symbol name",
));
}
Ok(())
}
fn validate_constant_identities(
graph: &ProgramGraph,
identities: &BTreeMap<GraphValueId, Digest>,
) -> Result<(), CompileError> {
let constants = graph
.values()
.iter()
.filter(|value| value.contract.lifetime == ValueLifetime::Constant)
.map(|value| value.id)
.collect::<BTreeSet<_>>();
if let Some(id) = constants.iter().find(|id| !identities.contains_key(*id)) {
return Err(failure(
CompilerFailureKind::MissingConstantIdentity,
format!("request.facts.constant_identities.{}", id.0),
"constant graph value has no verified content identity",
"supply one digest keyed by the constant GraphValueId",
));
}
if let Some(id) = identities.keys().find(|id| !constants.contains(id)) {
return Err(failure(
CompilerFailureKind::UnknownConstantIdentity,
format!("request.facts.constant_identities.{}", id.0),
"constant identity names a non-constant or missing graph value",
"remove stale identities and key constant content by GraphValueId",
));
}
Ok(())
}
fn build_abi(graph: &ProgramGraph) -> Result<ArtifactAbi, CompileError> {
let resources = graph
.values()
.iter()
.map(|value| {
let access = match value.contract.access.clone() {
BufferAccess::ReadOnly => AbiAccess::ReadOnly,
BufferAccess::WriteOnly => AbiAccess::WriteOnly,
BufferAccess::ReadWrite => AbiAccess::ReadWrite,
BufferAccess::Uniform => AbiAccess::Uniform,
unsupported => {
return Err(failure(
CompilerFailureKind::InvalidProgram,
format!("request.graph.values[{}].contract.access", value.id.0),
format!("access {unsupported:?} has no artifact ABI representation"),
"lower workgroup/private resources inside the node Program",
))
}
};
Ok(ResourceAbiRecord {
slot: value.id.0,
value: ArtifactValueId(value.id.0),
dtype: value.contract.dtype.clone(),
access,
})
})
.collect::<Result<Vec<_>, CompileError>>()?;
let entries = graph
.nodes()
.iter()
.map(|node| EntryAbiRecord {
node: ArtifactNodeId(node.id.0),
inputs: node
.inputs
.iter()
.map(|input| ArtifactValueId(input.value.0))
.collect(),
outputs: node
.outputs
.iter()
.map(|output| ArtifactValueId(output.0))
.collect(),
})
.collect();
Ok(ArtifactAbi { resources, entries })
}
fn ensure_node_dag(
count: usize,
dependencies: &[DependencyEdge],
code: CompilerFailureKind,
) -> Result<(), CompileError> {
let groups: Vec<_> = (0..count).map(|id| FusionGroupId(id as u32)).collect();
ensure_group_dag(count, dependencies, &groups, code)
}
fn ensure_group_dag(
count: usize,
dependencies: &[DependencyEdge],
node_groups: &[FusionGroupId],
code: CompilerFailureKind,
) -> Result<(), CompileError> {
group_stages_inner(count, dependencies, node_groups)
.map(|_| ())
.map_err(|_| {
failure(
code,
"artifact.dependencies",
"dependency graph contains a cycle",
"remove the cyclic semantic dependency",
)
})
}
fn group_stages(
count: usize,
dependencies: &[DependencyEdge],
node_groups: &[FusionGroupId],
) -> Result<Vec<u32>, CompileError> {
group_stages_inner(count, dependencies, node_groups).map_err(|_| {
failure(
CompilerFailureKind::DependencyCycle,
"artifact.dependencies",
"selected-plan dependency graph contains a cycle",
"fix compiler legality before plan selection",
)
})
}
fn group_stages_inner(
count: usize,
dependencies: &[DependencyEdge],
node_groups: &[FusionGroupId],
) -> Result<Vec<u32>, ()> {
let mut outgoing = vec![BTreeSet::<usize>::new(); count];
let mut indegree = vec![0usize; count];
for edge in dependencies {
let (DependencyEndpoint::Node(from), DependencyEndpoint::Node(to)) = (edge.from, edge.to)
else {
continue;
};
let from = node_groups[from.0 as usize].0 as usize;
let to = node_groups[to.0 as usize].0 as usize;
if from != to && outgoing[from].insert(to) {
indegree[to] += 1;
}
}
let mut ready: BTreeSet<usize> = indegree
.iter()
.enumerate()
.filter_map(|(index, degree)| (*degree == 0).then_some(index))
.collect();
let mut stage = vec![0u32; count];
let mut visited = 0usize;
while let Some(next) = ready.pop_first() {
visited += 1;
for successor in outgoing[next].iter().copied() {
stage[successor] = stage[successor].max(stage[next].checked_add(1).ok_or(())?);
indegree[successor] -= 1;
if indegree[successor] == 0 {
ready.insert(successor);
}
}
}
(visited == count).then_some(stage).ok_or(())
}
fn build_barriers(
dependencies: &[DependencyEdge],
node_groups: &[FusionGroupId],
stages: &[u32],
) -> Result<Vec<BarrierRecord>, CompileError> {
let max_stage = stages.iter().copied().max().unwrap_or(0);
let mut barriers = Vec::new();
for after_stage in 1..=max_stage {
let mut edge_ids = Vec::new();
for (index, edge) in dependencies.iter().enumerate() {
let (DependencyEndpoint::Node(from), DependencyEndpoint::Node(to)) =
(edge.from, edge.to)
else {
continue;
};
let from_stage = stages[node_groups[from.0 as usize].0 as usize];
let to_stage = stages[node_groups[to.0 as usize].0 as usize];
if from_stage < after_stage && to_stage == after_stage {
edge_ids.push(
u32::try_from(index).map_err(|_| {
overflow("artifact.dependencies", "edge identity exceeds u32")
})?,
);
}
}
barriers.push(BarrierRecord {
before_stage: after_stage - 1,
after_stage,
dependencies: edge_ids,
});
}
Ok(barriers)
}
fn build_materializations(
graph: &ProgramGraph,
node_groups: &[FusionGroupId],
stages: &[u32],
) -> Vec<MaterializationRecord> {
let mut records = Vec::new();
for value in graph.values() {
let Some(producer) = value.producer else {
continue;
};
let producer_node = ArtifactNodeId(producer.0);
let producer_group = node_groups[producer_node.0 as usize];
let producer_stage = stages[producer_group.0 as usize];
let cross_group = value.consumers.iter().any(|consumer| {
let consumer_node = ArtifactNodeId(consumer.0);
node_groups[consumer_node.0 as usize] != producer_group
});
let reason = match value.contract.lifetime {
ValueLifetime::Output => Some(MaterializationReason::Output),
ValueLifetime::Retained => Some(MaterializationReason::Retained),
_ if cross_group => Some(MaterializationReason::CrossGroupUse),
_ => None,
};
if let Some(reason) = reason {
records.push(MaterializationRecord {
value: ArtifactValueId(value.id.0),
producer: producer_group,
stage: producer_stage,
reason,
});
}
}
records.sort_by_key(|record| (record.value, record.reason as u8));
records
}
fn build_resources(
graph: &ProgramGraph,
bindings: &BTreeMap<String, u64>,
node_groups: &[FusionGroupId],
stages: &[u32],
) -> Result<(Vec<ResourceRecord>, ResourceEnvelope), CompileError> {
let final_stage = stages.iter().copied().max().unwrap_or(0);
let mut resources = Vec::with_capacity(graph.values().len());
for value in graph.values() {
let mut element_count = 1u64;
for dim in &value.contract.shape {
let extent = match dim {
ShapeDim::Known(extent) => *extent,
ShapeDim::Symbol(symbol) => bindings[symbol],
};
element_count = element_count.checked_mul(extent).ok_or_else(|| {
overflow(
format!("graph.values[{}].shape", value.name),
"shape element count exceeds u64",
)
})?;
}
let host_count = usize::try_from(element_count).map_err(|_| {
overflow(
format!("graph.values[{}].shape", value.name),
"shape element count exceeds addressable packed-size input",
)
})?;
let byte_count = value
.contract
.dtype
.packed_size_bytes(host_count)
.map_err(|message| overflow(format!("graph.values[{}].dtype", value.name), message))?
.ok_or_else(|| {
failure(
CompilerFailureKind::UnsizedResource,
format!("graph.values[{}].dtype", value.name),
"value representation has no fixed packed byte size",
"resolve the representation to a fixed-width typed value before compilation",
)
})?;
let byte_count = u64::try_from(byte_count).map_err(|_| {
overflow(
format!("graph.values[{}]", value.name),
"packed byte count exceeds u64",
)
})?;
let producer_stage = value.producer.map_or(0, |producer| {
stages[node_groups[producer.0 as usize].0 as usize]
});
let mut last_stage = value
.consumers
.iter()
.map(|consumer| stages[node_groups[consumer.0 as usize].0 as usize])
.max()
.unwrap_or(producer_stage);
if matches!(
value.contract.lifetime,
ValueLifetime::Output | ValueLifetime::Retained
) {
last_stage = last_stage.max(final_stage);
}
resources.push(ResourceRecord {
value: ArtifactValueId(value.id.0),
name: value.name.clone(),
element_count,
byte_count,
lifetime: match value.contract.lifetime {
ValueLifetime::Constant => ResourceLifetime::Constant,
ValueLifetime::Invocation => ResourceLifetime::Invocation,
ValueLifetime::Retained => ResourceLifetime::Retained,
ValueLifetime::Output => ResourceLifetime::Output,
},
first_stage: producer_stage,
last_stage,
});
}
resources.sort_by_key(|resource| resource.value);
let total_bytes = resources.iter().try_fold(0u64, |total, resource| {
total.checked_add(resource.byte_count).ok_or_else(|| {
overflow(
"artifact.resource_envelope.total_bytes",
"resource sum exceeds u64",
)
})
})?;
let mut peak_live_bytes = 0u64;
for stage in 0..=final_stage {
let live = resources
.iter()
.filter(|resource| resource.first_stage <= stage && stage <= resource.last_stage)
.try_fold(0u64, |total, resource| {
total.checked_add(resource.byte_count).ok_or_else(|| {
overflow(
"artifact.resource_envelope.peak_live_bytes",
"live resource sum exceeds u64",
)
})
})?;
peak_live_bytes = peak_live_bytes.max(live);
}
Ok((
resources,
ResourceEnvelope {
total_bytes,
peak_live_bytes,
},
))
}
fn encode_payload(payload: &ArtifactPayload) -> Result<Vec<u8>, CompileError> {
let body = serde_json::to_vec(payload).map_err(serialization_failure)?;
let body_len = u32::try_from(body.len()).map_err(|_| {
overflow(
"artifact.body",
"canonical body exceeds the u32 framing limit",
)
})?;
let digest = artifact_digest(payload.schema_version, &body);
let capacity = ARTIFACT_HEADER_BYTES
.checked_add(body.len())
.and_then(|len| len.checked_add(ARTIFACT_DIGEST_BYTES))
.ok_or_else(|| overflow("artifact", "encoded artifact length overflowed usize"))?;
let mut bytes = Vec::with_capacity(capacity);
bytes.extend_from_slice(ARTIFACT_MAGIC);
bytes.extend_from_slice(&payload.schema_version.to_le_bytes());
bytes.extend_from_slice(&body_len.to_le_bytes());
bytes.extend_from_slice(&body);
bytes.extend_from_slice(&digest.0);
Ok(bytes)
}
fn artifact_digest(version: u16, body: &[u8]) -> Digest {
let mut hasher = blake3::Hasher::new();
hasher.update(ARTIFACT_DIGEST_DOMAIN);
hasher.update(&version.to_le_bytes());
hasher.update(&(body.len() as u64).to_le_bytes());
hasher.update(body);
Digest(*hasher.finalize().as_bytes())
}
fn domain_digest(domain: &[u8], bytes: &[u8]) -> Digest {
let mut hasher = blake3::Hasher::new();
hasher.update(domain);
hasher.update(&(bytes.len() as u64).to_le_bytes());
hasher.update(bytes);
Digest(*hasher.finalize().as_bytes())
}
fn serialization_failure(error: serde_json::Error) -> CompileError {
failure(
CompilerFailureKind::MalformedArtifact,
"artifact.body",
error.to_string(),
"use values representable by the canonical artifact schema",
)
}
fn overflow(path: impl Into<String>, message: impl Into<String>) -> CompileError {
failure(
CompilerFailureKind::ResourceOverflow,
path,
message,
"reduce resolved extents or split the graph before compilation",
)
}
fn failure(
code: CompilerFailureKind,
path: impl Into<String>,
message: impl Into<String>,
fix: impl Into<String>,
) -> CompileError {
let stage = diagnostic_stage(code);
let retry = diagnostic_retry(code);
CompileError {
diagnostic: Diagnostic::error(code.as_str(), message.into())
.with_stage(stage)
.with_location(OpLocation::op("vyre-megakernel").with_path(path))
.with_fix(fix.into())
.with_retry(retry),
}
}