1use crate::fuse::{compress_dims, fuse_dims};
7use crate::inline::{InlineVec, OperandVec};
8use crate::{block, order, Result, StridedError};
9
10pub const SMALL_TENSOR_THRESHOLD: usize = 1024;
14
15#[derive(Debug)]
17pub struct KernelPlan {
18 #[cfg_attr(not(test), allow(dead_code))]
19 pub(crate) order: InlineVec<usize>, pub block: InlineVec<usize>,
21}
22
23#[cfg(test)]
30pub(crate) fn build_plan(
31 dims: &[usize],
32 strides_list: &[&[isize]],
33 dest_index: Option<usize>,
34 elem_size: usize,
35) -> KernelPlan {
36 let order = order::compute_order(dims, strides_list, dest_index);
37 let block = block::compute_block_sizes(dims, &order, strides_list, elem_size);
38 KernelPlan { order, block }
39}
40
41pub(crate) fn build_plan_fused(
53 dims: &[usize],
54 strides_list: &[&[isize]],
55 dest_index: Option<usize>,
56 elem_size: usize,
57) -> (InlineVec<usize>, OperandVec<InlineVec<isize>>, KernelPlan) {
58 let order = order::compute_order(dims, strides_list, dest_index);
60
61 let ordered_dims: InlineVec<usize> = order.iter().map(|&d| dims[d]).collect();
63 let ordered_strides: OperandVec<InlineVec<isize>> = strides_list
64 .iter()
65 .map(|strides| order.iter().map(|&d| strides[d]).collect())
66 .collect();
67 let ordered_strides_refs: OperandVec<&[isize]> =
68 ordered_strides.iter().map(|s| s.as_slice()).collect();
69
70 let fused_dims = fuse_dims(&ordered_dims, &ordered_strides_refs);
72
73 let (compressed_dims, compressed_strides) = compress_dims(&fused_dims, &ordered_strides_refs);
75 let compressed_strides_refs: OperandVec<&[isize]> =
76 compressed_strides.iter().map(|s| s.as_slice()).collect();
77
78 let identity: InlineVec<usize> = (0..compressed_dims.len()).collect();
80 let block = block::compute_block_sizes(
81 &compressed_dims,
82 &identity,
83 &compressed_strides_refs,
84 elem_size,
85 );
86
87 (
88 compressed_dims,
89 compressed_strides,
90 KernelPlan {
91 order: identity,
92 block,
93 },
94 )
95}
96
97pub(crate) fn build_plan_fused_small(
106 dims: &[usize],
107 strides_list: &[&[isize]],
108) -> (InlineVec<usize>, OperandVec<InlineVec<isize>>, KernelPlan) {
109 let fused = fuse_dims(dims, strides_list);
111 let (fused_dims, fused_strides) = compress_dims(&fused, strides_list);
112
113 let block = fused_dims.clone();
115 let identity: InlineVec<usize> = (0..fused_dims.len()).collect();
116
117 (
118 fused_dims,
119 fused_strides,
120 KernelPlan {
121 order: identity,
122 block,
123 },
124 )
125}
126
127#[cfg(test)]
137#[inline]
138pub(crate) fn for_each_inner_block<F>(
139 dims: &[usize],
140 plan: &KernelPlan,
141 strides_list: &[&[isize]],
142 mut f: F,
143) -> Result<()>
144where
145 F: FnMut(&[isize], usize, &[isize]) -> Result<()>,
146{
147 let rank = dims.len();
148 if rank == 0 {
149 let offsets: OperandVec<isize> = strides_list.iter().map(|_| 0isize).collect();
150 return f(&offsets, 1, &[]);
151 }
152
153 let ordered_dims: InlineVec<usize> = plan.order.iter().map(|&d| dims[d]).collect();
155 let ordered_blocks: InlineVec<usize> = plan.block.clone();
156
157 let mut ordered_strides: OperandVec<InlineVec<isize>> = OperandVec::new();
159 for strides in strides_list {
160 ordered_strides.push(plan.order.iter().map(|&d| strides[d]).collect());
161 }
162 let ordered_strides_refs: OperandVec<&[isize]> =
163 ordered_strides.iter().map(|s| s.as_slice()).collect();
164
165 let mut offsets: OperandVec<isize> = strides_list.iter().map(|_| 0isize).collect();
167
168 match rank {
170 1 => kernel_1d_inner(
171 &ordered_dims,
172 &ordered_blocks,
173 &ordered_strides_refs,
174 &mut offsets,
175 &mut f,
176 ),
177 2 => kernel_2d_inner(
178 &ordered_dims,
179 &ordered_blocks,
180 &ordered_strides_refs,
181 &mut offsets,
182 &mut f,
183 ),
184 3 => kernel_3d_inner(
185 &ordered_dims,
186 &ordered_blocks,
187 &ordered_strides_refs,
188 &mut offsets,
189 &mut f,
190 ),
191 4 => kernel_4d_inner(
192 &ordered_dims,
193 &ordered_blocks,
194 &ordered_strides_refs,
195 &mut offsets,
196 &mut f,
197 ),
198 5 => kernel_5d_inner(
199 &ordered_dims,
200 &ordered_blocks,
201 &ordered_strides_refs,
202 &mut offsets,
203 &mut f,
204 ),
205 6 => kernel_6d_inner(
206 &ordered_dims,
207 &ordered_blocks,
208 &ordered_strides_refs,
209 &mut offsets,
210 &mut f,
211 ),
212 7 => kernel_7d_inner(
213 &ordered_dims,
214 &ordered_blocks,
215 &ordered_strides_refs,
216 &mut offsets,
217 &mut f,
218 ),
219 8 => kernel_8d_inner(
220 &ordered_dims,
221 &ordered_blocks,
222 &ordered_strides_refs,
223 &mut offsets,
224 &mut f,
225 ),
226 _ => kernel_nd_inner_iterative(
227 &ordered_dims,
228 &ordered_blocks,
229 &ordered_strides_refs,
230 &mut offsets,
231 &mut f,
232 ),
233 }
234}
235
236macro_rules! elem_loops {
255 ($offsets:ident, $strides:ident, $f:ident, $blens:ident, $is:ident; $lv:literal) => {
259 for _i in 0..$blens[$lv] {
260 $f($offsets, $blens[0], &$is)?;
261 if _i + 1 < $blens[$lv] {
262 for (o, s) in $offsets.iter_mut().zip($strides.iter()) {
263 *o += s[$lv];
264 }
265 }
266 }
267 let _back = $blens[$lv].saturating_sub(1) as isize;
268 for (o, s) in $offsets.iter_mut().zip($strides.iter()) {
269 *o -= _back * s[$lv];
270 }
271 };
272 ($offsets:ident, $strides:ident, $f:ident, $blens:ident, $is:ident;
275 $lv:literal, $next:literal $(, $rest:literal)*) => {
276 for _i in 0..$blens[$lv] {
277 elem_loops!($offsets, $strides, $f, $blens, $is; $next $(, $rest)*);
278 if _i + 1 < $blens[$lv] {
279 for (o, s) in $offsets.iter_mut().zip($strides.iter()) {
280 *o += s[$lv];
281 }
282 }
283 }
284 let _back = $blens[$lv].saturating_sub(1) as isize;
285 for (o, s) in $offsets.iter_mut().zip($strides.iter()) {
286 *o -= _back * s[$lv];
287 }
288 };
289}
290
291macro_rules! block_loop {
298 ($dims:ident, $blocks:ident, $strides:ident, $offsets:ident, $f:ident,
300 $blens:ident, $is:ident; elem=[$($el:literal),+]; $lv0:literal; top=$top:literal) => {{
301 let mut _j = 0usize;
302 let mut _advanced = 0usize;
303 while _j < $dims[$lv0] {
304 $blens[$lv0] = $blocks[$lv0].max(1).min($dims[$lv0]).min($dims[$lv0] - _j);
305 elem_loops!($offsets, $strides, $f, $blens, $is; $($el),+);
306 _j += $blens[$lv0];
307 if _j < $dims[$lv0] {
308 for (o, s) in $offsets.iter_mut().zip($strides.iter()) {
309 *o += ($blens[$lv0] as isize) * s[$lv0];
310 }
311 _advanced = _j;
312 }
313 }
314 for (o, s) in $offsets.iter_mut().zip($strides.iter()) {
315 *o -= (_advanced as isize) * s[$lv0];
316 }
317 }};
318 ($dims:ident, $blocks:ident, $strides:ident, $offsets:ident, $f:ident,
320 $blens:ident, $is:ident; elem=[$($el:literal),+];
321 $lv:literal, $next:literal $(, $rest:literal)*; top=$top:literal) => {{
322 let mut _j = 0usize;
323 let mut _advanced = 0usize;
324 while _j < $dims[$lv] {
325 $blens[$lv] = $blocks[$lv].max(1).min($dims[$lv]).min($dims[$lv] - _j);
326 block_loop!($dims, $blocks, $strides, $offsets, $f, $blens, $is;
327 elem=[$($el),+]; $next $(, $rest)*; top=$top);
328 _j += $blens[$lv];
329 if _j < $dims[$lv] {
330 for (o, s) in $offsets.iter_mut().zip($strides.iter()) {
331 *o += ($blens[$lv] as isize) * s[$lv];
332 }
333 _advanced = _j;
334 }
335 }
336 for (o, s) in $offsets.iter_mut().zip($strides.iter()) {
337 *o -= (_advanced as isize) * s[$lv];
338 }
339 }};
340}
341
342macro_rules! make_kernel {
349 ($name:ident, rank=1) => {
351 #[inline]
352 fn $name<F>(
353 dims: &[usize],
354 blocks: &[usize],
355 strides: &[&[isize]],
356 offsets: &mut [isize],
357 f: &mut F,
358 ) -> Result<()>
359 where
360 F: FnMut(&[isize], usize, &[isize]) -> Result<()>,
361 {
362 let d0 = dims[0];
363 let b0 = blocks[0].max(1).min(d0);
364 let inner_strides: OperandVec<isize> = strides.iter().map(|s| s[0]).collect();
365
366 let mut j0 = 0usize;
367 let mut advanced = 0usize;
368 while j0 < d0 {
369 let blen0 = b0.min(d0 - j0);
370 f(offsets, blen0, &inner_strides)?;
371 j0 += blen0;
372 if j0 < d0 {
373 for (o, s) in offsets.iter_mut().zip(strides.iter()) {
374 *o += (blen0 as isize) * s[0];
375 }
376 advanced = j0;
377 }
378 }
379 for (o, s) in offsets.iter_mut().zip(strides.iter()) {
380 *o -= (advanced as isize) * s[0];
381 }
382 Ok(())
383 }
384 };
385 ($name:ident, rank=$rank:literal,
390 block=[$($blk:literal),+], elem=[$($el:literal),+], top=$top:literal) => {
391 #[inline]
392 fn $name<F>(
393 dims: &[usize],
394 blocks: &[usize],
395 strides: &[&[isize]],
396 offsets: &mut [isize],
397 f: &mut F,
398 ) -> Result<()>
399 where
400 F: FnMut(&[isize], usize, &[isize]) -> Result<()>,
401 {
402 let inner_strides: OperandVec<isize> = strides.iter().map(|s| s[0]).collect();
403 let mut blens = [0usize; $rank];
404 block_loop!(dims, blocks, strides, offsets, f, blens, inner_strides;
405 elem=[$($el),+]; $($blk),+; top=$top);
406 Ok(())
407 }
408 };
409}
410
411make_kernel!(kernel_1d_inner, rank = 1);
412make_kernel!(
413 kernel_2d_inner,
414 rank = 2,
415 block = [1, 0],
416 elem = [1],
417 top = 1
418);
419make_kernel!(
420 kernel_3d_inner,
421 rank = 3,
422 block = [2, 1, 0],
423 elem = [2, 1],
424 top = 2
425);
426make_kernel!(
427 kernel_4d_inner,
428 rank = 4,
429 block = [3, 2, 1, 0],
430 elem = [3, 2, 1],
431 top = 3
432);
433make_kernel!(
434 kernel_5d_inner,
435 rank = 5,
436 block = [4, 3, 2, 1, 0],
437 elem = [4, 3, 2, 1],
438 top = 4
439);
440make_kernel!(
441 kernel_6d_inner,
442 rank = 6,
443 block = [5, 4, 3, 2, 1, 0],
444 elem = [5, 4, 3, 2, 1],
445 top = 5
446);
447make_kernel!(
448 kernel_7d_inner,
449 rank = 7,
450 block = [6, 5, 4, 3, 2, 1, 0],
451 elem = [6, 5, 4, 3, 2, 1],
452 top = 6
453);
454make_kernel!(
455 kernel_8d_inner,
456 rank = 8,
457 block = [7, 6, 5, 4, 3, 2, 1, 0],
458 elem = [7, 6, 5, 4, 3, 2, 1],
459 top = 7
460);
461
462#[inline]
467fn kernel_nd_inner_iterative<F>(
468 dims: &[usize],
469 blocks: &[usize],
470 strides: &[&[isize]],
471 offsets: &mut [isize],
472 f: &mut F,
473) -> Result<()>
474where
475 F: FnMut(&[isize], usize, &[isize]) -> Result<()>,
476{
477 let rank = dims.len();
478 debug_assert!(rank >= 9);
479
480 let d0 = dims[0];
481 let b0 = blocks[0].max(1).min(d0);
482 let inner_strides: OperandVec<isize> = strides.iter().map(|s| s[0]).collect();
483
484 let mut idx: Vec<usize> = vec![0usize; rank];
486
487 if dims.contains(&0) {
491 return Ok(());
492 }
493 loop {
494 let mut j0 = 0usize;
496 let mut advanced = 0usize;
497 while j0 < d0 {
498 let blen0 = b0.min(d0 - j0);
499 f(offsets, blen0, &inner_strides)?;
500 j0 += blen0;
501 if j0 < d0 {
502 for (offset, s) in offsets.iter_mut().zip(strides.iter()) {
503 *offset += (blen0 as isize) * s[0];
504 }
505 advanced = j0;
506 }
507 }
508 for (offset, s) in offsets.iter_mut().zip(strides.iter()) {
509 *offset -= (advanced as isize) * s[0];
510 }
511
512 let mut level = 1usize;
514 loop {
515 if idx[level] + 1 < dims[level] {
516 idx[level] += 1;
517 for (offset, s) in offsets.iter_mut().zip(strides.iter()) {
518 *offset += s[level];
519 }
520 break;
521 }
522
523 let last = (dims[level] - 1) as isize;
524 idx[level] = 0;
525 for (offset, s) in offsets.iter_mut().zip(strides.iter()) {
526 *offset -= last * s[level];
527 }
528 level += 1;
529 if level == rank {
530 return Ok(());
531 }
532 }
533 }
534}
535
536#[inline]
556pub(crate) fn for_each_inner_block_preordered<S: AsRef<[isize]>, F>(
557 dims: &[usize],
558 blocks: &[usize],
559 strides: &[S],
560 initial_offsets: &[isize],
561 mut f: F,
562) -> Result<()>
563where
564 F: FnMut(&[isize], usize, &[isize]) -> Result<()>,
565{
566 let refs: OperandVec<&[isize]> = strides.iter().map(|s| s.as_ref()).collect();
567 preordered_dyn(dims, blocks, &refs, initial_offsets, &mut f)
568}
569
570type BlockFn<'a> = &'a mut dyn FnMut(&[isize], usize, &[isize]) -> Result<()>;
571
572#[inline(never)]
573fn preordered_dyn(
574 dims: &[usize],
575 blocks: &[usize],
576 strides: &[&[isize]],
577 initial_offsets: &[isize],
578 mut f: BlockFn<'_>,
579) -> Result<()> {
580 let rank = dims.len();
581 if rank == 0 {
582 return f(initial_offsets, 1, &[]);
583 }
584
585 let mut offsets: OperandVec<isize> = initial_offsets.iter().copied().collect();
587
588 match rank {
589 1 => kernel_1d_inner(dims, blocks, strides, &mut offsets, &mut f),
590 2 => kernel_2d_inner(dims, blocks, strides, &mut offsets, &mut f),
591 3 => kernel_3d_inner(dims, blocks, strides, &mut offsets, &mut f),
592 4 => kernel_4d_inner(dims, blocks, strides, &mut offsets, &mut f),
593 5 => kernel_5d_inner(dims, blocks, strides, &mut offsets, &mut f),
594 6 => kernel_6d_inner(dims, blocks, strides, &mut offsets, &mut f),
595 7 => kernel_7d_inner(dims, blocks, strides, &mut offsets, &mut f),
596 8 => kernel_8d_inner(dims, blocks, strides, &mut offsets, &mut f),
597 _ => kernel_nd_inner_iterative(dims, blocks, strides, &mut offsets, &mut f),
598 }
599}
600
601pub fn ensure_same_shape(a: &[usize], b: &[usize]) -> Result<()> {
618 if a.len() != b.len() {
619 return Err(crate::StridedError::RankMismatch(a.len(), b.len()));
620 }
621 if a != b {
622 return Err(crate::StridedError::ShapeMismatch(a.to_vec(), b.to_vec()));
623 }
624 Ok(())
625}
626
627#[derive(Copy, Clone, Debug, Eq, PartialEq)]
628pub(crate) enum ContiguousLayout {
629 RowMajor,
631 ColMajor,
633}
634
635pub(crate) fn contiguous_layout(dims: &[usize], strides: &[isize]) -> Option<ContiguousLayout> {
641 if dims.len() != strides.len() {
642 return None;
643 }
644 if dims.is_empty() {
645 return Some(ContiguousLayout::RowMajor);
646 }
647
648 let mut expected = 1isize;
650 let mut row_ok = true;
651 for (&dim, &stride) in dims.iter().rev().zip(strides.iter().rev()) {
652 if dim <= 1 {
653 continue;
654 }
655 if stride != expected {
656 row_ok = false;
657 break;
658 }
659 expected = expected.saturating_mul(dim as isize);
660 }
661 if row_ok {
662 return Some(ContiguousLayout::RowMajor);
663 }
664
665 let mut expected = 1isize;
667 for (&dim, &stride) in dims.iter().zip(strides.iter()) {
668 if dim <= 1 {
669 continue;
670 }
671 if stride != expected {
672 return None;
673 }
674 expected = expected.saturating_mul(dim as isize);
675 }
676 Some(ContiguousLayout::ColMajor)
677}
678
679pub(crate) fn total_len(dims: &[usize]) -> Result<usize> {
687 if dims.contains(&0) {
688 return Ok(0);
689 }
690 dims.iter()
691 .try_fold(1usize, |acc, &dim| acc.checked_mul(dim))
692 .ok_or(StridedError::OffsetOverflow)
693}
694
695#[inline]
703pub(crate) fn use_sequential_fast_path(total: usize) -> bool {
704 #[cfg(feature = "parallel")]
705 {
706 total <= crate::threading::MINTHREADLENGTH || crate::execution_policy::rayon_threads() <= 1
707 }
708 #[cfg(not(feature = "parallel"))]
709 {
710 let _ = total;
711 true
712 }
713}
714
715#[inline]
721pub(crate) fn same_contiguous_layout(
722 dims: &[usize],
723 strides_list: &[&[isize]],
724) -> Option<ContiguousLayout> {
725 let first = contiguous_layout(dims, strides_list.first()?)?;
726 for strides in &strides_list[1..] {
727 if contiguous_layout(dims, strides)? != first {
728 return None;
729 }
730 }
731 Some(first)
732}
733
734#[inline]
740pub(crate) fn sequential_contiguous_layout(
741 dims: &[usize],
742 strides_list: &[&[isize]],
743) -> Result<Option<ContiguousLayout>> {
744 if !use_sequential_fast_path(total_len(dims)?) {
745 return Ok(None);
746 }
747 Ok(same_contiguous_layout(dims, strides_list))
748}
749
750#[cfg(test)]
751#[path = "kernel/tests/tests.rs"]
752mod tests;