aleleba/lightcraft
1
0
Fork 0
mirror of https://github.com/storytold/lightcraft.git synced 2026-11-01 22:17:53 -06:00
Code Issues Packages Projects Releases Wiki Activity
Files
4e4ea5e3c3c63e654020ebd18e7d2b2b8da97964
lightcraft/crates/faces/src/net/shape.rs
T

501 lines
23 KiB
Rust
Raw Blame History

//! Bounded, allocation-free (apart from small shape lists) graph checking.
use super::{Data, Graph, HashMap, MAX_TENSOR, NetError, Node, Result, pads, pair, shape_err, unsupported};
use crate::graph::{Attr, Weight};
pub(super) fn elements(s: &[usize]) -> Result<usize> {
if s.len() > 8 || s.contains(&0) {
return shape_err("a tensor has a zero dimension or excessive rank");
}
s.iter()
.try_fold(1usize, |n, &d| n.checked_mul(d))
.filter(|&n| n <= MAX_TENSOR)
.ok_or_else(|| NetError::Shape("a tensor exceeds the element limit".into()))
}
pub(super) fn tensor(s: &[usize], n: usize, data: &[f32]) -> Result<()> {
if elements(s)? != n || data.iter().any(|v| !v.is_finite()) {
return shape_err("invalid tensor size or non-finite values");
}
Ok(())
}
pub(super) fn float(n: &Node, name: &str, default: f32) -> f32 {
match n.attrs.get(name) {
Some(Attr::Float(v)) => *v,
_ => default,
}
}
pub(super) fn weight<'a>(n: &Node, g: &'a Graph, i: usize) -> Result<&'a Weight> {
n.inputs.get(i).and_then(|s| g.weights.get(s)).ok_or_else(|| NetError::Unsupported(format!("{} needs constant input {i}", n.op)))
}
pub(super) fn floats<'a>(n: &Node, g: &'a Graph, i: usize) -> Result<&'a [f32]> {
match &weight(n, g, i)?.data {
Data::F32(v) => Ok(v),
_ => unsupported(format!("{} needs float32 weights", n.op)),
}
}
pub(super) fn check_graph(g: &Graph) -> Result<()> {
if g.inputs.len() != 1 || g.outputs.is_empty() || g.outputs.len() > 64 || g.nodes.len() > 20_000 || g.weights.len() > 20_000 {
return unsupported("one input and at most 64 outputs are required");
}
let mut known = std::collections::HashSet::new();
let mut total_weights = 0usize;
for (name, w) in &g.weights {
let count = w.dims.iter().try_fold(1usize, |n, &d| n.checked_mul(d));
let len = match &w.data {
Data::F32(v) => {
if v.iter().any(|v| !v.is_finite()) {
return shape_err("non-finite weights");
}
v.len()
}
Data::I64(v) => v.len(),
};
if name.is_empty()
|| name.len() > 4096
|| w.dims.len() > 8
|| w.dims.iter().any(|&d| d > MAX_TENSOR)
|| count != Some(len)
|| len > MAX_TENSOR
{
return shape_err("invalid weight name, shape or size");
}
total_weights = total_weights
.checked_add(len)
.filter(|&v| v <= MAX_TENSOR)
.ok_or_else(|| NetError::Shape("combined weights exceed 64 Mi elements".into()))?;
known.insert(name.as_str());
}
let input = g.inputs.first().ok_or_else(|| NetError::Shape("missing input".into()))?;
if input.0.is_empty() || !known.insert(&input.0) {
return shape_err("invalid or repeated input name");
}
for n in &g.nodes {
if n.inputs.len() > 64 || n.outputs.len() != 1 || n.outputs.first().is_none_or(|s| s.is_empty() || s.len() > 4096) {
return unsupported(format!("{} requires one named output", n.op));
}
if n.inputs.iter().any(|s| !s.is_empty() && !known.contains(s.as_str())) {
return shape_err(format!("{} reads an unknown or forward value", n.op));
}
if !known.insert(n.outputs.first().ok_or_else(|| NetError::Shape("missing output".into()))?) {
return shape_err("a graph value is defined twice");
}
}
if g.outputs.iter().any(|s| g.weights.contains_key(s)) {
return unsupported("constant graph outputs");
}
let mut output_names = std::collections::HashSet::new();
if g.outputs.iter().any(|s| !known.contains(s.as_str()) || !output_names.insert(s)) {
return shape_err("unknown or repeated graph output");
}
Ok(())
}
pub(super) fn check(n: &Node, g: &Graph) -> Result<()> {
let (min, max, attrs): (usize, usize, &[(&str, u8)]) = match n.op.as_str() {
"Conv" => (2, 3, &[("group", 0), ("kernel_shape", 2), ("strides", 2), ("pads", 2), ("dilations", 2), ("auto_pad", 3)]),
"BatchNormalization" => (5, 5, &[("epsilon", 1), ("momentum", 1), ("training_mode", 0)]),
"PRelu" | "Add" | "Sub" | "Mul" => (2, 2, &[]),
"Relu" | "Sigmoid" | "GlobalAveragePool" | "Identity" => (1, 1, &[]),
"Dropout" => (1, 1, &[("ratio", 1), ("seed", 0)]),
"Flatten" => (1, 1, &[("axis", 0)]),
"Gemm" => (2, 3, &[("alpha", 1), ("beta", 1), ("transA", 0), ("transB", 0)]),
"MaxPool" => {
(1, 1, &[("kernel_shape", 2), ("strides", 2), ("pads", 2), ("dilations", 2), ("ceil_mode", 0), ("auto_pad", 3), ("storage_order", 0)])
}
"Resize" => (
3,
4,
&[
("mode", 3),
("coordinate_transformation_mode", 3),
("nearest_mode", 3),
("cubic_coeff_a", 1),
("exclude_outside", 0),
("extrapolation_value", 1),
],
),
"Transpose" => (1, 1, &[("perm", 2)]),
"Reshape" => (2, 2, &[("allowzero", 0)]),
op => return unsupported(format!("the operator `{op}`")),
};
if !(min..=max).contains(&n.inputs.len()) || n.inputs.first().is_none_or(String::is_empty) {
return shape_err(format!("{} has invalid inputs", n.op));
}
for (name, value) in &n.attrs {
let actual = match value {
Attr::Int(_) => 0,
Attr::Float(v) if v.is_finite() => 1,
Attr::Ints(_) => 2,
Attr::Str(_) => 3,
_ => 255,
};
if attrs.iter().find(|(key, _)| *key == name).is_none_or(|(_, kind)| *kind != actual) {
return unsupported(format!("{} attribute `{name}` (or its type)", n.op));
}
}
for (i, name) in n.inputs.iter().enumerate() {
if let Some(w) = g.weights.get(name)
&& matches!(w.data, Data::I64(_))
&& !(n.op == "Reshape" && i == 1)
{
return unsupported(format!("{} needs float32 input {i}", n.op));
}
}
match n.op.as_str() {
"Conv" => {
let w = weight(n, g, 1)?;
floats(n, g, 1)?;
if w.dims.len() != 4 || w.dims.contains(&0) {
return shape_err("Conv needs nonempty 2-D weights");
}
if n.ints("kernel_shape").is_some_and(|v| {
v.iter().map(|&d| usize::try_from(d).ok()).collect::<Option<Vec<_>>>().as_deref() != Some(w.dims.get(2..).unwrap_or(&[]))
}) {
return shape_err("Conv kernel_shape differs from its weights");
}
if n.ints("dilations").is_some_and(|d| d != [1, 1]) || n.string("auto_pad").is_some_and(|s| s != "NOTSET") {
return unsupported("Conv dilation or automatic padding");
}
let (sh, sw) = pair(n, "strides", 1)?;
let (pt, pl, pb, pr) = pads(n)?;
if [sh, sw, pt, pl, pb, pr].iter().any(|&v| v > 65536) {
return shape_err("Conv stride or padding exceeds 65536");
}
if n.int("group").unwrap_or(1) < 1 {
return shape_err("Conv group must be positive");
}
if n.inputs.get(2).is_some_and(|s| !s.is_empty()) {
floats(n, g, 2)?;
}
}
"BatchNormalization" => {
if n.int("training_mode").unwrap_or(0) != 0 {
return unsupported("training BatchNormalization");
}
if float(n, "epsilon", 1e-5) <= 0.0 {
return shape_err("BatchNormalization epsilon must be positive");
}
for i in 1..5 {
floats(n, g, i)?;
}
}
"PRelu" => {
floats(n, g, 1)?;
}
"Gemm" => {
floats(n, g, 1)?;
if n.inputs.get(2).is_some_and(|s| !s.is_empty()) {
floats(n, g, 2)?;
}
}
"Dropout" if !(0.0..1.0).contains(&float(n, "ratio", 0.5)) => return shape_err("Dropout ratio is invalid"),
"MaxPool" => {
pair(n, "kernel_shape", 1)?;
pair(n, "strides", 1)?;
if n.ints("kernel_shape").is_none()
|| n.int("ceil_mode").unwrap_or(0) != 0
|| n.int("storage_order").unwrap_or(0) != 0
|| n.ints("dilations").is_some_and(|v| v != [1, 1])
|| pads(n)? != (0, 0, 0, 0)
|| n.string("auto_pad").is_some_and(|s| s != "NOTSET")
{
return unsupported("this MaxPool padding, dilation or ceil mode");
}
}
"Resize" => {
floats(n, g, 2)?;
if n.string("mode").unwrap_or("nearest") != "nearest" || n.inputs.get(3).is_some_and(|s| !s.is_empty()) {
return unsupported("Resize requires nearest constant scales");
}
if n.int("exclude_outside").unwrap_or(0) != 0 || float(n, "extrapolation_value", 0.0) != 0.0 {
return unsupported("Resize excludes or extrapolates outside values");
}
let coord = n.string("coordinate_transformation_mode").unwrap_or("half_pixel");
let nearest = n.string("nearest_mode").unwrap_or("round_prefer_floor");
if !matches!((coord, nearest), ("asymmetric", "floor") | ("half_pixel", "round_prefer_floor")) {
return unsupported("this Resize coordinate/nearest mode");
}
if n.inputs.get(1).is_some_and(|s| !s.is_empty()) && !floats(n, g, 1)?.is_empty() {
return unsupported("Resize with an ROI");
}
}
"Reshape" => {
if n.int("allowzero").unwrap_or(0) != 0 {
return unsupported("Reshape allowzero");
}
if !matches!(weight(n, g, 1)?.data, Data::I64(_)) {
return unsupported("Reshape needs a constant int64 shape");
}
}
_ => {}
}
Ok(())
}
pub(super) fn broadcast(a: &[usize], b: &[usize]) -> Result<Vec<usize>> {
elements(a)?;
elements(b)?;
let rank = a.len().max(b.len());
let mut out = vec![1; rank];
for (i, o) in out.iter_mut().enumerate() {
let aa = i.checked_sub(rank - a.len()).and_then(|j| a.get(j)).copied().unwrap_or(1);
let bb = i.checked_sub(rank - b.len()).and_then(|j| b.get(j)).copied().unwrap_or(1);
if aa != bb && aa != 1 && bb != 1 {
return shape_err("incompatible broadcast shapes");
}
*o = aa.max(bb);
}
elements(&out)?;
Ok(out)
}
pub(super) fn axis(n: &Node, rank: usize) -> Result<usize> {
let raw = n.int("axis").unwrap_or(1);
let raw = if raw < 0 { raw.checked_add(rank as i64).ok_or_else(|| NetError::Shape("axis overflow".into()))? } else { raw };
usize::try_from(raw).ok().filter(|&a| a <= rank).ok_or_else(|| NetError::Shape("Flatten axis is outside its rank".into()))
}
fn image(s: &[usize]) -> Result<(usize, usize, usize)> {
match s {
[1, c, h, w] => Ok((*c, *h, *w)),
_ => unsupported("a spatial operator needs [1,C,H,W]"),
}
}
pub(super) fn infer(g: &Graph, input: &[usize]) -> Result<Vec<Vec<usize>>> {
Ok(infer_all(g, input)?.1)
}
/// Each node's output shape, and the graph's outputs' shapes.
pub(super) fn infer_all(g: &Graph, input: &[usize]) -> Result<(Vec<Vec<usize>>, Vec<Vec<usize>>)> {
elements(input)?;
let mut per_node = Vec::with_capacity(g.nodes.len());
let mut shapes: HashMap<&str, Vec<usize>> = g.weights.iter().map(|(n, w)| (n.as_str(), w.dims.clone())).collect();
shapes.insert(g.inputs.first().ok_or_else(|| NetError::Shape("no input".into()))?.0.as_str(), input.to_vec());
let mut live = elements(input)?;
// Scratch grows independently and remains allocated across all subsequent layers.
let mut scratch_lower = 0;
let mut scratch_product = 0;
let mut last = HashMap::new();
for (i, n) in g.nodes.iter().enumerate() {
for s in &n.inputs {
last.insert(s.as_str(), i);
}
}
for s in &g.outputs {
last.insert(s.as_str(), usize::MAX);
}
for (i, n) in g.nodes.iter().enumerate() {
let a = n.inputs.first().and_then(|s| shapes.get(s.as_str())).ok_or_else(|| NetError::Shape("missing input shape".into()))?;
elements(a)?;
let out = match n.op.as_str() {
"Conv" => {
let (cin, h, w) = image(a)?;
let dims = &weight(n, g, 1)?.dims;
let [co, cg, kh, kw] = dims.as_slice() else {
return shape_err("Conv dimensions");
};
let groups = usize::try_from(n.int("group").unwrap_or(1)).map_err(|_| NetError::Shape("Conv group".into()))?;
if groups.checked_mul(*cg) != Some(cin) || !co.is_multiple_of(groups) {
return shape_err("Conv channels do not match weights/groups");
}
if n.inputs.get(2).is_some_and(|s| !s.is_empty()) && weight(n, g, 2)?.dims != [*co] {
return shape_err("Conv bias is not a channel vector");
}
let (sh, sw) = pair(n, "strides", 1)?;
let (pt, pl, pb, pr) = pads(n)?;
let padded = |size: usize, p: usize, q: usize, k: usize, s: usize| {
size.checked_add(p)
.and_then(|v| v.checked_add(q))
.and_then(|v| v.checked_sub(k))
.and_then(|v| (v / s).checked_add(1))
.ok_or_else(|| NetError::Shape("Conv padding overflow or oversized kernel".into()))
};
vec![1, *co, padded(h, pt, pb, *kh, sh)?, padded(w, pl, pr, *kw, sw)?]
}
"BatchNormalization" => {
let channels = *a.get(1).ok_or_else(|| NetError::Shape("BatchNormalization rank must be at least two".into()))?;
for j in 1..5 {
if weight(n, g, j)?.dims != [channels] {
return shape_err("BatchNormalization parameters must be channel vectors");
}
}
if floats(n, g, 4)?.iter().any(|v| *v < 0.0) {
return shape_err("BatchNormalization has negative variance");
}
let eps = float(n, "epsilon", 1e-5);
let (s, b, m, v) = (floats(n, g, 1)?, floats(n, g, 2)?, floats(n, g, 3)?, floats(n, g, 4)?);
for (((&s, &b), &m), &v) in s.iter().zip(b).zip(m).zip(v) {
let variance = v + eps;
let coefficient = s / variance.sqrt();
let offset = b - m * coefficient;
if !variance.is_finite() || !coefficient.is_finite() || !offset.is_finite() {
return shape_err("BatchNormalization coefficients overflow");
}
}
a.clone()
}
"PRelu" => {
if broadcast(a, &weight(n, g, 1)?.dims)? != *a {
return shape_err("PRelu slope would enlarge its input");
}
a.clone()
}
"Add" | "Sub" | "Mul" => {
let b = n.inputs.get(1).and_then(|s| shapes.get(s.as_str())).ok_or_else(|| NetError::Shape("missing second input".into()))?;
broadcast(a, b)?
}
"Gemm" => {
let b = &weight(n, g, 1)?.dims;
let [ar, ac] = a.as_slice() else {
return shape_err("Gemm needs matrix A");
};
let [br, bc] = b.as_slice() else {
return shape_err("Gemm needs matrix B");
};
let (m, k) = if n.int("transA").unwrap_or(0) != 0 { (*ac, *ar) } else { (*ar, *ac) };
let (bk, cols) = if n.int("transB").unwrap_or(0) != 0 { (*bc, *br) } else { (*br, *bc) };
if k != bk {
return shape_err("Gemm inner dimensions differ");
}
let s = vec![m, cols];
if n.inputs.get(2).is_some_and(|s| !s.is_empty()) && broadcast(&s, &weight(n, g, 2)?.dims)? != s {
return shape_err("Gemm bias would enlarge the output");
}
s
}
"Flatten" => {
let at = axis(n, a.len())?;
vec![
elements(a.get(..at).ok_or_else(|| NetError::Shape("bad axis".into()))?)?,
elements(a.get(at..).ok_or_else(|| NetError::Shape("bad axis".into()))?)?,
]
}
"GlobalAveragePool" => {
let (c, _, _) = image(a)?;
vec![1, c, 1, 1]
}
"MaxPool" => {
let (c, h, w) = image(a)?;
let (kh, kw) = pair(n, "kernel_shape", 1)?;
let (sh, sw) = pair(n, "strides", 1)?;
if h < kh || w < kw {
return shape_err("MaxPool kernel exceeds input");
}
vec![1, c, (h - kh) / sh + 1, (w - kw) / sw + 1]
}
"Resize" => {
let (c, h, w) = image(a)?;
let s = floats(n, g, 2)?;
if s.len() != 4
|| s.first() != Some(&1.0)
|| s.get(1) != Some(&1.0)
|| s.iter().any(|&v| !(1.0..=16.0).contains(&v) || v.fract() != 0.0)
{
return unsupported("Resize supports whole spatial scales from 1 to 16");
}
vec![
1,
c,
h.checked_mul(s.get(2).copied().unwrap_or(0.0) as usize).ok_or_else(|| NetError::Shape("Resize overflow".into()))?,
w.checked_mul(s.get(3).copied().unwrap_or(0.0) as usize).ok_or_else(|| NetError::Shape("Resize overflow".into()))?,
]
}
"Transpose" => {
let perm: Vec<i64> = n.ints("perm").map_or_else(|| (0..a.len() as i64).rev().collect(), Vec::from);
let mut seen = std::collections::HashSet::new();
let mut s = Vec::new();
if perm.len() != a.len() {
return shape_err("Transpose rank mismatch");
}
for d in perm {
let d = usize::try_from(d).map_err(|_| NetError::Shape("negative permutation".into()))?;
if !seen.insert(d) {
return shape_err("repeated permutation axis");
}
s.push(*a.get(d).ok_or_else(|| NetError::Shape("bad permutation axis".into()))?);
}
s
}
"Reshape" => {
let Data::I64(dims) = &weight(n, g, 1)?.data else {
return shape_err("Reshape datatype");
};
if dims.len() > 8 {
return shape_err("Reshape rank exceeds eight");
}
let total = elements(a)?;
let mut s = Vec::new();
let mut infer = None;
for (j, &v) in dims.iter().enumerate() {
s.push(match v {
-1 if infer.is_none() => {
infer = Some(j);
1
}
0 => *a.get(j).ok_or_else(|| NetError::Shape("Reshape copies missing dimension".into()))?,
v if v > 0 => usize::try_from(v).map_err(|_| NetError::Shape("Reshape dimension overflow".into()))?,
_ => return shape_err("invalid Reshape dimension"),
});
}
let known = elements(&s)?;
if let Some(j) = infer {
if !total.is_multiple_of(known) {
return shape_err("Reshape does not divide evenly");
}
*s.get_mut(j).ok_or_else(|| NetError::Shape("missing inferred dimension".into()))? = total / known;
}
if elements(&s)? != total {
return shape_err("Reshape changes element count");
}
s
}
_ => a.clone(),
};
let out_count = elements(&out)?;
let mut copied_constants = 0usize;
for (j, name) in n.inputs.iter().enumerate() {
// Execution borrows operator weights, but get(0) and binary get(1) materialize constants.
if (j == 0 || (j == 1 && matches!(n.op.as_str(), "Add" | "Sub" | "Mul")))
&& let Some(w) = g.weights.get(name)
{
copied_constants =
copied_constants.checked_add(elements(&w.dims)?).ok_or_else(|| NetError::Shape("constant copies overflow".into()))?;
}
}
if n.op == "Conv" {
let wd = &weight(n, g, 1)?.dims;
let [co, cg, kh, kw] = wd.as_slice() else {
return shape_err("Conv rank");
};
let block = 256.min(out.get(2).copied().unwrap_or(1) * out.get(3).copied().unwrap_or(1));
let lower = cg
.checked_mul(*kh)
.and_then(|v| v.checked_mul(*kw))
.and_then(|v| v.checked_mul(block))
.filter(|&v| v <= MAX_TENSOR)
.ok_or_else(|| NetError::Shape("convolution scratch exceeds the limit".into()))?;
scratch_lower = scratch_lower.max(lower);
scratch_product = scratch_product.max(co * block);
}
let scratch = scratch_lower + scratch_product;
if live
.checked_add(out_count)
.and_then(|v| v.checked_add(scratch))
.and_then(|v| v.checked_add(copied_constants))
.is_none_or(|v| v > 128 * 1024 * 1024)
{
return shape_err("tensors and scratch exceed the 512 MiB working limit");
}
live = live
.checked_add(out_count)
.filter(|&n| n <= 128 * 1024 * 1024)
.ok_or_else(|| NetError::Shape("live tensors exceed the 512 MiB working limit".into()))?;
per_node.push(out.clone());
shapes.insert(n.outputs.first().ok_or_else(|| NetError::Shape("no output".into()))?, out);
let mut removed = std::collections::HashSet::new();
for name in &n.inputs {
if last.get(name.as_str()) == Some(&i) && !g.weights.contains_key(name) && removed.insert(name) {
live = live.saturating_sub(shapes.get(name.as_str()).map_or(0, |s| elements(s).unwrap_or(0)));
}
}
}
let outputs =
g.outputs.iter().map(|s| shapes.get(s.as_str()).cloned().ok_or_else(|| NetError::Shape("missing output".into()))).collect::<Result<_>>()?;
Ok((per_node, outputs))
}
Reference in New Issue View Git Blame Copy Permalink