hrx/tensor/
ops.rs

1use super::{DType, DeviceTensor, TensorDesc};
2use crate::{
3    Error, Result,
4    inference::{Inference, ModelContext, PreparedModel},
5    loom::Specialization,
6    model::{Command, Dispatch, ModelFragment, ModelSession},
7    plan_cache::PlanCache,
8};
9use std::sync::Arc;
10
11/// Bounded shape-specialized device tensor operations in one shared context.
12pub struct TensorOps {
13    pub(super) context: ModelContext,
14    pub(super) linear: PlanCache<(TensorDesc, Vec<TensorDesc>, DType), PreparedModel>,
15    finite: PlanCache<TensorDesc, PreparedModel>,
16    gather: PlanCache<(TensorDesc, usize), PreparedModel>,
17}
18impl TensorOps {
19    /// Recordable row gather with directly bound source and U32 `[rows]`
20    /// indices. Output is host-visible so final records can be mapped without
21    /// a copy; GPU consumers can bind it directly. Out-of-range indices select
22    /// row zero, so callers must validate indices before publication.
23    pub fn gather_rows_fragment(&self, input: &TensorDesc, rows: usize) -> Result<ModelFragment> {
24        self.gather_fragment(input, rows, true)
25    }
26    fn gather_fragment(
27        &self,
28        input: &TensorDesc,
29        rows: usize,
30        mapped: bool,
31    ) -> Result<ModelFragment> {
32        if input.shape().len() != 2
33            || !input.is_contiguous()
34            || input.is_empty()
35            || input.bytes() > i32::MAX as usize
36            || rows == 0
37        {
38            return Err(Error::Message(
39                "gather requires nonempty contiguous rows".into(),
40            ));
41        }
42        let output = TensorDesc::new(input.dtype(), vec![rows, input.shape()[1]])?
43            .with_layout(input.layout())?;
44        if output.bytes() > i32::MAX as usize {
45            return Err(Error::Message(
46                "gather exceeds kernel indexing limits".into(),
47            ));
48        }
49        let ids_desc = TensorDesc::new(DType::U32, vec![rows])?;
50        if ids_desc.bytes() > i32::MAX as usize {
51            return Err(Error::Message(
52                "gather index storage exceeds kernel limits".into(),
53            ));
54        }
55        let mut model = ModelSession::in_context(&self.context)?;
56        let source = model.allocate(input.bytes())?;
57        let ids = model.allocate(ids_desc.bytes())?;
58        let destination = if mapped {
59            model.allocate_shared(output.bytes())?
60        } else {
61            model.allocate(output.bytes())?
62        };
63        let mut code = include_str!("gather.loom").to_owned();
64        for (name, value) in [
65            ("INPUT", input.bytes()),
66            ("LAST_INPUT", input.bytes() - 1),
67            ("ROWS", input.shape()[0]),
68            ("LAST_ROW", input.shape()[0] - 1),
69            ("ROW_BYTES", input.shape()[1] * input.dtype().bytes()),
70            ("BUCKET", rows),
71            ("LAST_INDEX", rows - 1),
72            ("OUTPUT", output.bytes()),
73            ("LAST_OUTPUT", output.bytes() - 1),
74            ("GRID", output.bytes().div_ceil(256)),
75        ] {
76            code = code.replace(&format!("@{name}@"), &value.to_string());
77        }
78        // The embedded kernel checks indices and bounds every byte access.
79        let kernel = unsafe { model.compile(&[(&code, Specialization::new("gather"))])? }[0];
80        unsafe {
81            model.freeze(&self.context)?.fragment(
82                &[Command::Dispatch(Dispatch::indices(
83                    kernel,
84                    [0],
85                    [output.bytes().div_ceil(256) as u32, 1, 1],
86                    vec![source.read(), ids.read(), destination.write()],
87                ))],
88                &[(source, input.clone()), (ids, ids_desc)],
89                &[(destination, output)],
90            )
91        }
92    }
93    /// Create lazy bounded caches, evicting only unleased and drained plans.
94    pub fn new(context: &ModelContext, capacity: usize) -> Result<Self> {
95        Ok(Self {
96            context: context.clone(),
97            linear: PlanCache::new(capacity, PreparedModel::is_idle)?,
98            finite: PlanCache::new(capacity, PreparedModel::is_idle)?,
99            gather: PlanCache::new(capacity, PreparedModel::is_idle)?,
100        })
101    }
102    /// Gather complete contiguous rows using checked host indices. Duplicates
103    /// and arbitrary ordering are allowed. Only indices are uploaded; row bytes
104    /// stay on device. Three private slots per power-of-two output-size bucket
105    /// bound preparation churn; the returned view retains its slot until dropped.
106    pub fn gather_rows(&self, input: &DeviceTensor, indices: &[usize]) -> Result<DeviceTensor> {
107        self.context.validate(input)?;
108        let desc = input.desc();
109        if desc.shape().len() != 2
110            || !desc.is_contiguous()
111            || desc.bytes() > i32::MAX as usize
112            || indices.iter().any(|&i| i >= desc.shape()[0])
113        {
114            return Err(Error::Message(
115                "gather requires contiguous rows and in-range indices".into(),
116            ));
117        }
118        let output = TensorDesc::new(desc.dtype(), vec![indices.len(), desc.shape()[1]])?
119            .with_layout(desc.layout())?;
120        if output.is_empty() {
121            input.completion().wait()?;
122            return self.context.allocate(output);
123        }
124        let bucket = indices
125            .len()
126            .checked_next_power_of_two()
127            .ok_or_else(|| Error::Message("gather size overflow".into()))?;
128        let plan = self.gather.get_or_prepare((desc.clone(), bucket), || {
129            self.gather_fragment(desc, bucket, false)?.prepare(3)
130        })?;
131        let bytes: Vec<_> = indices
132            .iter()
133            .copied()
134            .chain(std::iter::repeat(indices[0]))
135            .take(bucket)
136            .flat_map(|i| (i as u32).to_le_bytes())
137            .collect();
138        let ids = self
139            .context
140            .upload(TensorDesc::new(DType::U32, vec![bucket])?, &bytes)?;
141        let result = plan.submit(&[input.clone(), ids])?;
142        result.outputs()[0].view(0, output)
143    }
144    /// Reduce contiguous F32 to one I32 error flag: zero iff every value is
145    /// finite. Empty tensors are not accepted. Multi-pass bounded reductions
146    /// avoid atomics and leave input/output ownership with normal slot leases.
147    pub fn prepare_finite(&self, input: &TensorDesc) -> Result<Arc<PreparedModel>> {
148        self.finite
149            .get_or_prepare(input.clone(), || self.finite_fragment(input)?.prepare(3))
150    }
151    /// Record a finite-value reduction directly into a caller-owned graph.
152    /// The output is one I32 flag with the same contract as `prepare_finite`.
153    pub fn finite_fragment(&self, input: &TensorDesc) -> Result<ModelFragment> {
154        if input.dtype() != DType::F32
155            || !input.is_contiguous()
156            || input.is_empty()
157            || input.elements() > i32::MAX as usize
158        {
159            return Err(Error::Message(
160                "finite check requires nonempty contiguous F32".into(),
161            ));
162        }
163        {
164            let mut model = ModelSession::in_context(&self.context)?;
165            let src = model.allocate(input.bytes())?;
166            let mut current = src;
167            let mut count = input.elements();
168            let mut commands = Vec::new();
169            let mut first = true;
170            loop {
171                let out = count.div_ceil(256);
172                let dst = model.allocate(out * 4)?;
173                let ty = if first { "f32" } else { "i32" };
174                let bad = if first {
175                    "%abs = scalar.absf %value : f32\n%max = scalar.constant 3.402823466e38 : f32\n%finite = scalar.cmpf ole, %abs, %max : f32\n%bad = scf.select %finite, %fzero, %fone : f32"
176                } else {
177                    "%bad = scalar.uitofp %value : i32 to f32"
178                };
179                let mut source = include_str!("finite.loom")
180                    .replace("@TYPE@", ty)
181                    .replace("@CHECK@", bad);
182                for (name, value) in [
183                    ("GRID", out.div_ceil(64)),
184                    ("INPUT", count),
185                    ("OUTPUT", out),
186                    ("LAST_INPUT", count - 1),
187                    ("LAST_OUTPUT", out - 1),
188                ] {
189                    source = source.replace(&format!("@{name}@"), &value.to_string());
190                }
191                // Static bounds above cover every load, reduction flag and store.
192                let kernel =
193                    unsafe { model.compile(&[(&source, Specialization::new("finite"))])? }[0];
194                commands.push(Command::Dispatch(Dispatch::indices(
195                    kernel,
196                    [0],
197                    [out.div_ceil(64) as u32, 1, 1],
198                    vec![current.read(), dst.write()],
199                )));
200                current = dst;
201                if out == 1 {
202                    break;
203                }
204                count = out;
205                first = false;
206            }
207            unsafe {
208                model.freeze(&self.context)?.fragment(
209                    &commands,
210                    &[(src, input.clone())],
211                    &[(current, TensorDesc::new(DType::I32, vec![1])?)],
212                )
213            }
214        }
215    }
216    /// Return a resident non-finite flag without downloading the input tensor.
217    pub fn finite(&self, input: &DeviceTensor) -> Result<Inference> {
218        self.context.validate(input)?;
219        self.prepare_finite(input.desc())?
220            .submit(std::slice::from_ref(input))
221    }
222}
hrx/tensor/ops.rs

hrx/tensor/
ops.rs
