hrx/
image.rs

1//! Device-resident image transformations with bounded reusable plans.
2
3use crate::{
4    Error, Result,
5    inference::{Inference, ModelContext, PreparedModel},
6    loom::Specialization,
7    model::{Command, Dispatch, ModelFragment, ModelSession},
8    plan_cache::PlanCache,
9    tensor::{DType, DeviceTensor, Layout, TensorDesc},
10};
11use std::sync::Arc;
12mod affine;
13mod composite;
14mod normalize;
15mod resize;
16mod similarity;
17mod views;
18pub use affine::{RgbSampling, invert_affine};
19pub use composite::RgbComposite;
20pub use resize::RgbResize;
21pub use similarity::fit_similarity_2d;
22pub use views::RgbEncoding;
23
24type NormalizationKey = (TensorDesc, [u32; 3], [u32; 3]);
25
26/// GPU image operations sharing the caller's model context and compilation cache.
27pub struct ImageOps {
28    context: ModelContext,
29    patch_plans: PlanCache<(TensorDesc, usize), PreparedModel>,
30    normalize_plans: PlanCache<NormalizationKey, PreparedModel>,
31    affine_plans: PlanCache<(TensorDesc, usize, usize, usize, RgbSampling), PreparedModel>,
32    resize_plans: PlanCache<(TensorDesc, RgbResize), PreparedModel>,
33    similarity_plans: PlanCache<(TensorDesc, Vec<u32>), PreparedModel>,
34    view_plans: PlanCache<(TensorDesc, usize, RgbEncoding), PreparedModel>,
35    composite_plans: PlanCache<(TensorDesc, TensorDesc, RgbComposite), PreparedModel>,
36}
37impl ImageOps {
38    /// Create a bounded shape cache. Preparation is lazy, never per warm request.
39    pub fn new(context: &ModelContext, cache_capacity: usize) -> Result<Self> {
40        Ok(Self {
41            context: context.clone(),
42            patch_plans: PlanCache::new(cache_capacity, PreparedModel::is_idle)?,
43            normalize_plans: PlanCache::new(cache_capacity, PreparedModel::is_idle)?,
44            affine_plans: PlanCache::new(cache_capacity, PreparedModel::is_idle)?,
45            resize_plans: PlanCache::new(cache_capacity, PreparedModel::is_idle)?,
46            similarity_plans: PlanCache::new(cache_capacity, PreparedModel::is_idle)?,
47            view_plans: PlanCache::new(cache_capacity, PreparedModel::is_idle)?,
48            composite_plans: PlanCache::new(cache_capacity, PreparedModel::is_idle)?,
49        })
50    }
51
52    /// Prepare NCHW f32 patchification to `[batch, patches, channels*p*p]`.
53    /// Within a patch, values are in channel, row, column order. No arithmetic
54    /// changes the values, including non-finite payloads. Both axes must divide p.
55    pub fn prepare_patchify(&self, input: &TensorDesc, patch: usize) -> Result<Arc<PreparedModel>> {
56        self.patch_plans.get_or_prepare((input.clone(), patch), || {
57            self.patchify_fragment(input, patch)?.prepare(3)
58        })
59    }
60
61    /// Recordable patchification for direct composition with a model graph.
62    pub fn patchify_fragment(&self, input: &TensorDesc, patch: usize) -> Result<ModelFragment> {
63        if input.dtype() != DType::F32
64            || input.layout() != Layout::Nchw
65            || !input.is_contiguous()
66            || input.is_empty()
67        {
68            return Err(Error::Message(
69                "patchify requires nonempty contiguous NCHW f32".into(),
70            ));
71        }
72        let [batch, channels, height, width]: [usize; 4] = input
73            .shape()
74            .try_into()
75            .map_err(|_| Error::Message("patchify requires four axes".into()))?;
76        if patch == 0
77            || !height.is_multiple_of(patch)
78            || !width.is_multiple_of(patch)
79            || input.elements() > i32::MAX as usize
80        {
81            return Err(Error::Message("invalid patch or image size".into()));
82        }
83        let count = input.elements();
84        let cells = (height / patch) * (width / patch);
85        let patch_elements = channels * patch * patch;
86        let output = TensorDesc::new(DType::F32, vec![batch, cells, patch_elements])?;
87        let source = format!(
88            r#"
89amdgpu.target<gfx11-generic> @image_target {{subgroup_size = 32}}
90kernel.def target(@image_target) export("patchify") @patchify(%unused: index) {{
91  %one = index.constant 1 : index
92  %threads = index.constant 256 : index
93  %grid = index.constant {grid} : index
94  kernel.launch.config workgroups(%grid, %one, %one) workgroup_size(%threads, %one, %one) : index
95}} launch(%unused: index, %input: buffer, %output: buffer) {{
96  %zero = index.constant 0 : offset
97  %threads = index.constant 256 : index
98  %lane = kernel.workitem.id<x> : index
99  %group = kernel.workgroup.id<x> : index
100  %base = index.mul %group, %threads : index
101  %raw = index.add %base, %lane : index
102  %count = index.constant {count} : index
103  %p = index.constant {patch} : index
104  %pp = index.constant {pp} : index
105  %pe = index.constant {patch_elements} : index
106  %ch = index.constant {channels} : index
107  %pxs = index.constant {pxs} : index
108  %cells = index.constant {cells} : index
109  %h = index.constant {height} : index
110  %w = index.constant {width} : index
111  %inputg = buffer.assume.memory_space<global> %input : buffer
112  %outputg = buffer.assume.memory_space<global> %output : buffer
113  %src = buffer.view %inputg[%zero] : buffer -> view<{count}xf32>
114  %dst = buffer.view %outputg[%zero] : buffer -> view<{count}xf32>
115  %valid = index.cmp ult, %raw, %count : index
116  scf.if %valid {{
117    %i = index.assume %raw [range(%raw, 0, {count})] : index
118    %dx = index.rem %i, %p : index
119    %qdy = index.div %i, %p : index
120    %dy = index.rem %qdy, %p : index
121    %qc = index.div %i, %pp : index
122    %c = index.rem %qc, %ch : index
123    %cell = index.div %i, %pe : index
124    %image = index.div %cell, %cells : index
125    %incell = index.rem %cell, %cells : index
126    %px = index.rem %incell, %pxs : index
127    %py = index.div %incell, %pxs : index
128    %x0 = index.mul %px, %p : index
129    %x = index.add %x0, %dx : index
130    %y0 = index.mul %py, %p : index
131    %y = index.add %y0, %dy : index
132    %ic0 = index.mul %image, %ch : index
133    %ic = index.add %ic0, %c : index
134    %iy0 = index.mul %ic, %h : index
135    %iy = index.add %iy0, %y : index
136    %ix0 = index.mul %iy, %w : index
137    %ix = index.add %ix0, %x : index
138    %j = index.assume %ix [range(%ix, 0, {count})] : index
139    %value = view.load %src[%j] : view<{count}xf32> -> f32
140    view.store %value, %dst[%i] : f32, view<{count}xf32>
141  }}
142  kernel.return
143}}
144"#,
145            grid = count.div_ceil(256),
146            pp = patch * patch,
147            pxs = width / patch
148        );
149        let mut model = ModelSession::in_context(&self.context)?;
150        let src = model.allocate(input.bytes())?;
151        let dst = model.allocate(output.bytes())?;
152        // The generated kernel's scalar and memory extents are fixed above.
153        let kernel = unsafe { model.compile(&[(&source, Specialization::new("patchify"))])? }[0];
154        let definition = model.freeze(&self.context)?;
155        unsafe {
156            definition.fragment(
157                &[Command::Dispatch(Dispatch::indices(
158                    kernel,
159                    [0],
160                    [count.div_ceil(256) as u32, 1, 1],
161                    vec![src.read(), dst.write()],
162                ))],
163                &[(src, input.clone())],
164                &[(dst, output)],
165            )
166        }
167    }
168
169    /// Patchify a resident image batch without a host wait or readback.
170    pub fn patchify(&self, input: &DeviceTensor, patch: usize) -> Result<Inference> {
171        self.context.validate(input)?;
172        self.prepare_patchify(input.desc(), patch)?
173            .submit(std::slice::from_ref(input))
174    }
175}
hrx/image.rs

hrx/
image.rs
