hrx-rs 0.6.0

Owned GPU and NPU execution with Loom and IRON compilation
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149

use super::*;
use std::fmt::Write;

/// Least-squares orientation-preserving 2D similarity, mapping source points
/// to template points in FP32. No robust outlier rejection. Collinear points
/// are allowed if the fit has nonzero scale; collapsed/non-finite fits fail.
pub fn fit_similarity_2d(points: &[[f32; 2]], template: &[[f32; 2]]) -> Result<[[f32; 3]; 2]> {
    if points.len() < 2
        || points.len() != template.len()
        || !points
            .iter()
            .chain(template)
            .flatten()
            .all(|v| v.is_finite())
    {
        return Err(Error::Message(
            "similarity needs matching finite point sets".into(),
        ));
    }
    let mean = |points: &[[f32; 2]]| {
        [
            points.iter().map(|p| p[0]).sum::<f32>() / points.len() as f32,
            points.iter().map(|p| p[1]).sum::<f32>() / points.len() as f32,
        ]
    };
    let p = mean(points);
    let q = mean(template);
    let (mut den, mut a, mut b) = (0., 0., 0.);
    for (s, d) in points.iter().zip(template) {
        let (x, y, u, v) = (s[0] - p[0], s[1] - p[1], d[0] - q[0], d[1] - q[1]);
        den += x * x + y * y;
        a += x * u + y * v;
        b += x * v - y * u;
    }
    a /= den;
    b /= den;
    let result = [
        [a, -b, q[0] - a * p[0] + b * p[1]],
        [b, a, q[1] - b * p[0] - a * p[1]],
    ];
    invert_affine(result)?;
    Ok(result)
}

impl ImageOps {
    /// Prepare FP32 similarity fitting for `[batch,points,2]`. Positive
    /// non-overlapping strides allow direct landmark views of detection rows.
    /// Template coordinates are immutable shape-plan data (2..=64 points).
    /// Outputs: forward `[batch,2,3]`, inverse `[batch,2,3]`, status I32 `[batch]`.
    /// Status is zero for success and one for invalid/degenerate geometry; both
    /// matrices are NaN on failure. Consumers must preserve/check status, not
    /// treat a black crop from invalid geometry as a successful alignment.
    pub fn prepare_similarity_2d(
        &self,
        points: &TensorDesc,
        template: &[[f32; 2]],
    ) -> Result<Arc<PreparedModel>> {
        let shape = points.shape();
        if points.dtype() != DType::F32
            || points.layout() != Layout::General
            || shape.len() != 3
            || shape[2] != 2
            || shape[0] == 0
            || !(2..=64).contains(&template.len())
            || shape[1] != template.len()
            || points.bytes() / 4 > i32::MAX as usize
            || !template.iter().flatten().all(|v| v.is_finite())
        {
            return Err(Error::Message(
                "similarity needs F32 [batch,points,2] and a finite template".into(),
            ));
        }
        let key = (
            points.clone(),
            template.iter().flatten().map(|v| v.to_bits()).collect(),
        );
        self.similarity_plans.get_or_prepare(key, || {
            let batch = shape[0];
            let count = template.len();
            let elements = points.bytes() / 4;
            let matrices = TensorDesc::new(DType::F32, vec![batch, 2, 3])?;
            let status = TensorDesc::new(DType::I32, vec![batch])?;
            if matrices.elements() > i32::MAX as usize { return Err(Error::Message("similarity output too large".into())); }
            let q = [0,1].map(|axis| template.iter().map(|p| p[axis]).sum::<f32>() / count as f32);
            if !q.iter().all(|v| v.is_finite()) { return Err(Error::Message("template mean overflow".into())); }
            let mut body = String::new();
            for (i, point) in template.iter().enumerate() {
                for (axis, letter) in [(0,"x"),(1,"y")] {
                    let offset = i * points.strides()[1] + axis * points.strides()[2];
                    writeln!(body, r#"
    %off{letter}{i} = index.constant {offset} : index
    %raw{letter}{i} = index.add %inputbase, %off{letter}{i} : index
    %at{letter}{i} = index.assume %raw{letter}{i} [range(%raw{letter}{i}, 0, {last})] : index
    %{letter}{i} = view.load %src[%at{letter}{i}] : view<{elements}xf32> -> f32
    %abs{letter}{i} = scalar.absf %{letter}{i} : f32
    %finite{letter}{i} = scalar.cmpf ole, %abs{letter}{i}, %max : f32
    %sum{letter}{next} = scalar.addf %sum{letter}{i}, %{letter}{i} : f32
"#, last=elements-1, next=i+1).unwrap();
                }
                writeln!(body, "    %xyfinite{i} = scalar.andi %finitex{i}, %finitey{i} : i1\n    %finite{next} = scalar.andi %finite{i}, %xyfinite{i} : i1",next=i+1).unwrap();
                for (axis, letter) in [(0,"u"),(1,"v")] {
                    writeln!(body, "    %{letter}{i} = scalar.constant {:.9e} : f32",point[axis]-q[axis]).unwrap();
                }
            }
            writeln!(body,"    %px = scalar.divf %sumx{count}, %n : f32\n    %py = scalar.divf %sumy{count}, %n : f32").unwrap();
            for i in 0..count {
                writeln!(body, r#"
    %sx{i} = scalar.subf %x{i}, %px : f32
    %sy{i} = scalar.subf %y{i}, %py : f32
    %xx{i} = scalar.mulf %sx{i}, %sx{i} : f32
    %yy{i} = scalar.mulf %sy{i}, %sy{i} : f32
    %xy{i} = scalar.addf %xx{i}, %yy{i} : f32
    %den{next} = scalar.addf %den{i}, %xy{i} : f32
    %xu{i} = scalar.mulf %sx{i}, %u{i} : f32
    %yv{i} = scalar.mulf %sy{i}, %v{i} : f32
    %dot{i} = scalar.addf %xu{i}, %yv{i} : f32
    %dot{next}sum = scalar.addf %dot{i}sum, %dot{i} : f32
    %xv{i} = scalar.mulf %sx{i}, %v{i} : f32
    %yu{i} = scalar.mulf %sy{i}, %u{i} : f32
    %cross{i} = scalar.subf %xv{i}, %yu{i} : f32
    %cross{next}sum = scalar.addf %cross{i}sum, %cross{i} : f32
"#,next=i+1).unwrap();
            }
            let mut source = include_str!("similarity_2d.loom").replace("@BODY@", &body);
            for (name,value) in [("BATCH",batch),("BATCH_LAST",batch-1),("GRID",batch.div_ceil(64)),
                ("INPUT",elements),("STRIDE",points.strides()[0]),("COUNT",count),
                ("OUTPUT",batch*6),("OUTPUT_LAST",batch*6-1)] {
                source=source.replace(&format!("@{name}@"),&value.to_string());
            }
            source=source.replace("@QX@", &format!("{:.9e}",q[0])).replace("@QY@", &format!("{:.9e}",q[1]));
            let mut model = ModelSession::in_context(&self.context)?;
            let input = model.allocate(points.bytes())?;
            let forward = model.allocate(matrices.bytes())?;
            let inverse = model.allocate(matrices.bytes())?;
            let flags = model.allocate(status.bytes())?;
            let kernel = unsafe { model.compile(&[(&source, Specialization::new("similarity_2d"))])? }[0];
            unsafe { model.freeze(&self.context)?.prepare(&[Command::Dispatch(Dispatch::indices(
                kernel,[0],[batch.div_ceil(64) as u32,1,1],vec![input.read(),forward.write(),inverse.write(),flags.write()]))],
                &[(input,points.clone())],&[(forward,matrices.clone()),(inverse,matrices),(flags,status)],3) }
        })
    }

    /// Fit resident landmarks; no landmark or matrix readback is performed.
    pub fn similarity_2d(&self, points: &DeviceTensor, template: &[[f32; 2]]) -> Result<Inference> {
        self.context.validate(points)?;
        self.prepare_similarity_2d(points.desc(), template)?
            .submit(std::slice::from_ref(points))
    }
}