hrx/image/
views.rs

1use super::*;
2
3/// Numerical contracts for RGB model inputs and outputs.
4#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
5pub enum RgbEncoding {
6    /// Encode byte / 255; decode float * 255 without clipping or rounding.
7    Unit,
8    /// Encode byte / 127.5 - 1; decode round-even((clamp(x,-1,1)+1)*127.5).
9    Symmetric,
10}
11
12impl ImageOps {
13    /// Split packed U8 NHWC into interleaved F32 NCHW views. A factor of f
14    /// produces f² views per image: view (dy,dx) samples image[y*f+dy,x*f+dx].
15    /// This is a lossless layout permutation before the selected normalization,
16    /// not a resize. Both spatial dimensions must divide the factor.
17    pub fn prepare_encode_rgb_views(
18        &self,
19        input: &TensorDesc,
20        factor: usize,
21        encoding: RgbEncoding,
22    ) -> Result<Arc<PreparedModel>> {
23        if input.dtype() != DType::U8 || input.layout() != Layout::Nhwc {
24            return Err(Error::Message("RGB view encoding requires U8 NHWC".into()));
25        }
26        self.view_plans
27            .get_or_prepare((input.clone(), factor, encoding), || {
28                self.rgb_views_fragment(input, factor, encoding)?.prepare(3)
29            })
30    }
31
32    /// Reassemble interleaved F32 NCHW views as F32 NHWC byte-level RGB.
33    /// The view batch must be a multiple of factor². Unit values retain the
34    /// network's fractional output; Symmetric values use the documented clamp.
35    pub fn prepare_decode_rgb_views(
36        &self,
37        input: &TensorDesc,
38        factor: usize,
39        encoding: RgbEncoding,
40    ) -> Result<Arc<PreparedModel>> {
41        if input.dtype() != DType::F32 || input.layout() != Layout::Nchw {
42            return Err(Error::Message("RGB view decoding requires F32 NCHW".into()));
43        }
44        self.view_plans
45            .get_or_prepare((input.clone(), factor, encoding), || {
46                self.rgb_views_fragment(input, factor, encoding)?.prepare(3)
47            })
48    }
49
50    /// Record RGB view encoding with directly bound input pixels.
51    pub fn encode_rgb_views_fragment(
52        &self,
53        input: &TensorDesc,
54        factor: usize,
55        encoding: RgbEncoding,
56    ) -> Result<ModelFragment> {
57        if input.dtype() != DType::U8 || input.layout() != Layout::Nhwc {
58            return Err(Error::Message("RGB view encoding requires U8 NHWC".into()));
59        }
60        self.rgb_views_fragment(input, factor, encoding)
61    }
62
63    /// Record RGB view decoding with directly bound model output.
64    pub fn decode_rgb_views_fragment(
65        &self,
66        input: &TensorDesc,
67        factor: usize,
68        encoding: RgbEncoding,
69    ) -> Result<ModelFragment> {
70        if input.dtype() != DType::F32 || input.layout() != Layout::Nchw {
71            return Err(Error::Message("RGB view decoding requires F32 NCHW".into()));
72        }
73        self.rgb_views_fragment(input, factor, encoding)
74    }
75
76    fn rgb_views_fragment(
77        &self,
78        input: &TensorDesc,
79        factor: usize,
80        encoding: RgbEncoding,
81    ) -> Result<ModelFragment> {
82        if !input.is_contiguous()
83            || input.is_empty()
84            || input.shape().len() != 4
85            || input.elements() > i32::MAX as usize
86            || factor == 0
87        {
88            return Err(Error::Message("invalid RGB view dimensions".into()));
89        }
90        let ff = factor
91            .checked_mul(factor)
92            .ok_or_else(|| Error::Message("RGB view factor overflow".into()))?;
93        let encode = input.dtype() == DType::U8;
94        let s = input.shape();
95        let (batch, height, width) = if encode {
96            if s[3] != 3 || !s[1].is_multiple_of(factor) || !s[2].is_multiple_of(factor) {
97                return Err(Error::Message(
98                    "RGB axes must divide the view factor".into(),
99                ));
100            }
101            (s[0], s[1], s[2])
102        } else {
103            if s[1] != 3 || !s[0].is_multiple_of(ff) {
104                return Err(Error::Message(
105                    "RGB view batch must divide factor squared".into(),
106                ));
107            }
108            (s[0] / ff, s[2] * factor, s[3] * factor)
109        };
110        let vh = height / factor;
111        let vw = width / factor;
112        let output = if encode {
113            TensorDesc::new(DType::F32, vec![batch * ff, 3, vh, vw])?.with_layout(Layout::Nchw)?
114        } else {
115            TensorDesc::new(DType::F32, vec![batch, height, width, 3])?.with_layout(Layout::Nhwc)?
116        };
117        {
118            let count = input.elements();
119            let address = if encode {
120                format!(
121                    r#"
122    %x = index.rem %i, %vw : index
123    %yy = index.div %i, %vw : index
124    %y = index.rem %yy, %vh : index
125    %cc = index.div %i, %plane : index
126    %ch = index.rem %cc, %three : index
127    %view = index.div %cc, %three : index
128    %image = index.div %view, %ff : index
129    %part = index.rem %view, %ff : index
130    %dx = index.rem %part, %factor : index
131    %dy = index.div %part, %factor : index
132    %xf = index.mul %x, %factor : index
133    %xx = index.add %xf, %dx : index
134    %yf = index.mul %y, %factor : index
135    %yp = index.add %yf, %dy : index
136    %ih = index.mul %image, %height : index
137    %row = index.add %ih, %yp : index
138    %rw = index.mul %row, %width : index
139    %pixel = index.add %rw, %xx : index
140    %base = index.mul %pixel, %three : index
141    %addr = index.add %base, %ch : index
142    %at = index.assume %addr [range(%addr, 0, {count})] : index
143    %byte = view.load %src[%at] : view<{count}xi8> -> i8
144    %uint = scalar.extui %byte : i8 to i32
145    %value = scalar.uitofp %uint : i32 to f32
146"#
147                )
148            } else {
149                format!(
150                    r#"
151    %ch = index.rem %i, %three : index
152    %pixel = index.div %i, %three : index
153    %x = index.rem %pixel, %width : index
154    %yy = index.div %pixel, %width : index
155    %y = index.rem %yy, %height : index
156    %image = index.div %yy, %height : index
157    %dx = index.rem %x, %factor : index
158    %dy = index.rem %y, %factor : index
159    %xx = index.div %x, %factor : index
160    %yp = index.div %y, %factor : index
161    %df = index.mul %dy, %factor : index
162    %part = index.add %df, %dx : index
163    %im = index.mul %image, %ff : index
164    %view = index.add %im, %part : index
165    %vc = index.mul %view, %three : index
166    %channel = index.add %vc, %ch : index
167    %hh = index.mul %channel, %vh : index
168    %row = index.add %hh, %yp : index
169    %rw = index.mul %row, %vw : index
170    %addr = index.add %rw, %xx : index
171    %at = index.assume %addr [range(%addr, 0, {count})] : index
172    %value = view.load %src[%at] : view<{count}xf32> -> f32
173"#
174                )
175            };
176            let arithmetic = match (encode, encoding) {
177                (true, RgbEncoding::Unit) => {
178                    "%scale = scalar.constant 255.0 : f32\n%result = scalar.divf %value, %scale : f32"
179                }
180                (true, RgbEncoding::Symmetric) => {
181                    "%scale = scalar.constant 127.5 : f32\n%onef = scalar.constant 1.0 : f32\n%scaled = scalar.divf %value, %scale : f32\n%result = scalar.subf %scaled, %onef : f32"
182                }
183                (false, RgbEncoding::Unit) => {
184                    "%scale = scalar.constant 255.0 : f32\n%result = scalar.mulf %value, %scale : f32"
185                }
186                (false, RgbEncoding::Symmetric) => {
187                    "%scale = scalar.constant 127.5 : f32\n%onef = scalar.constant 1.0 : f32\n%negative = scalar.constant -1.0 : f32\n%low = scalar.maxnumf %value, %negative : f32\n%clamped = scalar.minnumf %low, %onef : f32\n%shifted = scalar.addf %clamped, %onef : f32\n%scaled = scalar.mulf %shifted, %scale : f32\n%result = scalar.roundevenf %scaled : f32"
188                }
189            };
190            let ty = if encode { "i8" } else { "f32" };
191            let source = format!(
192                r#"
193amdgpu.target<gfx11-generic> @views_target {{subgroup_size = 32}}
194kernel.def target(@views_target) export("rgb_views") @rgb_views(%unused: index) {{
195  %one = index.constant 1 : index
196  %threads = index.constant 256 : index
197  %grid = index.constant {grid} : index
198  kernel.launch.config workgroups(%grid, %one, %one) workgroup_size(%threads, %one, %one) : index
199}} launch(%unused: index, %input: buffer, %output: buffer) {{
200  %off = index.constant 0 : offset
201  %threads = index.constant 256 : index
202  %lane = kernel.workitem.id<x> : index
203  %group = kernel.workgroup.id<x> : index
204  %basei = index.mul %group, %threads : index
205  %raw = index.add %basei, %lane : index
206  %count = index.constant {count} : index
207  %three = index.constant 3 : index
208  %width = index.constant {width} : index
209  %height = index.constant {height} : index
210  %vw = index.constant {vw} : index
211  %vh = index.constant {vh} : index
212  %plane = index.constant {plane} : index
213  %factor = index.constant {factor} : index
214  %ff = index.constant {ff} : index
215  %ig = buffer.assume.memory_space<global> %input : buffer
216  %og = buffer.assume.memory_space<global> %output : buffer
217  %src = buffer.view %ig[%off] : buffer -> view<{count}x{ty}>
218  %dst = buffer.view %og[%off] : buffer -> view<{count}xf32>
219  %valid = index.cmp ult, %raw, %count : index
220  scf.if %valid {{
221    %i = index.assume %raw [range(%raw, 0, {count})] : index
222    {address}
223    {arithmetic}
224    view.store %result, %dst[%i] : f32, view<{count}xf32>
225  }}
226  kernel.return
227}}
228"#,
229                grid = count.div_ceil(256),
230                plane = vh * vw
231            );
232            let mut model = ModelSession::in_context(&self.context)?;
233            let src = model.allocate(input.bytes())?;
234            let dst = model.allocate(output.bytes())?;
235            // Permutations above are bijections within checked input/output extents.
236            let kernel =
237                unsafe { model.compile(&[(&source, Specialization::new("rgb_views"))])? }[0];
238            unsafe {
239                model.freeze(&self.context)?.fragment(
240                    &[Command::Dispatch(Dispatch::indices(
241                        kernel,
242                        [0],
243                        [count.div_ceil(256) as u32, 1, 1],
244                        vec![src.read(), dst.write()],
245                    ))],
246                    &[(src, input.clone())],
247                    &[(dst, output)],
248                )
249            }
250        }
251    }
252
253    /// Encode resident pixels without a host boundary.
254    pub fn encode_rgb_views(
255        &self,
256        input: &DeviceTensor,
257        factor: usize,
258        encoding: RgbEncoding,
259    ) -> Result<Inference> {
260        self.context.validate(input)?;
261        self.prepare_encode_rgb_views(input.desc(), factor, encoding)?
262            .submit(std::slice::from_ref(input))
263    }
264    /// Decode resident model views without a host boundary.
265    pub fn decode_rgb_views(
266        &self,
267        input: &DeviceTensor,
268        factor: usize,
269        encoding: RgbEncoding,
270    ) -> Result<Inference> {
271        self.context.validate(input)?;
272        self.prepare_decode_rgb_views(input.desc(), factor, encoding)?
273            .submit(std::slice::from_ref(input))
274    }
275}
hrx/image/views.rs

hrx/image/
views.rs
