1use std::{alloc::Layout, marker::PhantomData, ptr::NonNull};
81
82use diskann_utils::{Reborrow, ReborrowMut, strided::Strided, views::rowmajor};
83
84use super::matrix::{
85 Defaulted, LayoutError, Mat, MatMut, MatRef, NewCloned, NewMut, NewOwned, NewRef, Overflow,
86 Repr, ReprMut, ReprOwned, SliceError,
87};
88use crate::bits::{AsMutPtr, AsPtr, MutSlicePtr, SlicePtr};
89use crate::utils;
90
91#[inline]
93fn padded_ncols<const PACK: usize>(ncols: usize) -> usize {
94 ncols.next_multiple_of(PACK)
95}
96
97#[inline]
107fn compute_capacity<const GROUP: usize, const PACK: usize>(nrows: usize, ncols: usize) -> usize {
108 nrows.next_multiple_of(GROUP) * padded_ncols::<PACK>(ncols)
109}
110
111#[inline]
115fn checked_compute_capacity<const GROUP: usize, const PACK: usize>(
116 nrows: usize,
117 ncols: usize,
118) -> Option<usize> {
119 nrows
120 .checked_next_multiple_of(GROUP)?
121 .checked_mul(ncols.checked_next_multiple_of(PACK)?)
122}
123
124#[inline]
127fn linear_index<const GROUP: usize, const PACK: usize>(
128 row: usize,
129 col: usize,
130 ncols: usize,
131) -> usize {
132 let pncols = padded_ncols::<PACK>(ncols);
133 let block = row / GROUP;
134 let row_in_block = row % GROUP;
135 block * GROUP * pncols + (col / PACK) * GROUP * PACK + row_in_block * PACK + (col % PACK)
136}
137
138#[inline]
143fn col_offset<const GROUP: usize, const PACK: usize>(col: usize) -> usize {
144 (col / PACK) * GROUP * PACK + (col % PACK)
145}
146
147#[derive(Debug, Clone, Copy, PartialEq, Eq)]
152pub(crate) struct BlockTransposedRepr<T, const GROUP: usize, const PACK: usize = 1> {
153 nrows: usize,
154 ncols: usize,
155 _elem: PhantomData<T>,
156}
157
158impl<T: Copy, const GROUP: usize, const PACK: usize> BlockTransposedRepr<T, GROUP, PACK> {
159 const _ASSERTIONS: () = {
161 assert!(GROUP > 0, "group size GROUP must be positive");
162 assert!(PACK > 0, "packing factor PACK must be positive");
163 assert!(
164 GROUP.is_multiple_of(PACK),
165 "GROUP must be divisible by PACK"
166 );
167 };
168
169 pub fn new(nrows: usize, ncols: usize) -> Result<Self, Overflow> {
174 let () = Self::_ASSERTIONS;
175 let capacity = checked_compute_capacity::<GROUP, PACK>(nrows, ncols)
176 .ok_or_else(|| Overflow::for_type::<T>(nrows, ncols))?;
177 Overflow::check_byte_budget::<T>(capacity, nrows, ncols)?;
178 Ok(Self {
179 nrows,
180 ncols,
181 _elem: PhantomData,
182 })
183 }
184
185 #[inline]
189 fn storage_len(&self) -> usize {
190 compute_capacity::<GROUP, PACK>(self.nrows, self.ncols)
191 }
192
193 #[inline]
195 fn nrows(&self) -> usize {
196 self.nrows
197 }
198
199 #[inline]
201 pub fn ncols(&self) -> usize {
202 self.ncols
203 }
204
205 #[inline]
208 pub fn padded_ncols(&self) -> usize {
209 padded_ncols::<PACK>(self.ncols)
210 }
211
212 #[inline]
214 pub fn full_blocks(&self) -> usize {
215 self.nrows / GROUP
216 }
217
218 #[inline]
220 pub fn num_blocks(&self) -> usize {
221 self.nrows.div_ceil(GROUP)
222 }
223
224 #[inline]
226 pub fn remainder(&self) -> usize {
227 self.nrows % GROUP
228 }
229
230 #[inline]
235 pub fn padded_nrows(&self) -> usize {
236 self.num_blocks() * GROUP
237 }
238
239 #[inline]
241 fn block_stride(&self) -> usize {
242 GROUP * self.padded_ncols()
243 }
244
245 #[inline]
247 fn block_offset(&self, block: usize) -> usize {
248 block * self.block_stride()
249 }
250
251 fn check_slice(&self, slice: &[T]) -> Result<(), SliceError> {
253 let cap = self.storage_len();
254 if slice.len() != cap {
255 Err(SliceError::LengthMismatch {
256 expected: cap,
257 found: slice.len(),
258 })
259 } else {
260 Ok(())
261 }
262 }
263
264 unsafe fn box_to_mat(self, b: Box<[T]>) -> Mat<Self> {
270 debug_assert_eq!(b.len(), self.storage_len(), "safety contract violated");
271
272 let ptr = utils::box_into_nonnull(b).cast::<u8>();
273
274 unsafe { Mat::from_raw_parts(self, ptr) }
276 }
277}
278
279#[derive(Debug, Clone, Copy)]
288pub struct Row<'a, T, const GROUP: usize, const PACK: usize = 1> {
289 base: SlicePtr<'a, T>,
291 ncols: usize,
292}
293
294impl<T: Copy, const GROUP: usize, const PACK: usize> Row<'_, T, GROUP, PACK> {
295 #[inline]
297 pub fn len(&self) -> usize {
298 self.ncols
299 }
300
301 #[inline]
303 pub fn is_empty(&self) -> bool {
304 self.ncols == 0
305 }
306
307 #[inline]
309 pub fn get(&self, col: usize) -> Option<&T> {
310 if col < self.ncols {
311 Some(unsafe { &*self.base.as_ptr().add(col_offset::<GROUP, PACK>(col)) })
313 } else {
314 None
315 }
316 }
317
318 #[inline]
320 pub fn iter(&self) -> RowIter<'_, T, GROUP, PACK> {
321 RowIter {
322 base: self.base,
323 col: 0,
324 ncols: self.ncols,
325 }
326 }
327}
328
329impl<T: Copy, const GROUP: usize, const PACK: usize> std::ops::Index<usize>
330 for Row<'_, T, GROUP, PACK>
331{
332 type Output = T;
333
334 #[inline]
335 #[expect(clippy::panic)] fn index(&self, col: usize) -> &Self::Output {
337 self.get(col)
338 .unwrap_or_else(|| panic!("column index {col} out of bounds (ncols = {})", self.ncols))
339 }
340}
341
342#[derive(Debug, Clone)]
344pub struct RowIter<'a, T, const GROUP: usize, const PACK: usize = 1> {
345 base: SlicePtr<'a, T>,
346 col: usize,
347 ncols: usize,
348}
349
350impl<T: Copy, const GROUP: usize, const PACK: usize> Iterator for RowIter<'_, T, GROUP, PACK> {
351 type Item = T;
352
353 #[inline]
354 fn next(&mut self) -> Option<Self::Item> {
355 if self.col >= self.ncols {
356 return None;
357 }
358 let val = unsafe { *self.base.as_ptr().add(col_offset::<GROUP, PACK>(self.col)) };
360 self.col += 1;
361 Some(val)
362 }
363
364 #[inline]
365 fn size_hint(&self) -> (usize, Option<usize>) {
366 let remaining = self.ncols - self.col;
367 (remaining, Some(remaining))
368 }
369}
370
371impl<T: Copy, const GROUP: usize, const PACK: usize> ExactSizeIterator
372 for RowIter<'_, T, GROUP, PACK>
373{
374}
375impl<T: Copy, const GROUP: usize, const PACK: usize> std::iter::FusedIterator
376 for RowIter<'_, T, GROUP, PACK>
377{
378}
379
380#[derive(Debug)]
382pub struct RowMut<'a, T, const GROUP: usize, const PACK: usize = 1> {
383 base: MutSlicePtr<'a, T>,
384 ncols: usize,
385}
386
387impl<T: Copy, const GROUP: usize, const PACK: usize> RowMut<'_, T, GROUP, PACK> {
388 #[inline]
390 pub fn len(&self) -> usize {
391 self.ncols
392 }
393
394 #[inline]
396 pub fn is_empty(&self) -> bool {
397 self.ncols == 0
398 }
399
400 #[inline]
402 pub fn get(&self, col: usize) -> Option<&T> {
403 if col < self.ncols {
404 Some(unsafe { &*self.base.as_ptr().add(col_offset::<GROUP, PACK>(col)) })
406 } else {
407 None
408 }
409 }
410
411 #[inline]
413 pub fn get_mut(&mut self, col: usize) -> Option<&mut T> {
414 if col < self.ncols {
415 Some(unsafe { &mut *self.base.as_mut_ptr().add(col_offset::<GROUP, PACK>(col)) })
417 } else {
418 None
419 }
420 }
421
422 #[inline]
428 pub fn set(&mut self, col: usize, value: T) {
429 assert!(
430 col < self.ncols,
431 "column index {col} out of bounds (ncols = {})",
432 self.ncols
433 );
434 unsafe { *self.base.as_mut_ptr().add(col_offset::<GROUP, PACK>(col)) = value };
436 }
437}
438
439impl<T: Copy, const GROUP: usize, const PACK: usize> std::ops::Index<usize>
440 for RowMut<'_, T, GROUP, PACK>
441{
442 type Output = T;
443
444 #[inline]
445 #[expect(clippy::panic)] fn index(&self, col: usize) -> &Self::Output {
447 self.get(col)
448 .unwrap_or_else(|| panic!("column index {col} out of bounds (ncols = {})", self.ncols))
449 }
450}
451
452impl<T: Copy, const GROUP: usize, const PACK: usize> std::ops::IndexMut<usize>
453 for RowMut<'_, T, GROUP, PACK>
454{
455 #[inline]
456 #[expect(clippy::panic)] fn index_mut(&mut self, col: usize) -> &mut Self::Output {
458 let ncols = self.ncols;
459 self.get_mut(col)
460 .unwrap_or_else(|| panic!("column index {col} out of bounds (ncols = {ncols})"))
461 }
462}
463
464unsafe impl<T: Copy, const GROUP: usize, const PACK: usize> Repr
471 for BlockTransposedRepr<T, GROUP, PACK>
472{
473 type Row<'a>
474 = Row<'a, T, GROUP, PACK>
475 where
476 Self: 'a;
477
478 fn nrows(&self) -> usize {
479 self.nrows
480 }
481
482 fn layout(&self) -> Result<Layout, LayoutError> {
483 Ok(Layout::array::<T>(self.storage_len())?)
484 }
485
486 unsafe fn get_row<'a>(self, ptr: NonNull<u8>, i: usize) -> Self::Row<'a> {
487 debug_assert!(i < self.nrows);
488
489 if self.ncols == 0 {
492 return Row {
493 base: unsafe { SlicePtr::new_unchecked(NonNull::dangling()) },
496 ncols: 0,
497 };
498 }
499
500 let base_ptr = ptr.as_ptr().cast::<T>();
501 let offset = linear_index::<GROUP, PACK>(i, 0, self.ncols);
502
503 let row_base = unsafe { base_ptr.add(offset) };
506
507 Row {
508 base: unsafe { SlicePtr::new_unchecked(NonNull::new_unchecked(row_base)) },
511 ncols: self.ncols,
512 }
513 }
514}
515
516unsafe impl<T: Copy, const GROUP: usize, const PACK: usize> ReprMut
520 for BlockTransposedRepr<T, GROUP, PACK>
521{
522 type RowMut<'a>
523 = RowMut<'a, T, GROUP, PACK>
524 where
525 Self: 'a;
526
527 unsafe fn get_row_mut<'a>(self, ptr: NonNull<u8>, i: usize) -> Self::RowMut<'a> {
528 debug_assert!(i < self.nrows);
529
530 if self.ncols == 0 {
533 return RowMut {
534 base: unsafe { MutSlicePtr::new_unchecked(NonNull::dangling()) },
537 ncols: 0,
538 };
539 }
540
541 let base_ptr = ptr.as_ptr().cast::<T>();
542 let offset = linear_index::<GROUP, PACK>(i, 0, self.ncols);
543
544 let row_base = unsafe { base_ptr.add(offset) };
547
548 RowMut {
549 base: unsafe { MutSlicePtr::new_unchecked(NonNull::new_unchecked(row_base)) },
552 ncols: self.ncols,
553 }
554 }
555}
556
557unsafe impl<T: Copy, const GROUP: usize, const PACK: usize> ReprOwned
560 for BlockTransposedRepr<T, GROUP, PACK>
561{
562 unsafe fn drop(self, ptr: NonNull<u8>) {
563 unsafe {
565 let slice_ptr =
566 std::ptr::slice_from_raw_parts_mut(ptr.cast::<T>().as_ptr(), self.storage_len());
567 let _ = Box::from_raw(slice_ptr);
568 }
569 }
570}
571
572unsafe impl<T: Copy, const GROUP: usize, const PACK: usize> NewOwned<T>
578 for BlockTransposedRepr<T, GROUP, PACK>
579{
580 type Error = crate::error::Infallible;
581
582 fn new_owned(self, value: T) -> Result<Mat<Self>, Self::Error> {
583 let b: Box<[T]> = vec![value; self.storage_len()].into_boxed_slice();
584
585 Ok(unsafe { self.box_to_mat(b) })
587 }
588}
589
590unsafe impl<T: Copy + Default, const GROUP: usize, const PACK: usize> NewOwned<Defaulted>
592 for BlockTransposedRepr<T, GROUP, PACK>
593{
594 type Error = crate::error::Infallible;
595
596 fn new_owned(self, _: Defaulted) -> Result<Mat<Self>, Self::Error> {
597 self.new_owned(T::default())
598 }
599}
600
601impl<T: Copy, const GROUP: usize, const PACK: usize> NewCloned
602 for BlockTransposedRepr<T, GROUP, PACK>
603{
604 fn new_cloned(v: MatRef<'_, Self>) -> Mat<Self> {
605 let b: Box<[T]> = BlockTransposedRef::new(v).as_slice().into();
606
607 unsafe { v.repr().box_to_mat(b) }
610 }
611}
612
613unsafe impl<T: Copy, const GROUP: usize, const PACK: usize> NewRef<T>
615 for BlockTransposedRepr<T, GROUP, PACK>
616{
617 type Error = SliceError;
618
619 fn new_ref(self, data: &[T]) -> Result<MatRef<'_, Self>, Self::Error> {
620 self.check_slice(data)?;
621
622 Ok(unsafe { MatRef::from_raw_parts(self, utils::as_nonnull(data).cast::<u8>()) })
624 }
625}
626
627unsafe impl<T: Copy, const GROUP: usize, const PACK: usize> NewMut<T>
629 for BlockTransposedRepr<T, GROUP, PACK>
630{
631 type Error = SliceError;
632
633 fn new_mut(self, data: &mut [T]) -> Result<MatMut<'_, Self>, Self::Error> {
634 self.check_slice(data)?;
635
636 Ok(unsafe { MatMut::from_raw_parts(self, utils::as_nonnull_mut(data).cast::<u8>()) })
638 }
639}
640
641macro_rules! delegate_to_ref {
650 ($(#[$m:meta])* $vis:vis fn $name:ident(&self $(, $a:ident: $t:ty)*) $(-> $r:ty)?) => {
652 #[doc = concat!("See [`BlockTransposedRef::", stringify!($name), "`].")]
653 $(#[$m])*
654 #[inline]
655 $vis fn $name(&self $(, $a: $t)*) $(-> $r)? {
656 self.as_view().$name($($a),*)
657 }
658 };
659 ($(#[$m:meta])* unsafe $vis:vis fn $name:ident(&self $(, $a:ident: $t:ty)*) $(-> $r:ty)?) => {
661 #[doc = concat!("See [`BlockTransposedRef::", stringify!($name), "`].")]
662 $(#[$m])*
663 #[inline]
664 $vis unsafe fn $name(&self $(, $a: $t)*) $(-> $r)? {
665 unsafe { self.as_view().$name($($a),*) }
667 }
668 };
669}
670
671#[derive(Debug, Clone)]
689pub struct BlockTransposed<T: Copy, const GROUP: usize, const PACK: usize = 1> {
690 data: Mat<BlockTransposedRepr<T, GROUP, PACK>>,
691}
692
693#[derive(Debug, Clone, Copy)]
697pub struct BlockTransposedRef<'a, T: Copy, const GROUP: usize, const PACK: usize = 1> {
698 data: MatRef<'a, BlockTransposedRepr<T, GROUP, PACK>>,
699}
700
701pub struct BlockTransposedMut<'a, T: Copy, const GROUP: usize, const PACK: usize = 1> {
705 data: MatMut<'a, BlockTransposedRepr<T, GROUP, PACK>>,
706}
707
708impl<'a, T: Copy, const GROUP: usize, const PACK: usize> BlockTransposedRef<'a, T, GROUP, PACK> {
711 fn new(data: MatRef<'a, BlockTransposedRepr<T, GROUP, PACK>>) -> Self {
712 Self { data }
713 }
714
715 #[inline]
717 pub fn nrows(&self) -> usize {
718 self.data.repr().nrows()
719 }
720
721 #[inline]
723 pub fn ncols(&self) -> usize {
724 self.data.repr().ncols()
725 }
726
727 #[inline]
729 pub fn padded_ncols(&self) -> usize {
730 self.data.repr().padded_ncols()
731 }
732
733 pub const fn group_size(&self) -> usize {
735 GROUP
736 }
737
738 pub const fn const_group_size() -> usize {
740 GROUP
741 }
742
743 pub const fn pack_size(&self) -> usize {
745 PACK
746 }
747
748 #[inline]
750 pub fn full_blocks(&self) -> usize {
751 self.data.repr().full_blocks()
752 }
753
754 #[inline]
756 pub fn num_blocks(&self) -> usize {
757 self.data.repr().num_blocks()
758 }
759
760 #[inline]
763 pub fn remainder(&self) -> usize {
764 self.data.repr().remainder()
765 }
766
767 #[inline]
772 pub fn padded_nrows(&self) -> usize {
773 self.data.repr().padded_nrows()
774 }
775
776 #[inline]
778 pub fn as_ptr(&self) -> *const T {
779 self.data.as_raw_ptr().cast::<T>()
780 }
781
782 #[inline]
787 pub fn as_slice(&self) -> &'a [T] {
788 let len = self.data.repr().storage_len();
789 unsafe { std::slice::from_raw_parts(self.as_ptr(), len) }
791 }
792
793 #[inline]
805 pub unsafe fn block_ptr_unchecked(&self, block: usize) -> *const T {
806 debug_assert!(block < self.num_blocks());
807 unsafe { self.as_ptr().add(self.data.repr().block_offset(block)) }
809 }
810
811 #[expect(clippy::expect_used)]
821 pub fn block(&self, block: usize) -> rowmajor::Ref<'a, T> {
822 assert!(block < self.full_blocks());
823 let offset = self.data.repr().block_offset(block);
824 let stride = self.data.repr().block_stride();
825 let data: &[T] = unsafe { std::slice::from_raw_parts(self.as_ptr().add(offset), stride) };
828 rowmajor::Ref::try_from_data(data, self.padded_ncols() / PACK, GROUP * PACK)
829 .expect("base data should have been sized correctly")
830 }
831
832 #[expect(clippy::expect_used)]
838 pub fn remainder_block(&self) -> Option<rowmajor::Ref<'a, T>> {
839 if self.remainder() == 0 {
840 None
841 } else {
842 let offset = self.data.repr().block_offset(self.full_blocks());
843 let stride = self.data.repr().block_stride();
844 let data: &[T] =
847 unsafe { std::slice::from_raw_parts(self.as_ptr().add(offset), stride) };
848 Some(
849 rowmajor::Ref::try_from_data(data, self.padded_ncols() / PACK, GROUP * PACK)
850 .expect("base data should have been sized correctly"),
851 )
852 }
853 }
854
855 #[inline]
861 pub fn get_element(&self, row: usize, col: usize) -> T {
862 assert!(
863 row < self.nrows(),
864 "row {row} out of bounds (nrows = {})",
865 self.nrows()
866 );
867 assert!(
868 col < self.ncols(),
869 "col {col} out of bounds (ncols = {})",
870 self.ncols()
871 );
872 let idx = linear_index::<GROUP, PACK>(row, col, self.ncols());
873 unsafe { *self.as_ptr().add(idx) }
875 }
876
877 #[inline]
879 pub fn get_row(&self, i: usize) -> Option<Row<'_, T, GROUP, PACK>> {
880 self.data.get_row(i)
881 }
882}
883
884impl<'a, T: Copy, const GROUP: usize, const PACK: usize> BlockTransposedMut<'a, T, GROUP, PACK> {
887 fn new(data: MatMut<'a, BlockTransposedRepr<T, GROUP, PACK>>) -> Self {
888 Self { data }
889 }
890
891 #[inline]
893 pub fn as_view(&self) -> BlockTransposedRef<'_, T, GROUP, PACK> {
894 BlockTransposedRef::new(self.data.as_view())
895 }
896
897 delegate_to_ref!(pub fn nrows(&self) -> usize);
900 delegate_to_ref!(pub fn ncols(&self) -> usize);
901 delegate_to_ref!(pub fn padded_ncols(&self) -> usize);
902 delegate_to_ref!(pub fn full_blocks(&self) -> usize);
903 delegate_to_ref!(pub fn num_blocks(&self) -> usize);
904 delegate_to_ref!(pub fn remainder(&self) -> usize);
905 delegate_to_ref!(pub fn padded_nrows(&self) -> usize);
906 delegate_to_ref!(pub fn as_ptr(&self) -> *const T);
907 delegate_to_ref!(pub fn as_slice(&self) -> &[T]);
908 delegate_to_ref!(#[expect(clippy::missing_safety_doc)] unsafe pub fn block_ptr_unchecked(&self, block: usize) -> *const T);
909 delegate_to_ref!(pub fn block(&self, block: usize) -> rowmajor::Ref<'_, T>);
910 delegate_to_ref!(pub fn remainder_block(&self) -> Option<rowmajor::Ref<'_, T>>);
911 delegate_to_ref!(pub fn get_element(&self, row: usize, col: usize) -> T);
912
913 pub const fn group_size(&self) -> usize {
915 GROUP
916 }
917
918 pub const fn const_group_size() -> usize {
920 GROUP
921 }
922
923 pub const fn pack_size(&self) -> usize {
925 PACK
926 }
927
928 #[inline]
930 pub fn get_row(&self, i: usize) -> Option<Row<'_, T, GROUP, PACK>> {
931 self.data.get_row(i)
932 }
933
934 #[inline]
945 pub fn as_mut_slice(&mut self) -> &mut [T] {
946 self.reborrow_mut().mut_slice_inner()
947 }
948
949 fn mut_slice_inner(mut self) -> &'a mut [T] {
950 let len = self.data.repr().storage_len();
951 unsafe { std::slice::from_raw_parts_mut(self.data.as_raw_mut_ptr().cast::<T>(), len) }
953 }
954
955 pub fn block_mut(&mut self, block: usize) -> rowmajor::Mut<'_, T> {
961 self.reborrow_mut().block_mut_inner(block)
962 }
963
964 #[expect(clippy::expect_used)]
965 fn block_mut_inner(mut self, block: usize) -> rowmajor::Mut<'a, T> {
966 let repr = *self.data.repr();
967 assert!(block < repr.full_blocks());
968 let offset = repr.block_offset(block);
969 let stride = repr.block_stride();
970 let pncols = repr.padded_ncols();
971 let data: &mut [T] = unsafe {
973 std::slice::from_raw_parts_mut(
974 self.data.as_raw_mut_ptr().cast::<T>().add(offset),
975 stride,
976 )
977 };
978 rowmajor::Mut::try_from_data(data, pncols / PACK, GROUP * PACK)
979 .expect("base data should have been sized correctly")
980 }
981
982 pub fn remainder_block_mut(&mut self) -> Option<rowmajor::Mut<'_, T>> {
985 self.reborrow_mut().remainder_block_mut_inner()
986 }
987
988 #[expect(clippy::expect_used)]
989 fn remainder_block_mut_inner(mut self) -> Option<rowmajor::Mut<'a, T>> {
990 let repr = *self.data.repr();
991 if repr.remainder() == 0 {
992 None
993 } else {
994 let offset = repr.block_offset(repr.full_blocks());
995 let stride = repr.block_stride();
996 let pncols = repr.padded_ncols();
997 let data: &mut [T] = unsafe {
999 std::slice::from_raw_parts_mut(
1000 self.data.as_raw_mut_ptr().cast::<T>().add(offset),
1001 stride,
1002 )
1003 };
1004 Some(
1005 rowmajor::Mut::try_from_data(data, pncols / PACK, GROUP * PACK)
1006 .expect("base data should have been sized correctly"),
1007 )
1008 }
1009 }
1010
1011 #[inline]
1013 pub fn get_row_mut(&mut self, i: usize) -> Option<RowMut<'_, T, GROUP, PACK>> {
1014 self.data.get_row_mut(i)
1015 }
1016
1017 fn reborrow_mut(&mut self) -> BlockTransposedMut<'_, T, GROUP, PACK> {
1020 BlockTransposedMut::new(self.data.reborrow_mut())
1021 }
1022}
1023
1024impl<T: Copy, const GROUP: usize, const PACK: usize> BlockTransposed<T, GROUP, PACK> {
1027 pub fn as_view(&self) -> BlockTransposedRef<'_, T, GROUP, PACK> {
1029 BlockTransposedRef::new(self.data.as_view())
1030 }
1031
1032 pub fn as_view_mut(&mut self) -> BlockTransposedMut<'_, T, GROUP, PACK> {
1034 BlockTransposedMut::new(self.data.as_view_mut())
1035 }
1036
1037 delegate_to_ref!(pub fn nrows(&self) -> usize);
1040 delegate_to_ref!(pub fn ncols(&self) -> usize);
1041 delegate_to_ref!(pub fn padded_ncols(&self) -> usize);
1042 delegate_to_ref!(pub fn full_blocks(&self) -> usize);
1043 delegate_to_ref!(pub fn num_blocks(&self) -> usize);
1044 delegate_to_ref!(pub fn remainder(&self) -> usize);
1045 delegate_to_ref!(pub fn padded_nrows(&self) -> usize);
1046 delegate_to_ref!(pub fn as_ptr(&self) -> *const T);
1047 delegate_to_ref!(pub fn as_slice(&self) -> &[T]);
1048 delegate_to_ref!(#[expect(clippy::missing_safety_doc)] unsafe pub fn block_ptr_unchecked(&self, block: usize) -> *const T);
1049 delegate_to_ref!(pub fn block(&self, block: usize) -> rowmajor::Ref<'_, T>);
1050 delegate_to_ref!(pub fn remainder_block(&self) -> Option<rowmajor::Ref<'_, T>>);
1051 delegate_to_ref!(pub fn get_element(&self, row: usize, col: usize) -> T);
1052
1053 pub const fn group_size(&self) -> usize {
1055 GROUP
1056 }
1057
1058 pub const fn const_group_size() -> usize {
1060 GROUP
1061 }
1062
1063 pub const fn pack_size(&self) -> usize {
1065 PACK
1066 }
1067
1068 #[inline]
1070 pub fn get_row(&self, i: usize) -> Option<Row<'_, T, GROUP, PACK>> {
1071 self.data.get_row(i)
1072 }
1073
1074 #[inline]
1078 pub fn as_mut_slice(&mut self) -> &mut [T] {
1079 self.as_view_mut().mut_slice_inner()
1080 }
1081
1082 pub fn block_mut(&mut self, block: usize) -> rowmajor::Mut<'_, T> {
1084 self.as_view_mut().block_mut_inner(block)
1085 }
1086
1087 pub fn remainder_block_mut(&mut self) -> Option<rowmajor::Mut<'_, T>> {
1089 self.as_view_mut().remainder_block_mut_inner()
1090 }
1091
1092 #[inline]
1094 pub fn get_row_mut(&mut self, i: usize) -> Option<RowMut<'_, T, GROUP, PACK>> {
1095 self.data.get_row_mut(i)
1096 }
1097}
1098
1099impl<'this, T: Copy, const GROUP: usize, const PACK: usize> Reborrow<'this>
1102 for BlockTransposed<T, GROUP, PACK>
1103{
1104 type Target = BlockTransposedRef<'this, T, GROUP, PACK>;
1105
1106 #[inline]
1107 fn reborrow(&'this self) -> Self::Target {
1108 self.as_view()
1109 }
1110}
1111
1112impl<T: Copy + Default, const GROUP: usize, const PACK: usize> BlockTransposed<T, GROUP, PACK> {
1115 #[expect(clippy::expect_used)]
1121 pub fn new(nrows: usize, ncols: usize) -> Self {
1122 let repr = BlockTransposedRepr::<T, GROUP, PACK>::new(nrows, ncols)
1123 .expect("dimensions should not overflow");
1124 Self {
1125 data: Mat::new(repr, Defaulted).expect("infallible"),
1126 }
1127 }
1128
1129 pub fn try_new(nrows: usize, ncols: usize) -> Result<Self, Overflow> {
1131 let repr = BlockTransposedRepr::<T, GROUP, PACK>::new(nrows, ncols)?;
1132 Ok(Self {
1133 data: Mat::new(repr, Defaulted).expect("infallible"),
1134 })
1135 }
1136
1137 pub fn from_strided(v: Strided<'_, T>) -> Self {
1149 let nrows = v.nrows();
1150 let ncols = v.ncols();
1151 let mut mat = Self::new(nrows, ncols);
1152
1153 let repr = *mat.data.repr();
1154 let num_blocks = repr.num_blocks();
1155 let pncols = repr.padded_ncols();
1156 let num_col_groups = pncols / PACK;
1157
1158 let mut dst = mat.data.as_raw_mut_ptr().cast::<T>();
1162 for block in 0..num_blocks {
1163 let row_base = block * GROUP;
1164 for cg in 0..num_col_groups {
1165 let col_base = cg * PACK;
1166 for rib in 0..GROUP {
1167 let row = row_base + rib;
1168 if row < nrows {
1169 let src_row = unsafe { v.row_unchecked(row) };
1171 for p in 0..PACK {
1172 let col = col_base + p;
1173 if col < ncols {
1174 unsafe { *dst = *src_row.get_unchecked(col) };
1180 }
1181 dst = unsafe { dst.add(1) };
1185 }
1186 } else {
1187 dst = unsafe { dst.add(PACK) };
1190 }
1191 }
1192 }
1193 }
1194
1195 mat
1196 }
1197
1198 pub fn from_matrix_view(v: rowmajor::Ref<'_, T>) -> Self {
1200 Self::from_strided(v.into())
1201 }
1202}
1203
1204impl<T: Copy, const GROUP: usize, const PACK: usize> std::ops::Index<(usize, usize)>
1209 for BlockTransposed<T, GROUP, PACK>
1210{
1211 type Output = T;
1212
1213 #[inline]
1214 fn index(&self, (row, col): (usize, usize)) -> &Self::Output {
1215 assert!(row < self.nrows());
1216 assert!(col < self.ncols());
1217 let idx = linear_index::<GROUP, PACK>(row, col, self.ncols());
1218 unsafe { &*self.as_ptr().add(idx) }
1220 }
1221}
1222
1223#[cfg(test)]
1228mod tests {
1229 use diskann_utils::{
1244 lazy_format,
1245 views::rowmajor::{Matrix, MatrixMut},
1246 };
1247
1248 use super::*;
1249 use crate::utils::div_round_up;
1250
1251 fn gen_f32(i: usize) -> f32 {
1257 (i + 1) as f32
1258 }
1259 fn gen_i32(i: usize) -> i32 {
1260 (i + 1) as i32
1261 }
1262 fn gen_u8(i: usize) -> u8 {
1263 ((i % 255) + 1) as u8
1264 }
1265
1266 #[test]
1267 fn clone_has_independent_backing_allocation() {
1268 let mut data = rowmajor::Owned::from_element(5, 3, 0);
1269 data.as_mut_slice()
1270 .iter_mut()
1271 .enumerate()
1272 .for_each(|(i, value)| *value = (i + 1) as i32);
1273 let mut original = BlockTransposed::<i32, 4, 2>::from_matrix_view(data.as_view());
1274 let column_padding = linear_index::<4, 2>(0, 3, original.ncols());
1275 let row_padding = linear_index::<4, 2>(5, 0, original.ncols());
1276 let row_and_column_padding = linear_index::<4, 2>(5, 3, original.ncols());
1277 original.as_mut_slice()[column_padding] = -10;
1278 original.as_mut_slice()[row_padding] = -11;
1279 original.as_mut_slice()[row_and_column_padding] = -12;
1280
1281 let mut cloned = original.clone();
1282
1283 assert_eq!(cloned.as_slice(), original.as_slice());
1284 assert_eq!(cloned.as_slice()[column_padding], -10);
1285 assert_eq!(cloned.as_slice()[row_padding], -11);
1286 assert_eq!(cloned.as_slice()[row_and_column_padding], -12);
1287 assert_ne!(cloned.as_ptr(), original.as_ptr());
1288
1289 cloned.get_row_mut(0).unwrap()[0] = -1;
1290 assert_eq!(original[(0, 0)], 1);
1291 assert_eq!(cloned[(0, 0)], -1);
1292 }
1293
1294 fn test_full_api<
1307 T: Copy + Default + PartialEq + std::fmt::Debug + 'static,
1308 const GROUP: usize,
1309 const PACK: usize,
1310 >(
1311 nrows: usize,
1312 ncols: usize,
1313 gen_element: fn(usize) -> T,
1314 ) {
1315 let context = lazy_format!(
1316 "T={}, GROUP={}, PACK={}, nrows={}, ncols={}",
1317 std::any::type_name::<T>(),
1318 GROUP,
1319 PACK,
1320 nrows,
1321 ncols,
1322 );
1323
1324 let mut data = rowmajor::Owned::from_element(nrows, ncols, T::default());
1327 data.as_mut_slice()
1328 .iter_mut()
1329 .enumerate()
1330 .for_each(|(i, d)| *d = gen_element(i));
1331
1332 let mut transpose = BlockTransposed::<T, GROUP, PACK>::from_strided(data.as_view().into());
1333
1334 let expected_padded = div_round_up(ncols, PACK) * PACK;
1335 let expected_remainder = nrows % GROUP;
1336 let storage_len = transpose.as_slice().len();
1337
1338 assert_eq!(transpose.nrows(), nrows, "{}", context);
1341 assert_eq!(transpose.ncols(), ncols, "{}", context);
1342 assert_eq!(transpose.group_size(), GROUP, "{}", context);
1343 assert_eq!(
1344 BlockTransposed::<T, GROUP, PACK>::const_group_size(),
1345 GROUP,
1346 "{}",
1347 context
1348 );
1349 assert_eq!(transpose.pack_size(), PACK, "{}", context);
1350 assert_eq!(transpose.full_blocks(), nrows / GROUP, "{}", context);
1351 assert_eq!(
1352 transpose.num_blocks(),
1353 div_round_up(nrows, GROUP),
1354 "{}",
1355 context,
1356 );
1357 assert_eq!(transpose.remainder(), expected_remainder, "{}", context);
1358 assert_eq!(transpose.padded_ncols(), expected_padded, "{}", context);
1359
1360 for row in 0..nrows {
1363 for col in 0..ncols {
1364 assert_eq!(
1365 *data.element(row, col),
1366 transpose[(row, col)],
1367 "Index at ({}, {}) -- {}",
1368 row,
1369 col,
1370 context,
1371 );
1372 assert_eq!(
1373 *data.element(row, col),
1374 transpose.get_element(row, col),
1375 "get_element at ({}, {}) -- {}",
1376 row,
1377 col,
1378 context,
1379 );
1380 }
1381 }
1382
1383 let view = transpose.as_view();
1386 for row in 0..nrows {
1387 let row_view = view.get_row(row).unwrap();
1388 assert_eq!(row_view.len(), ncols, "{}", context);
1389 assert_eq!(row_view.is_empty(), ncols == 0, "{}", context);
1390 for col in 0..ncols {
1391 assert_eq!(
1392 *data.element(row, col),
1393 row_view[col],
1394 "row view at ({}, {}) -- {}",
1395 row,
1396 col,
1397 context,
1398 );
1399 }
1400 if ncols > 0 {
1402 assert_eq!(row_view.get(0), Some(data.element(row, 0)), "{}", context);
1403 }
1404 assert_eq!(row_view.get(ncols), None, "{}", context);
1405
1406 let iter = row_view.iter();
1408 assert_eq!(iter.len(), ncols, "{}", context);
1409 let (lo, hi) = iter.size_hint();
1410 assert_eq!(lo, ncols, "{}", context);
1411 assert_eq!(hi, Some(ncols), "{}", context);
1412
1413 let collected: Vec<T> = row_view.iter().collect();
1414 assert_eq!(collected.len(), ncols, "{}", context);
1415 assert_eq!(data.row(row), &*collected, "{}", context);
1416 }
1417 assert!(view.get_row(nrows).is_none(), "{}", context);
1419 let _ = view;
1420
1421 {
1424 let view = transpose.as_view();
1425 assert_eq!(view.nrows(), nrows, "{}", context);
1426 assert_eq!(view.ncols(), ncols, "{}", context);
1427 assert_eq!(view.padded_ncols(), expected_padded, "{}", context);
1428 assert_eq!(view.group_size(), GROUP, "{}", context);
1429 assert_eq!(
1430 BlockTransposedRef::<T, GROUP, PACK>::const_group_size(),
1431 GROUP,
1432 );
1433 assert_eq!(view.pack_size(), PACK, "{}", context);
1434 assert_eq!(view.full_blocks(), nrows / GROUP, "{}", context);
1435 assert_eq!(view.num_blocks(), div_round_up(nrows, GROUP), "{}", context,);
1436 assert_eq!(view.remainder(), expected_remainder, "{}", context);
1437 assert_eq!(view.as_ptr(), transpose.as_ptr(), "{}", context);
1438 assert_eq!(view.as_slice(), transpose.as_slice(), "{}", context);
1439
1440 for row in 0..nrows {
1441 for col in 0..ncols {
1442 assert_eq!(
1443 *data.element(row, col),
1444 view.get_element(row, col),
1445 "Ref get_element at ({}, {}) -- {}",
1446 row,
1447 col,
1448 context,
1449 );
1450 }
1451 let row_view = view.get_row(row).unwrap();
1452 for col in 0..ncols {
1453 assert_eq!(*data.element(row, col), row_view[col], "{}", context);
1454 }
1455 }
1456 assert!(view.get_row(nrows).is_none(), "{}", context);
1457 }
1458
1459 let expected_ptr = transpose.as_ptr();
1462 {
1463 let mut_view = transpose.as_view_mut();
1464 assert_eq!(mut_view.nrows(), nrows, "{}", context);
1465 assert_eq!(mut_view.ncols(), ncols, "{}", context);
1466 assert_eq!(mut_view.padded_ncols(), expected_padded, "{}", context);
1467 assert_eq!(mut_view.group_size(), GROUP, "{}", context);
1468 assert_eq!(
1469 BlockTransposedMut::<T, GROUP, PACK>::const_group_size(),
1470 GROUP,
1471 );
1472 assert_eq!(mut_view.pack_size(), PACK, "{}", context);
1473 assert_eq!(mut_view.full_blocks(), nrows / GROUP, "{}", context);
1474 assert_eq!(
1475 mut_view.num_blocks(),
1476 div_round_up(nrows, GROUP),
1477 "{}",
1478 context,
1479 );
1480 assert_eq!(mut_view.remainder(), expected_remainder, "{}", context);
1481 assert_eq!(mut_view.as_ptr(), expected_ptr, "{}", context);
1482 assert_eq!(mut_view.as_slice().len(), storage_len, "{}", context);
1483
1484 for row in 0..nrows {
1485 for col in 0..ncols {
1486 assert_eq!(
1487 *data.element(row, col),
1488 mut_view.get_element(row, col),
1489 "Mut get_element at ({}, {}) -- {}",
1490 row,
1491 col,
1492 context,
1493 );
1494 }
1495 let row_view = mut_view.get_row(row).unwrap();
1496 for col in 0..ncols {
1497 assert_eq!(*data.element(row, col), row_view[col], "{}", context);
1498 }
1499 }
1500 assert!(mut_view.get_row(nrows).is_none(), "{}", context);
1501 }
1502
1503 {
1506 let mut_view = transpose.as_view_mut();
1507 let ref_from_mut = mut_view.as_view();
1508 assert_eq!(ref_from_mut.nrows(), nrows, "{}", context);
1509 for row in 0..nrows {
1510 for col in 0..ncols {
1511 assert_eq!(
1512 *data.element(row, col),
1513 ref_from_mut.get_element(row, col),
1514 "{}",
1515 context,
1516 );
1517 }
1518 }
1519 }
1520
1521 {
1525 let mut mut_view = transpose.as_view_mut();
1526 assert_eq!(mut_view.as_mut_slice().len(), storage_len, "{}", context);
1527 }
1528 assert_eq!(transpose.as_mut_slice().len(), storage_len, "{}", context);
1530
1531 let expected_block_nrows = expected_padded / PACK;
1534 let expected_block_ncols = GROUP * PACK;
1535
1536 for b in 0..transpose.full_blocks() {
1537 let block_data: Vec<T>;
1538 let ptr: *const T;
1539 {
1540 let block = transpose.block(b);
1541 assert_eq!(block.nrows(), expected_block_nrows, "{}", context);
1542 assert_eq!(block.ncols(), expected_block_ncols, "{}", context);
1543
1544 ptr = unsafe { transpose.block_ptr_unchecked(b) };
1546 assert_eq!(ptr, block.as_slice().as_ptr(), "{}", context);
1547
1548 block_data = block.as_slice().to_vec();
1549 }
1550
1551 {
1553 let view = transpose.as_view();
1554 assert_eq!(view.block(b).as_slice(), &block_data[..], "{}", context);
1555 assert_eq!(unsafe { view.block_ptr_unchecked(b) }, ptr, "{}", context);
1557 }
1558
1559 {
1561 let mut_view = transpose.as_view_mut();
1562 assert_eq!(mut_view.block(b).as_slice(), &block_data[..], "{}", context);
1563 assert_eq!(
1564 unsafe { mut_view.block_ptr_unchecked(b) },
1566 ptr,
1567 "{}",
1568 context,
1569 );
1570 }
1571 }
1572
1573 if expected_remainder != 0 {
1575 let remainder_data: Vec<T>;
1576 let ptr: *const T;
1577 let fb = transpose.full_blocks();
1578 {
1579 let block = transpose.remainder_block().unwrap();
1580 assert_eq!(block.nrows(), expected_block_nrows, "{}", context);
1581 assert_eq!(block.ncols(), expected_block_ncols, "{}", context);
1582
1583 ptr = unsafe { transpose.block_ptr_unchecked(fb) };
1585 assert_eq!(ptr, block.as_slice().as_ptr(), "{}", context);
1586
1587 remainder_data = block.as_slice().to_vec();
1588 }
1589
1590 {
1592 let view = transpose.as_view();
1593 let ref_block = view.remainder_block().unwrap();
1594 assert_eq!(ref_block.as_slice(), &remainder_data[..], "{}", context);
1595 }
1596 {
1598 let mut_view = transpose.as_view_mut();
1599 let mut_block = mut_view.remainder_block().unwrap();
1600 assert_eq!(mut_block.as_slice(), &remainder_data[..], "{}", context);
1601 }
1602 } else {
1603 assert!(transpose.remainder_block().is_none(), "{}", context);
1604 {
1605 let view = transpose.as_view();
1606 assert!(view.remainder_block().is_none(), "{}", context);
1607 }
1608 {
1609 let mut_view = transpose.as_view_mut();
1610 assert!(mut_view.remainder_block().is_none(), "{}", context);
1611 }
1612 }
1613
1614 {
1617 let mut mut_view = transpose.as_view_mut();
1618 for b in 0..mut_view.full_blocks() {
1619 let block_mut = mut_view.block_mut(b);
1620 assert_eq!(block_mut.nrows(), expected_block_nrows, "{}", context);
1621 assert_eq!(block_mut.ncols(), expected_block_ncols, "{}", context);
1622 }
1623 if expected_remainder != 0 {
1624 let rem = mut_view.remainder_block_mut().unwrap();
1625 assert_eq!(rem.nrows(), expected_block_nrows, "{}", context);
1626 assert_eq!(rem.ncols(), expected_block_ncols, "{}", context);
1627 } else {
1628 assert!(mut_view.remainder_block_mut().is_none(), "{}", context);
1629 }
1630 }
1631
1632 for b in 0..transpose.full_blocks() {
1634 let block_mut = transpose.block_mut(b);
1635 assert_eq!(block_mut.nrows(), expected_block_nrows, "{}", context);
1636 assert_eq!(block_mut.ncols(), expected_block_ncols, "{}", context);
1637 }
1638 if expected_remainder != 0 {
1639 let rem = transpose.remainder_block_mut().unwrap();
1640 assert_eq!(rem.nrows(), expected_block_nrows, "{}", context);
1641 assert_eq!(rem.ncols(), expected_block_ncols, "{}", context);
1642 } else {
1643 assert!(transpose.remainder_block_mut().is_none(), "{}", context);
1644 }
1645
1646 {
1649 let mut mut_view = transpose.as_view_mut();
1650 for row in 0..nrows {
1651 let row_view = mut_view.get_row_mut(row).unwrap();
1652 assert_eq!(row_view.len(), ncols, "{}", context);
1653 assert_eq!(row_view.is_empty(), ncols == 0, "{}", context);
1654 for col in 0..ncols {
1655 assert_eq!(*data.element(row, col), row_view[col], "{}", context);
1656 }
1657 }
1658 assert!(mut_view.get_row_mut(nrows).is_none(), "{}", context);
1659 }
1660
1661 if nrows > 0 && ncols > 0 {
1664 {
1666 let view = transpose.as_view();
1667 let row = view.get_row(0).unwrap();
1668 assert_eq!(row.get(ncols), None, "{}", context);
1669 assert_eq!(row.get(usize::MAX), None, "{}", context);
1670 }
1671
1672 let row = transpose.get_row_mut(0).unwrap();
1674 assert_eq!(row.get(ncols), None, "{}", context);
1675
1676 let mut row = transpose.get_row_mut(0).unwrap();
1678 let sentinel = gen_element(usize::MAX / 2);
1679 let original = row[0];
1680 if let Some(v) = row.get_mut(0) {
1681 *v = sentinel;
1682 }
1683 assert_eq!(row.get_mut(ncols), None, "{}", context);
1684 let _ = row;
1686 assert_eq!(transpose.get_element(0, 0), sentinel, "{}", context);
1687 transpose.get_row_mut(0).unwrap().set(0, original);
1689 }
1690
1691 for b in 0..transpose.full_blocks() {
1694 transpose.block_mut(b).as_mut_slice().fill(T::default());
1695 }
1696 if transpose.remainder() != 0 {
1697 transpose
1698 .remainder_block_mut()
1699 .unwrap()
1700 .as_mut_slice()
1701 .fill(T::default());
1702 }
1703 assert!(
1704 transpose.as_slice().iter().all(|v| *v == T::default()),
1705 "not fully zeroed -- {}",
1706 context,
1707 );
1708
1709 let transpose = BlockTransposed::<T, GROUP, PACK>::from_strided(data.as_view().into());
1712 let raw = transpose.as_slice();
1713
1714 for row in 0..nrows {
1716 for col in ncols..expected_padded {
1717 let idx = linear_index::<GROUP, PACK>(row, col, ncols);
1718 assert_eq!(
1719 raw[idx],
1720 T::default(),
1721 "col padding at ({}, {}) -- {}",
1722 row,
1723 col,
1724 context,
1725 );
1726 }
1727 }
1728
1729 let padded_nrows = nrows.next_multiple_of(GROUP);
1731 for row in nrows..padded_nrows {
1732 for col in 0..expected_padded {
1733 let idx = linear_index::<GROUP, PACK>(row, col, ncols);
1734 assert_eq!(
1735 raw[idx],
1736 T::default(),
1737 "row padding at ({}, {}) -- {}",
1738 row,
1739 col,
1740 context,
1741 );
1742 }
1743 }
1744
1745 assert_eq!(
1748 transpose.as_view().padded_nrows(),
1749 padded_nrows,
1750 "padded_nrows() mismatch -- {}",
1751 context,
1752 );
1753
1754 if nrows > 0 && ncols > 0 {
1757 let via_matrix = BlockTransposed::<T, GROUP, PACK>::from_matrix_view(data.as_view());
1758 assert_eq!(via_matrix.as_slice(), transpose.as_slice(), "{}", context);
1759 }
1760 }
1761
1762 #[test]
1767 fn test_api_pack1_group16() {
1768 let rows: Vec<usize> = if cfg!(miri) {
1771 vec![0, 1, 15, 16, 17, 33]
1772 } else {
1773 (0..128).collect()
1774 };
1775 let cols: Vec<usize> = if cfg!(miri) {
1776 vec![0, 1, 2]
1777 } else {
1778 (0..5).collect()
1779 };
1780 for &nrows in &rows {
1781 for &ncols in &cols {
1782 test_full_api::<f32, 16, 1>(nrows, ncols, gen_f32);
1783 }
1784 }
1785 }
1786
1787 #[test]
1788 fn test_api_pack1_group8() {
1789 let rows: Vec<usize> = if cfg!(miri) {
1792 vec![0, 1, 7, 8, 9, 17]
1793 } else {
1794 (0..128).collect()
1795 };
1796 let cols: Vec<usize> = if cfg!(miri) {
1797 vec![0, 1, 2]
1798 } else {
1799 (0..5).collect()
1800 };
1801 for &nrows in &rows {
1802 for &ncols in &cols {
1803 test_full_api::<f32, 8, 1>(nrows, ncols, gen_f32);
1804 }
1805 }
1806 }
1807
1808 #[test]
1809 fn test_api_pack2() {
1810 let rows: Vec<usize> = if cfg!(miri) {
1813 vec![0, 1, 3, 4, 5, 7, 8, 9, 15, 16, 17]
1814 } else {
1815 (0..48).collect()
1816 };
1817 let cols: Vec<usize> = if cfg!(miri) {
1818 vec![0, 1, 2, 3, 4, 5]
1819 } else {
1820 (0..9).collect()
1821 };
1822 for &nrows in &rows {
1823 for &ncols in &cols {
1824 test_full_api::<f32, 4, 2>(nrows, ncols, gen_f32);
1825 test_full_api::<f32, 8, 2>(nrows, ncols, gen_f32);
1826 test_full_api::<f32, 16, 2>(nrows, ncols, gen_f32);
1827 }
1828 }
1829 }
1830
1831 #[test]
1832 fn test_api_pack4() {
1833 let rows: Vec<usize> = if cfg!(miri) {
1836 vec![0, 1, 3, 4, 5, 7, 8, 9, 15, 16, 17]
1837 } else {
1838 (0..48).collect()
1839 };
1840 let cols: Vec<usize> = if cfg!(miri) {
1841 vec![0, 1, 3, 4, 5, 8]
1842 } else {
1843 (0..9).collect()
1844 };
1845 for &nrows in &rows {
1846 for &ncols in &cols {
1847 test_full_api::<f32, 4, 4>(nrows, ncols, gen_f32);
1848 test_full_api::<f32, 8, 4>(nrows, ncols, gen_f32);
1849 test_full_api::<f32, 16, 4>(nrows, ncols, gen_f32);
1850 }
1851 }
1852 }
1853
1854 #[test]
1856 fn test_api_non_f32() {
1857 test_full_api::<i32, 4, 1>(10, 7, gen_i32);
1859 test_full_api::<i32, 8, 2>(12, 5, gen_i32);
1860
1861 test_full_api::<u8, 4, 2>(12, 5, gen_u8);
1863 test_full_api::<u8, 8, 1>(10, 7, gen_u8);
1864 }
1865
1866 fn test_block_layout_pack1<
1873 T: Copy + Default + PartialEq + std::fmt::Debug + 'static,
1874 const GROUP: usize,
1875 >(
1876 nrows: usize,
1877 ncols: usize,
1878 gen_element: fn(usize) -> T,
1879 ) {
1880 let mut data = rowmajor::Owned::from_element(nrows, ncols, T::default());
1881 data.as_mut_slice()
1882 .iter_mut()
1883 .enumerate()
1884 .for_each(|(i, d)| *d = gen_element(i));
1885
1886 let transpose = BlockTransposed::<T, GROUP, 1>::from_strided(data.as_view().into());
1887
1888 for b in 0..transpose.full_blocks() {
1890 let block = transpose.block(b);
1891 for i in 0..block.nrows() {
1892 for j in 0..block.ncols() {
1893 assert_eq!(
1894 *block.element(i, j),
1895 *data.element(GROUP * b + j, i),
1896 "block {} at ({}, {}) -- GROUP={}, nrows={}, ncols={}",
1897 b,
1898 i,
1899 j,
1900 GROUP,
1901 nrows,
1902 ncols,
1903 );
1904 }
1905 }
1906 }
1907
1908 if transpose.remainder() != 0 {
1910 let fb = transpose.full_blocks();
1911 let block = transpose.remainder_block().unwrap();
1912 for i in 0..block.nrows() {
1913 for j in 0..transpose.remainder() {
1914 assert_eq!(
1915 *block.element(i, j),
1916 *data.element(GROUP * fb + j, i),
1917 "remainder at ({}, {}) -- GROUP={}, nrows={}, ncols={}",
1918 i,
1919 j,
1920 GROUP,
1921 nrows,
1922 ncols,
1923 );
1924 }
1925 }
1926 }
1927 }
1928
1929 #[test]
1930 fn test_block_layout_pack1_group16() {
1931 let rows: Vec<usize> = if cfg!(miri) {
1932 vec![0, 1, 15, 16, 17, 33]
1933 } else {
1934 (0..128).collect()
1935 };
1936 let cols: Vec<usize> = if cfg!(miri) {
1937 vec![0, 1, 2]
1938 } else {
1939 (0..5).collect()
1940 };
1941 for &nrows in &rows {
1942 for &ncols in &cols {
1943 test_block_layout_pack1::<f32, 16>(nrows, ncols, gen_f32);
1944 }
1945 }
1946 }
1947
1948 #[test]
1949 fn test_block_layout_pack1_group8() {
1950 let rows: Vec<usize> = if cfg!(miri) {
1951 vec![0, 1, 7, 8, 9, 17]
1952 } else {
1953 (0..128).collect()
1954 };
1955 let cols: Vec<usize> = if cfg!(miri) {
1956 vec![0, 1, 2]
1957 } else {
1958 (0..5).collect()
1959 };
1960 for &nrows in &rows {
1961 for &ncols in &cols {
1962 test_block_layout_pack1::<f32, 8>(nrows, ncols, gen_f32);
1963 }
1964 }
1965 }
1966
1967 #[test]
1974 fn test_row_view_send_sync() {
1975 fn assert_send<T: Send>() {}
1976 fn assert_sync<T: Sync>() {}
1977
1978 assert_send::<Row<'_, f32, 16>>();
1979 assert_sync::<Row<'_, f32, 16>>();
1980 assert_send::<Row<'_, u8, 8, 2>>();
1981 assert_sync::<Row<'_, u8, 8, 2>>();
1982
1983 assert_send::<RowMut<'_, f32, 16>>();
1984 assert_sync::<RowMut<'_, f32, 16>>();
1985 assert_send::<RowMut<'_, i32, 4, 4>>();
1986 assert_sync::<RowMut<'_, i32, 4, 4>>();
1987 }
1988
1989 #[test]
1992 fn test_new_ref_and_new_mut() {
1993 let nrows = 5;
1994 let ncols = 3;
1995 let repr = BlockTransposedRepr::<f32, 4>::new(nrows, ncols).unwrap();
1996
1997 let mat = BlockTransposed::<f32, 4>::new(nrows, ncols);
1998 let raw: &[f32] = mat.as_slice();
1999
2000 let mat_ref = BlockTransposedRef::new(repr.new_ref(raw).unwrap());
2001 assert_eq!(mat_ref.nrows(), nrows);
2002 assert_eq!(mat_ref.ncols(), ncols);
2003 for row in 0..nrows {
2004 for col in 0..ncols {
2005 assert_eq!(mat_ref.get_element(row, col), mat.get_element(row, col));
2006 }
2007 }
2008
2009 let mut buf = raw.to_vec();
2010 let mat_mut = BlockTransposedMut::new(repr.new_mut(&mut buf).unwrap());
2011 assert_eq!(mat_mut.nrows(), nrows);
2012 assert_eq!(mat_mut.ncols(), ncols);
2013
2014 let mut short = vec![0.0_f32; 2];
2016 assert!(repr.new_ref(&short).is_err());
2017 assert!(repr.new_mut(&mut short).is_err());
2018 }
2019
2020 #[test]
2023 fn test_row_view_empty() {
2024 fn check_empty<const GROUP: usize, const PACK: usize>() {
2027 let mut mat = BlockTransposed::<f32, GROUP, PACK>::new(4, 0);
2028
2029 let view = mat.as_view();
2031 for i in 0..4 {
2032 let row = view.get_row(i).unwrap();
2033 assert!(row.is_empty());
2034 assert_eq!(row.len(), 0);
2035 assert_eq!(row.iter().count(), 0);
2036 }
2037
2038 for i in 0..4 {
2040 let row = mat.get_row_mut(i).unwrap();
2041 assert!(row.is_empty());
2042 assert_eq!(row.len(), 0);
2043 }
2044 }
2045
2046 check_empty::<16, 1>(); check_empty::<4, 2>(); check_empty::<4, 4>(); }
2050
2051 #[test]
2054 #[should_panic(expected = "column index 3 out of bounds")]
2055 fn test_row_view_index_oob() {
2056 let mat = BlockTransposed::<f32, 4>::new(4, 3);
2057 let view = mat.as_view();
2058 let row = view.get_row(0).unwrap();
2059 let _ = row[3];
2060 }
2061
2062 #[test]
2063 #[should_panic(expected = "column index 3 out of bounds")]
2064 fn test_row_view_mut_index_oob() {
2065 let mut mat = BlockTransposed::<f32, 4>::new(4, 3);
2066 let row = mat.get_row_mut(0).unwrap();
2067 let _ = row[3];
2068 }
2069
2070 #[test]
2071 #[should_panic(expected = "column index 3 out of bounds")]
2072 fn test_row_view_mut_index_mut_oob() {
2073 let mut mat = BlockTransposed::<f32, 4>::new(4, 3);
2074 let mut row = mat.get_row_mut(0).unwrap();
2075 row[3] = 1.0;
2076 }
2077
2078 #[test]
2079 #[should_panic(expected = "column index 3 out of bounds")]
2080 fn test_row_view_set_oob() {
2081 let mut mat = BlockTransposed::<f32, 4>::new(4, 3);
2082 let mut row = mat.get_row_mut(0).unwrap();
2083 row.set(3, 1.0);
2084 }
2085
2086 #[test]
2087 #[should_panic(expected = "row 4 out of bounds")]
2088 fn test_get_element_row_oob() {
2089 let mat = BlockTransposed::<f32, 4>::new(4, 3);
2090 mat.get_element(4, 0);
2091 }
2092
2093 #[test]
2094 #[should_panic(expected = "col 3 out of bounds")]
2095 fn test_get_element_col_oob() {
2096 let mat = BlockTransposed::<f32, 4>::new(4, 3);
2097 mat.get_element(0, 3);
2098 }
2099
2100 #[test]
2101 #[should_panic(expected = "assertion failed")]
2102 fn test_index_tuple_row_oob() {
2103 let mat = BlockTransposed::<f32, 4>::new(4, 3);
2104 let _ = mat[(4, 0)];
2105 }
2106
2107 #[test]
2108 #[should_panic(expected = "assertion failed")]
2109 fn test_index_tuple_col_oob() {
2110 let mat = BlockTransposed::<f32, 4>::new(4, 3);
2111 let _ = mat[(0, 3)];
2112 }
2113
2114 #[test]
2115 #[should_panic]
2116 fn test_block_oob() {
2117 let mat = BlockTransposed::<f32, 4>::new(4, 3);
2118 let _ = mat.block(1);
2119 }
2120
2121 #[test]
2122 #[should_panic]
2123 fn test_block_mut_oob() {
2124 let mut mat = BlockTransposed::<f32, 4>::new(4, 3);
2125 let _ = mat.block_mut(1);
2126 }
2127
2128 #[test]
2131 fn test_from_strided_nonunit_stride() {
2132 const GROUP: usize = 4;
2133 const PACK: usize = 2;
2134 let nrows = 5;
2135 let ncols = 3;
2136 let cstride = 8;
2137
2138 let required_len = (nrows - 1) * cstride + ncols;
2139 let mut flat = vec![0.0_f32; required_len];
2140 for row in 0..nrows {
2141 for col in 0..ncols {
2142 flat[row * cstride + col] = (row * 100 + col + 1) as f32;
2143 }
2144 }
2145
2146 let strided = Strided::try_from_data(&flat, nrows, ncols, cstride)
2147 .expect("should construct strided view");
2148 let transpose = BlockTransposed::<f32, GROUP, PACK>::from_strided(strided);
2149
2150 assert_eq!(transpose.nrows(), nrows);
2151 assert_eq!(transpose.ncols(), ncols);
2152
2153 for row in 0..nrows {
2154 for col in 0..ncols {
2155 let expected = (row * 100 + col + 1) as f32;
2156 assert_eq!(
2157 transpose[(row, col)],
2158 expected,
2159 "mismatch at ({}, {})",
2160 row,
2161 col,
2162 );
2163 }
2164 }
2165
2166 let padded_ncols = ncols.next_multiple_of(PACK);
2167 let raw: &[f32] = transpose.as_slice();
2168 for row in 0..nrows {
2169 for col in ncols..padded_ncols {
2170 let idx = linear_index::<GROUP, PACK>(row, col, ncols);
2171 assert_eq!(
2172 raw[idx], 0.0,
2173 "column-padding at ({}, {}) should be zero",
2174 row, col,
2175 );
2176 }
2177 }
2178 }
2179
2180 #[test]
2183 fn test_concurrent_row_mutation() {
2184 const GROUP: usize = 8;
2185 const PACK: usize = 2;
2186
2187 let (nrows, ncols, num_threads) = if cfg!(miri) { (8, 4, 2) } else { (64, 16, 4) };
2188
2189 let mut mat = BlockTransposed::<f32, GROUP, PACK>::new(nrows, ncols);
2190 let rows: Vec<RowMut<'_, f32, GROUP, PACK>> = mat.data.rows_mut().collect();
2191 let rows_per_thread = nrows / num_threads;
2192 let mut rows = rows.into_boxed_slice();
2193
2194 std::thread::scope(|s| {
2195 let mut remaining = &mut rows[..];
2196 for thread_id in 0..num_threads {
2197 let chunk_len = if thread_id == num_threads - 1 {
2198 remaining.len()
2199 } else {
2200 rows_per_thread
2201 };
2202 let (chunk, rest) = remaining.split_at_mut(chunk_len);
2203 remaining = rest;
2204 let start_row = thread_id * rows_per_thread;
2205
2206 s.spawn(move || {
2207 for (offset, row_view) in chunk.iter_mut().enumerate() {
2208 let row = start_row + offset;
2209 for col in 0..ncols {
2210 let value = (thread_id * 10000 + row * 100 + col) as f32;
2211 row_view.set(col, value);
2212 }
2213 }
2214 });
2215 }
2216 });
2217
2218 for row in 0..nrows {
2219 let thread_id = (row / rows_per_thread).min(num_threads - 1);
2220 for col in 0..ncols {
2221 let expected = (thread_id * 10000 + row * 100 + col) as f32;
2222 assert_eq!(
2223 mat.get_element(row, col),
2224 expected,
2225 "mismatch at ({}, {})",
2226 row,
2227 col,
2228 );
2229 }
2230 }
2231 }
2232}