Skip to content

linalg,core: zero the padding lanes the mmm kernels compute on #2596

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
czoli1976 wants to merge 1 commit into
base: main
Choose a base branch
from
+32 −25
Open

linalg,core: zero the padding lanes the mmm kernels compute on #2596

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 6 additions & 2 deletions core/src/ops/cnn/conv/im2col.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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::<PackedFormat>() {
(pf.single_panel_len(geometry.k), pf.alignment(), false)
Expand Down Expand Up @@ -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)(
Expand Down
39 changes: 16 additions & 23 deletions linalg/src/frame/mmm/scratch.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -381,11 +381,12 @@ impl<TI: LADatum> ScratchSpaceImpl<TI> {
.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::<TI>().add(down * ker.mr())
Expand All @@ -409,11 +410,9 @@ impl<TI: LADatum> ScratchSpaceImpl<TI> {
.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::<TI>().add(right * ker.nr())
Expand All @@ -435,11 +434,9 @@ impl<TI: LADatum> ScratchSpaceImpl<TI> {
.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::<TI>().add(down * ker.mr())
Expand All @@ -454,11 +451,9 @@ impl<TI: LADatum> ScratchSpaceImpl<TI> {
.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::<TI>().add(right * ker.nr())
Expand All @@ -473,9 +468,7 @@ impl<TI: LADatum> ScratchSpaceImpl<TI> {
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;
Expand Down
10 changes: 10 additions & 0 deletions linalg/src/frame/pack.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down