diff --git a/src/cgemm_kernel.rs b/src/cgemm_kernel.rs index a8d3362..a6572ef 100644 --- a/src/cgemm_kernel.rs +++ b/src/cgemm_kernel.rs @@ -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 } @@ -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 } @@ -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 } @@ -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 } diff --git a/src/dgemm_kernel.rs b/src/dgemm_kernel.rs index 0e1779a..16282a5 100644 --- a/src/dgemm_kernel.rs +++ b/src/dgemm_kernel.rs @@ -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 } @@ -120,11 +117,8 @@ impl GemmKernel for KernelFmaAvx2 { type MRTy = ::MRTy; type NRTy = ::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 } @@ -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 } @@ -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 } @@ -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 } diff --git a/src/gemm.rs b/src/gemm.rs index cebb8e8..3661369 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")] @@ -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_| { @@ -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); @@ -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 @@ -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 { diff --git a/src/kernel.rs b/src/kernel.rs index 190689a..4063690 100644 --- a/src/kernel.rs +++ b/src/kernel.rs @@ -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)] @@ -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, @@ -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..595967c 100644 --- a/src/sgemm_kernel.rs +++ b/src/sgemm_kernel.rs @@ -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 } @@ -121,11 +118,8 @@ impl GemmKernel for KernelFmaAvx2 { type MRTy = ::MRTy; type NRTy = ::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 } @@ -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 } @@ -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 } @@ -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 } @@ -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 } diff --git a/src/zgemm_kernel.rs b/src/zgemm_kernel.rs index 26df3b9..81d30dc 100644 --- a/src/zgemm_kernel.rs +++ b/src/zgemm_kernel.rs @@ -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 } @@ -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 } @@ -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 } @@ -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 }