1use std::{alloc::Layout, marker::PhantomData, ptr::NonNull};
81
82use diskann_utils::{
83 Reborrow, ReborrowMut,
84 strided::StridedView,
85 views::{MatrixView, MutMatrixView},
86};
87
88use super::matrix::{
89 Defaulted, LayoutError, Mat, MatMut, MatRef, NewCloned, NewMut, NewOwned, NewRef, Overflow,
90 Repr, ReprMut, ReprOwned, SliceError,
91};
92use crate::bits::{AsMutPtr, AsPtr, MutSlicePtr, SlicePtr};
93use crate::utils;
94
95#[inline]
97fn padded_ncols<const PACK: usize>(ncols: usize) -> usize {
98 ncols.next_multiple_of(PACK)
99}
100
101#[inline]
111fn compute_capacity<const GROUP: usize, const PACK: usize>(nrows: usize, ncols: usize) -> usize {
112 nrows.next_multiple_of(GROUP) * padded_ncols::<PACK>(ncols)
113}
114
115#[inline]
119fn checked_compute_capacity<const GROUP: usize, const PACK: usize>(
120 nrows: usize,
121 ncols: usize,
122) -> Option<usize> {
123 nrows
124 .checked_next_multiple_of(GROUP)?
125 .checked_mul(ncols.checked_next_multiple_of(PACK)?)
126}
127
128#[inline]
131fn linear_index<const GROUP: usize, const PACK: usize>(
132 row: usize,
133 col: usize,
134 ncols: usize,
135) -> usize {
136 let pncols = padded_ncols::<PACK>(ncols);
137 let block = row / GROUP;
138 let row_in_block = row % GROUP;
139 block * GROUP * pncols + (col / PACK) * GROUP * PACK + row_in_block * PACK + (col % PACK)
140}
141
142#[inline]
147fn col_offset<const GROUP: usize, const PACK: usize>(col: usize) -> usize {
148 (col / PACK) * GROUP * PACK + (col % PACK)
149}
150
151#[derive(Debug, Clone, Copy, PartialEq, Eq)]
156pub(crate) struct BlockTransposedRepr<T, const GROUP: usize, const PACK: usize = 1> {
157 nrows: usize,
158 ncols: usize,
159 _elem: PhantomData<T>,
160}
161
162impl<T: Copy, const GROUP: usize, const PACK: usize> BlockTransposedRepr<T, GROUP, PACK> {
163 const _ASSERTIONS: () = {
165 assert!(GROUP > 0, "group size GROUP must be positive");
166 assert!(PACK > 0, "packing factor PACK must be positive");
167 assert!(
168 GROUP.is_multiple_of(PACK),
169 "GROUP must be divisible by PACK"
170 );
171 };
172
173 pub fn new(nrows: usize, ncols: usize) -> Result<Self, Overflow> {
178 let () = Self::_ASSERTIONS;
179 let capacity = checked_compute_capacity::<GROUP, PACK>(nrows, ncols)
180 .ok_or_else(|| Overflow::for_type::<T>(nrows, ncols))?;
181 Overflow::check_byte_budget::<T>(capacity, nrows, ncols)?;
182 Ok(Self {
183 nrows,
184 ncols,
185 _elem: PhantomData,
186 })
187 }
188
189 #[inline]
193 fn storage_len(&self) -> usize {
194 compute_capacity::<GROUP, PACK>(self.nrows, self.ncols)
195 }
196
197 #[inline]
199 fn nrows(&self) -> usize {
200 self.nrows
201 }
202
203 #[inline]
205 pub fn ncols(&self) -> usize {
206 self.ncols
207 }
208
209 #[inline]
212 pub fn padded_ncols(&self) -> usize {
213 padded_ncols::<PACK>(self.ncols)
214 }
215
216 #[inline]
218 pub fn full_blocks(&self) -> usize {
219 self.nrows / GROUP
220 }
221
222 #[inline]
224 pub fn num_blocks(&self) -> usize {
225 self.nrows.div_ceil(GROUP)
226 }
227
228 #[inline]
230 pub fn remainder(&self) -> usize {
231 self.nrows % GROUP
232 }
233
234 #[inline]
239 pub fn padded_nrows(&self) -> usize {
240 self.num_blocks() * GROUP
241 }
242
243 #[inline]
245 fn block_stride(&self) -> usize {
246 GROUP * self.padded_ncols()
247 }
248
249 #[inline]
251 fn block_offset(&self, block: usize) -> usize {
252 block * self.block_stride()
253 }
254
255 fn check_slice(&self, slice: &[T]) -> Result<(), SliceError> {
257 let cap = self.storage_len();
258 if slice.len() != cap {
259 Err(SliceError::LengthMismatch {
260 expected: cap,
261 found: slice.len(),
262 })
263 } else {
264 Ok(())
265 }
266 }
267
268 unsafe fn box_to_mat(self, b: Box<[T]>) -> Mat<Self> {
274 debug_assert_eq!(b.len(), self.storage_len(), "safety contract violated");
275
276 let ptr = utils::box_into_nonnull(b).cast::<u8>();
277
278 unsafe { Mat::from_raw_parts(self, ptr) }
280 }
281}
282
283#[derive(Debug, Clone, Copy)]
292pub struct Row<'a, T, const GROUP: usize, const PACK: usize = 1> {
293 base: SlicePtr<'a, T>,
295 ncols: usize,
296}
297
298impl<T: Copy, const GROUP: usize, const PACK: usize> Row<'_, T, GROUP, PACK> {
299 #[inline]
301 pub fn len(&self) -> usize {
302 self.ncols
303 }
304
305 #[inline]
307 pub fn is_empty(&self) -> bool {
308 self.ncols == 0
309 }
310
311 #[inline]
313 pub fn get(&self, col: usize) -> Option<&T> {
314 if col < self.ncols {
315 Some(unsafe { &*self.base.as_ptr().add(col_offset::<GROUP, PACK>(col)) })
317 } else {
318 None
319 }
320 }
321
322 #[inline]
324 pub fn iter(&self) -> RowIter<'_, T, GROUP, PACK> {
325 RowIter {
326 base: self.base,
327 col: 0,
328 ncols: self.ncols,
329 }
330 }
331}
332
333impl<T: Copy, const GROUP: usize, const PACK: usize> std::ops::Index<usize>
334 for Row<'_, T, GROUP, PACK>
335{
336 type Output = T;
337
338 #[inline]
339 #[allow(clippy::panic)] fn index(&self, col: usize) -> &Self::Output {
341 self.get(col)
342 .unwrap_or_else(|| panic!("column index {col} out of bounds (ncols = {})", self.ncols))
343 }
344}
345
346#[derive(Debug, Clone)]
348pub struct RowIter<'a, T, const GROUP: usize, const PACK: usize = 1> {
349 base: SlicePtr<'a, T>,
350 col: usize,
351 ncols: usize,
352}
353
354impl<T: Copy, const GROUP: usize, const PACK: usize> Iterator for RowIter<'_, T, GROUP, PACK> {
355 type Item = T;
356
357 #[inline]
358 fn next(&mut self) -> Option<Self::Item> {
359 if self.col >= self.ncols {
360 return None;
361 }
362 let val = unsafe { *self.base.as_ptr().add(col_offset::<GROUP, PACK>(self.col)) };
364 self.col += 1;
365 Some(val)
366 }
367
368 #[inline]
369 fn size_hint(&self) -> (usize, Option<usize>) {
370 let remaining = self.ncols - self.col;
371 (remaining, Some(remaining))
372 }
373}
374
375impl<T: Copy, const GROUP: usize, const PACK: usize> ExactSizeIterator
376 for RowIter<'_, T, GROUP, PACK>
377{
378}
379impl<T: Copy, const GROUP: usize, const PACK: usize> std::iter::FusedIterator
380 for RowIter<'_, T, GROUP, PACK>
381{
382}
383
384#[derive(Debug)]
386pub struct RowMut<'a, T, const GROUP: usize, const PACK: usize = 1> {
387 base: MutSlicePtr<'a, T>,
388 ncols: usize,
389}
390
391impl<T: Copy, const GROUP: usize, const PACK: usize> RowMut<'_, T, GROUP, PACK> {
392 #[inline]
394 pub fn len(&self) -> usize {
395 self.ncols
396 }
397
398 #[inline]
400 pub fn is_empty(&self) -> bool {
401 self.ncols == 0
402 }
403
404 #[inline]
406 pub fn get(&self, col: usize) -> Option<&T> {
407 if col < self.ncols {
408 Some(unsafe { &*self.base.as_ptr().add(col_offset::<GROUP, PACK>(col)) })
410 } else {
411 None
412 }
413 }
414
415 #[inline]
417 pub fn get_mut(&mut self, col: usize) -> Option<&mut T> {
418 if col < self.ncols {
419 Some(unsafe { &mut *self.base.as_mut_ptr().add(col_offset::<GROUP, PACK>(col)) })
421 } else {
422 None
423 }
424 }
425
426 #[inline]
432 pub fn set(&mut self, col: usize, value: T) {
433 assert!(
434 col < self.ncols,
435 "column index {col} out of bounds (ncols = {})",
436 self.ncols
437 );
438 unsafe { *self.base.as_mut_ptr().add(col_offset::<GROUP, PACK>(col)) = value };
440 }
441}
442
443impl<T: Copy, const GROUP: usize, const PACK: usize> std::ops::Index<usize>
444 for RowMut<'_, T, GROUP, PACK>
445{
446 type Output = T;
447
448 #[inline]
449 #[allow(clippy::panic)] fn index(&self, col: usize) -> &Self::Output {
451 self.get(col)
452 .unwrap_or_else(|| panic!("column index {col} out of bounds (ncols = {})", self.ncols))
453 }
454}
455
456impl<T: Copy, const GROUP: usize, const PACK: usize> std::ops::IndexMut<usize>
457 for RowMut<'_, T, GROUP, PACK>
458{
459 #[inline]
460 #[allow(clippy::panic)] fn index_mut(&mut self, col: usize) -> &mut Self::Output {
462 let ncols = self.ncols;
463 self.get_mut(col)
464 .unwrap_or_else(|| panic!("column index {col} out of bounds (ncols = {ncols})"))
465 }
466}
467
468unsafe impl<T: Copy, const GROUP: usize, const PACK: usize> Repr
475 for BlockTransposedRepr<T, GROUP, PACK>
476{
477 type Row<'a>
478 = Row<'a, T, GROUP, PACK>
479 where
480 Self: 'a;
481
482 fn nrows(&self) -> usize {
483 self.nrows
484 }
485
486 fn layout(&self) -> Result<Layout, LayoutError> {
487 Ok(Layout::array::<T>(self.storage_len())?)
488 }
489
490 unsafe fn get_row<'a>(self, ptr: NonNull<u8>, i: usize) -> Self::Row<'a> {
491 debug_assert!(i < self.nrows);
492
493 if self.ncols == 0 {
496 return Row {
497 base: unsafe { SlicePtr::new_unchecked(NonNull::dangling()) },
500 ncols: 0,
501 };
502 }
503
504 let base_ptr = ptr.as_ptr().cast::<T>();
505 let offset = linear_index::<GROUP, PACK>(i, 0, self.ncols);
506
507 let row_base = unsafe { base_ptr.add(offset) };
510
511 Row {
512 base: unsafe { SlicePtr::new_unchecked(NonNull::new_unchecked(row_base)) },
515 ncols: self.ncols,
516 }
517 }
518}
519
520unsafe impl<T: Copy, const GROUP: usize, const PACK: usize> ReprMut
524 for BlockTransposedRepr<T, GROUP, PACK>
525{
526 type RowMut<'a>
527 = RowMut<'a, T, GROUP, PACK>
528 where
529 Self: 'a;
530
531 unsafe fn get_row_mut<'a>(self, ptr: NonNull<u8>, i: usize) -> Self::RowMut<'a> {
532 debug_assert!(i < self.nrows);
533
534 if self.ncols == 0 {
537 return RowMut {
538 base: unsafe { MutSlicePtr::new_unchecked(NonNull::dangling()) },
541 ncols: 0,
542 };
543 }
544
545 let base_ptr = ptr.as_ptr().cast::<T>();
546 let offset = linear_index::<GROUP, PACK>(i, 0, self.ncols);
547
548 let row_base = unsafe { base_ptr.add(offset) };
551
552 RowMut {
553 base: unsafe { MutSlicePtr::new_unchecked(NonNull::new_unchecked(row_base)) },
556 ncols: self.ncols,
557 }
558 }
559}
560
561unsafe impl<T: Copy, const GROUP: usize, const PACK: usize> ReprOwned
564 for BlockTransposedRepr<T, GROUP, PACK>
565{
566 unsafe fn drop(self, ptr: NonNull<u8>) {
567 unsafe {
569 let slice_ptr =
570 std::ptr::slice_from_raw_parts_mut(ptr.cast::<T>().as_ptr(), self.storage_len());
571 let _ = Box::from_raw(slice_ptr);
572 }
573 }
574}
575
576unsafe impl<T: Copy, const GROUP: usize, const PACK: usize> NewOwned<T>
582 for BlockTransposedRepr<T, GROUP, PACK>
583{
584 type Error = crate::error::Infallible;
585
586 fn new_owned(self, value: T) -> Result<Mat<Self>, Self::Error> {
587 let b: Box<[T]> = vec![value; self.storage_len()].into_boxed_slice();
588
589 Ok(unsafe { self.box_to_mat(b) })
591 }
592}
593
594unsafe impl<T: Copy + Default, const GROUP: usize, const PACK: usize> NewOwned<Defaulted>
596 for BlockTransposedRepr<T, GROUP, PACK>
597{
598 type Error = crate::error::Infallible;
599
600 fn new_owned(self, _: Defaulted) -> Result<Mat<Self>, Self::Error> {
601 self.new_owned(T::default())
602 }
603}
604
605impl<T: Copy, const GROUP: usize, const PACK: usize> NewCloned
606 for BlockTransposedRepr<T, GROUP, PACK>
607{
608 fn new_cloned(v: MatRef<'_, Self>) -> Mat<Self> {
609 let b: Box<[T]> = BlockTransposedRef::new(v).as_slice().into();
610
611 unsafe { v.repr().box_to_mat(b) }
614 }
615}
616
617unsafe impl<T: Copy, const GROUP: usize, const PACK: usize> NewRef<T>
619 for BlockTransposedRepr<T, GROUP, PACK>
620{
621 type Error = SliceError;
622
623 fn new_ref(self, data: &[T]) -> Result<MatRef<'_, Self>, Self::Error> {
624 self.check_slice(data)?;
625
626 Ok(unsafe { MatRef::from_raw_parts(self, utils::as_nonnull(data).cast::<u8>()) })
628 }
629}
630
631unsafe impl<T: Copy, const GROUP: usize, const PACK: usize> NewMut<T>
633 for BlockTransposedRepr<T, GROUP, PACK>
634{
635 type Error = SliceError;
636
637 fn new_mut(self, data: &mut [T]) -> Result<MatMut<'_, Self>, Self::Error> {
638 self.check_slice(data)?;
639
640 Ok(unsafe { MatMut::from_raw_parts(self, utils::as_nonnull_mut(data).cast::<u8>()) })
642 }
643}
644
645macro_rules! delegate_to_ref {
654 ($(#[$m:meta])* $vis:vis fn $name:ident(&self $(, $a:ident: $t:ty)*) $(-> $r:ty)?) => {
656 #[doc = concat!("See [`BlockTransposedRef::", stringify!($name), "`].")]
657 $(#[$m])*
658 #[inline]
659 $vis fn $name(&self $(, $a: $t)*) $(-> $r)? {
660 self.as_view().$name($($a),*)
661 }
662 };
663 ($(#[$m:meta])* unsafe $vis:vis fn $name:ident(&self $(, $a:ident: $t:ty)*) $(-> $r:ty)?) => {
665 #[doc = concat!("See [`BlockTransposedRef::", stringify!($name), "`].")]
666 $(#[$m])*
667 #[inline]
668 $vis unsafe fn $name(&self $(, $a: $t)*) $(-> $r)? {
669 unsafe { self.as_view().$name($($a),*) }
671 }
672 };
673}
674
675#[derive(Debug, Clone)]
693pub struct BlockTransposed<T: Copy, const GROUP: usize, const PACK: usize = 1> {
694 data: Mat<BlockTransposedRepr<T, GROUP, PACK>>,
695}
696
697#[derive(Debug, Clone, Copy)]
701pub struct BlockTransposedRef<'a, T: Copy, const GROUP: usize, const PACK: usize = 1> {
702 data: MatRef<'a, BlockTransposedRepr<T, GROUP, PACK>>,
703}
704
705pub struct BlockTransposedMut<'a, T: Copy, const GROUP: usize, const PACK: usize = 1> {
709 data: MatMut<'a, BlockTransposedRepr<T, GROUP, PACK>>,
710}
711
712impl<'a, T: Copy, const GROUP: usize, const PACK: usize> BlockTransposedRef<'a, T, GROUP, PACK> {
715 fn new(data: MatRef<'a, BlockTransposedRepr<T, GROUP, PACK>>) -> Self {
716 Self { data }
717 }
718
719 #[inline]
721 pub fn nrows(&self) -> usize {
722 self.data.repr().nrows()
723 }
724
725 #[inline]
727 pub fn ncols(&self) -> usize {
728 self.data.repr().ncols()
729 }
730
731 #[inline]
733 pub fn padded_ncols(&self) -> usize {
734 self.data.repr().padded_ncols()
735 }
736
737 pub const fn group_size(&self) -> usize {
739 GROUP
740 }
741
742 pub const fn const_group_size() -> usize {
744 GROUP
745 }
746
747 pub const fn pack_size(&self) -> usize {
749 PACK
750 }
751
752 #[inline]
754 pub fn full_blocks(&self) -> usize {
755 self.data.repr().full_blocks()
756 }
757
758 #[inline]
760 pub fn num_blocks(&self) -> usize {
761 self.data.repr().num_blocks()
762 }
763
764 #[inline]
767 pub fn remainder(&self) -> usize {
768 self.data.repr().remainder()
769 }
770
771 #[inline]
776 pub fn padded_nrows(&self) -> usize {
777 self.data.repr().padded_nrows()
778 }
779
780 #[inline]
782 pub fn as_ptr(&self) -> *const T {
783 self.data.as_raw_ptr().cast::<T>()
784 }
785
786 #[inline]
791 pub fn as_slice(&self) -> &'a [T] {
792 let len = self.data.repr().storage_len();
793 unsafe { std::slice::from_raw_parts(self.as_ptr(), len) }
795 }
796
797 #[inline]
809 pub unsafe fn block_ptr_unchecked(&self, block: usize) -> *const T {
810 debug_assert!(block < self.num_blocks());
811 unsafe { self.as_ptr().add(self.data.repr().block_offset(block)) }
813 }
814
815 #[allow(clippy::expect_used)]
825 pub fn block(&self, block: usize) -> MatrixView<'a, T> {
826 assert!(block < self.full_blocks());
827 let offset = self.data.repr().block_offset(block);
828 let stride = self.data.repr().block_stride();
829 let data: &[T] = unsafe { std::slice::from_raw_parts(self.as_ptr().add(offset), stride) };
832 MatrixView::try_from(data, self.padded_ncols() / PACK, GROUP * PACK)
833 .expect("base data should have been sized correctly")
834 }
835
836 #[allow(clippy::expect_used)]
842 pub fn remainder_block(&self) -> Option<MatrixView<'a, T>> {
843 if self.remainder() == 0 {
844 None
845 } else {
846 let offset = self.data.repr().block_offset(self.full_blocks());
847 let stride = self.data.repr().block_stride();
848 let data: &[T] =
851 unsafe { std::slice::from_raw_parts(self.as_ptr().add(offset), stride) };
852 Some(
853 MatrixView::try_from(data, self.padded_ncols() / PACK, GROUP * PACK)
854 .expect("base data should have been sized correctly"),
855 )
856 }
857 }
858
859 #[inline]
865 pub fn get_element(&self, row: usize, col: usize) -> T {
866 assert!(
867 row < self.nrows(),
868 "row {row} out of bounds (nrows = {})",
869 self.nrows()
870 );
871 assert!(
872 col < self.ncols(),
873 "col {col} out of bounds (ncols = {})",
874 self.ncols()
875 );
876 let idx = linear_index::<GROUP, PACK>(row, col, self.ncols());
877 unsafe { *self.as_ptr().add(idx) }
879 }
880
881 #[inline]
883 pub fn get_row(&self, i: usize) -> Option<Row<'_, T, GROUP, PACK>> {
884 self.data.get_row(i)
885 }
886}
887
888impl<'a, T: Copy, const GROUP: usize, const PACK: usize> BlockTransposedMut<'a, T, GROUP, PACK> {
891 fn new(data: MatMut<'a, BlockTransposedRepr<T, GROUP, PACK>>) -> Self {
892 Self { data }
893 }
894
895 #[inline]
897 pub fn as_view(&self) -> BlockTransposedRef<'_, T, GROUP, PACK> {
898 BlockTransposedRef::new(self.data.as_view())
899 }
900
901 delegate_to_ref!(pub fn nrows(&self) -> usize);
904 delegate_to_ref!(pub fn ncols(&self) -> usize);
905 delegate_to_ref!(pub fn padded_ncols(&self) -> usize);
906 delegate_to_ref!(pub fn full_blocks(&self) -> usize);
907 delegate_to_ref!(pub fn num_blocks(&self) -> usize);
908 delegate_to_ref!(pub fn remainder(&self) -> usize);
909 delegate_to_ref!(pub fn padded_nrows(&self) -> usize);
910 delegate_to_ref!(pub fn as_ptr(&self) -> *const T);
911 delegate_to_ref!(pub fn as_slice(&self) -> &[T]);
912 delegate_to_ref!(#[allow(clippy::missing_safety_doc)] unsafe pub fn block_ptr_unchecked(&self, block: usize) -> *const T);
913 delegate_to_ref!(#[allow(clippy::expect_used)] pub fn block(&self, block: usize) -> MatrixView<'_, T>);
914 delegate_to_ref!(#[allow(clippy::expect_used)] pub fn remainder_block(&self) -> Option<MatrixView<'_, T>>);
915 delegate_to_ref!(pub fn get_element(&self, row: usize, col: usize) -> T);
916
917 pub const fn group_size(&self) -> usize {
919 GROUP
920 }
921
922 pub const fn const_group_size() -> usize {
924 GROUP
925 }
926
927 pub const fn pack_size(&self) -> usize {
929 PACK
930 }
931
932 #[inline]
934 pub fn get_row(&self, i: usize) -> Option<Row<'_, T, GROUP, PACK>> {
935 self.data.get_row(i)
936 }
937
938 #[inline]
949 pub fn as_mut_slice(&mut self) -> &mut [T] {
950 self.reborrow_mut().mut_slice_inner()
951 }
952
953 fn mut_slice_inner(mut self) -> &'a mut [T] {
954 let len = self.data.repr().storage_len();
955 unsafe { std::slice::from_raw_parts_mut(self.data.as_raw_mut_ptr().cast::<T>(), len) }
957 }
958
959 #[allow(clippy::expect_used)]
965 pub fn block_mut(&mut self, block: usize) -> MutMatrixView<'_, T> {
966 self.reborrow_mut().block_mut_inner(block)
967 }
968
969 #[allow(clippy::expect_used)]
970 fn block_mut_inner(mut self, block: usize) -> MutMatrixView<'a, T> {
971 let repr = *self.data.repr();
972 assert!(block < repr.full_blocks());
973 let offset = repr.block_offset(block);
974 let stride = repr.block_stride();
975 let pncols = repr.padded_ncols();
976 let data: &mut [T] = unsafe {
978 std::slice::from_raw_parts_mut(
979 self.data.as_raw_mut_ptr().cast::<T>().add(offset),
980 stride,
981 )
982 };
983 MutMatrixView::try_from(data, pncols / PACK, GROUP * PACK)
984 .expect("base data should have been sized correctly")
985 }
986
987 #[allow(clippy::expect_used)]
990 pub fn remainder_block_mut(&mut self) -> Option<MutMatrixView<'_, T>> {
991 self.reborrow_mut().remainder_block_mut_inner()
992 }
993
994 #[allow(clippy::expect_used)]
995 fn remainder_block_mut_inner(mut self) -> Option<MutMatrixView<'a, T>> {
996 let repr = *self.data.repr();
997 if repr.remainder() == 0 {
998 None
999 } else {
1000 let offset = repr.block_offset(repr.full_blocks());
1001 let stride = repr.block_stride();
1002 let pncols = repr.padded_ncols();
1003 let data: &mut [T] = unsafe {
1005 std::slice::from_raw_parts_mut(
1006 self.data.as_raw_mut_ptr().cast::<T>().add(offset),
1007 stride,
1008 )
1009 };
1010 Some(
1011 MutMatrixView::try_from(data, pncols / PACK, GROUP * PACK)
1012 .expect("base data should have been sized correctly"),
1013 )
1014 }
1015 }
1016
1017 #[inline]
1019 pub fn get_row_mut(&mut self, i: usize) -> Option<RowMut<'_, T, GROUP, PACK>> {
1020 self.data.get_row_mut(i)
1021 }
1022
1023 fn reborrow_mut(&mut self) -> BlockTransposedMut<'_, T, GROUP, PACK> {
1026 BlockTransposedMut::new(self.data.reborrow_mut())
1027 }
1028}
1029
1030impl<T: Copy, const GROUP: usize, const PACK: usize> BlockTransposed<T, GROUP, PACK> {
1033 pub fn as_view(&self) -> BlockTransposedRef<'_, T, GROUP, PACK> {
1035 BlockTransposedRef::new(self.data.as_view())
1036 }
1037
1038 pub fn as_view_mut(&mut self) -> BlockTransposedMut<'_, T, GROUP, PACK> {
1040 BlockTransposedMut::new(self.data.as_view_mut())
1041 }
1042
1043 delegate_to_ref!(pub fn nrows(&self) -> usize);
1046 delegate_to_ref!(pub fn ncols(&self) -> usize);
1047 delegate_to_ref!(pub fn padded_ncols(&self) -> usize);
1048 delegate_to_ref!(pub fn full_blocks(&self) -> usize);
1049 delegate_to_ref!(pub fn num_blocks(&self) -> usize);
1050 delegate_to_ref!(pub fn remainder(&self) -> usize);
1051 delegate_to_ref!(pub fn padded_nrows(&self) -> usize);
1052 delegate_to_ref!(pub fn as_ptr(&self) -> *const T);
1053 delegate_to_ref!(pub fn as_slice(&self) -> &[T]);
1054 delegate_to_ref!(#[allow(clippy::missing_safety_doc)] unsafe pub fn block_ptr_unchecked(&self, block: usize) -> *const T);
1055 delegate_to_ref!(#[allow(clippy::expect_used)] pub fn block(&self, block: usize) -> MatrixView<'_, T>);
1056 delegate_to_ref!(#[allow(clippy::expect_used)] pub fn remainder_block(&self) -> Option<MatrixView<'_, T>>);
1057 delegate_to_ref!(pub fn get_element(&self, row: usize, col: usize) -> T);
1058
1059 pub const fn group_size(&self) -> usize {
1061 GROUP
1062 }
1063
1064 pub const fn const_group_size() -> usize {
1066 GROUP
1067 }
1068
1069 pub const fn pack_size(&self) -> usize {
1071 PACK
1072 }
1073
1074 #[inline]
1076 pub fn get_row(&self, i: usize) -> Option<Row<'_, T, GROUP, PACK>> {
1077 self.data.get_row(i)
1078 }
1079
1080 #[inline]
1084 pub fn as_mut_slice(&mut self) -> &mut [T] {
1085 self.as_view_mut().mut_slice_inner()
1086 }
1087
1088 #[allow(clippy::expect_used)]
1090 pub fn block_mut(&mut self, block: usize) -> MutMatrixView<'_, T> {
1091 self.as_view_mut().block_mut_inner(block)
1092 }
1093
1094 #[allow(clippy::expect_used)]
1096 pub fn remainder_block_mut(&mut self) -> Option<MutMatrixView<'_, T>> {
1097 self.as_view_mut().remainder_block_mut_inner()
1098 }
1099
1100 #[inline]
1102 pub fn get_row_mut(&mut self, i: usize) -> Option<RowMut<'_, T, GROUP, PACK>> {
1103 self.data.get_row_mut(i)
1104 }
1105}
1106
1107impl<'this, T: Copy, const GROUP: usize, const PACK: usize> Reborrow<'this>
1110 for BlockTransposed<T, GROUP, PACK>
1111{
1112 type Target = BlockTransposedRef<'this, T, GROUP, PACK>;
1113
1114 #[inline]
1115 fn reborrow(&'this self) -> Self::Target {
1116 self.as_view()
1117 }
1118}
1119
1120impl<T: Copy + Default, const GROUP: usize, const PACK: usize> BlockTransposed<T, GROUP, PACK> {
1123 #[allow(clippy::expect_used)]
1129 pub fn new(nrows: usize, ncols: usize) -> Self {
1130 let repr = BlockTransposedRepr::<T, GROUP, PACK>::new(nrows, ncols)
1131 .expect("dimensions should not overflow");
1132 Self {
1133 data: Mat::new(repr, Defaulted).expect("infallible"),
1134 }
1135 }
1136
1137 pub fn try_new(nrows: usize, ncols: usize) -> Result<Self, Overflow> {
1139 let repr = BlockTransposedRepr::<T, GROUP, PACK>::new(nrows, ncols)?;
1140 Ok(Self {
1141 data: Mat::new(repr, Defaulted).expect("infallible"),
1142 })
1143 }
1144
1145 pub fn from_strided(v: StridedView<'_, T>) -> Self {
1157 let nrows = v.nrows();
1158 let ncols = v.ncols();
1159 let mut mat = Self::new(nrows, ncols);
1160
1161 let repr = *mat.data.repr();
1162 let num_blocks = repr.num_blocks();
1163 let pncols = repr.padded_ncols();
1164 let num_col_groups = pncols / PACK;
1165
1166 let mut dst = mat.data.as_raw_mut_ptr().cast::<T>();
1170 for block in 0..num_blocks {
1171 let row_base = block * GROUP;
1172 for cg in 0..num_col_groups {
1173 let col_base = cg * PACK;
1174 for rib in 0..GROUP {
1175 let row = row_base + rib;
1176 if row < nrows {
1177 let src_row = unsafe { v.get_row_unchecked(row) };
1179 for p in 0..PACK {
1180 let col = col_base + p;
1181 if col < ncols {
1182 unsafe { *dst = *src_row.get_unchecked(col) };
1188 }
1189 dst = unsafe { dst.add(1) };
1193 }
1194 } else {
1195 dst = unsafe { dst.add(PACK) };
1198 }
1199 }
1200 }
1201 }
1202
1203 mat
1204 }
1205
1206 pub fn from_matrix_view(v: MatrixView<'_, T>) -> Self {
1208 Self::from_strided(v.into())
1209 }
1210}
1211
1212impl<T: Copy, const GROUP: usize, const PACK: usize> std::ops::Index<(usize, usize)>
1217 for BlockTransposed<T, GROUP, PACK>
1218{
1219 type Output = T;
1220
1221 #[inline]
1222 fn index(&self, (row, col): (usize, usize)) -> &Self::Output {
1223 assert!(row < self.nrows());
1224 assert!(col < self.ncols());
1225 let idx = linear_index::<GROUP, PACK>(row, col, self.ncols());
1226 unsafe { &*self.as_ptr().add(idx) }
1228 }
1229}
1230
1231#[cfg(test)]
1236mod tests {
1237 use diskann_utils::{lazy_format, views::Matrix};
1252
1253 use super::*;
1254 use crate::utils::div_round_up;
1255
1256 fn gen_f32(i: usize) -> f32 {
1262 (i + 1) as f32
1263 }
1264 fn gen_i32(i: usize) -> i32 {
1265 (i + 1) as i32
1266 }
1267 fn gen_u8(i: usize) -> u8 {
1268 ((i % 255) + 1) as u8
1269 }
1270
1271 #[test]
1272 fn clone_has_independent_backing_allocation() {
1273 let mut data = Matrix::new(0, 5, 3);
1274 data.as_mut_slice()
1275 .iter_mut()
1276 .enumerate()
1277 .for_each(|(i, value)| *value = (i + 1) as i32);
1278 let mut original = BlockTransposed::<i32, 4, 2>::from_matrix_view(data.as_view());
1279 let column_padding = linear_index::<4, 2>(0, 3, original.ncols());
1280 let row_padding = linear_index::<4, 2>(5, 0, original.ncols());
1281 let row_and_column_padding = linear_index::<4, 2>(5, 3, original.ncols());
1282 original.as_mut_slice()[column_padding] = -10;
1283 original.as_mut_slice()[row_padding] = -11;
1284 original.as_mut_slice()[row_and_column_padding] = -12;
1285
1286 let mut cloned = original.clone();
1287
1288 assert_eq!(cloned.as_slice(), original.as_slice());
1289 assert_eq!(cloned.as_slice()[column_padding], -10);
1290 assert_eq!(cloned.as_slice()[row_padding], -11);
1291 assert_eq!(cloned.as_slice()[row_and_column_padding], -12);
1292 assert_ne!(cloned.as_ptr(), original.as_ptr());
1293
1294 cloned.get_row_mut(0).unwrap()[0] = -1;
1295 assert_eq!(original[(0, 0)], 1);
1296 assert_eq!(cloned[(0, 0)], -1);
1297 }
1298
1299 fn test_full_api<
1312 T: Copy + Default + PartialEq + std::fmt::Debug + 'static,
1313 const GROUP: usize,
1314 const PACK: usize,
1315 >(
1316 nrows: usize,
1317 ncols: usize,
1318 gen_element: fn(usize) -> T,
1319 ) {
1320 let context = lazy_format!(
1321 "T={}, GROUP={}, PACK={}, nrows={}, ncols={}",
1322 std::any::type_name::<T>(),
1323 GROUP,
1324 PACK,
1325 nrows,
1326 ncols,
1327 );
1328
1329 let mut data = Matrix::new(T::default(), nrows, ncols);
1332 data.as_mut_slice()
1333 .iter_mut()
1334 .enumerate()
1335 .for_each(|(i, d)| *d = gen_element(i));
1336
1337 let mut transpose = BlockTransposed::<T, GROUP, PACK>::from_strided(data.as_view().into());
1338
1339 let expected_padded = div_round_up(ncols, PACK) * PACK;
1340 let expected_remainder = nrows % GROUP;
1341 let storage_len = transpose.as_slice().len();
1342
1343 assert_eq!(transpose.nrows(), nrows, "{}", context);
1346 assert_eq!(transpose.ncols(), ncols, "{}", context);
1347 assert_eq!(transpose.group_size(), GROUP, "{}", context);
1348 assert_eq!(
1349 BlockTransposed::<T, GROUP, PACK>::const_group_size(),
1350 GROUP,
1351 "{}",
1352 context
1353 );
1354 assert_eq!(transpose.pack_size(), PACK, "{}", context);
1355 assert_eq!(transpose.full_blocks(), nrows / GROUP, "{}", context);
1356 assert_eq!(
1357 transpose.num_blocks(),
1358 div_round_up(nrows, GROUP),
1359 "{}",
1360 context,
1361 );
1362 assert_eq!(transpose.remainder(), expected_remainder, "{}", context);
1363 assert_eq!(transpose.padded_ncols(), expected_padded, "{}", context);
1364
1365 for row in 0..nrows {
1368 for col in 0..ncols {
1369 assert_eq!(
1370 data[(row, col)],
1371 transpose[(row, col)],
1372 "Index at ({}, {}) -- {}",
1373 row,
1374 col,
1375 context,
1376 );
1377 assert_eq!(
1378 data[(row, col)],
1379 transpose.get_element(row, col),
1380 "get_element at ({}, {}) -- {}",
1381 row,
1382 col,
1383 context,
1384 );
1385 }
1386 }
1387
1388 let view = transpose.as_view();
1391 for row in 0..nrows {
1392 let row_view = view.get_row(row).unwrap();
1393 assert_eq!(row_view.len(), ncols, "{}", context);
1394 assert_eq!(row_view.is_empty(), ncols == 0, "{}", context);
1395 for col in 0..ncols {
1396 assert_eq!(
1397 data[(row, col)],
1398 row_view[col],
1399 "row view at ({}, {}) -- {}",
1400 row,
1401 col,
1402 context,
1403 );
1404 }
1405 if ncols > 0 {
1407 assert_eq!(row_view.get(0), Some(&data[(row, 0)]), "{}", context);
1408 }
1409 assert_eq!(row_view.get(ncols), None, "{}", context);
1410
1411 let iter = row_view.iter();
1413 assert_eq!(iter.len(), ncols, "{}", context);
1414 let (lo, hi) = iter.size_hint();
1415 assert_eq!(lo, ncols, "{}", context);
1416 assert_eq!(hi, Some(ncols), "{}", context);
1417
1418 let collected: Vec<T> = row_view.iter().collect();
1419 assert_eq!(collected.len(), ncols, "{}", context);
1420 for col in 0..ncols {
1421 assert_eq!(data[(row, col)], collected[col], "{}", context);
1422 }
1423 }
1424 assert!(view.get_row(nrows).is_none(), "{}", context);
1426 let _ = view;
1427
1428 {
1431 let view = transpose.as_view();
1432 assert_eq!(view.nrows(), nrows, "{}", context);
1433 assert_eq!(view.ncols(), ncols, "{}", context);
1434 assert_eq!(view.padded_ncols(), expected_padded, "{}", context);
1435 assert_eq!(view.group_size(), GROUP, "{}", context);
1436 assert_eq!(
1437 BlockTransposedRef::<T, GROUP, PACK>::const_group_size(),
1438 GROUP,
1439 );
1440 assert_eq!(view.pack_size(), PACK, "{}", context);
1441 assert_eq!(view.full_blocks(), nrows / GROUP, "{}", context);
1442 assert_eq!(view.num_blocks(), div_round_up(nrows, GROUP), "{}", context,);
1443 assert_eq!(view.remainder(), expected_remainder, "{}", context);
1444 assert_eq!(view.as_ptr(), transpose.as_ptr(), "{}", context);
1445 assert_eq!(view.as_slice(), transpose.as_slice(), "{}", context);
1446
1447 for row in 0..nrows {
1448 for col in 0..ncols {
1449 assert_eq!(
1450 data[(row, col)],
1451 view.get_element(row, col),
1452 "Ref get_element at ({}, {}) -- {}",
1453 row,
1454 col,
1455 context,
1456 );
1457 }
1458 let row_view = view.get_row(row).unwrap();
1459 for col in 0..ncols {
1460 assert_eq!(data[(row, col)], row_view[col], "{}", context);
1461 }
1462 }
1463 assert!(view.get_row(nrows).is_none(), "{}", context);
1464 }
1465
1466 let expected_ptr = transpose.as_ptr();
1469 {
1470 let mut_view = transpose.as_view_mut();
1471 assert_eq!(mut_view.nrows(), nrows, "{}", context);
1472 assert_eq!(mut_view.ncols(), ncols, "{}", context);
1473 assert_eq!(mut_view.padded_ncols(), expected_padded, "{}", context);
1474 assert_eq!(mut_view.group_size(), GROUP, "{}", context);
1475 assert_eq!(
1476 BlockTransposedMut::<T, GROUP, PACK>::const_group_size(),
1477 GROUP,
1478 );
1479 assert_eq!(mut_view.pack_size(), PACK, "{}", context);
1480 assert_eq!(mut_view.full_blocks(), nrows / GROUP, "{}", context);
1481 assert_eq!(
1482 mut_view.num_blocks(),
1483 div_round_up(nrows, GROUP),
1484 "{}",
1485 context,
1486 );
1487 assert_eq!(mut_view.remainder(), expected_remainder, "{}", context);
1488 assert_eq!(mut_view.as_ptr(), expected_ptr, "{}", context);
1489 assert_eq!(mut_view.as_slice().len(), storage_len, "{}", context);
1490
1491 for row in 0..nrows {
1492 for col in 0..ncols {
1493 assert_eq!(
1494 data[(row, col)],
1495 mut_view.get_element(row, col),
1496 "Mut get_element at ({}, {}) -- {}",
1497 row,
1498 col,
1499 context,
1500 );
1501 }
1502 let row_view = mut_view.get_row(row).unwrap();
1503 for col in 0..ncols {
1504 assert_eq!(data[(row, col)], row_view[col], "{}", context);
1505 }
1506 }
1507 assert!(mut_view.get_row(nrows).is_none(), "{}", context);
1508 }
1509
1510 {
1513 let mut_view = transpose.as_view_mut();
1514 let ref_from_mut = mut_view.as_view();
1515 assert_eq!(ref_from_mut.nrows(), nrows, "{}", context);
1516 for row in 0..nrows {
1517 for col in 0..ncols {
1518 assert_eq!(
1519 data[(row, col)],
1520 ref_from_mut.get_element(row, col),
1521 "{}",
1522 context,
1523 );
1524 }
1525 }
1526 }
1527
1528 {
1532 let mut mut_view = transpose.as_view_mut();
1533 assert_eq!(mut_view.as_mut_slice().len(), storage_len, "{}", context);
1534 }
1535 assert_eq!(transpose.as_mut_slice().len(), storage_len, "{}", context);
1537
1538 let expected_block_nrows = expected_padded / PACK;
1541 let expected_block_ncols = GROUP * PACK;
1542
1543 for b in 0..transpose.full_blocks() {
1544 let block_data: Vec<T>;
1545 let ptr: *const T;
1546 {
1547 let block = transpose.block(b);
1548 assert_eq!(block.nrows(), expected_block_nrows, "{}", context);
1549 assert_eq!(block.ncols(), expected_block_ncols, "{}", context);
1550
1551 ptr = unsafe { transpose.block_ptr_unchecked(b) };
1553 assert_eq!(ptr, block.as_slice().as_ptr(), "{}", context);
1554
1555 block_data = block.as_slice().to_vec();
1556 }
1557
1558 {
1560 let view = transpose.as_view();
1561 assert_eq!(view.block(b).as_slice(), &block_data[..], "{}", context);
1562 assert_eq!(unsafe { view.block_ptr_unchecked(b) }, ptr, "{}", context);
1564 }
1565
1566 {
1568 let mut_view = transpose.as_view_mut();
1569 assert_eq!(mut_view.block(b).as_slice(), &block_data[..], "{}", context);
1570 assert_eq!(
1571 unsafe { mut_view.block_ptr_unchecked(b) },
1573 ptr,
1574 "{}",
1575 context,
1576 );
1577 }
1578 }
1579
1580 if expected_remainder != 0 {
1582 let remainder_data: Vec<T>;
1583 let ptr: *const T;
1584 let fb = transpose.full_blocks();
1585 {
1586 let block = transpose.remainder_block().unwrap();
1587 assert_eq!(block.nrows(), expected_block_nrows, "{}", context);
1588 assert_eq!(block.ncols(), expected_block_ncols, "{}", context);
1589
1590 ptr = unsafe { transpose.block_ptr_unchecked(fb) };
1592 assert_eq!(ptr, block.as_slice().as_ptr(), "{}", context);
1593
1594 remainder_data = block.as_slice().to_vec();
1595 }
1596
1597 {
1599 let view = transpose.as_view();
1600 let ref_block = view.remainder_block().unwrap();
1601 assert_eq!(ref_block.as_slice(), &remainder_data[..], "{}", context);
1602 }
1603 {
1605 let mut_view = transpose.as_view_mut();
1606 let mut_block = mut_view.remainder_block().unwrap();
1607 assert_eq!(mut_block.as_slice(), &remainder_data[..], "{}", context);
1608 }
1609 } else {
1610 assert!(transpose.remainder_block().is_none(), "{}", context);
1611 {
1612 let view = transpose.as_view();
1613 assert!(view.remainder_block().is_none(), "{}", context);
1614 }
1615 {
1616 let mut_view = transpose.as_view_mut();
1617 assert!(mut_view.remainder_block().is_none(), "{}", context);
1618 }
1619 }
1620
1621 {
1624 let mut mut_view = transpose.as_view_mut();
1625 for b in 0..mut_view.full_blocks() {
1626 let block_mut = mut_view.block_mut(b);
1627 assert_eq!(block_mut.nrows(), expected_block_nrows, "{}", context);
1628 assert_eq!(block_mut.ncols(), expected_block_ncols, "{}", context);
1629 }
1630 if expected_remainder != 0 {
1631 let rem = mut_view.remainder_block_mut().unwrap();
1632 assert_eq!(rem.nrows(), expected_block_nrows, "{}", context);
1633 assert_eq!(rem.ncols(), expected_block_ncols, "{}", context);
1634 } else {
1635 assert!(mut_view.remainder_block_mut().is_none(), "{}", context);
1636 }
1637 }
1638
1639 for b in 0..transpose.full_blocks() {
1641 let block_mut = transpose.block_mut(b);
1642 assert_eq!(block_mut.nrows(), expected_block_nrows, "{}", context);
1643 assert_eq!(block_mut.ncols(), expected_block_ncols, "{}", context);
1644 }
1645 if expected_remainder != 0 {
1646 let rem = transpose.remainder_block_mut().unwrap();
1647 assert_eq!(rem.nrows(), expected_block_nrows, "{}", context);
1648 assert_eq!(rem.ncols(), expected_block_ncols, "{}", context);
1649 } else {
1650 assert!(transpose.remainder_block_mut().is_none(), "{}", context);
1651 }
1652
1653 {
1656 let mut mut_view = transpose.as_view_mut();
1657 for row in 0..nrows {
1658 let row_view = mut_view.get_row_mut(row).unwrap();
1659 assert_eq!(row_view.len(), ncols, "{}", context);
1660 assert_eq!(row_view.is_empty(), ncols == 0, "{}", context);
1661 for col in 0..ncols {
1662 assert_eq!(data[(row, col)], row_view[col], "{}", context);
1663 }
1664 }
1665 assert!(mut_view.get_row_mut(nrows).is_none(), "{}", context);
1666 }
1667
1668 if nrows > 0 && ncols > 0 {
1671 {
1673 let view = transpose.as_view();
1674 let row = view.get_row(0).unwrap();
1675 assert_eq!(row.get(ncols), None, "{}", context);
1676 assert_eq!(row.get(usize::MAX), None, "{}", context);
1677 }
1678
1679 let row = transpose.get_row_mut(0).unwrap();
1681 assert_eq!(row.get(ncols), None, "{}", context);
1682
1683 let mut row = transpose.get_row_mut(0).unwrap();
1685 let sentinel = gen_element(usize::MAX / 2);
1686 let original = row[0];
1687 if let Some(v) = row.get_mut(0) {
1688 *v = sentinel;
1689 }
1690 assert_eq!(row.get_mut(ncols), None, "{}", context);
1691 let _ = row;
1693 assert_eq!(transpose.get_element(0, 0), sentinel, "{}", context);
1694 transpose.get_row_mut(0).unwrap().set(0, original);
1696 }
1697
1698 for b in 0..transpose.full_blocks() {
1701 transpose.block_mut(b).as_mut_slice().fill(T::default());
1702 }
1703 if transpose.remainder() != 0 {
1704 transpose
1705 .remainder_block_mut()
1706 .unwrap()
1707 .as_mut_slice()
1708 .fill(T::default());
1709 }
1710 assert!(
1711 transpose.as_slice().iter().all(|v| *v == T::default()),
1712 "not fully zeroed -- {}",
1713 context,
1714 );
1715
1716 let transpose = BlockTransposed::<T, GROUP, PACK>::from_strided(data.as_view().into());
1719 let raw = transpose.as_slice();
1720
1721 for row in 0..nrows {
1723 for col in ncols..expected_padded {
1724 let idx = linear_index::<GROUP, PACK>(row, col, ncols);
1725 assert_eq!(
1726 raw[idx],
1727 T::default(),
1728 "col padding at ({}, {}) -- {}",
1729 row,
1730 col,
1731 context,
1732 );
1733 }
1734 }
1735
1736 let padded_nrows = nrows.next_multiple_of(GROUP);
1738 for row in nrows..padded_nrows {
1739 for col in 0..expected_padded {
1740 let idx = linear_index::<GROUP, PACK>(row, col, ncols);
1741 assert_eq!(
1742 raw[idx],
1743 T::default(),
1744 "row padding at ({}, {}) -- {}",
1745 row,
1746 col,
1747 context,
1748 );
1749 }
1750 }
1751
1752 assert_eq!(
1755 transpose.as_view().padded_nrows(),
1756 padded_nrows,
1757 "padded_nrows() mismatch -- {}",
1758 context,
1759 );
1760
1761 if nrows > 0 && ncols > 0 {
1764 let via_matrix = BlockTransposed::<T, GROUP, PACK>::from_matrix_view(data.as_view());
1765 assert_eq!(via_matrix.as_slice(), transpose.as_slice(), "{}", context);
1766 }
1767 }
1768
1769 #[test]
1774 fn test_api_pack1_group16() {
1775 let rows: Vec<usize> = if cfg!(miri) {
1778 vec![0, 1, 15, 16, 17, 33]
1779 } else {
1780 (0..128).collect()
1781 };
1782 let cols: Vec<usize> = if cfg!(miri) {
1783 vec![0, 1, 2]
1784 } else {
1785 (0..5).collect()
1786 };
1787 for &nrows in &rows {
1788 for &ncols in &cols {
1789 test_full_api::<f32, 16, 1>(nrows, ncols, gen_f32);
1790 }
1791 }
1792 }
1793
1794 #[test]
1795 fn test_api_pack1_group8() {
1796 let rows: Vec<usize> = if cfg!(miri) {
1799 vec![0, 1, 7, 8, 9, 17]
1800 } else {
1801 (0..128).collect()
1802 };
1803 let cols: Vec<usize> = if cfg!(miri) {
1804 vec![0, 1, 2]
1805 } else {
1806 (0..5).collect()
1807 };
1808 for &nrows in &rows {
1809 for &ncols in &cols {
1810 test_full_api::<f32, 8, 1>(nrows, ncols, gen_f32);
1811 }
1812 }
1813 }
1814
1815 #[test]
1816 fn test_api_pack2() {
1817 let rows: Vec<usize> = if cfg!(miri) {
1820 vec![0, 1, 3, 4, 5, 7, 8, 9, 15, 16, 17]
1821 } else {
1822 (0..48).collect()
1823 };
1824 let cols: Vec<usize> = if cfg!(miri) {
1825 vec![0, 1, 2, 3, 4, 5]
1826 } else {
1827 (0..9).collect()
1828 };
1829 for &nrows in &rows {
1830 for &ncols in &cols {
1831 test_full_api::<f32, 4, 2>(nrows, ncols, gen_f32);
1832 test_full_api::<f32, 8, 2>(nrows, ncols, gen_f32);
1833 test_full_api::<f32, 16, 2>(nrows, ncols, gen_f32);
1834 }
1835 }
1836 }
1837
1838 #[test]
1839 fn test_api_pack4() {
1840 let rows: Vec<usize> = if cfg!(miri) {
1843 vec![0, 1, 3, 4, 5, 7, 8, 9, 15, 16, 17]
1844 } else {
1845 (0..48).collect()
1846 };
1847 let cols: Vec<usize> = if cfg!(miri) {
1848 vec![0, 1, 3, 4, 5, 8]
1849 } else {
1850 (0..9).collect()
1851 };
1852 for &nrows in &rows {
1853 for &ncols in &cols {
1854 test_full_api::<f32, 4, 4>(nrows, ncols, gen_f32);
1855 test_full_api::<f32, 8, 4>(nrows, ncols, gen_f32);
1856 test_full_api::<f32, 16, 4>(nrows, ncols, gen_f32);
1857 }
1858 }
1859 }
1860
1861 #[test]
1863 fn test_api_non_f32() {
1864 test_full_api::<i32, 4, 1>(10, 7, gen_i32);
1866 test_full_api::<i32, 8, 2>(12, 5, gen_i32);
1867
1868 test_full_api::<u8, 4, 2>(12, 5, gen_u8);
1870 test_full_api::<u8, 8, 1>(10, 7, gen_u8);
1871 }
1872
1873 fn test_block_layout_pack1<
1880 T: Copy + Default + PartialEq + std::fmt::Debug + 'static,
1881 const GROUP: usize,
1882 >(
1883 nrows: usize,
1884 ncols: usize,
1885 gen_element: fn(usize) -> T,
1886 ) {
1887 let mut data = Matrix::new(T::default(), nrows, ncols);
1888 data.as_mut_slice()
1889 .iter_mut()
1890 .enumerate()
1891 .for_each(|(i, d)| *d = gen_element(i));
1892
1893 let transpose = BlockTransposed::<T, GROUP, 1>::from_strided(data.as_view().into());
1894
1895 for b in 0..transpose.full_blocks() {
1897 let block = transpose.block(b);
1898 for i in 0..block.nrows() {
1899 for j in 0..block.ncols() {
1900 assert_eq!(
1901 block[(i, j)],
1902 data[(GROUP * b + j, i)],
1903 "block {} at ({}, {}) -- GROUP={}, nrows={}, ncols={}",
1904 b,
1905 i,
1906 j,
1907 GROUP,
1908 nrows,
1909 ncols,
1910 );
1911 }
1912 }
1913 }
1914
1915 if transpose.remainder() != 0 {
1917 let fb = transpose.full_blocks();
1918 let block = transpose.remainder_block().unwrap();
1919 for i in 0..block.nrows() {
1920 for j in 0..transpose.remainder() {
1921 assert_eq!(
1922 block[(i, j)],
1923 data[(GROUP * fb + j, i)],
1924 "remainder at ({}, {}) -- GROUP={}, nrows={}, ncols={}",
1925 i,
1926 j,
1927 GROUP,
1928 nrows,
1929 ncols,
1930 );
1931 }
1932 }
1933 }
1934 }
1935
1936 #[test]
1937 fn test_block_layout_pack1_group16() {
1938 let rows: Vec<usize> = if cfg!(miri) {
1939 vec![0, 1, 15, 16, 17, 33]
1940 } else {
1941 (0..128).collect()
1942 };
1943 let cols: Vec<usize> = if cfg!(miri) {
1944 vec![0, 1, 2]
1945 } else {
1946 (0..5).collect()
1947 };
1948 for &nrows in &rows {
1949 for &ncols in &cols {
1950 test_block_layout_pack1::<f32, 16>(nrows, ncols, gen_f32);
1951 }
1952 }
1953 }
1954
1955 #[test]
1956 fn test_block_layout_pack1_group8() {
1957 let rows: Vec<usize> = if cfg!(miri) {
1958 vec![0, 1, 7, 8, 9, 17]
1959 } else {
1960 (0..128).collect()
1961 };
1962 let cols: Vec<usize> = if cfg!(miri) {
1963 vec![0, 1, 2]
1964 } else {
1965 (0..5).collect()
1966 };
1967 for &nrows in &rows {
1968 for &ncols in &cols {
1969 test_block_layout_pack1::<f32, 8>(nrows, ncols, gen_f32);
1970 }
1971 }
1972 }
1973
1974 #[test]
1981 fn test_row_view_send_sync() {
1982 fn assert_send<T: Send>() {}
1983 fn assert_sync<T: Sync>() {}
1984
1985 assert_send::<Row<'_, f32, 16>>();
1986 assert_sync::<Row<'_, f32, 16>>();
1987 assert_send::<Row<'_, u8, 8, 2>>();
1988 assert_sync::<Row<'_, u8, 8, 2>>();
1989
1990 assert_send::<RowMut<'_, f32, 16>>();
1991 assert_sync::<RowMut<'_, f32, 16>>();
1992 assert_send::<RowMut<'_, i32, 4, 4>>();
1993 assert_sync::<RowMut<'_, i32, 4, 4>>();
1994 }
1995
1996 #[test]
1999 fn test_new_ref_and_new_mut() {
2000 let nrows = 5;
2001 let ncols = 3;
2002 let repr = BlockTransposedRepr::<f32, 4>::new(nrows, ncols).unwrap();
2003
2004 let mat = BlockTransposed::<f32, 4>::new(nrows, ncols);
2005 let raw: &[f32] = mat.as_slice();
2006
2007 let mat_ref = BlockTransposedRef::new(repr.new_ref(raw).unwrap());
2008 assert_eq!(mat_ref.nrows(), nrows);
2009 assert_eq!(mat_ref.ncols(), ncols);
2010 for row in 0..nrows {
2011 for col in 0..ncols {
2012 assert_eq!(mat_ref.get_element(row, col), mat.get_element(row, col));
2013 }
2014 }
2015
2016 let mut buf = raw.to_vec();
2017 let mat_mut = BlockTransposedMut::new(repr.new_mut(&mut buf).unwrap());
2018 assert_eq!(mat_mut.nrows(), nrows);
2019 assert_eq!(mat_mut.ncols(), ncols);
2020
2021 let mut short = vec![0.0_f32; 2];
2023 assert!(repr.new_ref(&short).is_err());
2024 assert!(repr.new_mut(&mut short).is_err());
2025 }
2026
2027 #[test]
2030 fn test_row_view_empty() {
2031 fn check_empty<const GROUP: usize, const PACK: usize>() {
2034 let mut mat = BlockTransposed::<f32, GROUP, PACK>::new(4, 0);
2035
2036 let view = mat.as_view();
2038 for i in 0..4 {
2039 let row = view.get_row(i).unwrap();
2040 assert!(row.is_empty());
2041 assert_eq!(row.len(), 0);
2042 assert_eq!(row.iter().count(), 0);
2043 }
2044
2045 for i in 0..4 {
2047 let row = mat.get_row_mut(i).unwrap();
2048 assert!(row.is_empty());
2049 assert_eq!(row.len(), 0);
2050 }
2051 }
2052
2053 check_empty::<16, 1>(); check_empty::<4, 2>(); check_empty::<4, 4>(); }
2057
2058 #[test]
2061 #[should_panic(expected = "column index 3 out of bounds")]
2062 fn test_row_view_index_oob() {
2063 let mat = BlockTransposed::<f32, 4>::new(4, 3);
2064 let view = mat.as_view();
2065 let row = view.get_row(0).unwrap();
2066 let _ = row[3];
2067 }
2068
2069 #[test]
2070 #[should_panic(expected = "column index 3 out of bounds")]
2071 fn test_row_view_mut_index_oob() {
2072 let mut mat = BlockTransposed::<f32, 4>::new(4, 3);
2073 let row = mat.get_row_mut(0).unwrap();
2074 let _ = row[3];
2075 }
2076
2077 #[test]
2078 #[should_panic(expected = "column index 3 out of bounds")]
2079 fn test_row_view_mut_index_mut_oob() {
2080 let mut mat = BlockTransposed::<f32, 4>::new(4, 3);
2081 let mut row = mat.get_row_mut(0).unwrap();
2082 row[3] = 1.0;
2083 }
2084
2085 #[test]
2086 #[should_panic(expected = "column index 3 out of bounds")]
2087 fn test_row_view_set_oob() {
2088 let mut mat = BlockTransposed::<f32, 4>::new(4, 3);
2089 let mut row = mat.get_row_mut(0).unwrap();
2090 row.set(3, 1.0);
2091 }
2092
2093 #[test]
2094 #[should_panic(expected = "row 4 out of bounds")]
2095 fn test_get_element_row_oob() {
2096 let mat = BlockTransposed::<f32, 4>::new(4, 3);
2097 mat.get_element(4, 0);
2098 }
2099
2100 #[test]
2101 #[should_panic(expected = "col 3 out of bounds")]
2102 fn test_get_element_col_oob() {
2103 let mat = BlockTransposed::<f32, 4>::new(4, 3);
2104 mat.get_element(0, 3);
2105 }
2106
2107 #[test]
2108 #[should_panic(expected = "assertion failed")]
2109 fn test_index_tuple_row_oob() {
2110 let mat = BlockTransposed::<f32, 4>::new(4, 3);
2111 let _ = mat[(4, 0)];
2112 }
2113
2114 #[test]
2115 #[should_panic(expected = "assertion failed")]
2116 fn test_index_tuple_col_oob() {
2117 let mat = BlockTransposed::<f32, 4>::new(4, 3);
2118 let _ = mat[(0, 3)];
2119 }
2120
2121 #[test]
2122 #[should_panic]
2123 fn test_block_oob() {
2124 let mat = BlockTransposed::<f32, 4>::new(4, 3);
2125 let _ = mat.block(1);
2126 }
2127
2128 #[test]
2129 #[should_panic]
2130 fn test_block_mut_oob() {
2131 let mut mat = BlockTransposed::<f32, 4>::new(4, 3);
2132 let _ = mat.block_mut(1);
2133 }
2134
2135 #[test]
2138 fn test_from_strided_nonunit_stride() {
2139 use diskann_utils::strided::StridedView;
2140
2141 const GROUP: usize = 4;
2142 const PACK: usize = 2;
2143 let nrows = 5;
2144 let ncols = 3;
2145 let cstride = 8;
2146
2147 let required_len = (nrows - 1) * cstride + ncols;
2148 let mut flat = vec![0.0_f32; required_len];
2149 for row in 0..nrows {
2150 for col in 0..ncols {
2151 flat[row * cstride + col] = (row * 100 + col + 1) as f32;
2152 }
2153 }
2154
2155 let strided = StridedView::try_shrink_from(&flat, nrows, ncols, cstride)
2156 .expect("should construct strided view");
2157 let transpose = BlockTransposed::<f32, GROUP, PACK>::from_strided(strided);
2158
2159 assert_eq!(transpose.nrows(), nrows);
2160 assert_eq!(transpose.ncols(), ncols);
2161
2162 for row in 0..nrows {
2163 for col in 0..ncols {
2164 let expected = (row * 100 + col + 1) as f32;
2165 assert_eq!(
2166 transpose[(row, col)],
2167 expected,
2168 "mismatch at ({}, {})",
2169 row,
2170 col,
2171 );
2172 }
2173 }
2174
2175 let padded_ncols = ncols.next_multiple_of(PACK);
2176 let raw: &[f32] = transpose.as_slice();
2177 for row in 0..nrows {
2178 for col in ncols..padded_ncols {
2179 let idx = linear_index::<GROUP, PACK>(row, col, ncols);
2180 assert_eq!(
2181 raw[idx], 0.0,
2182 "column-padding at ({}, {}) should be zero",
2183 row, col,
2184 );
2185 }
2186 }
2187 }
2188
2189 #[test]
2192 fn test_concurrent_row_mutation() {
2193 const GROUP: usize = 8;
2194 const PACK: usize = 2;
2195
2196 let (nrows, ncols, num_threads) = if cfg!(miri) { (8, 4, 2) } else { (64, 16, 4) };
2197
2198 let mut mat = BlockTransposed::<f32, GROUP, PACK>::new(nrows, ncols);
2199 let rows: Vec<RowMut<'_, f32, GROUP, PACK>> = mat.data.rows_mut().collect();
2200 let rows_per_thread = nrows / num_threads;
2201 let mut rows = rows.into_boxed_slice();
2202
2203 std::thread::scope(|s| {
2204 let mut remaining = &mut rows[..];
2205 for thread_id in 0..num_threads {
2206 let chunk_len = if thread_id == num_threads - 1 {
2207 remaining.len()
2208 } else {
2209 rows_per_thread
2210 };
2211 let (chunk, rest) = remaining.split_at_mut(chunk_len);
2212 remaining = rest;
2213 let start_row = thread_id * rows_per_thread;
2214
2215 s.spawn(move || {
2216 for (offset, row_view) in chunk.iter_mut().enumerate() {
2217 let row = start_row + offset;
2218 for col in 0..ncols {
2219 let value = (thread_id * 10000 + row * 100 + col) as f32;
2220 row_view.set(col, value);
2221 }
2222 }
2223 });
2224 }
2225 });
2226
2227 for row in 0..nrows {
2228 let thread_id = (row / rows_per_thread).min(num_threads - 1);
2229 for col in 0..ncols {
2230 let expected = (thread_id * 10000 + row * 100 + col) as f32;
2231 assert_eq!(
2232 mat.get_element(row, col),
2233 expected,
2234 "mismatch at ({}, {})",
2235 row,
2236 col,
2237 );
2238 }
2239 }
2240 }
2241}