fuel_backend_contract/
storage.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2//! Type-erased tensor storage wrapper (Phase 7.5 work item G fix-up).
3//!
4//! `Storage` was originally `fuel-core::Storage` and lived alongside the
5//! eager-dispatch op methods that mutate it. Phase 7.5 work item G moved
6//! the struct and the methods that depend only on backend-contract traits
7//! here so that `fuel_graph::Graph` can own a NodeId-keyed map of
8//! `Arc<RwLock<Storage>>` slots without inverting the dependency graph.
9//!
10//! As of B0.3 (the `fuel-backend-contract` cut) it lives in this crate —
11//! it IS the type-erased `Box<dyn DynBackendStorage>` handle that carries a
12//! backend's storage. The bundle *data* it references (`OutputView`,
13//! `OutputViewSpec`, `compose_bundle`) stays in [`fuel_ir::storage`].
14//!
15//! `Storage::device()` returns `Arc<dyn DynBackendDevice>` rather than
16//! the `Device` wrapper (which still lives in fuel-core); callers wrap as
17//! needed.
18//!
19//! ## Multi-output bundles (Option C, Session 1)
20//!
21//! A `Storage` may optionally carry [`OutputView`] side-table entries
22//! describing how the inner byte buffer is partitioned into multiple
23//! logically independent outputs. A multi-output op allocates one
24//! bundled `Storage` whose `bundle` slice describes per-slot dtype/shape/
25//! layout at byte offsets within the same underlying buffer. The
26//! backend trait surface ([`DynBackendStorage`]) is intentionally
27//! unchanged — backends still produce single typed-byte buffers; bundle
28//! metadata lives only on the `Storage` newtype.
29
30use crate::dyn_backend::{DynBackendDevice, DynBackendStorage};
31use crate::inplace_op::{InplaceOp1, InplaceOp2, InplaceOp3};
32use fuel_ir::op::{self, BinaryOp, CmpOp, ReduceOp};
33use fuel_ir::storage::{OutputView, OutputViewSpec, compose_bundle};
34use fuel_ir::stype::SType;
35use fuel_ir::{DType, Error, HostBuffer, Layout, Result, Scalar, Shape, conv};
36use std::sync::Arc;
37
38/// Allocate a bundled [`Storage`] on `device` covering all of `specs`.
39///
40/// Sizes the underlying allocation to `total_bytes / primary_size`
41/// elements of the primary (slot 0) dtype, rounding up so the buffer
42/// is guaranteed to hold every slot's bytes. Zero-initialised — the
43/// kernel can overwrite each slot on first write.
44///
45/// Returns a [`Storage`] whose [`Storage::bundle()`] is the resolved
46/// slot-view list and whose primary dtype equals slot 0's dtype.
47pub fn allocate_bundled_storage(
48    device: &dyn DynBackendDevice,
49    specs: &[OutputViewSpec],
50) -> Result<Storage> {
51    let (total_bytes, views) = compose_bundle(specs)?;
52    let primary_dtype = specs[0].dtype;
53    let primary_size = primary_dtype.size_in_bytes().max(1);
54    // Round up so the allocation holds every slot's bytes even when
55    // the bundle's total isn't a clean multiple of the primary dtype's
56    // size (e.g. F32 primary + F64 secondary at an odd boundary).
57    let flat_elems = total_bytes.div_ceil(primary_size).max(1);
58    let inner = device.zeros_impl_dyn(&Shape::from_dims(&[flat_elems]), primary_dtype)?;
59    Storage::from_dyn_bundled(inner, Arc::from(views.into_boxed_slice()))
60}
61
62/// Owns a typed contiguous buffer on one device. The boxed
63/// `DynBackendStorage` is the actual byte holder; `Storage` is a thin
64/// wrapper that gives the eager-dispatch op methods (matmul, conv,
65/// unary, binary, …) somewhere to live.
66///
67/// Optionally carries a `bundle` side-table describing how the inner
68/// buffer is partitioned into multiple logically independent outputs.
69/// See [`OutputView`] and the module-level "Multi-output bundles" doc
70/// for the contract.
71///
72/// We do not implement `Clone` because cloning storage may fail on
73/// out-of-memory; use [`Self::try_clone`] for the fallible version.
74#[derive(Debug)]
75pub struct Storage {
76    pub(crate) inner: Box<dyn DynBackendStorage>,
77    /// `None` for single-output storage (today's default). `Some(_)`
78    /// for multi-output bundles: a shared Arc'd slice of per-slot
79    /// [`OutputView`] entries, one per logical output. `Op::View`
80    /// nodes share this Arc so the bundle stays alive as long as any
81    /// view holds a reference.
82    pub(crate) bundle: Option<Arc<[OutputView]>>,
83    /// How the bytes are ENCODED (orthogonal to the logical element type
84    /// reported by `inner.dtype_dyn()`). Empty = plain dense dtype,
85    /// byte-identical to pre-SType behaviour. v1: PRIMARY storage only —
86    /// bundle slots keep `dtype` only (per-slot SType is a future
87    /// addition). See [`fuel_ir::SType`] and `docs/specs/storage-encoding.md`.
88    pub(crate) stype: SType,
89}
90
91impl Storage {
92    /// Construct storage from any concrete `DynBackendStorage` implementor.
93    ///
94    /// This is the backend-agnostic entry point — backends provide a type
95    /// implementing `DynBackendStorage`, and `Storage::new` boxes it.
96    /// Single-output (no bundle metadata); use
97    /// [`Storage::from_dyn_bundled`] for the multi-output case.
98    pub fn new<B: DynBackendStorage + 'static>(b: B) -> Self {
99        Storage {
100            inner: Box::new(b),
101            bundle: None,
102            stype: SType::default(),
103        }
104    }
105
106    /// Wrap an already-boxed `dyn DynBackendStorage`. Used by callers
107    /// (notably the quantized fast-paths) that produce a `Box<dyn ..>`
108    /// directly from trait dispatch. Single-output; use
109    /// [`Storage::from_dyn_bundled`] for the multi-output case.
110    pub fn from_dyn(b: Box<dyn DynBackendStorage>) -> Self {
111        Storage {
112            inner: b,
113            bundle: None,
114            stype: SType::default(),
115        }
116    }
117
118    /// Wrap an already-boxed `dyn DynBackendStorage` together with a
119    /// per-slot bundle side-table. Used by multi-output op authors:
120    /// the backend allocates one bundled byte buffer (whose
121    /// `dtype_dyn()` is the primary/slot-0 dtype), and the caller
122    /// attaches the slot metadata here.
123    ///
124    /// Validates at construction time that
125    /// `bundle[0].dtype == inner.dtype_dyn()` so callers can trust
126    /// `Storage::dtype()` ≡ `Storage::primary_dtype()` ≡
127    /// `Storage::slot_dtype(0)` on a bundled storage.
128    pub fn from_dyn_bundled(
129        b: Box<dyn DynBackendStorage>,
130        bundle: Arc<[OutputView]>,
131    ) -> Result<Self> {
132        if bundle.is_empty() {
133            return Err(Error::Msg(
134                "Storage::from_dyn_bundled: bundle slice must be non-empty".into(),
135            )
136            .bt());
137        }
138        let primary = bundle[0].dtype;
139        let inner_dtype = b.dtype_dyn();
140        if primary != inner_dtype {
141            return Err(Error::Msg(format!(
142                "Storage::from_dyn_bundled: slot 0 dtype {primary:?} must \
143                 match inner backend dtype {inner_dtype:?}",
144            ))
145            .bt());
146        }
147        Ok(Storage {
148            inner: b,
149            bundle: Some(bundle),
150            stype: SType::default(),
151        })
152    }
153
154    /// Attach a bundle side-table to an existing single-output
155    /// `Storage`. The dtype check from [`Storage::from_dyn_bundled`]
156    /// applies. Panics in `debug_assertions` mode if a bundle is
157    /// already attached — re-bundling silently is a contract bug.
158    pub fn with_bundle(mut self, bundle: Arc<[OutputView]>) -> Result<Self> {
159        debug_assert!(
160            self.bundle.is_none(),
161            "Storage::with_bundle: bundle already attached",
162        );
163        if bundle.is_empty() {
164            return Err(
165                Error::Msg("Storage::with_bundle: bundle slice must be non-empty".into()).bt(),
166            );
167        }
168        let primary = bundle[0].dtype;
169        let inner_dtype = self.inner.dtype_dyn();
170        if primary != inner_dtype {
171            return Err(Error::Msg(format!(
172                "Storage::with_bundle: slot 0 dtype {primary:?} must \
173                 match inner backend dtype {inner_dtype:?}",
174            ))
175            .bt());
176        }
177        self.bundle = Some(bundle);
178        Ok(self)
179    }
180
181    /// Whether this storage carries multi-output bundle metadata.
182    pub fn is_bundled(&self) -> bool {
183        self.bundle.is_some()
184    }
185
186    /// Number of logical output slots in this storage. `1` for
187    /// single-output storage; for a bundled storage, the number of
188    /// [`OutputView`] entries.
189    pub fn slot_count(&self) -> usize {
190        match &self.bundle {
191            Some(b) => b.len(),
192            None => 1,
193        }
194    }
195
196    /// Borrow the full bundle slice, or `None` if this is a
197    /// single-output storage.
198    pub fn bundle(&self) -> Option<&[OutputView]> {
199        self.bundle.as_deref()
200    }
201
202    /// Clone the `Arc<[OutputView]>` handle. Used by `Op::View`'s
203    /// realization path: a View's output storage shares this Arc so
204    /// the bundle stays alive as long as any view holds a reference.
205    pub fn bundle_arc(&self) -> Option<Arc<[OutputView]>> {
206        self.bundle.clone()
207    }
208
209    /// Per-slot view for `idx`. Returns `None` for out-of-range
210    /// indices or for single-output storage with `idx != 0`. For
211    /// single-output storage with `idx == 0` this returns `None` as
212    /// well — the slot's dtype/shape live in the (`Storage::dtype`,
213    /// caller-supplied Layout) pair, NOT in a synthetic `OutputView`.
214    pub fn slot_view(&self, idx: usize) -> Option<&OutputView> {
215        self.bundle.as_deref().and_then(|b| b.get(idx))
216    }
217
218    /// Primary slot's dtype — for a bundled storage this is slot 0's
219    /// dtype (enforced equal to the inner backend dtype at construction
220    /// time); for single-output storage this is just the inner dtype.
221    /// Use [`Self::slot_dtype`] when the consumer is talking about a
222    /// specific bundle slot, not the primary one.
223    pub fn primary_dtype(&self) -> DType {
224        self.inner.dtype_dyn()
225    }
226
227    /// Dtype of a specific bundle slot. Returns `None` for
228    /// out-of-range indices or for single-output storage. For
229    /// single-output storage, callers should use [`Self::dtype`]
230    /// directly.
231    pub fn slot_dtype(&self, idx: usize) -> Option<DType> {
232        self.slot_view(idx).map(|v| v.dtype)
233    }
234
235    /// Borrow the inner storage as a `DynBackendStorage` trait object.
236    ///
237    /// Backends that need to peel back to their concrete storage type can
238    /// downcast via `storage.as_dyn().as_any().downcast_ref::<MyStorage>()`.
239    pub fn as_dyn(&self) -> &dyn DynBackendStorage {
240        &*self.inner
241    }
242
243    /// Mutable variant of [`Self::as_dyn`].
244    pub fn as_dyn_mut(&mut self) -> &mut dyn DynBackendStorage {
245        &mut *self.inner
246    }
247
248    /// Downcast the inner storage to a concrete backend type.
249    pub fn downcast_ref<T: 'static>(&self) -> Option<&T> {
250        self.inner.as_any().downcast_ref::<T>()
251    }
252
253    /// Mutable variant of [`downcast_ref`](Self::downcast_ref).
254    pub fn downcast_mut<T: 'static>(&mut self) -> Option<&mut T> {
255        self.inner.as_any_mut().downcast_mut::<T>()
256    }
257
258    pub fn try_clone(&self, layout: &Layout) -> Result<Self> {
259        // Preserve the encoding scheme across a clone (v1: cloning a plain
260        // storage stays plain since `from_dyn` defaults empty; cloning an
261        // *encoded* storage must carry its `stype` forward).
262        Ok(Storage::from_dyn(self.inner.try_clone_dyn(layout)?).with_stype(self.stype.clone()))
263    }
264
265    /// Return an `Arc` to the owning device as a trait object.
266    /// fuel-core wraps this in its `Device` newtype; other consumers
267    /// (fuel-graph, tests) can use the trait object directly.
268    pub fn device(&self) -> Arc<dyn DynBackendDevice> {
269        self.inner.device_arc_dyn()
270    }
271
272    pub fn dtype(&self) -> DType {
273        self.inner.dtype_dyn()
274    }
275
276    /// Attach an encoding scheme to this storage (consuming builder).
277    /// Does not touch the bytes or the logical dtype — only describes HOW
278    /// the bytes are encoded. See [`fuel_ir::SType`].
279    pub fn with_stype(mut self, stype: SType) -> Self {
280        self.stype = stype;
281        self
282    }
283
284    /// The encoding scheme. Empty = plain dense dtype.
285    pub fn stype(&self) -> &SType {
286        &self.stype
287    }
288
289    /// Pre-G this method consulted `Device::same_device` for the Metal
290    /// pointer-identity check; post-G it goes through the
291    /// `DynBackendDevice::same_device_dyn` trait method which has the
292    /// same semantics.
293    pub fn same_device(&self, rhs: &Self, op: &'static str) -> Result<()> {
294        let lhs_dev = self.inner.device_dyn();
295        let rhs_dev = rhs.inner.device_dyn();
296        let lhs_loc = lhs_dev.location_dyn();
297        let rhs_loc = rhs_dev.location_dyn();
298        let same = if matches!(lhs_loc, fuel_ir::DeviceLocation::Metal { .. }) {
299            // On metal, require physical identity (matches pre-G behaviour).
300            lhs_dev.same_device_dyn(rhs_dev)
301        } else {
302            lhs_loc == rhs_loc
303        };
304        if !same {
305            Err(Error::DeviceMismatchBinaryOp {
306                lhs: lhs_loc,
307                rhs: rhs_loc,
308                op,
309            }
310            .bt())
311        } else {
312            Ok(())
313        }
314    }
315
316    pub fn same_dtype(&self, rhs: &Self, op: &'static str) -> Result<()> {
317        let lhs = self.dtype();
318        let rhs = rhs.dtype();
319        if lhs != rhs {
320            Err(Error::DTypeMismatchBinaryOp { lhs, rhs, op }.bt())
321        } else {
322            Ok(())
323        }
324    }
325
326    pub fn const_set(&mut self, v: Scalar, l: &Layout) -> Result<()> {
327        self.inner.const_set_dyn(v, l)
328    }
329
330    pub fn affine(&self, layout: &Layout, mul: f64, add: f64) -> Result<Self> {
331        Ok(Storage::from_dyn(self.inner.affine_dyn(layout, mul, add)?))
332    }
333
334    pub fn powf(&self, layout: &Layout, e: f64) -> Result<Self> {
335        Ok(Storage::from_dyn(self.inner.powf_dyn(layout, e)?))
336    }
337
338    pub fn elu(&self, layout: &Layout, alpha: f64) -> Result<Self> {
339        Ok(Storage::from_dyn(self.inner.elu_dyn(layout, alpha)?))
340    }
341
342    pub fn cmp(
343        &self,
344        op: CmpOp,
345        rhs: &Self,
346        lhs_layout: &Layout,
347        rhs_layout: &Layout,
348    ) -> Result<Self> {
349        Ok(Storage::from_dyn(self.inner.cmp_dyn(
350            op,
351            &*rhs.inner,
352            lhs_layout,
353            rhs_layout,
354        )?))
355    }
356
357    pub fn reduce_op(&self, op: ReduceOp, layout: &Layout, reduce_dims: &[usize]) -> Result<Self> {
358        Ok(Storage::from_dyn(self.inner.reduce_op_dyn(
359            op,
360            layout,
361            reduce_dims,
362        )?))
363    }
364
365    pub fn to_dtype(&self, layout: &Layout, dtype: DType) -> Result<Self> {
366        Ok(Storage::from_dyn(self.inner.to_dtype_dyn(layout, dtype)?))
367    }
368
369    pub fn to_cpu_storage(&self) -> Result<HostBuffer> {
370        self.inner.to_host_buffer_dyn()
371    }
372
373    pub fn inplace_op1(&mut self, l: &Layout, c: &dyn InplaceOp1) -> Result<()> {
374        c.fwd(&mut *self.inner, l)
375    }
376
377    pub fn inplace_op2(
378        &mut self,
379        l1: &Layout,
380        t2: &Self,
381        l2: &Layout,
382        c: &dyn InplaceOp2,
383    ) -> Result<()> {
384        self.same_device(t2, c.name())?;
385        c.fwd(&mut *self.inner, l1, &*t2.inner, l2)
386    }
387
388    pub fn inplace_op3(
389        &mut self,
390        l1: &Layout,
391        t2: &Self,
392        l2: &Layout,
393        t3: &Self,
394        l3: &Layout,
395        c: &dyn InplaceOp3,
396    ) -> Result<()> {
397        self.same_device(t2, c.name())?;
398        self.same_device(t3, c.name())?;
399        c.fwd(&mut *self.inner, l1, &*t2.inner, l2, &*t3.inner, l3)
400    }
401
402    // -----------------------------------------------------------------------
403    // Unary / Binary dispatch
404    // -----------------------------------------------------------------------
405
406    pub fn unary_impl<B: op::UnaryOpT>(&self, layout: &Layout) -> Result<Self> {
407        let op = op::UnaryOp::from_name(B::NAME)
408            .ok_or_else(|| Error::Msg(format!("unknown unary op '{}'", B::NAME)))?;
409        Ok(Storage::from_dyn(self.inner.unary_op_dyn(layout, op)?))
410    }
411
412    pub fn binary_impl<B: op::BinaryOpT>(
413        &self,
414        rhs: &Self,
415        lhs_layout: &Layout,
416        rhs_layout: &Layout,
417    ) -> Result<Self> {
418        self.same_device(rhs, B::NAME)?;
419        self.same_dtype(rhs, B::NAME)?;
420        let op = BinaryOp::from_name(B::NAME)
421            .ok_or_else(|| Error::Msg(format!("unknown binary op '{}'", B::NAME)))?;
422        Ok(Storage::from_dyn(self.inner.binary_op_dyn(
423            &*rhs.inner,
424            lhs_layout,
425            rhs_layout,
426            op,
427        )?))
428    }
429
430    // -----------------------------------------------------------------------
431    // Convolutions, pooling, upsampling
432    // -----------------------------------------------------------------------
433
434    pub fn conv1d(
435        &self,
436        l: &Layout,
437        kernel: &Self,
438        kernel_l: &Layout,
439        params: &conv::ParamsConv1D,
440    ) -> Result<Self> {
441        self.same_device(kernel, "conv1d")?;
442        self.same_dtype(kernel, "conv1d")?;
443        Ok(Storage::from_dyn(self.inner.conv1d_dyn(
444            l,
445            &*kernel.inner,
446            kernel_l,
447            params,
448        )?))
449    }
450
451    pub fn conv_transpose1d(
452        &self,
453        l: &Layout,
454        kernel: &Self,
455        kernel_l: &Layout,
456        params: &conv::ParamsConvTranspose1D,
457    ) -> Result<Self> {
458        self.same_device(kernel, "conv-transpose1d")?;
459        self.same_dtype(kernel, "conv-transpose1d")?;
460        Ok(Storage::from_dyn(self.inner.conv_transpose1d_dyn(
461            l,
462            &*kernel.inner,
463            kernel_l,
464            params,
465        )?))
466    }
467
468    pub fn conv2d(
469        &self,
470        l: &Layout,
471        kernel: &Self,
472        kernel_l: &Layout,
473        params: &conv::ParamsConv2D,
474    ) -> Result<Self> {
475        self.same_device(kernel, "conv2d")?;
476        self.same_dtype(kernel, "conv2d")?;
477        Ok(Storage::from_dyn(self.inner.conv2d_dyn(
478            l,
479            &*kernel.inner,
480            kernel_l,
481            params,
482        )?))
483    }
484
485    pub fn conv_transpose2d(
486        &self,
487        l: &Layout,
488        kernel: &Self,
489        kernel_l: &Layout,
490        params: &conv::ParamsConvTranspose2D,
491    ) -> Result<Self> {
492        self.same_device(kernel, "conv_transpose2d")?;
493        self.same_dtype(kernel, "conv_transpose2d")?;
494        Ok(Storage::from_dyn(self.inner.conv_transpose2d_dyn(
495            l,
496            &*kernel.inner,
497            kernel_l,
498            params,
499        )?))
500    }
501
502    pub fn avg_pool2d(
503        &self,
504        layout: &Layout,
505        kernel_size: (usize, usize),
506        stride: (usize, usize),
507    ) -> Result<Self> {
508        Ok(Storage::from_dyn(self.inner.avg_pool2d_dyn(
509            layout,
510            kernel_size,
511            stride,
512        )?))
513    }
514
515    pub fn max_pool2d(
516        &self,
517        layout: &Layout,
518        kernel_size: (usize, usize),
519        stride: (usize, usize),
520    ) -> Result<Self> {
521        Ok(Storage::from_dyn(self.inner.max_pool2d_dyn(
522            layout,
523            kernel_size,
524            stride,
525        )?))
526    }
527
528    pub fn upsample_nearest1d(&self, layout: &Layout, sz: usize) -> Result<Self> {
529        Ok(Storage::from_dyn(
530            self.inner.upsample_nearest1d_dyn(layout, sz)?,
531        ))
532    }
533
534    pub fn upsample_nearest2d(&self, layout: &Layout, h: usize, w: usize) -> Result<Self> {
535        Ok(Storage::from_dyn(
536            self.inner.upsample_nearest2d_dyn(layout, h, w)?,
537        ))
538    }
539
540    pub fn upsample_bilinear2d(
541        &self,
542        layout: &Layout,
543        h: usize,
544        w: usize,
545        align_corners: bool,
546        scale_h: Option<f64>,
547        scale_w: Option<f64>,
548    ) -> Result<Self> {
549        Ok(Storage::from_dyn(self.inner.upsample_bilinear2d_dyn(
550            layout,
551            h,
552            w,
553            align_corners,
554            scale_h,
555            scale_w,
556        )?))
557    }
558
559    // -----------------------------------------------------------------------
560    // Gather / Scatter / Index
561    // -----------------------------------------------------------------------
562
563    pub fn where_cond(
564        &self,
565        layout: &Layout,
566        t: &Self,
567        layout_t: &Layout,
568        f: &Self,
569        layout_f: &Layout,
570    ) -> Result<Self> {
571        self.same_device(t, "where")?;
572        self.same_device(f, "where")?;
573        t.same_dtype(f, "where")?;
574        Ok(Storage::from_dyn(self.inner.where_cond_dyn(
575            layout, &*t.inner, layout_t, &*f.inner, layout_f,
576        )?))
577    }
578
579    pub fn gather(&self, l: &Layout, indexes: &Self, indexes_l: &Layout, d: usize) -> Result<Self> {
580        self.same_device(indexes, "index-add")?;
581        Ok(Storage::from_dyn(self.inner.gather_dyn(
582            l,
583            &*indexes.inner,
584            indexes_l,
585            d,
586        )?))
587    }
588
589    pub fn scatter_set(
590        &mut self,
591        l: &Layout,
592        indexes: &Self,
593        indexes_l: &Layout,
594        source: &Self,
595        source_l: &Layout,
596        d: usize,
597    ) -> Result<()> {
598        self.same_device(indexes, "scatter-set")?;
599        self.same_device(source, "scatter-set")?;
600        self.inner
601            .scatter_set_dyn(l, &*source.inner, source_l, &*indexes.inner, indexes_l, d)
602    }
603
604    pub fn scatter_add(
605        &mut self,
606        l: &Layout,
607        indexes: &Self,
608        indexes_l: &Layout,
609        source: &Self,
610        source_l: &Layout,
611        d: usize,
612    ) -> Result<()> {
613        self.same_device(indexes, "scatter-add")?;
614        self.same_device(source, "scatter-add")?;
615        self.inner
616            .scatter_add_set_dyn(l, &*source.inner, source_l, &*indexes.inner, indexes_l, d)
617    }
618
619    pub fn index_add(
620        &self,
621        l: &Layout,
622        indexes: &Self,
623        indexes_l: &Layout,
624        source: &Self,
625        source_l: &Layout,
626        d: usize,
627    ) -> Result<Self> {
628        self.same_device(indexes, "index-add")?;
629        self.same_device(source, "index-add")?;
630        Ok(Storage::from_dyn(self.inner.index_add_dyn(
631            l,
632            &*indexes.inner,
633            indexes_l,
634            &*source.inner,
635            source_l,
636            d,
637        )?))
638    }
639
640    pub fn index_select(
641        &self,
642        rhs: &Self,
643        lhs_l: &Layout,
644        rhs_l: &Layout,
645        d: usize,
646    ) -> Result<Self> {
647        self.same_device(rhs, "index-select")?;
648        Ok(Storage::from_dyn(self.inner.index_select_dyn(
649            &*rhs.inner,
650            lhs_l,
651            rhs_l,
652            d,
653        )?))
654    }
655
656    // -----------------------------------------------------------------------
657    // Matmul and copy
658    // -----------------------------------------------------------------------
659
660    pub fn matmul(
661        &self,
662        rhs: &Self,
663        bmnk: (usize, usize, usize, usize),
664        lhs_layout: &Layout,
665        rhs_layout: &Layout,
666    ) -> Result<Self> {
667        self.same_device(rhs, "matmul")?;
668        self.same_dtype(rhs, "matmul")?;
669        Ok(Storage::from_dyn(self.inner.matmul_dyn(
670            &*rhs.inner,
671            bmnk,
672            lhs_layout,
673            rhs_layout,
674        )?))
675    }
676
677    /// `self`, the source, can be strided whereas `dst` is contiguous.
678    pub fn copy_strided_src(
679        &self,
680        dst: &mut Self,
681        dst_offset: usize,
682        src_l: &Layout,
683    ) -> Result<()> {
684        self.inner
685            .copy_strided_src_dyn(&mut *dst.inner, dst_offset, src_l)
686    }
687
688    pub fn copy2d(
689        &self,
690        dst: &mut Self,
691        d1: usize,
692        d2: usize,
693        src_s: usize,
694        dst_s: usize,
695        src_o: usize,
696        dst_o: usize,
697    ) -> Result<()> {
698        self.inner
699            .copy2d_dyn(&mut *dst.inner, d1, d2, src_s, dst_s, src_o, dst_o)
700    }
701}
702
703#[cfg(test)]
704mod multi_output_specs {
705    use super::*;
706
707    /// Helper: contiguous F32 spec of the given shape.
708    fn f32_spec(dims: &[usize]) -> OutputViewSpec {
709        OutputViewSpec::contiguous(DType::F32, Shape::from_dims(dims))
710    }
711
712    /// Helper: contiguous F64 spec of the given shape.
713    fn f64_spec(dims: &[usize]) -> OutputViewSpec {
714        OutputViewSpec::contiguous(DType::F64, Shape::from_dims(dims))
715    }
716
717    /// compose_bundle composes a single-slot spec into one OutputView
718    /// with byte_offset 0 and total_bytes equal to the slot's
719    /// footprint.
720    #[test]
721    fn compose_bundle_single_slot() {
722        let specs = vec![f32_spec(&[2, 3])];
723        let (total, views) = compose_bundle(&specs).expect("single slot composes");
724        assert_eq!(views.len(), 1);
725        assert_eq!(views[0].byte_offset, 0);
726        assert_eq!(views[0].len_elements, 6);
727        assert_eq!(views[0].dtype, DType::F32);
728        assert_eq!(total, 24); // 6 * 4 bytes
729    }
730
731    /// compose_bundle stacks slots with per-slot dtype alignment.
732    /// Slot 0 = F32[6] = 24 bytes (aligned to 4).
733    /// Slot 1 = F64[3] = 24 bytes (aligned to 8). cursor at 24 is
734    /// already aligned to 8, so byte_offset = 24.
735    #[test]
736    fn compose_bundle_two_slot_aligned() {
737        let specs = vec![f32_spec(&[2, 3]), f64_spec(&[3])];
738        let (total, views) = compose_bundle(&specs).expect("aligned compose");
739        assert_eq!(views.len(), 2);
740        assert_eq!(views[0].byte_offset, 0);
741        assert_eq!(views[0].len_elements, 6);
742        assert_eq!(views[1].byte_offset, 24);
743        assert_eq!(views[1].len_elements, 3);
744        assert_eq!(views[1].dtype, DType::F64);
745        assert_eq!(total, 48); // 24 (slot 0) + 24 (slot 1)
746    }
747
748    /// compose_bundle pads when slot 1's alignment requires it.
749    /// Slot 0 = F32[1] = 4 bytes, cursor at 4.
750    /// Slot 1 = F64[1], alignment 8; 4 % 8 = 4, pad to 8. byte_offset = 8.
751    #[test]
752    fn compose_bundle_pads_for_alignment() {
753        let specs = vec![f32_spec(&[1]), f64_spec(&[1])];
754        let (total, views) = compose_bundle(&specs).expect("padded compose");
755        assert_eq!(views[0].byte_offset, 0);
756        assert_eq!(views[1].byte_offset, 8); // padded from 4 to next-8-multiple
757        assert_eq!(total, 16); // 8 (start of slot 1) + 8 (slot 1)
758    }
759
760    /// compose_bundle rejects an empty spec list.
761    #[test]
762    fn compose_bundle_rejects_empty() {
763        let err = compose_bundle(&[]).expect_err("empty spec list must error");
764        assert!(format!("{err}").contains("non-empty"));
765    }
766
767    /// compose_bundle rejects a spec whose layout.shape() disagrees
768    /// with its declared shape (mirrors with_bundle's invariant).
769    #[test]
770    fn compose_bundle_rejects_shape_layout_mismatch() {
771        let s = Shape::from_dims(&[2, 3]);
772        let bogus_layout = Layout::contiguous(Shape::from_dims(&[3, 2]));
773        let bad = OutputViewSpec {
774            dtype: DType::F32,
775            shape: s,
776            layout: bogus_layout,
777            name: None,
778        };
779        let err = compose_bundle(&[bad]).expect_err("shape/layout mismatch must error");
780        assert!(format!("{err}").contains("disagrees"));
781    }
782
783    /// OutputViewSpec::contiguous wires the default layout correctly.
784    #[test]
785    fn output_view_spec_contiguous_helper() {
786        let s = f32_spec(&[4, 5]);
787        assert_eq!(s.dtype, DType::F32);
788        assert_eq!(s.shape, Shape::from_dims(&[4, 5]));
789        assert_eq!(s.layout.shape(), &Shape::from_dims(&[4, 5]));
790        assert_eq!(s.elem_count(), 20);
791        assert_eq!(s.len_bytes(), 80);
792        assert!(s.name.is_none());
793    }
794}
795
796#[cfg(test)]
797mod stype_attach {
798    //! Step 2 born-red coverage for the trait-object `Storage`'s `stype`
799    //! field. There is no concrete `DynBackendStorage` impl in this crate
800    //! (backends live downstream and would cycle as a dev-dep), so the test
801    //! uses a minimal mock whose only real method is `try_clone_dyn`.
802    use super::*;
803    use fuel_ir::op::UnaryOp;
804    use std::any::Any;
805
806    /// Minimal `DynBackendStorage`: an empty unit struct. Only
807    /// `try_clone_dyn` / `dtype_dyn` / `as_any*` are exercised; every other
808    /// method panics (never called by these tests).
809    #[derive(Debug)]
810    struct MockStorage;
811
812    #[allow(unused_variables)]
813    impl DynBackendStorage for MockStorage {
814        fn try_clone_dyn(&self, _layout: &Layout) -> Result<Box<dyn DynBackendStorage>> {
815            Ok(Box::new(MockStorage))
816        }
817        fn dtype_dyn(&self) -> DType {
818            DType::F32
819        }
820        fn as_any(&self) -> &dyn Any {
821            self
822        }
823        fn as_any_mut(&mut self) -> &mut dyn Any {
824            self
825        }
826
827        fn device_dyn(&self) -> &dyn DynBackendDevice {
828            unimplemented!()
829        }
830        fn device_arc_dyn(&self) -> Arc<dyn DynBackendDevice> {
831            unimplemented!()
832        }
833        fn to_host_buffer_dyn(&self) -> Result<HostBuffer> {
834            unimplemented!()
835        }
836        fn affine_dyn(&self, l: &Layout, mul: f64, add: f64) -> Result<Box<dyn DynBackendStorage>> {
837            unimplemented!()
838        }
839        fn powf_dyn(&self, l: &Layout, e: f64) -> Result<Box<dyn DynBackendStorage>> {
840            unimplemented!()
841        }
842        fn elu_dyn(&self, l: &Layout, alpha: f64) -> Result<Box<dyn DynBackendStorage>> {
843            unimplemented!()
844        }
845        fn reduce_op_dyn(
846            &self,
847            op: ReduceOp,
848            l: &Layout,
849            axes: &[usize],
850        ) -> Result<Box<dyn DynBackendStorage>> {
851            unimplemented!()
852        }
853        fn cmp_dyn(
854            &self,
855            op: CmpOp,
856            rhs: &dyn DynBackendStorage,
857            ll: &Layout,
858            rl: &Layout,
859        ) -> Result<Box<dyn DynBackendStorage>> {
860            unimplemented!()
861        }
862        fn to_dtype_dyn(&self, l: &Layout, dtype: DType) -> Result<Box<dyn DynBackendStorage>> {
863            unimplemented!()
864        }
865        fn unary_op_dyn(&self, l: &Layout, op: UnaryOp) -> Result<Box<dyn DynBackendStorage>> {
866            unimplemented!()
867        }
868        fn binary_op_dyn(
869            &self,
870            rhs: &dyn DynBackendStorage,
871            ll: &Layout,
872            rl: &Layout,
873            op: BinaryOp,
874        ) -> Result<Box<dyn DynBackendStorage>> {
875            unimplemented!()
876        }
877        fn where_cond_dyn(
878            &self,
879            cl: &Layout,
880            t: &dyn DynBackendStorage,
881            tl: &Layout,
882            f: &dyn DynBackendStorage,
883            fl: &Layout,
884        ) -> Result<Box<dyn DynBackendStorage>> {
885            unimplemented!()
886        }
887        fn conv1d_dyn(
888            &self,
889            l: &Layout,
890            k: &dyn DynBackendStorage,
891            kl: &Layout,
892            p: &conv::ParamsConv1D,
893        ) -> Result<Box<dyn DynBackendStorage>> {
894            unimplemented!()
895        }
896        fn conv_transpose1d_dyn(
897            &self,
898            l: &Layout,
899            k: &dyn DynBackendStorage,
900            kl: &Layout,
901            p: &conv::ParamsConvTranspose1D,
902        ) -> Result<Box<dyn DynBackendStorage>> {
903            unimplemented!()
904        }
905        fn conv2d_dyn(
906            &self,
907            l: &Layout,
908            k: &dyn DynBackendStorage,
909            kl: &Layout,
910            p: &conv::ParamsConv2D,
911        ) -> Result<Box<dyn DynBackendStorage>> {
912            unimplemented!()
913        }
914        fn conv_transpose2d_dyn(
915            &self,
916            l: &Layout,
917            k: &dyn DynBackendStorage,
918            kl: &Layout,
919            p: &conv::ParamsConvTranspose2D,
920        ) -> Result<Box<dyn DynBackendStorage>> {
921            unimplemented!()
922        }
923        fn avg_pool2d_dyn(
924            &self,
925            l: &Layout,
926            k: (usize, usize),
927            s: (usize, usize),
928        ) -> Result<Box<dyn DynBackendStorage>> {
929            unimplemented!()
930        }
931        fn max_pool2d_dyn(
932            &self,
933            l: &Layout,
934            k: (usize, usize),
935            s: (usize, usize),
936        ) -> Result<Box<dyn DynBackendStorage>> {
937            unimplemented!()
938        }
939        fn upsample_nearest1d_dyn(
940            &self,
941            l: &Layout,
942            t: usize,
943        ) -> Result<Box<dyn DynBackendStorage>> {
944            unimplemented!()
945        }
946        fn upsample_nearest2d_dyn(
947            &self,
948            l: &Layout,
949            h: usize,
950            w: usize,
951        ) -> Result<Box<dyn DynBackendStorage>> {
952            unimplemented!()
953        }
954        fn upsample_bilinear2d_dyn(
955            &self,
956            l: &Layout,
957            h: usize,
958            w: usize,
959            ac: bool,
960            sh: Option<f64>,
961            sw: Option<f64>,
962        ) -> Result<Box<dyn DynBackendStorage>> {
963            unimplemented!()
964        }
965        fn gather_dyn(
966            &self,
967            sl: &Layout,
968            ids: &dyn DynBackendStorage,
969            il: &Layout,
970            dim: usize,
971        ) -> Result<Box<dyn DynBackendStorage>> {
972            unimplemented!()
973        }
974        fn scatter_set_dyn(
975            &mut self,
976            sl: &Layout,
977            src: &dyn DynBackendStorage,
978            srl: &Layout,
979            ids: &dyn DynBackendStorage,
980            il: &Layout,
981            dim: usize,
982        ) -> Result<()> {
983            unimplemented!()
984        }
985        fn scatter_add_set_dyn(
986            &mut self,
987            sl: &Layout,
988            src: &dyn DynBackendStorage,
989            srl: &Layout,
990            ids: &dyn DynBackendStorage,
991            il: &Layout,
992            dim: usize,
993        ) -> Result<()> {
994            unimplemented!()
995        }
996        fn index_select_dyn(
997            &self,
998            ids: &dyn DynBackendStorage,
999            sl: &Layout,
1000            il: &Layout,
1001            dim: usize,
1002        ) -> Result<Box<dyn DynBackendStorage>> {
1003            unimplemented!()
1004        }
1005        fn index_add_dyn(
1006            &self,
1007            sl: &Layout,
1008            ids: &dyn DynBackendStorage,
1009            il: &Layout,
1010            src: &dyn DynBackendStorage,
1011            srl: &Layout,
1012            dim: usize,
1013        ) -> Result<Box<dyn DynBackendStorage>> {
1014            unimplemented!()
1015        }
1016        fn matmul_dyn(
1017            &self,
1018            rhs: &dyn DynBackendStorage,
1019            bmnk: (usize, usize, usize, usize),
1020            ll: &Layout,
1021            rl: &Layout,
1022        ) -> Result<Box<dyn DynBackendStorage>> {
1023            unimplemented!()
1024        }
1025        fn copy_strided_src_dyn(
1026            &self,
1027            dst: &mut dyn DynBackendStorage,
1028            off: usize,
1029            sl: &Layout,
1030        ) -> Result<()> {
1031            unimplemented!()
1032        }
1033        fn copy2d_dyn(
1034            &self,
1035            dst: &mut dyn DynBackendStorage,
1036            d1: usize,
1037            d2: usize,
1038            ss1: usize,
1039            ds1: usize,
1040            so: usize,
1041            dofs: usize,
1042        ) -> Result<()> {
1043            unimplemented!()
1044        }
1045        fn const_set_dyn(&mut self, value: Scalar, l: &Layout) -> Result<()> {
1046            unimplemented!()
1047        }
1048    }
1049
1050    /// Born-red: a freshly constructed trait-object Storage carries a plain
1051    /// (empty) SType by default.
1052    #[test]
1053    fn trait_object_storage_defaults_plain() {
1054        let s = Storage::new(MockStorage);
1055        assert!(
1056            s.stype().is_plain(),
1057            "default Storage must carry a plain SType"
1058        );
1059        assert_eq!(s.stype().layers().len(), 0);
1060    }
1061
1062    /// Born-red: `try_clone` carries an attached `stype` forward (the cheap,
1063    /// correct v1 choice — cloning an encoded storage preserves its scheme).
1064    #[test]
1065    fn try_clone_preserves_stype() {
1066        use fuel_ir::GgmlDType;
1067        use fuel_ir::stype::Encoding;
1068        let s = Storage::new(MockStorage).with_stype(SType::from_layer(Encoding::GgmlBlock {
1069            ggml_dtype: GgmlDType::Q4_0,
1070        }));
1071        let layout = Layout::contiguous(Shape::from_dims(&[4]));
1072        let cloned = s.try_clone(&layout).expect("clone");
1073        assert_eq!(cloned.stype(), s.stype(), "try_clone must preserve stype");
1074        assert!(!cloned.stype().is_plain());
1075    }
1076}
fuel_backend_contract/storage.rs

fuel_backend_contract/
storage.rs
