From 8e4c86c163222f7e3cdb437d078791526733615d Mon Sep 17 00:00:00 2001 From: Ulrik Sverdrup Date: Tue, 14 Jul 2026 16:30:14 +0200 Subject: [PATCH 1/3] gemm: Reduce monomorphizations of masked copy Use more specific generic params rather than the whole GemmKernel --- src/gemm.rs | 16 +++++++--------- 1 file changed, 7 insertions(+), 9 deletions(-) diff --git a/src/gemm.rs b/src/gemm.rs index cebb8e8..bf69c1d 100644 --- a/src/gemm.rs +++ b/src/gemm.rs @@ -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")] @@ -524,23 +524,21 @@ unsafe fn masked_kernel(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(beta: T, - c: *mut T, rsc: isize, csc: isize, - rows: usize, cols: usize, - mask_buf: &[T]) - where K: GemmKernel, T: Element, +unsafe fn c_to_masked_ab_beta_c(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 { From 5280dbfc8640e82b1b6c20a2d20bb842e40dc714 Mon Sep 17 00:00:00 2001 From: Ulrik Sverdrup Date: Tue, 14 Jul 2026 16:30:14 +0200 Subject: [PATCH 2/3] kernel: convert alignment to associated const --- src/cgemm_kernel.rs | 12 +++--------- src/dgemm_kernel.rs | 15 ++++----------- src/gemm.rs | 10 +++++----- src/kernel.rs | 4 ++-- src/sgemm_kernel.rs | 18 +++++------------- src/zgemm_kernel.rs | 12 +++--------- 6 files changed, 22 insertions(+), 49 deletions(-) diff --git a/src/cgemm_kernel.rs b/src/cgemm_kernel.rs index a8d3362..052fd96 100644 --- a/src/cgemm_kernel.rs +++ b/src/cgemm_kernel.rs @@ -65,8 +65,7 @@ impl GemmKernel for KernelAvx512 { type MRTy = U8; type NRTy = U4; - #[inline(always)] - fn align_to() -> usize { 32 } + const ALIGNMENT: usize = 32; #[inline(always)] fn always_masked() -> bool { KernelFallback::always_masked() } @@ -99,8 +98,7 @@ impl GemmKernel for KernelAvx2 { type MRTy = U4; type NRTy = U4; - #[inline(always)] - fn align_to() -> usize { 32 } + const ALIGNMENT: usize = 32; #[inline(always)] fn always_masked() -> bool { KernelFallback::always_masked() } @@ -133,8 +131,7 @@ impl GemmKernel for KernelNeon { type MRTy = U4; type NRTy = U2; - #[inline(always)] - fn align_to() -> usize { 16 } + const ALIGNMENT: usize = 16; #[inline(always)] fn always_masked() -> bool { KernelFallback::always_masked() } @@ -166,9 +163,6 @@ impl GemmKernel for KernelFallback { type MRTy = U4; type NRTy = U2; - #[inline(always)] - fn align_to() -> usize { 0 } - #[inline(always)] fn always_masked() -> bool { true } diff --git a/src/dgemm_kernel.rs b/src/dgemm_kernel.rs index 0e1779a..a949d5f 100644 --- a/src/dgemm_kernel.rs +++ b/src/dgemm_kernel.rs @@ -85,8 +85,7 @@ impl GemmKernel for KernelAvx { type MRTy = U8; type NRTy = U4; - #[inline(always)] - fn align_to() -> usize { 32 } + const ALIGNMENT: usize = 32; #[inline(always)] fn always_masked() -> bool { false } @@ -120,8 +119,7 @@ impl GemmKernel for KernelFmaAvx2 { type MRTy = ::MRTy; type NRTy = ::NRTy; - #[inline(always)] - fn align_to() -> usize { KernelAvx::align_to() } + const ALIGNMENT: usize = KernelAvx::ALIGNMENT; #[inline(always)] fn always_masked() -> bool { KernelAvx::always_masked() } @@ -172,8 +170,7 @@ impl GemmKernel for KernelAvx512 { type MRTy = U8; type NRTy = U8; - #[inline(always)] - fn align_to() -> usize { 64 } + const ALIGNMENT: usize = 64; #[inline(always)] fn always_masked() -> bool { false } @@ -220,8 +217,7 @@ impl GemmKernel for KernelNeon { type MRTy = U8; type NRTy = U4; - #[inline(always)] - fn align_to() -> usize { 32 } + const ALIGNMENT: usize = 32; #[inline(always)] fn always_masked() -> bool { false } @@ -251,9 +247,6 @@ impl GemmKernel for KernelFallback { type MRTy = U4; type NRTy = U4; - #[inline(always)] - fn align_to() -> usize { 0 } - #[inline(always)] fn always_masked() -> bool { true } diff --git a/src/gemm.rs b/src/gemm.rs index bf69c1d..eb2fb67 100644 --- a/src/gemm.rs +++ b/src/gemm.rs @@ -245,10 +245,10 @@ fn ensure_kernel_params() assert!(mr > 0 && mr <= KERNEL_MAX_MR); assert!(nr > 0 && nr <= KERNEL_MAX_NR); assert!(mr * nr * size_of::() <= 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::() * min(mr, nr); - assert!(K::align_to() <= max_align); + assert!(K::ALIGNMENT <= max_align); assert!(K::MR <= K::mc()); assert!(K::mc() <= K::kc()); @@ -405,7 +405,7 @@ unsafe fn gemm_packed(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::() <= KERNEL_MAX_SIZE && K::align_to() <= KERNEL_MAX_ALIGN); + assert!(mr * nr * size_of::() <= KERNEL_MAX_SIZE && K::ALIGNMENT <= KERNEL_MAX_ALIGN); #[cfg(not(feature = "std"))] let mut mask_buf = MaskBuffer { buffer: [0; MASK_BUF_SIZE] }; @@ -424,7 +424,7 @@ unsafe fn gemm_packed(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::()) }) .for_each(move |_tp, mask_buf, l2, nr_| { @@ -483,7 +483,7 @@ unsafe fn make_packing_buffer(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 diff --git a/src/kernel.rs b/src/kernel.rs index 190689a..3ccac89 100644 --- a/src/kernel.rs +++ b/src/kernel.rs @@ -24,7 +24,7 @@ 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; @@ -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) } } diff --git a/src/sgemm_kernel.rs b/src/sgemm_kernel.rs index 6b27cda..3c69970 100644 --- a/src/sgemm_kernel.rs +++ b/src/sgemm_kernel.rs @@ -89,8 +89,7 @@ impl GemmKernel for KernelAvx { type MRTy = U8; type NRTy = U8; - #[inline(always)] - fn align_to() -> usize { 32 } + const ALIGNMENT: usize = 32; #[inline(always)] fn always_masked() -> bool { false } @@ -121,8 +120,7 @@ impl GemmKernel for KernelFmaAvx2 { type MRTy = ::MRTy; type NRTy = ::NRTy; - #[inline(always)] - fn align_to() -> usize { KernelAvx::align_to() } + const ALIGNMENT: usize = KernelAvx::ALIGNMENT; #[inline(always)] fn always_masked() -> bool { KernelAvx::always_masked() } @@ -169,8 +167,7 @@ impl GemmKernel for KernelAvx512 { type MRTy = U16; type NRTy = U16; - #[inline(always)] - fn align_to() -> usize { 64 } + const ALIGNMENT: usize = 64; #[inline(always)] fn always_masked() -> bool { false } @@ -217,8 +214,7 @@ impl GemmKernel for KernelNeon { type MRTy = U8; type NRTy = U8; - #[inline(always)] - fn align_to() -> usize { 32 } + const ALIGNMENT: usize = 32; #[inline(always)] fn always_masked() -> bool { false } @@ -248,9 +244,6 @@ impl GemmKernel for KernelFallback { type MRTy = U8; type NRTy = U4; - #[inline(always)] - fn align_to() -> usize { 0 } - #[inline(always)] fn always_masked() -> bool { true } @@ -280,8 +273,7 @@ impl GemmKernel for KernelWasmSimd { type MRTy = U8; type NRTy = U8; - #[inline(always)] - fn align_to() -> usize { 16 } + const ALIGNMENT: usize = 16; #[inline(always)] fn always_masked() -> bool { false } diff --git a/src/zgemm_kernel.rs b/src/zgemm_kernel.rs index 26df3b9..72577fe 100644 --- a/src/zgemm_kernel.rs +++ b/src/zgemm_kernel.rs @@ -66,8 +66,7 @@ impl GemmKernel for KernelAvx512 { type MRTy = U4; type NRTy = U4; - #[inline(always)] - fn align_to() -> usize { 32 } + const ALIGNMENT: usize = 32; #[inline(always)] fn always_masked() -> bool { KernelFallback::always_masked() } @@ -100,8 +99,7 @@ impl GemmKernel for KernelAvx2 { type MRTy = U4; type NRTy = U2; - #[inline(always)] - fn align_to() -> usize { 32 } + const ALIGNMENT: usize = 32; #[inline(always)] fn always_masked() -> bool { KernelFallback::always_masked() } @@ -134,8 +132,7 @@ impl GemmKernel for KernelNeon { type MRTy = U4; type NRTy = U2; - #[inline(always)] - fn align_to() -> usize { 16 } + const ALIGNMENT: usize = 32; #[inline(always)] fn always_masked() -> bool { KernelFallback::always_masked() } @@ -167,9 +164,6 @@ impl GemmKernel for KernelFallback { type MRTy = U4; type NRTy = U2; - #[inline(always)] - fn align_to() -> usize { 0 } - #[inline(always)] fn always_masked() -> bool { true } From 3be284c5f123cd5cc0a8d802bcc9a88861ef4834 Mon Sep 17 00:00:00 2001 From: Ulrik Sverdrup Date: Tue, 14 Jul 2026 16:30:14 +0200 Subject: [PATCH 3/3] kernel: Convert always masked to associated const --- src/cgemm_kernel.rs | 15 ++++----------- src/dgemm_kernel.rs | 19 +++++-------------- src/gemm.rs | 2 +- src/kernel.rs | 6 +++--- src/sgemm_kernel.rs | 23 ++++++----------------- src/zgemm_kernel.rs | 15 ++++----------- 6 files changed, 23 insertions(+), 57 deletions(-) diff --git a/src/cgemm_kernel.rs b/src/cgemm_kernel.rs index 052fd96..a6572ef 100644 --- a/src/cgemm_kernel.rs +++ b/src/cgemm_kernel.rs @@ -66,9 +66,7 @@ impl GemmKernel for KernelAvx512 { type NRTy = U4; const ALIGNMENT: usize = 32; - - #[inline(always)] - fn always_masked() -> bool { KernelFallback::always_masked() } + const ALWAYS_MASKED: bool = KernelFallback::ALWAYS_MASKED; #[inline(always)] fn nc() -> usize { archparam::C_NC } @@ -99,9 +97,7 @@ impl GemmKernel for KernelAvx2 { type NRTy = U4; const ALIGNMENT: usize = 32; - - #[inline(always)] - fn always_masked() -> bool { KernelFallback::always_masked() } + const ALWAYS_MASKED: bool = KernelFallback::ALWAYS_MASKED; #[inline(always)] fn nc() -> usize { archparam::C_NC } @@ -132,9 +128,7 @@ impl GemmKernel for KernelNeon { type NRTy = U2; const ALIGNMENT: usize = 16; - - #[inline(always)] - fn always_masked() -> bool { KernelFallback::always_masked() } + const ALWAYS_MASKED: bool = KernelFallback::ALWAYS_MASKED; #[inline(always)] fn nc() -> usize { archparam::C_NC } @@ -163,8 +157,7 @@ impl GemmKernel for KernelFallback { type MRTy = U4; type NRTy = U2; - #[inline(always)] - fn always_masked() -> bool { true } + const ALWAYS_MASKED: bool = true; #[inline(always)] fn nc() -> usize { archparam::C_NC } diff --git a/src/dgemm_kernel.rs b/src/dgemm_kernel.rs index a949d5f..16282a5 100644 --- a/src/dgemm_kernel.rs +++ b/src/dgemm_kernel.rs @@ -86,9 +86,7 @@ impl GemmKernel for KernelAvx { type NRTy = U4; const ALIGNMENT: usize = 32; - - #[inline(always)] - fn always_masked() -> bool { false } + const ALWAYS_MASKED: bool = false; #[inline(always)] fn nc() -> usize { archparam::D_NC } @@ -120,9 +118,7 @@ impl GemmKernel for KernelFmaAvx2 { type NRTy = ::NRTy; const ALIGNMENT: usize = KernelAvx::ALIGNMENT; - - #[inline(always)] - fn always_masked() -> bool { KernelAvx::always_masked() } + const ALWAYS_MASKED: bool = KernelAvx::ALWAYS_MASKED; #[inline(always)] fn nc() -> usize { archparam::D_NC } @@ -171,9 +167,7 @@ impl GemmKernel for KernelAvx512 { type NRTy = U8; const ALIGNMENT: usize = 64; - - #[inline(always)] - fn always_masked() -> bool { false } + const ALWAYS_MASKED: bool = false; #[inline(always)] fn nc() -> usize { archparam::D_NC } @@ -218,9 +212,7 @@ impl GemmKernel for KernelNeon { type NRTy = U4; const ALIGNMENT: usize = 32; - - #[inline(always)] - fn always_masked() -> bool { false } + const ALWAYS_MASKED: bool = false; #[inline(always)] fn nc() -> usize { archparam::S_NC } @@ -247,8 +239,7 @@ impl GemmKernel for KernelFallback { type MRTy = U4; type NRTy = U4; - #[inline(always)] - fn always_masked() -> bool { true } + const ALWAYS_MASKED: bool = true; #[inline(always)] fn nc() -> usize { archparam::D_NC } diff --git a/src/gemm.rs b/src/gemm.rs index eb2fb67..3661369 100644 --- a/src/gemm.rs +++ b/src/gemm.rs @@ -439,7 +439,7 @@ unsafe fn gemm_packed(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); diff --git a/src/kernel.rs b/src/kernel.rs index 3ccac89..4063690 100644 --- a/src/kernel.rs +++ b/src/kernel.rs @@ -27,7 +27,7 @@ pub(crate) trait GemmKernel { 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)] @@ -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, diff --git a/src/sgemm_kernel.rs b/src/sgemm_kernel.rs index 3c69970..595967c 100644 --- a/src/sgemm_kernel.rs +++ b/src/sgemm_kernel.rs @@ -90,9 +90,7 @@ impl GemmKernel for KernelAvx { type NRTy = U8; const ALIGNMENT: usize = 32; - - #[inline(always)] - fn always_masked() -> bool { false } + const ALWAYS_MASKED: bool = false; #[inline(always)] fn nc() -> usize { archparam::S_NC } @@ -121,9 +119,7 @@ impl GemmKernel for KernelFmaAvx2 { type NRTy = ::NRTy; const ALIGNMENT: usize = KernelAvx::ALIGNMENT; - - #[inline(always)] - fn always_masked() -> bool { KernelAvx::always_masked() } + const ALWAYS_MASKED: bool = KernelAvx::ALWAYS_MASKED; #[inline(always)] fn nc() -> usize { archparam::S_NC } @@ -168,9 +164,7 @@ impl GemmKernel for KernelAvx512 { type NRTy = U16; const ALIGNMENT: usize = 64; - - #[inline(always)] - fn always_masked() -> bool { false } + const ALWAYS_MASKED: bool = false; #[inline(always)] fn nc() -> usize { archparam::S_NC } @@ -215,9 +209,7 @@ impl GemmKernel for KernelNeon { type NRTy = U8; const ALIGNMENT: usize = 32; - - #[inline(always)] - fn always_masked() -> bool { false } + const ALWAYS_MASKED: bool = false; #[inline(always)] fn nc() -> usize { archparam::S_NC } @@ -244,8 +236,7 @@ impl GemmKernel for KernelFallback { type MRTy = U8; type NRTy = U4; - #[inline(always)] - fn always_masked() -> bool { true } + const ALWAYS_MASKED: bool = true; #[inline(always)] fn nc() -> usize { archparam::S_NC } @@ -274,9 +265,7 @@ impl GemmKernel for KernelWasmSimd { type NRTy = U8; const ALIGNMENT: usize = 16; - - #[inline(always)] - fn always_masked() -> bool { false } + const ALWAYS_MASKED: bool = false; #[inline(always)] fn nc() -> usize { archparam::S_NC } diff --git a/src/zgemm_kernel.rs b/src/zgemm_kernel.rs index 72577fe..81d30dc 100644 --- a/src/zgemm_kernel.rs +++ b/src/zgemm_kernel.rs @@ -67,9 +67,7 @@ impl GemmKernel for KernelAvx512 { type NRTy = U4; const ALIGNMENT: usize = 32; - - #[inline(always)] - fn always_masked() -> bool { KernelFallback::always_masked() } + const ALWAYS_MASKED: bool = KernelFallback::ALWAYS_MASKED; #[inline(always)] fn nc() -> usize { archparam::Z_NC } @@ -100,9 +98,7 @@ impl GemmKernel for KernelAvx2 { type NRTy = U2; const ALIGNMENT: usize = 32; - - #[inline(always)] - fn always_masked() -> bool { KernelFallback::always_masked() } + const ALWAYS_MASKED: bool = KernelFallback::ALWAYS_MASKED; #[inline(always)] fn nc() -> usize { archparam::Z_NC } @@ -133,9 +129,7 @@ impl GemmKernel for KernelNeon { type NRTy = U2; const ALIGNMENT: usize = 32; - - #[inline(always)] - fn always_masked() -> bool { KernelFallback::always_masked() } + const ALWAYS_MASKED: bool = KernelFallback::ALWAYS_MASKED; #[inline(always)] fn nc() -> usize { archparam::Z_NC } @@ -164,8 +158,7 @@ impl GemmKernel for KernelFallback { type MRTy = U4; type NRTy = U2; - #[inline(always)] - fn always_masked() -> bool { true } + const ALWAYS_MASKED: bool = true; #[inline(always)] fn nc() -> usize { archparam::Z_NC }