Skip to content

Small simplifications, use associated consts #108

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

Merged
bluss merged 3 commits into from
Jul 14, 2026
Merged

Small simplifications, use associated consts #108

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
8e4c86c
gemm: Reduce monomorphizations of masked copy
bluss Jul 14, 2026
5280dbf
kernel: convert alignment to associated const
bluss Jul 14, 2026
3be284c
kernel: Convert always masked to associated const
bluss Jul 14, 2026
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
27 changes: 7 additions & 20 deletions src/cgemm_kernel.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -65,11 +65,8 @@ impl GemmKernel for KernelAvx512 {
type MRTy = U8;
type NRTy = U4;

#[inline(always)]
fn align_to() -> usize { 32 }

#[inline(always)]
fn always_masked() -> bool { KernelFallback::always_masked() }
const ALIGNMENT: usize = 32;
const ALWAYS_MASKED: bool = KernelFallback::ALWAYS_MASKED;

#[inline(always)]
fn nc() -> usize { archparam::C_NC }
Expand Down Expand Up @@ -99,11 +96,8 @@ impl GemmKernel for KernelAvx2 {
type MRTy = U4;
type NRTy = U4;

#[inline(always)]
fn align_to() -> usize { 32 }

#[inline(always)]
fn always_masked() -> bool { KernelFallback::always_masked() }
const ALIGNMENT: usize = 32;
const ALWAYS_MASKED: bool = KernelFallback::ALWAYS_MASKED;

#[inline(always)]
fn nc() -> usize { archparam::C_NC }
Expand Down Expand Up @@ -133,11 +127,8 @@ impl GemmKernel for KernelNeon {
type MRTy = U4;
type NRTy = U2;

#[inline(always)]
fn align_to() -> usize { 16 }

#[inline(always)]
fn always_masked() -> bool { KernelFallback::always_masked() }
const ALIGNMENT: usize = 16;
const ALWAYS_MASKED: bool = KernelFallback::ALWAYS_MASKED;

#[inline(always)]
fn nc() -> usize { archparam::C_NC }
Expand Down Expand Up @@ -166,11 +157,7 @@ impl GemmKernel for KernelFallback {
type MRTy = U4;
type NRTy = U2;

#[inline(always)]
fn align_to() -> usize { 0 }

#[inline(always)]
fn always_masked() -> bool { true }
const ALWAYS_MASKED: bool = true;

#[inline(always)]
fn nc() -> usize { archparam::C_NC }
Expand Down
34 changes: 9 additions & 25 deletions src/dgemm_kernel.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -85,11 +85,8 @@ impl GemmKernel for KernelAvx {
type MRTy = U8;
type NRTy = U4;

#[inline(always)]
fn align_to() -> usize { 32 }

#[inline(always)]
fn always_masked() -> bool { false }
const ALIGNMENT: usize = 32;
const ALWAYS_MASKED: bool = false;

#[inline(always)]
fn nc() -> usize { archparam::D_NC }
Expand Down Expand Up @@ -120,11 +117,8 @@ impl GemmKernel for KernelFmaAvx2 {
type MRTy = <KernelAvx as GemmKernel>::MRTy;
type NRTy = <KernelAvx as GemmKernel>::NRTy;

#[inline(always)]
fn align_to() -> usize { KernelAvx::align_to() }

#[inline(always)]
fn always_masked() -> bool { KernelAvx::always_masked() }
const ALIGNMENT: usize = KernelAvx::ALIGNMENT;
const ALWAYS_MASKED: bool = KernelAvx::ALWAYS_MASKED;

#[inline(always)]
fn nc() -> usize { archparam::D_NC }
Expand Down Expand Up @@ -172,11 +166,8 @@ impl GemmKernel for KernelAvx512 {
type MRTy = U8;
type NRTy = U8;

#[inline(always)]
fn align_to() -> usize { 64 }

#[inline(always)]
fn always_masked() -> bool { false }
const ALIGNMENT: usize = 64;
const ALWAYS_MASKED: bool = false;

#[inline(always)]
fn nc() -> usize { archparam::D_NC }
Expand Down Expand Up @@ -220,11 +211,8 @@ impl GemmKernel for KernelNeon {
type MRTy = U8;
type NRTy = U4;

#[inline(always)]
fn align_to() -> usize { 32 }

#[inline(always)]
fn always_masked() -> bool { false }
const ALIGNMENT: usize = 32;
const ALWAYS_MASKED: bool = false;

#[inline(always)]
fn nc() -> usize { archparam::S_NC }
Expand All @@ -251,11 +239,7 @@ impl GemmKernel for KernelFallback {
type MRTy = U4;
type NRTy = U4;

#[inline(always)]
fn align_to() -> usize { 0 }

#[inline(always)]
fn always_masked() -> bool { true }
const ALWAYS_MASKED: bool = true;

#[inline(always)]
fn nc() -> usize { archparam::D_NC }
Expand Down
28 changes: 13 additions & 15 deletions src/gemm.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@ use crate::ptr::Ptr;
use crate::util::range_chunk;
use crate::util::round_up_to;

use crate::kernel::Element;
use crate::kernel::{ConstNum, Element};
use crate::kernel::GemmKernel;
use crate::kernel::GemmSelect;
#[cfg(feature = "cgemm")]
Expand Down Expand Up @@ -245,10 +245,10 @@ fn ensure_kernel_params<K>()
assert!(mr > 0 && mr <= KERNEL_MAX_MR);
assert!(nr > 0 && nr <= KERNEL_MAX_NR);
assert!(mr * nr * size_of::<K::Elem>() <= KERNEL_MAX_SIZE);
assert!(K::align_to() <= KERNEL_MAX_ALIGN);
assert!(K::ALIGNMENT <= KERNEL_MAX_ALIGN);
// one row/col of the kernel is limiting the max align we can provide
let max_align = size_of::<K::Elem>() * min(mr, nr);
assert!(K::align_to() <= max_align);
assert!(K::ALIGNMENT <= max_align);

assert!(K::MR <= K::mc());
assert!(K::mc() <= K::kc());
Expand Down Expand Up @@ -405,7 +405,7 @@ unsafe fn gemm_packed<K>(nc: usize, kc: usize, mc: usize,
let mr = K::MR;
let nr = K::NR;
// check for the mask buffer that fits 8 x 8 f32 and 8 x 4 f64 kernels and alignment
assert!(mr * nr * size_of::<K::Elem>() <= KERNEL_MAX_SIZE && K::align_to() <= KERNEL_MAX_ALIGN);
assert!(mr * nr * size_of::<K::Elem>() <= KERNEL_MAX_SIZE && K::ALIGNMENT <= KERNEL_MAX_ALIGN);

#[cfg(not(feature = "std"))]
let mut mask_buf = MaskBuffer { buffer: [0; MASK_BUF_SIZE] };
Expand All @@ -424,7 +424,7 @@ unsafe fn gemm_packed<K>(nc: usize, kc: usize, mc: usize,
{
ptr = MASK_BUF.with(|buf| (*buf.get()).buffer.as_mut_ptr());
}
ptr = align_ptr(K::align_to(), ptr);
ptr = align_ptr(K::ALIGNMENT, ptr);
slice::from_raw_parts_mut(ptr as *mut K::Elem, KERNEL_MAX_SIZE / size_of::<K::Elem>())
})
.for_each(move |_tp, mask_buf, l2, nr_| {
Expand All @@ -439,7 +439,7 @@ unsafe fn gemm_packed<K>(nc: usize, kc: usize, mc: usize,
// GEMM KERNEL
// NOTE: For the rust kernels, it performs better to simply
// always use the masked kernel function!
if K::always_masked() || nr_ < nr || mr_ < mr {
if K::ALWAYS_MASKED || nr_ < nr || mr_ < mr {
masked_kernel::<_, K>(kc, alpha, app.ptr(), bpp.ptr(),
beta, c.ptr(), rsc, csc,
mr_, nr_, mask_buf);
Expand Down Expand Up @@ -483,7 +483,7 @@ unsafe fn make_packing_buffer<K>(m: usize, k: usize, n: usize, na: usize)
nelem, apack_size, bpack_size,
m,k,n, na);

(Alloc::new(nelem, K::align_to()), apack_size, bpack_size)
(Alloc::new(nelem, K::ALIGNMENT), apack_size, bpack_size)
}

/// offset the ptr forwards to align to a specific byte count
Expand Down Expand Up @@ -524,23 +524,21 @@ unsafe fn masked_kernel<T, K>(k: usize, alpha: T,
{
// use column major order for `mask_buf`
K::kernel(k, alpha, a, b, T::zero(), mask_buf.as_mut_ptr(), 1, K::MR as isize);
c_to_masked_ab_beta_c::<_, K>(beta, c, rsc, csc, rows, cols, &*mask_buf);
c_to_masked_ab_beta_c::<_, K::MRTy, K::NRTy>(beta, c, rsc, csc, rows, cols, &*mask_buf);
}

/// Copy output in `mask_buf` to the actual c matrix
///
/// C ← M + βC where M is the `mask_buf`
#[inline]
unsafe fn c_to_masked_ab_beta_c<T, K>(beta: T,
c: *mut T, rsc: isize, csc: isize,
rows: usize, cols: usize,
mask_buf: &[T])
where K: GemmKernel<Elem=T>, T: Element,
unsafe fn c_to_masked_ab_beta_c<T, MR, NR>(beta: T, c: *mut T, rsc: isize, csc: isize,
rows: usize, cols: usize, mask_buf: &[T])
where T: Element, MR: ConstNum, NR: ConstNum,
{
// note: use separate function here with `&T` argument for mask buf,
// so that the compiler sees that `c` and `mask_buf` never alias.
let mr = K::MR;
let nr = K::NR;
let mr = MR::VALUE;
let nr = NR::VALUE;
let mut ab = mask_buf.as_ptr();
for j in 0..nr {
for i in 0..mr {
Expand Down
10 changes: 5 additions & 5 deletions src/kernel.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,10 +24,10 @@ pub(crate) trait GemmKernel {
type NRTy: ConstNum;

/// align inputs to this
fn align_to() -> usize;
const ALIGNMENT: usize = 0;

/// Whether to always use the masked wrapper around the kernel.
fn always_masked() -> bool;
const ALWAYS_MASKED: bool = true;

// These should ideally be tuned per kernel and per microarch
#[inline(always)]
Expand Down Expand Up @@ -81,8 +81,8 @@ pub(crate) trait GemmKernel {
/// read from c, its value is to be treated as if it was zero.
///
/// When masked, the kernel is always called with β=0 but α is passed
/// as usual. (This is only useful information if you return `true` from
/// `always_masked`.)
/// as usual. (This is only useful information if you have `true` in
/// `ALWAYS_MASKED`.)
unsafe fn kernel(
k: usize,
alpha: Self::Elem,
Expand Down Expand Up @@ -223,7 +223,7 @@ pub(crate) mod test {
K::Elem: Copy,
{
unsafe {
Alloc::new(n, K::align_to()).init_with(elt)
Alloc::new(n, K::ALIGNMENT).init_with(elt)
}
}

Expand Down
41 changes: 11 additions & 30 deletions src/sgemm_kernel.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -89,11 +89,8 @@ impl GemmKernel for KernelAvx {
type MRTy = U8;
type NRTy = U8;

#[inline(always)]
fn align_to() -> usize { 32 }

#[inline(always)]
fn always_masked() -> bool { false }
const ALIGNMENT: usize = 32;
const ALWAYS_MASKED: bool = false;

#[inline(always)]
fn nc() -> usize { archparam::S_NC }
Expand Down Expand Up @@ -121,11 +118,8 @@ impl GemmKernel for KernelFmaAvx2 {
type MRTy = <KernelAvx as GemmKernel>::MRTy;
type NRTy = <KernelAvx as GemmKernel>::NRTy;

#[inline(always)]
fn align_to() -> usize { KernelAvx::align_to() }

#[inline(always)]
fn always_masked() -> bool { KernelAvx::always_masked() }
const ALIGNMENT: usize = KernelAvx::ALIGNMENT;
const ALWAYS_MASKED: bool = KernelAvx::ALWAYS_MASKED;

#[inline(always)]
fn nc() -> usize { archparam::S_NC }
Expand Down Expand Up @@ -169,11 +163,8 @@ impl GemmKernel for KernelAvx512 {
type MRTy = U16;
type NRTy = U16;

#[inline(always)]
fn align_to() -> usize { 64 }

#[inline(always)]
fn always_masked() -> bool { false }
const ALIGNMENT: usize = 64;
const ALWAYS_MASKED: bool = false;

#[inline(always)]
fn nc() -> usize { archparam::S_NC }
Expand Down Expand Up @@ -217,11 +208,8 @@ impl GemmKernel for KernelNeon {
type MRTy = U8;
type NRTy = U8;

#[inline(always)]
fn align_to() -> usize { 32 }

#[inline(always)]
fn always_masked() -> bool { false }
const ALIGNMENT: usize = 32;
const ALWAYS_MASKED: bool = false;

#[inline(always)]
fn nc() -> usize { archparam::S_NC }
Expand All @@ -248,11 +236,7 @@ impl GemmKernel for KernelFallback {
type MRTy = U8;
type NRTy = U4;

#[inline(always)]
fn align_to() -> usize { 0 }

#[inline(always)]
fn always_masked() -> bool { true }
const ALWAYS_MASKED: bool = true;

#[inline(always)]
fn nc() -> usize { archparam::S_NC }
Expand Down Expand Up @@ -280,11 +264,8 @@ impl GemmKernel for KernelWasmSimd {
type MRTy = U8;
type NRTy = U8;

#[inline(always)]
fn align_to() -> usize { 16 }

#[inline(always)]
fn always_masked() -> bool { false }
const ALIGNMENT: usize = 16;
const ALWAYS_MASKED: bool = false;

#[inline(always)]
fn nc() -> usize { archparam::S_NC }
Expand Down
27 changes: 7 additions & 20 deletions src/zgemm_kernel.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -66,11 +66,8 @@ impl GemmKernel for KernelAvx512 {
type MRTy = U4;
type NRTy = U4;

#[inline(always)]
fn align_to() -> usize { 32 }

#[inline(always)]
fn always_masked() -> bool { KernelFallback::always_masked() }
const ALIGNMENT: usize = 32;
const ALWAYS_MASKED: bool = KernelFallback::ALWAYS_MASKED;

#[inline(always)]
fn nc() -> usize { archparam::Z_NC }
Expand Down Expand Up @@ -100,11 +97,8 @@ impl GemmKernel for KernelAvx2 {
type MRTy = U4;
type NRTy = U2;

#[inline(always)]
fn align_to() -> usize { 32 }

#[inline(always)]
fn always_masked() -> bool { KernelFallback::always_masked() }
const ALIGNMENT: usize = 32;
const ALWAYS_MASKED: bool = KernelFallback::ALWAYS_MASKED;

#[inline(always)]
fn nc() -> usize { archparam::Z_NC }
Expand Down Expand Up @@ -134,11 +128,8 @@ impl GemmKernel for KernelNeon {
type MRTy = U4;
type NRTy = U2;

#[inline(always)]
fn align_to() -> usize { 16 }

#[inline(always)]
fn always_masked() -> bool { KernelFallback::always_masked() }
const ALIGNMENT: usize = 32;
const ALWAYS_MASKED: bool = KernelFallback::ALWAYS_MASKED;

#[inline(always)]
fn nc() -> usize { archparam::Z_NC }
Expand Down Expand Up @@ -167,11 +158,7 @@ impl GemmKernel for KernelFallback {
type MRTy = U4;
type NRTy = U2;

#[inline(always)]
fn align_to() -> usize { 0 }

#[inline(always)]
fn always_masked() -> bool { true }
const ALWAYS_MASKED: bool = true;

#[inline(always)]
fn nc() -> usize { archparam::Z_NC }
Expand Down
Loading