diff --git a/src/cgemm_common.rs b/src/cgemm_common.rs index 59d37d5..1dcb6bb 100644 --- a/src/cgemm_common.rs +++ b/src/cgemm_common.rs @@ -13,6 +13,7 @@ use rawpointer::PointerExt; use crate::kernel::Element; use crate::kernel::ConstNum; +use crate::packing::PackSlice; #[cfg(feature = "std")] macro_rules! fmuladd { @@ -108,14 +109,14 @@ macro_rules! kernel_fallback_impl_complex { macro_rules! pack_methods { () => { #[inline] - unsafe fn pack_mr(kc: usize, mc: usize, pack: &mut [Self::Elem], + unsafe fn pack_mr(kc: usize, mc: usize, pack: PackSlice, a: *const Self::Elem, rsa: isize, csa: isize) { pack_complex::(kc, mc, pack, a, rsa, csa) } #[inline] - unsafe fn pack_nr(kc: usize, mc: usize, pack: &mut [Self::Elem], + unsafe fn pack_nr(kc: usize, mc: usize, pack: PackSlice, a: *const Self::Elem, rsa: isize, csa: isize) { pack_complex::(kc, mc, pack, a, rsa, csa) @@ -137,14 +138,14 @@ macro_rules! pack_methods { /// qy q_ q_ q_ .. (y = 2 * MR) /// ... /// ] -pub(crate) unsafe fn pack_complex(kc: usize, mc: usize, pack: &mut [T], +pub(crate) unsafe fn pack_complex(kc: usize, mc: usize, pack: PackSlice, a: *const T, rsa: isize, csa: isize) where MR: ConstNum, T: Element, TReal: Element, { // use pointers as pointer to TReal - let pack = pack.as_mut_ptr() as *mut TReal; + let pack = pack.as_mut_ptr().cast::(); let areal = a as *const TReal; let aimag = areal.add(1); diff --git a/src/cgemm_kernel.rs b/src/cgemm_kernel.rs index cdd9aaf..795704e 100644 --- a/src/cgemm_kernel.rs +++ b/src/cgemm_kernel.rs @@ -11,6 +11,7 @@ use crate::kernel::GemmSelect; use crate::kernel::{U2, U4, c32, Element, c32_mul as mul}; use crate::archparam; use crate::cgemm_common::pack_complex; +use crate::packing::PackSlice; #[cfg(any(target_arch="x86", target_arch="x86_64"))] struct KernelAvx2; diff --git a/src/dgemm_kernel.rs b/src/dgemm_kernel.rs index a5c3eb9..aa46150 100644 --- a/src/dgemm_kernel.rs +++ b/src/dgemm_kernel.rs @@ -19,6 +19,9 @@ use core::arch::x86_64::*; #[cfg(any(target_arch="x86", target_arch="x86_64"))] use crate::x86::{FusedMulAdd, AvxMulAdd, DMultiplyAdd}; +#[cfg(any(target_arch="x86", target_arch="x86_64"))] +use crate::packing::PackSlice; + #[cfg(any(target_arch="x86", target_arch="x86_64"))] struct KernelAvx; #[cfg(any(target_arch="x86", target_arch="x86_64"))] @@ -166,7 +169,7 @@ impl GemmKernel for KernelFmaAvx2 { fn mc() -> usize { archparam::D_MC } #[inline] - unsafe fn pack_mr(kc: usize, mc: usize, pack: &mut [Self::Elem], + unsafe fn pack_mr(kc: usize, mc: usize, pack: PackSlice, a: *const Self::Elem, rsa: isize, csa: isize) { // safety: Avx2 is enabled @@ -174,7 +177,7 @@ impl GemmKernel for KernelFmaAvx2 { } #[inline] - unsafe fn pack_nr(kc: usize, mc: usize, pack: &mut [Self::Elem], + unsafe fn pack_nr(kc: usize, mc: usize, pack: PackSlice, a: *const Self::Elem, rsa: isize, csa: isize) { // safety: Avx2 is enabled diff --git a/src/gemm.rs b/src/gemm.rs index 63fc230..fa2284a 100644 --- a/src/gemm.rs +++ b/src/gemm.rs @@ -24,6 +24,7 @@ use crate::kernel::GemmSelect; #[cfg(feature = "cgemm")] use crate::kernel::{c32, c64}; use crate::threading::{get_thread_pool, ThreadPoolCtx, LoopThreadConfig}; +use crate::packing::PackSlice; use crate::sgemm_kernel; use crate::dgemm_kernel; #[cfg(feature = "cgemm")] @@ -300,7 +301,7 @@ unsafe fn gemm_loop( let a = a.stride_offset(csa, kkc * l4); // Pack B -> B~ - K::pack_nr(kc, nc, slice::from_raw_parts_mut(bpp.ptr(), bp_size), + K::pack_nr(kc, nc, PackSlice::from(bpp.ptr(), bp_size), b.ptr(), csb, rsb); // First time writing to C, use user's `beta`, else accumulate @@ -320,7 +321,7 @@ unsafe fn gemm_loop( let c = c.stride_offset(rsc, kmc * l3); // Pack A -> A~ - K::pack_mr(kc, mc, slice::from_raw_parts_mut(app.ptr(), ap_size), + K::pack_mr(kc, mc, PackSlice::from(app.ptr(), ap_size), a.ptr(), rsa, csa); // LOOP 2 and 1 diff --git a/src/kernel.rs b/src/kernel.rs index 66fb910..5ea4e38 100644 --- a/src/kernel.rs +++ b/src/kernel.rs @@ -8,6 +8,7 @@ use crate::archparam; use crate::packing::pack; +use crate::packing::PackSlice; /// General matrix multiply kernel pub(crate) trait GemmKernel { @@ -43,7 +44,7 @@ pub(crate) trait GemmKernel { /// Override only if the default packing function does not /// use the right layout. #[inline] - unsafe fn pack_mr(kc: usize, mc: usize, pack_buf: &mut [Self::Elem], + unsafe fn pack_mr(kc: usize, mc: usize, pack_buf: PackSlice, a: *const Self::Elem, rsa: isize, csa: isize) { pack::(kc, mc, pack_buf, a, rsa, csa) @@ -56,7 +57,7 @@ pub(crate) trait GemmKernel { /// Override only if the default packing function does not /// use the right layout. #[inline] - unsafe fn pack_nr(kc: usize, mc: usize, pack_buf: &mut [Self::Elem], + unsafe fn pack_nr(kc: usize, mc: usize, pack_buf: PackSlice, a: *const Self::Elem, rsa: isize, csa: isize) { pack::(kc, mc, pack_buf, a, rsa, csa) @@ -304,6 +305,7 @@ pub(crate) mod test { TReal: Element + fmt::Debug + PartialEq, { use crate::cgemm_common::pack_complex; + use crate::packing::PackSlice; const K: usize = 16; let mr = K::MR; @@ -333,8 +335,8 @@ pub(crate) mod test { // unlike test_a_kernel, we need custom packing for these kernels unsafe { - pack_complex::(K, mr, &mut apack[..], a.ptr_mut(), 1, mr as isize); - pack_complex::(nr, K, &mut bpack[..], b.ptr_mut(), nr as isize, 1); + pack_complex::(K, mr, PackSlice::from_slice(&mut *apack), a.ptr_mut(), 1, mr as isize); + pack_complex::(nr, K, PackSlice::from_slice(&mut *bpack), b.ptr_mut(), nr as isize, 1); } let mut c = vec![T::zero(); mr * nr]; diff --git a/src/packing.rs b/src/packing.rs index a6b4be2..167222d 100644 --- a/src/packing.rs +++ b/src/packing.rs @@ -8,11 +8,43 @@ use rawpointer::PointerExt; +use core::mem::MaybeUninit; use core::ptr::copy_nonoverlapping; use crate::kernel::ConstNum; use crate::kernel::Element; +/// An internal wrapper around a valid mutable pointer to uninitialized memory +#[repr(transparent)] +pub(crate) struct PackSlice<'a, T>(&'a mut [MaybeUninit]); + +impl<'a, T> PackSlice<'a, T> { + #[inline(always)] + pub(crate) fn as_mut_ptr(self) -> *mut T { + self.0.as_mut_ptr().cast::() + } + + /// Create a slice that points to uninitialized memory + /// + /// # Safety + /// + /// The caller must ensure the pointer is non-null and valid for &mut access to the whole + /// slice. And the slice is not used past the memory's lifetime (should be used as temporary only). + #[inline] + pub(crate) unsafe fn from<'b>(ptr: *mut T, size: usize) -> PackSlice<'b, T> { + PackSlice(core::slice::from_raw_parts_mut(ptr.cast::>(), size)) + } + + #[cfg(all(test, feature = "cgemm"))] + #[inline] + pub(crate) fn from_slice(slice: &'_ mut [T]) -> PackSlice<'_, T> { + // SAFETY: follows from mutable slice invariants + unsafe { + PackSlice::from(slice.as_mut_ptr(), slice.len()) + } + } +} + /// Pack matrix into `pack` /// /// + kc: length of the micropanel @@ -26,7 +58,7 @@ use crate::kernel::Element; // If one of pack and a is of a reference type, it gets a noalias annotation which // gives benefits to optimization. The packing buffer is contiguous so it can be passed as a slice // here. -pub(crate) unsafe fn pack(kc: usize, mc: usize, pack: &mut [T], +pub(crate) unsafe fn pack(kc: usize, mc: usize, pack: PackSlice, a: *const T, rsa: isize, csa: isize) where T: Element, MR: ConstNum, @@ -38,7 +70,7 @@ pub(crate) unsafe fn pack(kc: usize, mc: usize, pack: &mut [T], /// Safety: Requires AVX2 #[cfg(any(target_arch="x86", target_arch="x86_64"))] #[target_feature(enable="avx2")] -pub(crate) unsafe fn pack_avx2(kc: usize, mc: usize, pack: &mut [T], +pub(crate) unsafe fn pack_avx2(kc: usize, mc: usize, pack: PackSlice, a: *const T, rsa: isize, csa: isize) where T: Element, MR: ConstNum, @@ -50,7 +82,7 @@ pub(crate) unsafe fn pack_avx2(kc: usize, mc: usize, pack: &mut [T], /// /// Uses inline(always) so that it can be instantiated for different target features. #[inline(always)] -unsafe fn pack_impl(kc: usize, mc: usize, pack: &mut [T], +unsafe fn pack_impl(kc: usize, mc: usize, pack: PackSlice, a: *const T, rsa: isize, csa: isize) where T: Element, MR: ConstNum, diff --git a/src/sgemm_kernel.rs b/src/sgemm_kernel.rs index 305fd07..7030c50 100644 --- a/src/sgemm_kernel.rs +++ b/src/sgemm_kernel.rs @@ -18,6 +18,9 @@ use core::arch::x86_64::*; #[cfg(any(target_arch="x86", target_arch="x86_64"))] use crate::x86::{FusedMulAdd, AvxMulAdd, SMultiplyAdd}; +#[cfg(any(target_arch="x86", target_arch="x86_64"))] +use crate::packing::PackSlice; + #[cfg(any(target_arch="x86", target_arch="x86_64"))] struct KernelAvx; #[cfg(any(target_arch="x86", target_arch="x86_64"))] @@ -162,7 +165,7 @@ impl GemmKernel for KernelFmaAvx2 { fn mc() -> usize { archparam::S_MC } #[inline] - unsafe fn pack_mr(kc: usize, mc: usize, pack: &mut [Self::Elem], + unsafe fn pack_mr(kc: usize, mc: usize, pack: PackSlice, a: *const Self::Elem, rsa: isize, csa: isize) { // safety: Avx2 is enabled @@ -170,7 +173,7 @@ impl GemmKernel for KernelFmaAvx2 { } #[inline] - unsafe fn pack_nr(kc: usize, mc: usize, pack: &mut [Self::Elem], + unsafe fn pack_nr(kc: usize, mc: usize, pack: PackSlice, a: *const Self::Elem, rsa: isize, csa: isize) { // safety: Avx2 is enabled diff --git a/src/zgemm_kernel.rs b/src/zgemm_kernel.rs index d7bcb66..a689a63 100644 --- a/src/zgemm_kernel.rs +++ b/src/zgemm_kernel.rs @@ -11,6 +11,7 @@ use crate::kernel::GemmSelect; use crate::kernel::{U2, U4, c64, Element, c64_mul as mul}; use crate::archparam; use crate::cgemm_common::pack_complex; +use crate::packing::PackSlice; #[cfg(any(target_arch="x86", target_arch="x86_64"))] struct KernelAvx2;