Skip to main content

polydat_core/compile/
simd_plan.rs

1// Copyright 2024-2026 Jonathan Shook
2// SPDX-License-Identifier: Apache-2.0
3
4//! Static type and node-variant qualification for scalar-flow SIMD promotion.
5//!
6//! This module does not rewrite a graph. It establishes the production
7//! boundary a rewrite must pass before it may form a packet plan: an explicit
8//! scalar-node contract, a supported scalar/register lane shape, pure nodes,
9//! and exact register-typed I/O. Final Cranelift compilation remains the
10//! authoritative lowering probe.
11
12use std::fmt;
13
14use crate::ast::{Lifecycle, PolydatNode, PortType, Purity};
15
16/// Interpretation of scalar values sharing one physical register lane shape.
17#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
18pub enum SimdLaneKind {
19    /// Unsigned integer lanes.
20    UnsignedInteger,
21    /// Signed integer lanes.
22    SignedInteger,
23    /// Floating-point lanes.
24    Float,
25}
26
27/// A scalar type's fixed 128-bit Polydat register representation.
28#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
29pub struct SimdTypeShape {
30    /// The scalar type.
31    pub scalar: PortType,
32    /// The register type that carries a packet of it.
33    pub register: PortType,
34    /// How the lanes are interpreted.
35    pub lane_kind: SimdLaneKind,
36    /// Bits per lane.
37    pub lane_bits: u8,
38    /// Lanes per register.
39    pub lanes: u8,
40}
41
42impl SimdTypeShape {
43    /// Bits per packet: lanes times lane bits.
44    pub const fn packet_bits(self) -> u16 {
45        self.lane_bits as u16 * self.lanes as u16
46    }
47}
48
49/// Physically representable fixed-width shapes for the first promotion tier.
50///
51/// Unsigned integer types deliberately share the corresponding `RegI*x*`
52/// physical shape. Add/sub/mul and raw bitwise operations are bit-identical;
53/// signedness remains semantic metadata for comparisons, division, shifts,
54/// and conversion recipes. A returned shape is necessary, not sufficient:
55/// every member node still needs explicit variant metadata and a successful
56/// whole-cone backend compilation.
57pub const fn promotable_type_shape(scalar: PortType) -> Option<SimdTypeShape> {
58    use PortType::*;
59    use SimdLaneKind::*;
60
61    let shape = match scalar {
62        U8 => SimdTypeShape {
63            scalar,
64            register: RegI8x16,
65            lane_kind: UnsignedInteger,
66            lane_bits: 8,
67            lanes: 16,
68        },
69        I8 => SimdTypeShape {
70            scalar,
71            register: RegI8x16,
72            lane_kind: SignedInteger,
73            lane_bits: 8,
74            lanes: 16,
75        },
76        U16 => SimdTypeShape {
77            scalar,
78            register: RegI16x8,
79            lane_kind: UnsignedInteger,
80            lane_bits: 16,
81            lanes: 8,
82        },
83        I16 => SimdTypeShape {
84            scalar,
85            register: RegI16x8,
86            lane_kind: SignedInteger,
87            lane_bits: 16,
88            lanes: 8,
89        },
90        U32 => SimdTypeShape {
91            scalar,
92            register: RegI32x4,
93            lane_kind: UnsignedInteger,
94            lane_bits: 32,
95            lanes: 4,
96        },
97        I32 => SimdTypeShape {
98            scalar,
99            register: RegI32x4,
100            lane_kind: SignedInteger,
101            lane_bits: 32,
102            lanes: 4,
103        },
104        U64 => SimdTypeShape {
105            scalar,
106            register: RegI64x2,
107            lane_kind: UnsignedInteger,
108            lane_bits: 64,
109            lanes: 2,
110        },
111        I64 => SimdTypeShape {
112            scalar,
113            register: RegI64x2,
114            lane_kind: SignedInteger,
115            lane_bits: 64,
116            lanes: 2,
117        },
118        F32 => SimdTypeShape {
119            scalar,
120            register: RegF32x4,
121            lane_kind: Float,
122            lane_bits: 32,
123            lanes: 4,
124        },
125        F64 => SimdTypeShape {
126            scalar,
127            register: RegF64x2,
128            lane_kind: Float,
129            lane_bits: 64,
130            lanes: 2,
131        },
132        // F16X8 exists in Polydat's type plane, but the installed Cranelift
133        // x64 path is not a production arithmetic lowering. It enters this
134        // table only after a backend probe and operation catalog justify it.
135        F16 | U128 | I128 | Bool | Str | Bytes | Json | Ext | Handle | VecF32 | VecF64 | VecF16
136        | VecI8 | VecI16 | VecI32 | VecI64 | Reg128 | RegI8x16 | RegI16x8 | RegI32x4 | RegI64x2
137        | RegF16x8 | RegF32x4 | RegF64x2 => return None,
138    };
139    Some(shape)
140}
141
142/// A scalar node whose declared register variant has passed static shape and
143/// semantic validation. It still needs whole-cone Cranelift compilation.
144#[derive(Clone, Debug, PartialEq, Eq)]
145pub struct ValidatedSimdVariant {
146    /// The scalar node's name.
147    pub scalar_node: String,
148    /// The register node that computes a packet of it.
149    pub vector_node: &'static str,
150    /// The lane shape.
151    pub shape: SimdTypeShape,
152    /// Wire inputs the scalar node takes.
153    pub wire_inputs: u8,
154}
155
156#[derive(Clone, Debug, PartialEq, Eq)]
157/// Why a scalar node's declared register variant was not accepted.
158pub enum SimdVariantError {
159    /// The node declares no register variant.
160    Undeclared,
161    /// The scalar node is not pure.
162    ScalarNodeNotPure,
163    /// The variant is not declared exact.
164    VariantNotExact,
165    /// The variant is not declared total.
166    VariantNotTotal,
167    /// The variant is not declared lane-independent.
168    VariantNotLaneIndependent,
169    /// The scalar type has no fixed register shape.
170    UnsupportedScalarShape(PortType),
171    /// The scalar node's inputs and output are not all of one type.
172    ScalarSignatureNotUniform,
173    /// The scalar node's ports are not all cycle-lifecycle.
174    ScalarLifecycleNotCycle,
175    /// The register node is not registered.
176    VectorNodeUnavailable(String),
177    /// The register node is not pure.
178    VectorNodeNotPure,
179    /// The register node's signature does not match the scalar node's shape.
180    VectorSignatureMismatch,
181    /// The register node's ports are not all cycle-lifecycle.
182    VectorLifecycleNotCycle,
183}
184
185impl fmt::Display for SimdVariantError {
186    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
187        match self {
188            Self::Undeclared => f.write_str("scalar node declares no SIMD variant"),
189            Self::ScalarNodeNotPure => f.write_str("scalar node is not pure"),
190            Self::VariantNotExact => f.write_str("SIMD variant is not exact"),
191            Self::VariantNotTotal => f.write_str("SIMD variant is not total"),
192            Self::VariantNotLaneIndependent => f.write_str("SIMD variant is not lane-independent"),
193            Self::UnsupportedScalarShape(t) => write!(f, "unsupported scalar SIMD shape {t}"),
194            Self::ScalarSignatureNotUniform => {
195                f.write_str("scalar node inputs and output do not share one scalar lane type")
196            }
197            Self::ScalarLifecycleNotCycle => {
198                f.write_str("scalar node has an init-lifecycle data port")
199            }
200            Self::VectorNodeUnavailable(e) => write!(f, "register variant unavailable: {e}"),
201            Self::VectorNodeNotPure => f.write_str("register variant is not pure"),
202            Self::VectorSignatureMismatch => {
203                f.write_str("register variant signature does not match the scalar lane shape")
204            }
205            Self::VectorLifecycleNotCycle => {
206                f.write_str("register variant has an init-lifecycle data port")
207            }
208        }
209    }
210}
211
212impl std::error::Error for SimdVariantError {}
213
214/// Validate an instantiated scalar node and its declared register node.
215///
216/// Tier 1 currently accepts uniform element-wise operations: all scalar wire
217/// inputs and the sole output have one type, and the vector variant has the
218/// same arity under that type's register shape. Broadcast-vs-varying input
219/// roles are a property of the eventual cone plan, not the node variant.
220pub fn validate_simd_variant(
221    scalar: &dyn PolydatNode,
222) -> Result<ValidatedSimdVariant, SimdVariantError> {
223    let variant = scalar.simd_variant().ok_or(SimdVariantError::Undeclared)?;
224    if scalar.purity() != Purity::Pure {
225        return Err(SimdVariantError::ScalarNodeNotPure);
226    }
227    if !variant.exact {
228        return Err(SimdVariantError::VariantNotExact);
229    }
230    if !variant.total {
231        return Err(SimdVariantError::VariantNotTotal);
232    }
233    if !variant.lane_independent {
234        return Err(SimdVariantError::VariantNotLaneIndependent);
235    }
236
237    let scalar_meta = scalar.meta();
238    let [scalar_output] = scalar_meta.outs.as_slice() else {
239        return Err(SimdVariantError::ScalarSignatureNotUniform);
240    };
241    let shape = promotable_type_shape(scalar_output.typ)
242        .ok_or(SimdVariantError::UnsupportedScalarShape(scalar_output.typ))?;
243    let scalar_inputs = scalar_meta.wire_inputs();
244    if scalar_inputs.is_empty() || scalar_inputs.iter().any(|port| port.typ != shape.scalar) {
245        return Err(SimdVariantError::ScalarSignatureNotUniform);
246    }
247    if scalar_output.lifecycle != Lifecycle::Cycle
248        || scalar_inputs
249            .iter()
250            .any(|port| port.lifecycle != Lifecycle::Cycle)
251    {
252        return Err(SimdVariantError::ScalarLifecycleNotCycle);
253    }
254
255    let wires = vec![crate::compile::assembly::WireRef::input("__simd_probe"); scalar_inputs.len()];
256    let wire_types = vec![shape.register; scalar_inputs.len()];
257    let vector = crate::dsl::factory::build_node(variant.vector_node, &wires, &wire_types, &[])
258        .map_err(|e| SimdVariantError::VectorNodeUnavailable(e.to_string()))?;
259    if vector.purity() != Purity::Pure {
260        return Err(SimdVariantError::VectorNodeNotPure);
261    }
262    let vector_meta = vector.meta();
263    let [vector_output] = vector_meta.outs.as_slice() else {
264        return Err(SimdVariantError::VectorSignatureMismatch);
265    };
266    let vector_inputs = vector_meta.wire_inputs();
267    if vector_output.typ != shape.register
268        || vector_inputs.len() != scalar_inputs.len()
269        || vector_inputs.iter().any(|port| port.typ != shape.register)
270    {
271        return Err(SimdVariantError::VectorSignatureMismatch);
272    }
273    if vector_output.lifecycle != Lifecycle::Cycle
274        || vector_inputs
275            .iter()
276            .any(|port| port.lifecycle != Lifecycle::Cycle)
277    {
278        return Err(SimdVariantError::VectorLifecycleNotCycle);
279    }
280
281    Ok(ValidatedSimdVariant {
282        scalar_node: scalar_meta.name.clone(),
283        vector_node: variant.vector_node,
284        shape,
285        wire_inputs: scalar_inputs.len() as u8,
286    })
287}
288
289#[cfg(test)]
290mod tests {
291    use super::*;
292
293    #[test]
294    fn common_numeric_shapes_cover_one_128_bit_word() {
295        for ty in [
296            PortType::U8,
297            PortType::I8,
298            PortType::U16,
299            PortType::I16,
300            PortType::U32,
301            PortType::I32,
302            PortType::U64,
303            PortType::I64,
304            PortType::F32,
305            PortType::F64,
306        ] {
307            assert_eq!(promotable_type_shape(ty).unwrap().packet_bits(), 128);
308        }
309        assert!(promotable_type_shape(PortType::F16).is_none());
310        assert!(promotable_type_shape(PortType::U128).is_none());
311        assert!(promotable_type_shape(PortType::Str).is_none());
312    }
313}