From 7e545597c53a5ece6b8b590547df561d39f8169c Mon Sep 17 00:00:00 2001 From: ckristian Date: Thu, 6 Aug 2026 11:53:08 +0200 Subject: [PATCH] linalg,core: zero the padding lanes the mmm kernels compute on MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Packed-operand buffers (pack_tensor/pack_tensor_view, im2col's PackedFormat panels) and the border-tile scratch buffers were zeroed only under debug_assertions, assuming arithmetic on discarded padding lanes is free. Uninitialized bytes there frequently decode to denormal floats, and every kernel op touching such a lane takes a microcode assist — measured 14x on partial tiles, and unfixable downstream under WASM where FTZ cannot be set. Zero the partial last panel at pack time and the border scratch buffers unconditionally (also fixing AddRowColProducts' col-tail zeroing hitting the rows buffer). Co-Authored-By: Claude Fable 5 --- core/src/ops/cnn/conv/im2col.rs | 8 +++++-- linalg/src/frame/mmm/scratch.rs | 39 ++++++++++++++------------------- linalg/src/frame/pack.rs | 10 +++++++++ 3 files changed, 32 insertions(+), 25 deletions(-) diff --git a/core/src/ops/cnn/conv/im2col.rs b/core/src/ops/cnn/conv/im2col.rs index 59148c8392..c1ca64c5af 100644 --- a/core/src/ops/cnn/conv/im2col.rs +++ b/core/src/ops/cnn/conv/im2col.rs @@ -157,8 +157,10 @@ impl EvalOp for Im2Col { let r = geometry.out_format.r(); // Buffer geometry. zero_init for PackedI8K4: the K=4-inner writer skips // the K-padding lanes (k..k_aligned), which SMOPA accumulates — they must - // be 0. PackedFormat has no K padding; its mn-padding maps to discarded - // output rows, so uninitialized is fine (matches prior behaviour). + // be 0. PackedFormat has no K padding; its mn-padding lanes are computed + // on (then discarded) by the kernel, so the partial last panel still + // needs zeroing — garbage there decodes to denormals that stall the fp + // pipeline. Done after allocation below. let (single_panel_len, buf_align, zero_init) = if let Some(pf) = geometry.out_format.downcast_ref::() { (pf.single_panel_len(geometry.k), pf.alignment(), false) @@ -186,6 +188,8 @@ impl EvalOp for Im2Col { )?; if zero_init { data.as_bytes_mut().fill(0); + } else if n % r != 0 { + data.as_bytes_mut()[(n / r) * panel_bytes..].fill(0); } if n > 0 { dispatch_copy_by_size!(Patcher::patch(dt)( diff --git a/linalg/src/frame/mmm/scratch.rs b/linalg/src/frame/mmm/scratch.rs index 3d3c373f3f..07e48d9ce5 100644 --- a/linalg/src/frame/mmm/scratch.rs +++ b/linalg/src/frame/mmm/scratch.rs @@ -381,11 +381,12 @@ impl ScratchSpaceImpl { .get_unchecked(..m_remnant), ); } - if cfg!(debug_assertions) { - buf.get_unchecked_mut(m_remnant..) - .iter_mut() - .for_each(|x| *x = TI::zero()); - } + // 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::().add(down * ker.mr()) @@ -409,11 +410,9 @@ impl ScratchSpaceImpl { .get_unchecked(..n_remnant), ); } - if cfg!(debug_assertions) { - buf.get_unchecked_mut(n_remnant..) - .iter_mut() - .for_each(|x| *x = TI::zero()); - } + buf.get_unchecked_mut(n_remnant..) + .iter_mut() + .for_each(|x| *x = TI::zero()); buf.as_ptr() } else { v.as_ptr_unchecked::().add(right * ker.nr()) @@ -435,11 +434,9 @@ impl ScratchSpaceImpl { .get_unchecked(down * ker.mr()..) .get_unchecked(..m_remnant), ); - if cfg!(debug_assertions) { - r.get_unchecked_mut(m_remnant..) - .iter_mut() - .for_each(|x| *x = TI::zero()); - } + r.get_unchecked_mut(m_remnant..) + .iter_mut() + .for_each(|x| *x = TI::zero()); r.as_ptr() } else { rows.as_ptr_unchecked::().add(down * ker.mr()) @@ -454,11 +451,9 @@ impl ScratchSpaceImpl { .get_unchecked(right * ker.nr()..) .get_unchecked(..n_remnant), ); - if cfg!(debug_assertions) { - r.get_unchecked_mut(n_remnant..) - .iter_mut() - .for_each(|x| *x = TI::zero()); - } + c.get_unchecked_mut(n_remnant..) + .iter_mut() + .for_each(|x| *x = TI::zero()); c.as_ptr() } else { cols.as_ptr_unchecked::().add(right * ker.nr()) @@ -473,9 +468,7 @@ impl ScratchSpaceImpl { 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()); - if cfg!(debug_assertions) { - tmp_d_tile.iter_mut().for_each(|t| *t = TI::zero()); - } + 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; diff --git a/linalg/src/frame/pack.rs b/linalg/src/frame/pack.rs index fbdfc053d6..b19dda4c20 100644 --- a/linalg/src/frame/pack.rs +++ b/linalg/src/frame/pack.rs @@ -203,6 +203,11 @@ impl PackedFormat { ); if cfg!(debug_assertions) { packed.as_bytes_mut().fill(0u8); + } else if mn % self.r != 0 { + // The kernel computes on the last panel's padding lanes before + // their results are discarded; garbage bytes there decode to + // denormals and stall the fp pipeline. Zero the partial panel. + packed.as_bytes_mut()[(mn / self.r) * panel_bytes..].fill(0u8); } dispatch_copy!(Self::pack_t(t.datum_type())( self, @@ -246,6 +251,11 @@ impl PackedFormat { ); if cfg!(debug_assertions) { packed.as_bytes_mut().fill(0u8); + } else if mn % self.r != 0 { + // The kernel computes on the last panel's padding lanes before + // their results are discarded; garbage bytes there decode to + // denormals and stall the fp pipeline. Zero the partial panel. + packed.as_bytes_mut()[(mn / self.r) * panel_bytes..].fill(0u8); } dispatch_copy!(Self::pack_t(t.datum_type())( self,