hrx/image/
normalize.rs

1use super::*;
2
3impl ImageOps {
4    /// Prepare packed RGB U8 NHWC to F32 NCHW conversion, computing
5    /// `(pixel / 255 - mean[channel]) / std[channel]` in that order.
6    pub fn prepare_normalize_rgb(
7        &self,
8        input: &TensorDesc,
9        mean: [f32; 3],
10        std: [f32; 3],
11    ) -> Result<Arc<PreparedModel>> {
12        self.normalize_plans.get_or_prepare(
13            (input.clone(), mean.map(f32::to_bits), std.map(f32::to_bits)),
14            || self.normalize_rgb_fragment(input, mean, std)?.prepare(3),
15        )
16    }
17
18    /// Recordable normalization with no slot pool or host staging. Connect its
19    /// output directly to a model fragment in a caller-owned graph.
20    pub fn normalize_rgb_fragment(
21        &self,
22        input: &TensorDesc,
23        mean: [f32; 3],
24        std: [f32; 3],
25    ) -> Result<ModelFragment> {
26        if input.dtype() != DType::U8
27            || input.layout() != Layout::Nhwc
28            || !input.is_contiguous()
29            || input.is_empty()
30            || input.shape().last() != Some(&3)
31            || input.elements() > i32::MAX as usize
32            || mean.iter().any(|v| !v.is_finite())
33            || std.iter().any(|v| !v.is_finite() || *v <= 0.)
34        {
35            return Err(Error::Message(
36                "normalize requires nonempty NHWC RGB and finite means/positive deviations".into(),
37            ));
38        }
39        let [batch, height, width, _]: [usize; 4] = input
40            .shape()
41            .try_into()
42            .map_err(|_| Error::Message("expected NHWC input".into()))?;
43        let count = input.elements();
44        let output = TensorDesc::new(DType::F32, vec![batch, 3, height, width])?
45            .with_layout(Layout::Nchw)?;
46        let source = format!(
47            r#"
48amdgpu.target<gfx11-generic> @normalize_target {{subgroup_size = 32}}
49kernel.def target(@normalize_target) export("normalize_rgb") @normalize_rgb(%unused: index) {{
50  %one = index.constant 1 : index
51  %threads = index.constant 256 : index
52  %grid = index.constant {grid} : index
53  kernel.launch.config workgroups(%grid, %one, %one) workgroup_size(%threads, %one, %one) : index
54}} launch(%unused: index, %input: buffer, %output: buffer) {{
55  %off = index.constant 0 : offset
56  %zero = index.constant 0 : index
57  %one = index.constant 1 : index
58  %three = index.constant 3 : index
59  %plane = index.constant {plane} : index
60  %image_size = index.constant {image_size} : index
61  %count = index.constant {count} : index
62  %threads = index.constant 256 : index
63  %lane = kernel.workitem.id<x> : index
64  %group = kernel.workgroup.id<x> : index
65  %base = index.mul %group, %threads : index
66  %raw = index.add %base, %lane : index
67  %ig = buffer.assume.memory_space<global> %input : buffer
68  %og = buffer.assume.memory_space<global> %output : buffer
69  %iv = buffer.view %ig[%off] : buffer -> view<{count}xi8>
70  %ov = buffer.view %og[%off] : buffer -> view<{count}xf32>
71  %valid = index.cmp ult, %raw, %count : index
72  scf.if %valid {{
73    %i = index.assume %raw [range(%raw, 0, {count})] : index
74    %image = index.div %i, %image_size : index
75    %pixel = index.rem %i, %plane : index
76    %channel0 = index.div %i, %plane : index
77    %channel = index.rem %channel0, %three : index
78    %ip0 = index.mul %image, %plane : index
79    %ip = index.add %ip0, %pixel : index
80    %ic0 = index.mul %ip, %three : index
81    %ic = index.add %ic0, %channel : index
82    %at = index.assume %ic [range(%ic, 0, {count})] : index
83    %byte = view.load %iv[%at] : view<{count}xi8> -> i8
84    %uint = scalar.extui %byte : i8 to i32
85    %value = scalar.uitofp %uint : i32 to f32
86    %scale = scalar.constant 255.0 : f32
87    %unit = scalar.divf %value, %scale : f32
88    %m0 = scalar.constant {m0:.9e} : f32
89    %m1 = scalar.constant {m1:.9e} : f32
90    %m2 = scalar.constant {m2:.9e} : f32
91    %s0 = scalar.constant {s0:.9e} : f32
92    %s1 = scalar.constant {s1:.9e} : f32
93    %s2 = scalar.constant {s2:.9e} : f32
94    %red = index.cmp eq, %channel, %zero : index
95    %green = index.cmp eq, %channel, %one : index
96    %m12 = scf.select %green, %m1, %m2 : f32
97    %m = scf.select %red, %m0, %m12 : f32
98    %s12 = scf.select %green, %s1, %s2 : f32
99    %s = scf.select %red, %s0, %s12 : f32
100    %centered = scalar.subf %unit, %m : f32
101    %normalized = scalar.divf %centered, %s : f32
102    view.store %normalized, %ov[%i] : f32, view<{count}xf32>
103  }}
104  kernel.return
105}}
106"#,
107            grid = count.div_ceil(256),
108            plane = height * width,
109            image_size = height * width * 3,
110            m0 = mean[0],
111            m1 = mean[1],
112            m2 = mean[2],
113            s0 = std[0],
114            s1 = std[1],
115            s2 = std[2]
116        );
117        let mut model = ModelSession::in_context(&self.context)?;
118        let src = model.allocate(input.bytes())?;
119        let dst = model.allocate(output.bytes())?;
120        // The generated source's accesses are bounded to these descriptors.
121        let kernel =
122            unsafe { model.compile(&[(&source, Specialization::new("normalize_rgb"))])? }[0];
123        unsafe {
124            model.freeze(&self.context)?.fragment(
125                &[Command::Dispatch(Dispatch::indices(
126                    kernel,
127                    [0],
128                    [count.div_ceil(256) as u32, 1, 1],
129                    vec![src.read(), dst.write()],
130                ))],
131                &[(src, input.clone())],
132                &[(dst, output)],
133            )
134        }
135    }
136
137    /// Normalize resident RGB pixels and change layout without a host boundary.
138    pub fn normalize_rgb(
139        &self,
140        input: &DeviceTensor,
141        mean: [f32; 3],
142        std: [f32; 3],
143    ) -> Result<Inference> {
144        self.context.validate(input)?;
145        self.prepare_normalize_rgb(input.desc(), mean, std)?
146            .submit(std::slice::from_ref(input))
147    }
148}
hrx/image/normalize.rs

hrx/image/
normalize.rs
