hrx/loom/model/
definition.rs

1use super::*;
2use crate::{
3    execution::{BindingContract, BufferView, GpuKernel, Graph, KernelContract, MemoryPlacement},
4    inference::{InferenceGraph, ModelContext, PreparedModel},
5    tensor::{DeviceTensor, TensorDesc},
6};
7
8#[derive(Clone)]
9enum Storage {
10    Weight(crate::execution::Buffer),
11    Scratch(usize, Placement),
12}
13
14/// Immutable model code and weights in a shared context. Shape-specific plans
15/// allocate private scratch and graph storage for each bounded inference slot.
16pub struct ModelDefinition {
17    id: u64,
18    context: ModelContext,
19    allocations: Vec<Storage>,
20    kernels: Vec<Kernel>,
21}
22
23/// Validated, shape-specific model commands with shared immutable weights.
24/// Record into a caller-owned graph to connect stages without copies, or build
25/// a standalone bounded inference pool with [`Self::prepare`].
26pub struct ModelFragment {
27    context: ModelContext,
28    allocations: Vec<Storage>,
29    extents: Vec<usize>,
30    commands: Vec<(Option<GpuKernel>, Vec<Region>, Option<u8>)>,
31    inputs: Vec<(Region, TensorDesc)>,
32    outputs: Vec<(Region, TensorDesc)>,
33    reuse_private_scratch: bool,
34}
35
36impl ModelSession {
37    /// Transfer code and immutable weights to a coordinated context. All raw
38    /// graphs are destroyed after synchronization, before adopting allocations.
39    /// Scratch and host staging are discarded; prepared slots initialize their
40    /// own scratch. Values needed across inferences must use [`Self::weight`].
41    pub fn freeze(mut self, context: &ModelContext) -> Result<ModelDefinition> {
42        self.synchronize()?;
43        if self.stream.device_id()
44            != Device::open(context.runtime().gpu()?.index())?
45                .stream()?
46                .device_id()
47        {
48            return Err(Error::Message("model belongs to another device".into()));
49        }
50        self.graphs.clear();
51        let allocations = self
52            .allocations
53            .drain(..)
54            .map(|allocation| {
55                if allocation.immutable {
56                    // The session exclusively owns the allocation. Upload has been
57                    // drained above; all old native graphs have been destroyed.
58                    let buffer = unsafe {
59                        context
60                            .runtime()
61                            .adopt_gpu_buffer(allocation.buffer, MemoryPlacement::GpuLocal)?
62                    };
63                    Ok(Storage::Weight(buffer))
64                } else {
65                    Ok(Storage::Scratch(
66                        allocation.buffer.bytes(),
67                        allocation.placement,
68                    ))
69                }
70            })
71            .collect::<Result<Vec<_>>>()?;
72        Ok(ModelDefinition {
73            id: self.id,
74            context: context.clone(),
75            allocations,
76            kernels: self.kernels,
77        })
78    }
79}
80
81impl ModelDefinition {
82    /// Shared scheduling and allocation domain.
83    pub fn context(&self) -> &ModelContext {
84        &self.context
85    }
86
87    /// Evaluate a loaded kernel's compiler launch program during graph preparation.
88    pub fn launch_config(
89        &self,
90        id: KernelId,
91        workload_bits: &[u64],
92    ) -> Result<crate::loom::LaunchConfig> {
93        if id.session != self.id {
94            return Err(Error::Message(
95                "model kernel belongs to another definition".into(),
96            ));
97        }
98        self.kernels
99            .get(id.index)
100            .ok_or_else(|| Error::Message("invalid model kernel".into()))?
101            .launch_config(workload_bits)
102    }
103
104    /// Physical immutable weight bytes, counted once regardless of plan or slot count.
105    pub fn weight_bytes(&self) -> usize {
106        self.allocations
107            .iter()
108            .map(|storage| match storage {
109                Storage::Weight(buffer) => buffer.len(),
110                Storage::Scratch(..) => 0,
111            })
112            .sum()
113    }
114
115    fn validate_region(&self, region: Region) -> Result<()> {
116        if region.session != self.id {
117            return Err(Error::Message(
118                "model region belongs to another definition".into(),
119            ));
120        }
121        let bytes = match self.allocations.get(region.allocation) {
122            Some(Storage::Weight(buffer)) => buffer.len(),
123            Some(Storage::Scratch(bytes, _)) => *bytes,
124            None => return Err(Error::Message("invalid model allocation".into())),
125        };
126        if region
127            .offset
128            .checked_add(region.bytes)
129            .is_none_or(|end| end > bytes)
130        {
131            return Err(Error::Message("model region exceeds allocation".into()));
132        }
133        Ok(())
134    }
135
136    /// Prepare a shape with shared weights and fixed-address private slots.
137    /// Input/output descriptors select leading bytes of their corresponding
138    /// regions. Commands must restrict accesses to the selected shape.
139    ///
140    /// # Safety
141    /// Dispatch bindings and scalar arguments must exactly describe the trusted
142    /// compiled code, including accessible extents and every read or write.
143    /// Writable argument aliasing must be safe for that kernel's execution.
144    pub unsafe fn prepare(
145        &self,
146        commands: &[Command],
147        inputs: &[(Region, TensorDesc)],
148        outputs: &[(Region, TensorDesc)],
149        capacity: usize,
150    ) -> Result<PreparedModel> {
151        // Safety: forwarded unchanged from the caller.
152        unsafe { self.fragment(commands, inputs, outputs)? }.prepare(capacity)
153    }
154
155    /// Validate a shape's commands once, without allocating inference slots.
156    /// The fragment owns shared weights and code independently of this definition.
157    ///
158    /// # Safety
159    /// Dispatch contracts, accessible extents and writable aliases must satisfy
160    /// the same requirements as [`Self::prepare`].
161    pub unsafe fn fragment(
162        &self,
163        commands: &[Command],
164        inputs: &[(Region, TensorDesc)],
165        outputs: &[(Region, TensorDesc)],
166    ) -> Result<ModelFragment> {
167        for (region, desc) in inputs.iter().chain(outputs) {
168            self.validate_region(*region)?;
169            if desc.is_empty() || desc.bytes() > region.bytes {
170                return Err(Error::Message(
171                    "model IO descriptor exceeds its region or is empty".into(),
172                ));
173            }
174            if matches!(self.allocations[region.allocation], Storage::Weight(_)) {
175                return Err(Error::Message(
176                    "model IO cannot alias immutable weights".into(),
177                ));
178            }
179        }
180        let mut prepared = Vec::with_capacity(commands.len());
181        for command in commands {
182            match command {
183                Command::Fill { region, value } => {
184                    self.validate_region(*region)?;
185                    if matches!(self.allocations[region.allocation], Storage::Weight(_)) {
186                        return Err(Error::Message("cannot fill immutable weights".into()));
187                    }
188                    prepared.push((None, vec![*region], Some(*value)));
189                }
190                Command::Dispatch(dispatch) => {
191                    if dispatch.kernel.session != self.id {
192                        return Err(Error::Message(
193                            "kernel belongs to another definition".into(),
194                        ));
195                    }
196                    let kernel = self
197                        .kernels
198                        .get(dispatch.kernel.index)
199                        .ok_or_else(|| Error::Message("invalid model kernel".into()))?;
200                    let mut bindings = Vec::with_capacity(dispatch.bindings.len());
201                    for binding in &dispatch.bindings {
202                        self.validate_region(binding.region)?;
203                        if binding.access != Access::Read
204                            && matches!(
205                                self.allocations[binding.region.allocation],
206                                Storage::Weight(_)
207                            )
208                        {
209                            return Err(Error::Message("cannot write immutable weights".into()));
210                        }
211                        bindings.push(BindingContract {
212                            bytes: binding.region.bytes,
213                            alignment: 1,
214                            access: binding.access,
215                            layout: "model region".into(),
216                        });
217                    }
218                    let constants = match &dispatch.arguments {
219                        Arguments::Indices(values) => Constants::indices(kernel, values)?,
220                        Arguments::Packed(constants) => *constants.clone(),
221                    };
222                    // Safety: inherited from this method's caller.
223                    let kernel = unsafe {
224                        self.context.runtime().adopt_gpu_kernel(
225                            kernel.clone(),
226                            dispatch.grid,
227                            kernel.info().workgroup_size,
228                            KernelContract {
229                                bindings,
230                                constants: constants.as_bytes().to_vec(),
231                            },
232                        )?
233                    };
234                    prepared.push((
235                        Some(kernel),
236                        dispatch.bindings.iter().map(|b| b.region).collect(),
237                        None,
238                    ));
239                }
240            }
241        }
242        let mut extents = vec![0; self.allocations.len()];
243        for (_, regions, _) in &prepared {
244            for region in regions {
245                extents[region.allocation] =
246                    extents[region.allocation].max(region.offset + region.bytes);
247            }
248        }
249        for (region, desc) in inputs.iter().chain(outputs) {
250            extents[region.allocation] =
251                extents[region.allocation].max(region.offset + desc.bytes());
252        }
253        Ok(ModelFragment {
254            context: self.context.clone(),
255            allocations: self.allocations.clone(),
256            extents,
257            commands: prepared,
258            inputs: inputs.to_vec(),
259            outputs: outputs.to_vec(),
260            reuse_private_scratch: false,
261        })
262    }
263}
264
265impl ModelFragment {
266    /// Reuse private temporaries across fragments recorded into the same graph.
267    /// Input/output allocations and immutable weights are never recycled. Each
268    /// inference slot still owns a separate workspace. Graph hazards order reads
269    /// before a later fragment overwrites that storage, potentially serializing
270    /// otherwise independent branches in exchange for lower peak memory.
271    ///
272    /// # Safety
273    /// Every private scratch byte read by this fragment must first be written
274    /// by this fragment on every execution, including padding and conditional
275    /// paths. Neither initial zeroes nor another fragment's values may be used.
276    /// All scratch accesses must already be declared by the kernel contracts.
277    pub unsafe fn reuse_private_scratch(mut self) -> Self {
278        self.reuse_private_scratch = true;
279        self
280    }
281    /// Allocate independent fixed-address slots, exposing their actual model IO.
282    pub fn prepare(&self, capacity: usize) -> Result<PreparedModel> {
283        PreparedModel::prepare(&self.context, capacity, |context| {
284            let mut graph = context.runtime().graph();
285            let (inputs, outputs) = self.append(&mut graph, None)?;
286            Ok(InferenceGraph {
287                inputs,
288                outputs,
289                graph: graph.prepare()?,
290            })
291        })
292    }
293
294    /// Append model commands directly to an unprepared graph. Inputs are bound
295    /// without a device copy; outputs and scratch are allocated for this graph.
296    /// Adjacent stages can consume the returned tensors before graph preparation.
297    /// No pool, host staging, submission or intermediate synchronization is added.
298    ///
299    /// Inputs must describe the complete accessed regions of their allocations.
300    /// Partial aliases which would require a copy are rejected. Initialization
301    /// producers are drained here; this is a preparation-time operation.
302    pub fn record(&self, graph: &mut Graph, inputs: &[DeviceTensor]) -> Result<Vec<DeviceTensor>> {
303        Ok(self.append(graph, Some(inputs))?.1)
304    }
305
306    fn append(
307        &self,
308        graph: &mut Graph,
309        supplied: Option<&[DeviceTensor]>,
310    ) -> Result<(Vec<DeviceTensor>, Vec<DeviceTensor>)> {
311        graph.validate_runtime(self.context.runtime())?;
312        let mut buffers: Vec<Option<(usize, BufferView)>> = vec![None; self.allocations.len()];
313        if let Some(inputs) = supplied {
314            if inputs.len() != self.inputs.len() {
315                return Err(Error::Message("fragment input count mismatch".into()));
316            }
317            for ((region, desc), tensor) in self.inputs.iter().zip(inputs) {
318                self.context.validate(tensor)?;
319                if tensor.desc() != desc {
320                    return Err(Error::Message("fragment input descriptor mismatch".into()));
321                }
322                tensor.completion().wait()?;
323                let binding = tensor.binding().expect("nonempty fragment IO");
324                if buffers.iter().enumerate().any(|(index, old)| {
325                    index != region.allocation && old.as_ref().is_some_and(|(_, old)| old.overlaps(&binding))
326                }) || self.allocations.iter().any(|storage| {
327                    matches!(storage, Storage::Weight(weight) if weight.view().overlaps(&binding))
328                }) {
329                    return Err(Error::Message("fragment remapping introduces an undeclared alias".into()));
330                }
331                if let Some((base, old)) = &buffers[region.allocation] {
332                    if *base != region.offset || !old.same_region(&binding) {
333                        return Err(Error::Message(
334                            "fragment inputs alias one allocation inconsistently".into(),
335                        ));
336                    }
337                } else {
338                    buffers[region.allocation] = Some((region.offset, binding));
339                }
340            }
341        }
342        if self.reuse_private_scratch {
343            let private: Vec<_> = self
344                .allocations
345                .iter()
346                .enumerate()
347                .filter_map(|(index, storage)| match storage {
348                    Storage::Scratch(_, placement)
349                        if self.extents[index] > 0
350                            && !self
351                                .inputs
352                                .iter()
353                                .chain(&self.outputs)
354                                .any(|(region, _)| region.allocation == index) =>
355                    {
356                        Some((
357                            index,
358                            (self.extents[index], matches!(placement, Placement::Shared)),
359                        ))
360                    }
361                    _ => None,
362                })
363                .collect();
364            let requests: Vec<_> = private.iter().map(|(_, request)| *request).collect();
365            for ((index, _), view) in private.into_iter().zip(graph.fragment_scratch(&requests)?) {
366                buffers[index] = Some((0, view));
367            }
368        }
369        for (index, storage) in self.allocations.iter().enumerate() {
370            if buffers[index].is_none() {
371                buffers[index] = match storage {
372                    Storage::Weight(buffer) => Some((0, buffer.view())),
373                    Storage::Scratch(..) if self.extents[index] == 0 => None,
374                    Storage::Scratch(_, placement) => Some((
375                        0,
376                        self.context
377                            .runtime()
378                            .allocate(
379                                self.extents[index],
380                                match placement {
381                                    Placement::Device => MemoryPlacement::GpuLocal,
382                                    Placement::Shared => MemoryPlacement::HostVisible,
383                                },
384                            )?
385                            .view(),
386                    )),
387                };
388            }
389        }
390        let view = |region: Region| -> Result<BufferView> {
391            let (base, buffer) = buffers[region.allocation]
392                .as_ref()
393                .expect("referenced allocation");
394            let start = region
395                .offset
396                .checked_sub(*base)
397                .ok_or_else(|| Error::Message("fragment access precedes bound input".into()))?;
398            buffer.slice(start..start + region.bytes)
399        };
400        let tensors = |io: &[(Region, TensorDesc)]| {
401            io.iter()
402                .map(|(region, desc)| {
403                    self.context.tensor(
404                        desc.clone(),
405                        view(region.slice(0, desc.bytes())?)?,
406                        crate::Completion::ready(),
407                    )
408                })
409                .collect::<Result<Vec<_>>>()
410        };
411        let inputs = tensors(&self.inputs)?;
412        let outputs = tensors(&self.outputs)?;
413        // Validate every binding before mutating the caller's graph.
414        let bindings = self
415            .commands
416            .iter()
417            .map(|(_, regions, _)| {
418                regions
419                    .iter()
420                    .map(|&region| view(region))
421                    .collect::<Result<Vec<_>>>()
422            })
423            .collect::<Result<Vec<_>>>()?;
424        for ((kernel, _, fill), bindings) in self.commands.iter().zip(bindings) {
425            if let Some(value) = fill {
426                graph.fill(bindings[0].clone(), *value)?;
427            } else if let Some(kernel) = kernel {
428                // Safety: fragment construction validated the caller's trusted
429                // contracts. Remapping preserves region offsets and extents.
430                unsafe {
431                    graph.gpu_aliasing(kernel, &bindings)?;
432                }
433            }
434        }
435        Ok((inputs, outputs))
436    }
437}
hrx/loom/model/definition.rs

hrx/loom/model/
definition.rs
