Use MaybeUninit for uninitialized packing buffer memory

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<Self::Elem>,
                           a: *const Self::Elem, rsa: isize, csa: isize)
         {
             pack_complex::<Self::MRTy, T, TReal>(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<Self::Elem>,
                         a: *const Self::Elem, rsa: isize, csa: isize)
         {
             pack_complex::<Self::NRTy, T, TReal>(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<MR, T, TReal>(kc: usize, mc: usize, pack: &mut [T],
+pub(crate) unsafe fn pack_complex<MR, T, TReal>(kc: usize, mc: usize, pack: PackSlice<T>,
                                                 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::<TReal>();
     let areal = a as *const TReal;
     let aimag = areal.add(1);
 

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;

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,15 +169,15 @@ 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<Self::Elem>,
                       a: *const Self::Elem, rsa: isize, csa: isize)
     {
         // safety: Avx2 is enabled
         crate::packing::pack_avx2::<Self::MRTy, T>(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<Self::Elem>,
                       a: *const Self::Elem, rsa: isize, csa: isize)
     {
         // safety: Avx2 is enabled

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<K>(
             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<K>(
                     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

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<Self::Elem>,
                       a: *const Self::Elem, rsa: isize, csa: isize)
     {
         pack::<Self::MRTy, _>(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<Self::Elem>,
                       a: *const Self::Elem, rsa: isize, csa: isize)
     {
         pack::<Self::NRTy, _>(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::MRTy, T, TReal>(K, mr, &mut apack[..], a.ptr_mut(), 1, mr as isize);
-            pack_complex::<K::NRTy, T, TReal>(nr, K, &mut bpack[..], b.ptr_mut(), nr as isize, 1);
+            pack_complex::<K::MRTy, T, TReal>(K, mr, PackSlice::from_slice(&mut *apack), a.ptr_mut(), 1, mr as isize);
+            pack_complex::<K::NRTy, T, TReal>(nr, K, PackSlice::from_slice(&mut *bpack), b.ptr_mut(), nr as isize, 1);
         }
 
         let mut c = vec![T::zero(); mr * nr];

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<T>]);
+
+impl<'a, T> PackSlice<'a, T> {
+    #[inline(always)]
+    pub(crate) fn as_mut_ptr(self) -> *mut T {
+        self.0.as_mut_ptr().cast::<T>()
+    }
+
+    /// 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::<MaybeUninit<T>>(), 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<MR, T>(kc: usize, mc: usize, pack: &mut [T],
+pub(crate) unsafe fn pack<MR, T>(kc: usize, mc: usize, pack: PackSlice<T>,
                                  a: *const T, rsa: isize, csa: isize)
     where T: Element,
           MR: ConstNum,
@@ -38,7 +70,7 @@ pub(crate) unsafe fn pack<MR, T>(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<MR, T>(kc: usize, mc: usize, pack: &mut [T],
+pub(crate) unsafe fn pack_avx2<MR, T>(kc: usize, mc: usize, pack: PackSlice<T>,
                                      a: *const T, rsa: isize, csa: isize)
     where T: Element,
           MR: ConstNum,
@@ -50,7 +82,7 @@ pub(crate) unsafe fn pack_avx2<MR, T>(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<MR, T>(kc: usize, mc: usize, pack: &mut [T],
+unsafe fn pack_impl<MR, T>(kc: usize, mc: usize, pack: PackSlice<T>,
                            a: *const T, rsa: isize, csa: isize)
     where T: Element,
           MR: ConstNum,

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,15 +165,15 @@ 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<Self::Elem>,
                       a: *const Self::Elem, rsa: isize, csa: isize)
     {
         // safety: Avx2 is enabled
         crate::packing::pack_avx2::<Self::MRTy, T>(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<Self::Elem>,
                       a: *const Self::Elem, rsa: isize, csa: isize)
     {
         // safety: Avx2 is enabled

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;