tract_linalg/frame/mmm/

scratch.rs

use super::{FusedKerSpec, FusedSpec, MatMatMulKer, OutputStoreKer};
use crate::{BinOp, LADatum};
use downcast_rs::{Downcast, impl_downcast};
use std::cell::RefCell;
use std::fmt::Debug;
use std::sync::atomic::AtomicUsize;
use tract_data::internal::num_integer::Integer;
use tract_data::internal::*;

static GENERATION: AtomicUsize = AtomicUsize::new(1);

thread_local! {
    static TLS: RefCell<TLSScratch> = Default::default();
}

#[derive(Default, Debug)]
pub(crate) struct TLSScratch {
    generation: usize,
    blob: Blob,
    ker_specs_16: Vec<FusedKerSpec<f16>>,
    ker_specs_32: Vec<FusedKerSpec<f32>>,
    ker_specs_64: Vec<FusedKerSpec<f64>>,
}

impl TLSScratch {
    #[allow(unknown_lints, clippy::missing_transmute_annotations)]
    fn ker_specs<TI: LADatum>(&mut self) -> &mut Vec<FusedKerSpec<TI>> {
        unsafe {
            if TI::datum_type() == f32::datum_type() || TI::datum_type() == i32::datum_type() {
                std::mem::transmute(&mut self.ker_specs_32)
            } else if TI::datum_type() == f16::datum_type() {
                std::mem::transmute(&mut self.ker_specs_16)
            } else if TI::datum_type() == f64::datum_type() {
                std::mem::transmute(&mut self.ker_specs_64)
            } else {
                todo!();
            }
        }
    }

    fn sync<TI: LADatum>(&mut self, scratch: &ScratchSpaceImpl<TI>) {
        if self.generation == scratch.generation {
            return;
        }
        let ker_specs = self.ker_specs::<TI>();
        ker_specs.clear();
        ker_specs.extend_from_slice(&scratch.ker_specs);

        unsafe {
            self.blob.ensure_size_and_align(scratch.blob_size, scratch.blob_align);

            for LocDependent { loc, ker_spec, .. } in &scratch.loc_dependent {
                #[allow(clippy::single_match)]
                if matches!(scratch.ker_specs[*ker_spec], FusedKerSpec::AddMatMul { .. }) {
                    let scratch = &mut *(self.blob.as_ptr().add(*loc) as *mut AddMatMulTemp);
                    scratch.panel_a_id = usize::MAX;
                    scratch.panel_b_id = usize::MAX;
                };
            }
        }
        self.generation = scratch.generation;
    }
}

pub trait ScratchSpace: Downcast + Send {}
impl_downcast!(ScratchSpace);

#[derive(Debug, Default)]
pub struct ScratchSpaceImpl<TI: LADatum> {
    generation: usize,
    blob_size: usize,
    blob_align: usize,
    ker_specs: Vec<FusedKerSpec<TI>>,
    loc_dependent: TVec<LocDependent>,
    valid_down_tiles: usize,
    remnant_down: usize,
    valid_right_tiles: usize,
    remnant_right: usize,
}

#[derive(Debug, new)]
struct LocDependent {
    spec: usize,
    ker_spec: usize,
    // offset for the location dependent structure
    loc: usize,
    // offset of its associated dynamic-size buffers
    buffer_a: Option<usize>,
    buffer_b: Option<usize>,
}

impl<TI: LADatum> ScratchSpace for ScratchSpaceImpl<TI> {}
unsafe impl<TI: LADatum> Send for ScratchSpaceImpl<TI> {}

#[derive(Debug)]
struct AddMatMulTemp {
    ptr_a: *const u8,
    panel_a_id: usize,
    ptr_b: *const u8,
    panel_b_id: usize,
}

/// True when this border tile can store its live rows straight into C.
///
/// Only `sme_mmm_f32_32x32` understands the packing hint, and only for a
/// full N panel whose columns are contiguous. The tile must also contain a
/// matmul: the hint travels in that op, and without it the kernel stores
/// all 32 rows. A count of 16 or less also skips the bottom-tile FMOPAs;
/// 17..31 only shortens the store.
fn partial_m_direct<TI: LADatum>(
    ker: &impl MatMatMulKer<Acc = TI>,
    specs: &[FusedSpec],
    m_remnant: usize,
    n_remnant: usize,
) -> bool {
    if ker.name() != "sme_mmm_f32_32x32" {
        return false;
    }
    if m_remnant == 0
        || m_remnant >= ker.mr()
        || n_remnant != ker.nr()
        || !ker.stores_row_major_tile()
    {
        return false;
    }
    let mut saw_store = false;
    let mut saw_matmul = false;
    for spec in specs {
        match spec {
            FusedSpec::Store(store) => {
                saw_store = true;
                if store.col_byte_stride != store.item_size as isize || store.item_size != 4 {
                    return false;
                }
            }
            FusedSpec::AddMatMul { .. } => saw_matmul = true,
            _ => {}
        }
    }
    saw_store && saw_matmul
}

impl<TI: LADatum> ScratchSpaceImpl<TI> {
    pub unsafe fn prepare(
        &mut self,
        ker: &impl MatMatMulKer<Acc = TI>,
        m: usize,
        n: usize,
        specs: &[FusedSpec],
    ) -> TractResult<()> {
        use FusedKerSpec as FKS;
        use FusedSpec as FS;
        self.ker_specs.clear();
        self.loc_dependent.clear();
        self.ker_specs.reserve(specs.len() + 2);
        self.ker_specs.push(FusedKerSpec::Clear);
        self.valid_down_tiles = m / ker.mr();
        self.remnant_down = m % ker.mr();
        self.valid_right_tiles = n / ker.nr();
        self.remnant_right = n % ker.nr();
        let mut offset = 0;
        let mut align = std::mem::size_of::<*const ()>();
        fn ld(spec: usize, uspec: usize, loc: usize) -> LocDependent {
            LocDependent { spec, ker_spec: uspec, loc, buffer_a: None, buffer_b: None }
        }
        for (ix, spec) in specs.iter().enumerate() {
            offset = offset.next_multiple_of(&align);
            let ker_spec = match spec {
                FS::BinScalar(t, op) => match op {
                    BinOp::Min => FKS::ScalarMin(*t.try_as_plain_ram()?.to_scalar()?),
                    BinOp::Max => FKS::ScalarMax(*t.try_as_plain_ram()?.to_scalar()?),
                    BinOp::Mul => FKS::ScalarMul(*t.try_as_plain_ram()?.to_scalar()?),
                    BinOp::Add => FKS::ScalarAdd(*t.try_as_plain_ram()?.to_scalar()?),
                    BinOp::Sub => FKS::ScalarSub(*t.try_as_plain_ram()?.to_scalar()?),
                    BinOp::SubF => FKS::ScalarSubF(*t.try_as_plain_ram()?.to_scalar()?),
                },
                FS::ShiftLeft(s) => FKS::ShiftLeft(*s),
                FS::RoundingShiftRight(s, rp) => FKS::RoundingShiftRight(*s, *rp),
                FS::QScale(s, rp, m) => FKS::QScale(*s, *rp, *m),
                FS::BinPerRow(_, _) => {
                    self.loc_dependent.push(ld(ix, self.ker_specs.len(), offset));
                    offset += TI::datum_type().size_of() * ker.mr();
                    FusedKerSpec::Done
                }
                FS::BinPerCol(_, _) => {
                    self.loc_dependent.push(ld(ix, self.ker_specs.len(), offset));
                    offset += TI::datum_type().size_of() * ker.nr();
                    FusedKerSpec::Done
                }
                FS::AddRowColProducts(_, _) => {
                    self.loc_dependent.push(ld(ix, self.ker_specs.len(), offset));
                    offset += TI::datum_type().size_of() * (ker.mr() + ker.nr());
                    FusedKerSpec::Done
                }
                FS::AddUnicast(_) => {
                    self.loc_dependent.push(ld(ix, self.ker_specs.len(), offset));
                    offset += TI::datum_type().size_of() * ker.mr() * ker.nr();
                    FusedKerSpec::Done
                }
                FS::Store(store) => {
                    // Only worth a row-major tile when the destination is itself
                    // n-contiguous: set_from_tile then copies without transposing
                    // and the kernel's aligned bulk-store path applies.
                    let row_major = ker.stores_row_major_tile()
                        && store.col_byte_stride == store.item_size as isize;
                    let tile_bytes = if row_major {
                        // 128-align the row-major store tile, with rows padded to
                        // 128 bytes, so the kernel can hit its aligned bulk-store
                        // path (e.g. Apple AMX stz-direct).
                        align = align.lcm(&128);
                        offset = Integer::next_multiple_of(&offset, &128);
                        Integer::next_multiple_of(&(store.item_size * ker.nr()), &128) * ker.mr()
                    } else {
                        store.item_size * ker.mr() * ker.nr()
                    };
                    self.loc_dependent.push(ld(ix, self.ker_specs.len(), offset));
                    offset += tile_bytes;
                    FusedKerSpec::Done
                }
                FS::LeakyRelu(t) => FKS::LeakyRelu(*t.try_as_plain_ram()?.to_scalar()?),
                FS::AddMatMul { a, b, packing } => {
                    let mut ld = ld(ix, self.ker_specs.len(), offset);
                    offset += std::mem::size_of::<AddMatMulTemp>();
                    if let Some(tmp) = a.scratch_panel_buffer_layout() {
                        align = tmp.align().lcm(&align);
                        offset = Integer::next_multiple_of(&offset, &tmp.align());
                        ld.buffer_a = Some(offset);
                        offset += tmp.size();
                    }
                    if let Some(tmp) = b.scratch_panel_buffer_layout() {
                        align = tmp.align().lcm(&align);
                        offset = Integer::next_multiple_of(&offset, &tmp.align());
                        ld.buffer_b = Some(offset);
                        offset += tmp.size();
                    }
                    self.loc_dependent.push(ld);
                    FusedKerSpec::AddMatMul {
                        k: 0,
                        pa: std::ptr::null(),
                        pb: std::ptr::null(),
                        packing: *packing,
                    }
                }
            };
            self.ker_specs.push(ker_spec);
        }
        self.ker_specs.push(FKS::Done);
        self.blob_size = offset;
        self.blob_align = align;

        self.generation = GENERATION.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
        Ok(())
    }

    pub unsafe fn run(
        &self,
        ker: &impl MatMatMulKer<Acc = TI>,
        specs: &[FusedSpec],
        down: usize,
        right: usize,
    ) -> TractResult<()> {
        // Per-tile entry: enter the TLS scope (does sync once) then run a single
        // tile. Single-threaded callers should prefer `run_in_tls_scope`+
        // `run_one_tile` to amortise the TLS borrow + sync over many tiles.
        unsafe {
            self.run_in_tls_scope(|this, tls| this.run_one_tile(ker, specs, tls, down, right))
        }
    }

    /// Borrow the per-thread scratch blob for a single MMM call and `sync` it
    /// once. The closure is invoked once with a mutable reference to the TLS
    /// scratch and to `self`. Used by single-threaded matmul drivers to avoid
    /// re-entering TLS / re-running `sync` per tile.
    pub(crate) unsafe fn run_in_tls_scope<F, R>(&self, f: F) -> R
    where
        F: FnOnce(&Self, &mut TLSScratch) -> R,
    {
        TLS.with_borrow_mut(|tls| {
            tls.sync(self);
            f(self, tls)
        })
    }

    /// Run a single tile against an already-borrowed TLS scratch. Caller is
    /// responsible for entering `run_in_tls_scope` first (so `sync` has run).
    #[inline(always)]
    pub(crate) unsafe fn run_one_tile(
        &self,
        ker: &impl MatMatMulKer<Acc = TI>,
        specs: &[FusedSpec],
        tls: &mut TLSScratch,
        down: usize,
        right: usize,
    ) -> TractResult<()> {
        unsafe {
            if down < self.valid_down_tiles && right < self.valid_right_tiles {
                self.for_valid_tile(ker, specs, tls, down, right)?;
                let err = ker.kernel(tls.ker_specs());
                ensure!(err == 0, "Kernel {} returned error {err}", ker.name());
            } else {
                let remnant_down =
                    if down < self.valid_down_tiles { ker.mr() } else { self.remnant_down };
                let remnant_right =
                    if right < self.valid_right_tiles { ker.nr() } else { self.remnant_right };
                let direct = self.for_border_tile(
                    ker,
                    specs,
                    tls,
                    down,
                    right,
                    remnant_down,
                    remnant_right,
                )?;
                let err = ker.kernel(tls.ker_specs());
                ensure!(err == 0, "Kernel {} returned error {err}", ker.name());
                if !direct {
                    self.postprocess_tile(specs, tls, down, right, remnant_down, remnant_right)?;
                }
            }
            Ok(())
        }
    }

    #[inline(always)]
    unsafe fn for_valid_tile(
        &self,
        ker: &impl MatMatMulKer<Acc = TI>,
        specs: &[FusedSpec],
        tls: &mut TLSScratch,
        down: usize,
        right: usize,
    ) -> TractResult<()> {
        unsafe {
            use FusedKerSpec as FKS;
            use FusedSpec as FS;
            let ScratchSpaceImpl { ker_specs, loc_dependent, .. } = self;
            debug_assert!(specs.len() + 2 == ker_specs.len());
            for LocDependent { spec, ker_spec, loc, buffer_a, buffer_b } in loc_dependent {
                let spec = specs.get_unchecked(*spec);
                let it = match spec {
                    FS::BinPerRow(v, op) => {
                        let v = v.as_ptr_unchecked::<TI>().add(down * ker.mr());
                        match op {
                            BinOp::Min => FKS::PerRowMin(v),
                            BinOp::Max => FKS::PerRowMax(v),
                            BinOp::Add => FKS::PerRowAdd(v),
                            BinOp::Mul => FKS::PerRowMul(v),
                            BinOp::Sub => FKS::PerRowSub(v),
                            BinOp::SubF => FKS::PerRowSubF(v),
                        }
                    }
                    FS::BinPerCol(v, op) => {
                        let v = v.as_ptr_unchecked::<TI>().add(right * ker.nr());
                        match op {
                            BinOp::Min => FKS::PerColMin(v),
                            BinOp::Max => FKS::PerColMax(v),
                            BinOp::Add => FKS::PerColAdd(v),
                            BinOp::Mul => FKS::PerColMul(v),
                            BinOp::Sub => FKS::PerColSub(v),
                            BinOp::SubF => FKS::PerColSubF(v),
                        }
                    }
                    FS::AddRowColProducts(rows, cols) => {
                        let row_ptr = rows.as_ptr_unchecked::<TI>().add(down * ker.mr());
                        let col_ptr = cols.as_ptr_unchecked::<TI>().add(right * ker.nr());
                        FKS::AddRowColProducts(row_ptr, col_ptr)
                    }
                    FS::AddUnicast(store) => FKS::AddUnicast(store.tile_c(down, right)),
                    FS::Store(c_store) => FKS::Store(c_store.tile_c(down, right)),
                    FS::AddMatMul { a, b, packing } => {
                        let scratch = (tls.blob.as_mut_ptr().add(*loc) as *mut AddMatMulTemp)
                            .as_mut()
                            .unwrap();
                        if scratch.panel_a_id != down {
                            scratch.ptr_a = a.panel_bytes(
                                down,
                                buffer_a.map(|o| tls.blob.as_mut_ptr().add(o)),
                            )?;
                            scratch.panel_a_id = down;
                        }
                        if scratch.panel_b_id != right {
                            scratch.ptr_b = b.panel_bytes(
                                right,
                                buffer_b.map(|o| tls.blob.as_mut_ptr().add(o)),
                            )?;
                            scratch.panel_b_id = right;
                        }
                        FKS::AddMatMul {
                            k: b.k(),
                            pa: scratch.ptr_a,
                            pb: scratch.ptr_b,
                            packing: *packing,
                        }
                    }
                    _ => std::hint::unreachable_unchecked(),
                };
                *tls.ker_specs().get_unchecked_mut(*ker_spec) = it;
            }
            Ok(())
        }
    }

    #[inline(never)]
    #[allow(clippy::too_many_arguments)]
    unsafe fn for_border_tile(
        &self,
        ker: &impl MatMatMulKer<Acc = TI>,
        specs: &[FusedSpec],
        tls: &mut TLSScratch,
        down: usize,
        right: usize,
        m_remnant: usize,
        n_remnant: usize,
    ) -> TractResult<bool> {
        unsafe {
            use FusedKerSpec as FKS;
            use FusedSpec as FS;
            // sme_mmm_f32_32x32 reads a live-M count from packing bits 8..15
            // and, on the horizontal store, writes that many rows. The count
            // is set only together with a direct store, so the scratch copy
            // is skipped for the same tile. Every other kernel and every
            // short-N or non-contiguous store stays on the scratch path.
            let direct = partial_m_direct(ker, specs, m_remnant, n_remnant);
            for LocDependent { spec, ker_spec: uspec, loc, buffer_a, buffer_b } in
                &self.loc_dependent
            {
                let loc = tls.blob.as_mut_ptr().add(*loc);
                let spec = specs.get_unchecked(*spec);
                let it = match spec {
                    FS::BinPerRow(v, op) => {
                        let buf = std::slice::from_raw_parts_mut(loc as *mut TI, ker.mr());
                        let ptr = if m_remnant < ker.mr() {
                            if m_remnant > 0 {
                                buf.get_unchecked_mut(..m_remnant).copy_from_slice(
                                    v.as_slice_unchecked()
                                        .get_unchecked(down * ker.mr()..)
                                        .get_unchecked(..m_remnant),
                                );
                            }
                            // The kernel computes on the tail lanes before their
                            // results are discarded; garbage there decodes to
                            // denormals and stalls the fp pipeline. Zero them.
                            buf.get_unchecked_mut(m_remnant..)
                                .iter_mut()
                                .for_each(|x| *x = TI::zero());
                            buf.as_ptr()
                        } else {
                            v.as_ptr_unchecked::<TI>().add(down * ker.mr())
                        };
                        match op {
                            BinOp::Min => FKS::PerRowMin(ptr),
                            BinOp::Max => FKS::PerRowMax(ptr),
                            BinOp::Add => FKS::PerRowAdd(ptr),
                            BinOp::Mul => FKS::PerRowMul(ptr),
                            BinOp::Sub => FKS::PerRowSub(ptr),
                            BinOp::SubF => FKS::PerRowSubF(ptr),
                        }
                    }
                    FS::BinPerCol(v, op) => {
                        let buf = std::slice::from_raw_parts_mut(loc as *mut TI, ker.nr());
                        let ptr = if n_remnant < ker.nr() {
                            if n_remnant > 0 {
                                buf.get_unchecked_mut(..n_remnant).copy_from_slice(
                                    v.as_slice_unchecked()
                                        .get_unchecked(right * ker.nr()..)
                                        .get_unchecked(..n_remnant),
                                );
                            }
                            buf.get_unchecked_mut(n_remnant..)
                                .iter_mut()
                                .for_each(|x| *x = TI::zero());
                            buf.as_ptr()
                        } else {
                            v.as_ptr_unchecked::<TI>().add(right * ker.nr())
                        };
                        match op {
                            BinOp::Min => FKS::PerColMin(ptr),
                            BinOp::Max => FKS::PerColMax(ptr),
                            BinOp::Add => FKS::PerColAdd(ptr),
                            BinOp::Mul => FKS::PerColMul(ptr),
                            BinOp::Sub => FKS::PerColSub(ptr),
                            BinOp::SubF => FKS::PerColSubF(ptr),
                        }
                    }
                    FS::AddRowColProducts(rows, cols) => {
                        let r = std::slice::from_raw_parts_mut(loc as *mut TI, ker.mr());
                        let row_ptr = if m_remnant < ker.mr() {
                            r.get_unchecked_mut(..m_remnant).copy_from_slice(
                                rows.as_slice_unchecked()
                                    .get_unchecked(down * ker.mr()..)
                                    .get_unchecked(..m_remnant),
                            );
                            r.get_unchecked_mut(m_remnant..)
                                .iter_mut()
                                .for_each(|x| *x = TI::zero());
                            r.as_ptr()
                        } else {
                            rows.as_ptr_unchecked::<TI>().add(down * ker.mr())
                        };
                        let c = std::slice::from_raw_parts_mut(
                            (loc as *mut TI).add(ker.mr()),
                            ker.nr(),
                        );
                        let col_ptr = if n_remnant < ker.nr() {
                            c.get_unchecked_mut(..n_remnant).copy_from_slice(
                                cols.as_slice_unchecked()
                                    .get_unchecked(right * ker.nr()..)
                                    .get_unchecked(..n_remnant),
                            );
                            c.get_unchecked_mut(n_remnant..)
                                .iter_mut()
                                .for_each(|x| *x = TI::zero());
                            c.as_ptr()
                        } else {
                            cols.as_ptr_unchecked::<TI>().add(right * ker.nr())
                        };
                        FKS::AddRowColProducts(row_ptr, col_ptr)
                    }
                    FS::AddUnicast(store) => {
                        let row_byte_stride = store.row_byte_stride;
                        let col_byte_stride = store.col_byte_stride;
                        let tile_offset = row_byte_stride * down as isize * ker.mr() as isize
                            + col_byte_stride * right as isize * ker.nr() as isize;
                        let tile_ptr = store.ptr.offset(tile_offset);
                        let tmp_d_tile =
                            std::slice::from_raw_parts_mut(loc as *mut TI, ker.mr() * ker.nr());
                        tmp_d_tile.iter_mut().for_each(|t| *t = TI::zero());
                        for r in 0..m_remnant as isize {
                            for c in 0..n_remnant as isize {
                                let inner_offset = c * col_byte_stride + r * row_byte_stride;
                                if inner_offset + tile_offset
                                    < (store.item_size * store.item_count) as isize
                                {
                                    *tmp_d_tile
                                        .get_unchecked_mut(r as usize + c as usize * ker.mr()) =
                                        *(tile_ptr.offset(inner_offset) as *const TI);
                                }
                            }
                        }
                        FKS::AddUnicast(OutputStoreKer {
                            ptr: tmp_d_tile.as_ptr() as _,
                            row_byte_stride: std::mem::size_of::<TI>() as isize,
                            col_byte_stride: (std::mem::size_of::<TI>() * ker.mr()) as isize,
                            item_size: std::mem::size_of::<TI>(),
                        })
                    }
                    FS::Store(c_store) => {
                        if direct {
                            FKS::Store(c_store.tile_c(down, right))
                        } else {
                            let row_major = ker.stores_row_major_tile()
                                && c_store.col_byte_stride == c_store.item_size as isize;
                            let (row_byte_stride, col_byte_stride) = if row_major {
                                // Pad the row stride to 128 bytes so it stays aligned
                                // for the kernel's bulk-store path.
                                let row = Integer::next_multiple_of(
                                    &(c_store.item_size * ker.nr()),
                                    &128,
                                );
                                (row as isize, c_store.item_size as isize)
                            } else {
                                (
                                    c_store.item_size as isize,
                                    (c_store.item_size * ker.mr()) as isize,
                                )
                            };
                            let tmpc = OutputStoreKer {
                                ptr: loc as _,
                                item_size: c_store.item_size,
                                row_byte_stride,
                                col_byte_stride,
                            };
                            FKS::Store(tmpc)
                        }
                    }
                    FS::AddMatMul { a, b, packing } => {
                        let scratch = (loc as *mut AddMatMulTemp).as_mut().unwrap();
                        if scratch.panel_a_id != down {
                            scratch.ptr_a = a.panel_bytes(
                                down,
                                buffer_a.map(|o| tls.blob.as_mut_ptr().add(o)),
                            )?;
                            scratch.panel_a_id = down;
                        }
                        if scratch.panel_b_id != right {
                            scratch.ptr_b = b.panel_bytes(
                                right,
                                buffer_b.map(|o| tls.blob.as_mut_ptr().add(o)),
                            )?;
                            scratch.panel_b_id = right;
                        }
                        let packing = if direct { *packing | (m_remnant << 8) } else { *packing };
                        FKS::AddMatMul { k: b.k(), pa: scratch.ptr_a, pb: scratch.ptr_b, packing }
                    }
                    _ => std::hint::unreachable_unchecked(),
                };
                *tls.ker_specs().get_unchecked_mut(*uspec) = it;
            }
            Ok(direct)
        }
    }

    #[inline]
    pub fn uspecs(&self) -> &[FusedKerSpec<TI>] {
        &self.ker_specs
    }

    unsafe fn postprocess_tile(
        &self,
        specs: &[FusedSpec],
        tls: &mut TLSScratch,
        down: usize,
        right: usize,
        m_remnant: usize,
        n_remnant: usize,
    ) -> TractResult<()>
    where
        TI: LADatum,
    {
        unsafe {
            for LocDependent { spec, ker_spec: uspec, .. } in self.loc_dependent.iter() {
                let spec = specs.get_unchecked(*spec);
                let ker_spec = tls.ker_specs::<TI>().get_unchecked(*uspec);
                if let (FusedSpec::Store(c_store), FusedKerSpec::Store(tmp)) = (spec, ker_spec) {
                    c_store.set_from_tile(down, right, m_remnant, n_remnant, tmp)
                }
            }
            Ok(())
        }
    }
}