From 5e3fce18f6830aa9b563b5ecee3869609e46619e Mon Sep 17 00:00:00 2001 From: GiggleLiu Date: Sun, 12 Apr 2026 15:22:07 +0800 Subject: [PATCH 1/6] perf(cpu): add contraction layout analysis --- src/backend/cpu/contract.rs | 334 ++++++++++++++++++++++++++++++++---- 1 file changed, 296 insertions(+), 38 deletions(-) diff --git a/src/backend/cpu/contract.rs b/src/backend/cpu/contract.rs index 0fb1766..46ed2de 100644 --- a/src/backend/cpu/contract.rs +++ b/src/backend/cpu/contract.rs @@ -78,6 +78,214 @@ pub(super) fn compute_permutation( target.iter().map(|m| mode_position(current, *m)).collect() } +#[derive(Debug, Clone, PartialEq, Eq)] +enum MaterializationPlan { + NoCopy, + MakeContiguous, + Permute { perm: Vec, shape: Vec }, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +struct ContractionLayoutPlan { + batch_modes: Vec, + left_modes: Vec, + right_modes: Vec, + contracted_modes: Vec, + left_materialization: MaterializationPlan, + right_materialization: MaterializationPlan, + output_perm: Option>, + batch_size: usize, + left_size: usize, + right_size: usize, + contract_size: usize, +} + +fn physical_axis_order(strides: &[usize]) -> Vec { + let mut axes: Vec = (0..strides.len()).collect(); + axes.sort_by_key(|&axis| (strides[axis], axis)); + axes +} + +fn is_flattenable_group( + shape: &[usize], + strides: &[usize], + axes: &[usize], + require_unit_base: bool, +) -> bool { + if axes.is_empty() { + return true; + } + + let mut expected = if require_unit_base { 1 } else { strides[axes[0]] }; + if require_unit_base && strides[axes[0]] != 1 { + return false; + } + + for &axis in axes { + if strides[axis] != expected { + return false; + } + expected = expected.saturating_mul(shape[axis].max(1)); + } + + true +} + +fn is_permutation_like(shape: &[usize], strides: &[usize], physical_order: &[usize]) -> bool { + let mut expected = 1usize; + for &axis in physical_order { + if strides[axis] != expected { + return false; + } + expected = expected.saturating_mul(shape[axis].max(1)); + } + true +} + +fn analyze_operand_materialization( + shape: &[usize], + strides: &[usize], + modes: &[i32], + batch_modes: &[i32], + row_modes: &[i32], + col_modes: &[i32], +) -> MaterializationPlan { + let batch_axes: Vec = batch_modes + .iter() + .map(|&mode| mode_position(modes, mode)) + .collect(); + let row_axes: Vec = row_modes.iter().map(|&mode| mode_position(modes, mode)).collect(); + let col_axes: Vec = col_modes.iter().map(|&mode| mode_position(modes, mode)).collect(); + let physical_order = physical_axis_order(strides); + + let mut no_copy_orders = vec![batch_axes + .iter() + .chain(row_axes.iter()) + .chain(col_axes.iter()) + .copied() + .collect::>()]; + if row_axes != col_axes { + no_copy_orders.push( + batch_axes + .iter() + .chain(col_axes.iter()) + .chain(row_axes.iter()) + .copied() + .collect(), + ); + } + + let batch_len = batch_axes.len(); + let row_len = row_axes.len(); + for order in &no_copy_orders { + let rows = &order[batch_len..batch_len + row_len]; + let cols = &order[batch_len + row_len..]; + if *order == physical_order + && is_flattenable_group(shape, strides, &batch_axes, !batch_axes.is_empty()) + && is_flattenable_group(shape, strides, rows, false) + && is_flattenable_group(shape, strides, cols, false) + { + return MaterializationPlan::NoCopy; + } + } + + if is_permutation_like(shape, strides, &physical_order) { + let target_axes: Vec = batch_axes + .iter() + .chain(row_axes.iter()) + .chain(col_axes.iter()) + .copied() + .collect(); + let perm: Vec = target_axes + .iter() + .map(|axis| { + physical_order + .iter() + .position(|physical_axis| physical_axis == axis) + .expect("axis must exist in physical order") + }) + .collect(); + let physical_shape: Vec = physical_order.iter().map(|&axis| shape[axis]).collect(); + return MaterializationPlan::Permute { + perm, + shape: physical_shape, + }; + } + + MaterializationPlan::MakeContiguous +} + +fn analyze_contraction_layout( + shape_a: &[usize], + strides_a: &[usize], + modes_a: &[i32], + shape_b: &[usize], + strides_b: &[usize], + modes_b: &[i32], + modes_c: &[i32], +) -> ContractionLayoutPlan { + let (batch_modes, left_candidates, right_candidates, contracted_modes) = + classify_modes(modes_a, modes_b, modes_c); + let output_set: HashSet = modes_c.iter().copied().collect(); + let left_modes: Vec = left_candidates + .into_iter() + .filter(|mode| output_set.contains(mode)) + .collect(); + let right_modes: Vec = right_candidates + .into_iter() + .filter(|mode| output_set.contains(mode)) + .collect(); + + let left_materialization = analyze_operand_materialization( + shape_a, + strides_a, + modes_a, + &batch_modes, + &left_modes, + &contracted_modes, + ); + let right_materialization = analyze_operand_materialization( + shape_b, + strides_b, + modes_b, + &batch_modes, + &contracted_modes, + &right_modes, + ); + + let current_output: Vec = left_modes + .iter() + .chain(right_modes.iter()) + .chain(batch_modes.iter()) + .copied() + .collect(); + let output_perm = (current_output != modes_c).then(|| { + modes_c + .iter() + .map(|mode| { + current_output + .iter() + .position(|current_mode| current_mode == mode) + .expect("output mode must exist in current output") + }) + .collect() + }); + + ContractionLayoutPlan { + batch_size: product_of_dims(&batch_modes, modes_a, shape_a), + left_size: product_of_dims(&left_modes, modes_a, shape_a), + right_size: product_of_dims(&right_modes, modes_b, shape_b), + contract_size: product_of_dims(&contracted_modes, modes_a, shape_a), + batch_modes, + left_modes, + right_modes, + contracted_modes, + left_materialization, + right_materialization, + output_perm, + } +} + use crate::algebra::Algebra; use crate::backend::Cpu; use crate::tensor::compute_contiguous_strides; @@ -165,7 +373,7 @@ where let b_contig = ensure_contiguous(b, shape_b, strides_b); // 2. Classify modes - let (batch, left, right, contracted) = classify_modes(modes_a, modes_b, modes_c); + let (_batch, left, right, _contracted) = classify_modes(modes_a, modes_b, modes_c); // 3. Handle trace modes: modes in only one input that are NOT in the output. // GEMM can only contract modes shared by both inputs. Single-input modes @@ -193,69 +401,80 @@ where (b_contig, shape_b.to_vec(), modes_b.to_vec()) }; - // Free modes (left/right modes that ARE in the output) - let left_free: Vec = left.iter().filter(|m| c_set.contains(m)).copied().collect(); - let right_free: Vec = right - .iter() - .filter(|m| c_set.contains(m)) - .copied() - .collect(); - - // 4. Compute dimension sizes (using reduced inputs) - let batch_size = product_of_dims(&batch, &a_modes, &a_shape); - let left_size = product_of_dims(&left_free, &a_modes, &a_shape); - let right_size = product_of_dims(&right_free, &b_modes, &b_shape); - let contract_size = product_of_dims(&contracted, &a_modes, &a_shape); - - // 5. Permute A to [left_free, contracted, batch] - batch LAST for correct memory layout - let a_perm = compute_permutation(&a_modes, &left_free, &contracted, &batch); + let a_layout_strides = compute_contiguous_strides(&a_shape); + let b_layout_strides = compute_contiguous_strides(&b_shape); + let plan = analyze_contraction_layout( + &a_shape, + &a_layout_strides, + &a_modes, + &b_shape, + &b_layout_strides, + &b_modes, + modes_c, + ); + + // 4. Permute A to [left_free, contracted, batch] - batch LAST for current GEMM layout + let a_perm = compute_permutation( + &a_modes, + &plan.left_modes, + &plan.contracted_modes, + &plan.batch_modes, + ); let a_permuted = permute_data(&a_data, &a_shape, &a_perm); - // 6. Permute B to [contracted, right_free, batch] - batch LAST - let b_perm = compute_permutation(&b_modes, &contracted, &right_free, &batch); + // 5. Permute B to [contracted, right_free, batch] - batch LAST + let b_perm = compute_permutation( + &b_modes, + &plan.contracted_modes, + &plan.right_modes, + &plan.batch_modes, + ); let b_permuted = permute_data(&b_data, &b_shape, &b_perm); - // 7. Call GEMM - let c_data = if batch.is_empty() { + // 6. Call GEMM + let c_data = if plan.batch_modes.is_empty() { cpu.gemm_internal::( &a_permuted, - left_size, - contract_size, + plan.left_size, + plan.contract_size, &b_permuted, - right_size, + plan.right_size, ) } else { cpu.gemm_batched_internal::( &a_permuted, - batch_size, - left_size, - contract_size, + plan.batch_size, + plan.left_size, + plan.contract_size, &b_permuted, - right_size, + plan.right_size, ) }; - // 8. Permute result to output order + // 7. Permute result to output order // Result is in [left_free, right_free, batch] order - let current_order: Vec = left_free + let current_order: Vec = plan + .left_modes .iter() - .chain(right_free.iter()) - .chain(batch.iter()) + .chain(plan.right_modes.iter()) + .chain(plan.batch_modes.iter()) .copied() .collect(); - if current_order == modes_c { + if plan.output_perm.is_none() { c_data } else { let c_shape_current: Vec = current_order .iter() .map(|&m| shape_c[mode_position(modes_c, m)]) .collect(); - let out_perm: Vec = modes_c - .iter() - .map(|m| current_order.iter().position(|x| x == m).unwrap()) - .collect(); - permute_data(&c_data, &c_shape_current, &out_perm) + permute_data( + &c_data, + &c_shape_current, + plan.output_perm + .as_ref() + .expect("output permutation must exist"), + ) } } @@ -481,4 +700,43 @@ mod tests { let perm = compute_permutation(&[0, 1, 2], &[0], &[2], &[1]); assert_eq!(perm, vec![0, 2, 1]); } + + #[test] + fn test_analyze_contraction_layout_detects_copy_free_matmul() { + let plan = analyze_contraction_layout( + &[2, 3], + &[1, 2], + &[0, 1], + &[3, 4], + &[1, 3], + &[1, 2], + &[0, 2], + ); + + assert!(plan.batch_modes.is_empty()); + assert_eq!(plan.left_modes, vec![0]); + assert_eq!(plan.right_modes, vec![2]); + assert_eq!(plan.contracted_modes, vec![1]); + assert!(matches!(plan.left_materialization, MaterializationPlan::NoCopy)); + assert!(matches!(plan.right_materialization, MaterializationPlan::NoCopy)); + } + + #[test] + fn test_analyze_contraction_layout_marks_single_side_materialization() { + let plan = analyze_contraction_layout( + &[2, 2, 2], + &[1, 2, 4], + &[0, 1, 2], + &[2, 2, 2], + &[2, 1, 4], + &[0, 2, 3], + &[0, 1, 3], + ); + + assert!(matches!(plan.left_materialization, MaterializationPlan::NoCopy)); + assert!(matches!( + plan.right_materialization, + MaterializationPlan::Permute { .. } + )); + } } From 253953aea7b624232720dbc3d0423e0c99b6442e Mon Sep 17 00:00:00 2001 From: GiggleLiu Date: Sun, 12 Apr 2026 15:37:40 +0800 Subject: [PATCH 2/6] perf(cpu): add standard strided gemm fast path --- src/backend/cpu/contract.rs | 173 +++++++++++++++++++++++++------ src/backend/cpu/mod.rs | 169 ++++++++++++++++++++++++++++++ tests/suites/backend_contract.rs | 38 +++++++ 3 files changed, 346 insertions(+), 34 deletions(-) diff --git a/src/backend/cpu/contract.rs b/src/backend/cpu/contract.rs index 46ed2de..e7220fa 100644 --- a/src/backend/cpu/contract.rs +++ b/src/backend/cpu/contract.rs @@ -286,6 +286,7 @@ fn analyze_contraction_layout( } } +use super::MatrixLayout; use crate::algebra::Algebra; use crate::backend::Cpu; use crate::tensor::compute_contiguous_strides; @@ -350,6 +351,67 @@ where (result, new_shape, new_modes) } +fn group_base_stride(modes: &[i32], strides: &[usize], group_modes: &[i32]) -> isize { + group_modes + .first() + .map(|&mode| strides[mode_position(modes, mode)] as isize) + .unwrap_or(0) +} + +fn matrix_layout_from_operand<'a, T>( + data: &'a [T], + shape: &[usize], + strides: &[usize], + modes: &[i32], + row_modes: &[i32], + col_modes: &[i32], +) -> MatrixLayout<'a, T> { + MatrixLayout { + data, + rows: product_of_dims(row_modes, modes, shape), + cols: product_of_dims(col_modes, modes, shape), + row_stride: group_base_stride(modes, strides, row_modes), + col_stride: group_base_stride(modes, strides, col_modes), + } +} + +fn materialize_matrix_operand( + data: &[T], + shape: &[usize], + strides: &[usize], + modes: &[i32], + row_modes: &[i32], + col_modes: &[i32], +) -> Vec { + let contiguous = ensure_contiguous(data, shape, strides); + let perm = compute_permutation(modes, row_modes, col_modes, &[]); + permute_data(&contiguous, shape, &perm) +} + +fn finalize_contraction_output( + c_data: Vec, + plan: &ContractionLayoutPlan, + shape_c: &[usize], + modes_c: &[i32], +) -> Vec { + if let Some(output_perm) = &plan.output_perm { + let current_order: Vec = plan + .left_modes + .iter() + .chain(plan.right_modes.iter()) + .chain(plan.batch_modes.iter()) + .copied() + .collect(); + let c_shape_current: Vec = current_order + .iter() + .map(|&mode| shape_c[mode_position(modes_c, mode)]) + .collect(); + permute_data(&c_data, &c_shape_current, output_perm) + } else { + c_data + } +} + /// Execute tensor contraction on CPU via reshape→GEMM→reshape. #[allow(clippy::too_many_arguments)] pub(super) fn contract( @@ -368,14 +430,10 @@ pub(super) fn contract( where A::Scalar: crate::algebra::Scalar, { - // 1. Make inputs contiguous if needed - let a_contig = ensure_contiguous(a, shape_a, strides_a); - let b_contig = ensure_contiguous(b, shape_b, strides_b); - - // 2. Classify modes + // 1. Classify modes let (_batch, left, right, _contracted) = classify_modes(modes_a, modes_b, modes_c); - // 3. Handle trace modes: modes in only one input that are NOT in the output. + // 2. Handle trace modes: modes in only one input that are NOT in the output. // GEMM can only contract modes shared by both inputs. Single-input modes // not in the output must be summed over (traced) before GEMM. let c_set: HashSet = modes_c.iter().copied().collect(); @@ -390,6 +448,77 @@ where .copied() .collect(); + if left_trace.is_empty() && right_trace.is_empty() { + let plan = + analyze_contraction_layout(shape_a, strides_a, modes_a, shape_b, strides_b, modes_b, modes_c); + if plan.batch_modes.is_empty() { + let left_nocopy = matches!(plan.left_materialization, MaterializationPlan::NoCopy); + let right_nocopy = matches!(plan.right_materialization, MaterializationPlan::NoCopy); + if left_nocopy || right_nocopy { + let left_materialized; + let a_layout = if left_nocopy { + matrix_layout_from_operand( + a, + shape_a, + strides_a, + modes_a, + &plan.left_modes, + &plan.contracted_modes, + ) + } else { + left_materialized = materialize_matrix_operand( + a, + shape_a, + strides_a, + modes_a, + &plan.left_modes, + &plan.contracted_modes, + ); + MatrixLayout::column_major( + &left_materialized, + plan.left_size, + plan.contract_size, + ) + }; + + let right_materialized; + let b_layout = if right_nocopy { + matrix_layout_from_operand( + b, + shape_b, + strides_b, + modes_b, + &plan.contracted_modes, + &plan.right_modes, + ) + } else { + right_materialized = materialize_matrix_operand( + b, + shape_b, + strides_b, + modes_b, + &plan.contracted_modes, + &plan.right_modes, + ); + MatrixLayout::column_major( + &right_materialized, + plan.contract_size, + plan.right_size, + ) + }; + + if let Some(c_data) = cpu.gemm_standard_layout_internal::(a_layout, b_layout) { + return finalize_contraction_output(c_data, &plan, shape_c, modes_c); + } + } + } + } + + // 3. Make inputs contiguous if needed for the generic materialized path. + let a_contig = ensure_contiguous(a, shape_a, strides_a); + let b_contig = ensure_contiguous(b, shape_b, strides_b); + + // 4. Reduce trace modes before generic GEMM. let (a_data, a_shape, a_modes) = if !left_trace.is_empty() { reduce_trace_modes::(&a_contig, shape_a, modes_a, &left_trace) } else { @@ -413,7 +542,7 @@ where modes_c, ); - // 4. Permute A to [left_free, contracted, batch] - batch LAST for current GEMM layout + // 5. Permute A to [left_free, contracted, batch] - batch LAST for current GEMM layout let a_perm = compute_permutation( &a_modes, &plan.left_modes, @@ -422,7 +551,7 @@ where ); let a_permuted = permute_data(&a_data, &a_shape, &a_perm); - // 5. Permute B to [contracted, right_free, batch] - batch LAST + // 6. Permute B to [contracted, right_free, batch] - batch LAST let b_perm = compute_permutation( &b_modes, &plan.contracted_modes, @@ -431,7 +560,7 @@ where ); let b_permuted = permute_data(&b_data, &b_shape, &b_perm); - // 6. Call GEMM + // 7. Call GEMM let c_data = if plan.batch_modes.is_empty() { cpu.gemm_internal::( &a_permuted, @@ -451,31 +580,7 @@ where ) }; - // 7. Permute result to output order - // Result is in [left_free, right_free, batch] order - let current_order: Vec = plan - .left_modes - .iter() - .chain(plan.right_modes.iter()) - .chain(plan.batch_modes.iter()) - .copied() - .collect(); - - if plan.output_perm.is_none() { - c_data - } else { - let c_shape_current: Vec = current_order - .iter() - .map(|&m| shape_c[mode_position(modes_c, m)]) - .collect(); - permute_data( - &c_data, - &c_shape_current, - plan.output_perm - .as_ref() - .expect("output permutation must exist"), - ) - } + finalize_contraction_output(c_data, &plan, shape_c, modes_c) } /// Ensure data is contiguous (copy if strided). diff --git a/src/backend/cpu/mod.rs b/src/backend/cpu/mod.rs index 2d60ba5..d3f4a8e 100644 --- a/src/backend/cpu/mod.rs +++ b/src/backend/cpu/mod.rs @@ -10,7 +10,128 @@ use std::any::TypeId; #[derive(Clone, Debug, Default)] pub struct Cpu; +#[derive(Clone, Copy)] +pub(crate) struct MatrixLayout<'a, T> { + pub data: &'a [T], + pub rows: usize, + pub cols: usize, + pub row_stride: isize, + pub col_stride: isize, +} + +impl<'a, T> MatrixLayout<'a, T> { + pub(crate) fn column_major(data: &'a [T], rows: usize, cols: usize) -> Self { + Self { + data, + rows, + cols, + row_stride: 1, + col_stride: rows as isize, + } + } + + #[cfg(test)] + pub(crate) fn column_major_transposed(data: &'a [T], rows: usize, cols: usize) -> Self { + Self { + data, + rows, + cols, + row_stride: cols as isize, + col_stride: 1, + } + } +} + +fn layout_offset_bounds(layout: &MatrixLayout<'_, T>) -> (isize, isize) { + let row_extent = if layout.rows == 0 { + 0 + } else { + (layout.rows as isize - 1) * layout.row_stride + }; + let col_extent = if layout.cols == 0 { + 0 + } else { + (layout.cols as isize - 1) * layout.col_stride + }; + let offsets = [0, row_extent, col_extent, row_extent + col_extent]; + ( + *offsets.iter().min().expect("offset bounds must exist"), + *offsets.iter().max().expect("offset bounds must exist"), + ) +} + +fn faer_mat_ref<'a, T>(layout: MatrixLayout<'a, T>) -> faer::MatRef<'a, T> { + let (min_offset, max_offset) = layout_offset_bounds(&layout); + if layout.rows > 0 && layout.cols > 0 { + assert!(!layout.data.is_empty(), "matrix layout requires backing storage"); + assert!( + max_offset >= min_offset, + "matrix layout offsets must be ordered" + ); + assert!( + ((max_offset - min_offset) as usize) < layout.data.len(), + "matrix layout exceeds backing storage" + ); + } + + let ptr = unsafe { layout.data.as_ptr().offset(-min_offset) }; + unsafe { + faer::MatRef::from_raw_parts( + ptr, + layout.rows, + layout.cols, + layout.row_stride, + layout.col_stride, + ) + } +} + impl Cpu { + pub(crate) fn gemm_standard_layout_internal( + &self, + a: MatrixLayout<'_, A::Scalar>, + b: MatrixLayout<'_, A::Scalar>, + ) -> Option> { + if TypeId::of::() == TypeId::of::>() { + let a_f32 = MatrixLayout { + data: unsafe { std::mem::transmute::<&[A::Scalar], &[f32]>(a.data) }, + rows: a.rows, + cols: a.cols, + row_stride: a.row_stride, + col_stride: a.col_stride, + }; + let b_f32 = MatrixLayout { + data: unsafe { std::mem::transmute::<&[A::Scalar], &[f32]>(b.data) }, + rows: b.rows, + cols: b.cols, + row_stride: b.row_stride, + col_stride: b.col_stride, + }; + let result = faer_gemm_f32_layout(a_f32, b_f32); + return Some(unsafe { std::mem::transmute::, Vec>(result) }); + } + if TypeId::of::() == TypeId::of::>() { + let a_f64 = MatrixLayout { + data: unsafe { std::mem::transmute::<&[A::Scalar], &[f64]>(a.data) }, + rows: a.rows, + cols: a.cols, + row_stride: a.row_stride, + col_stride: a.col_stride, + }; + let b_f64 = MatrixLayout { + data: unsafe { std::mem::transmute::<&[A::Scalar], &[f64]>(b.data) }, + rows: b.rows, + cols: b.cols, + row_stride: b.row_stride, + col_stride: b.col_stride, + }; + let result = faer_gemm_f64_layout(a_f64, b_f64); + return Some(unsafe { std::mem::transmute::, Vec>(result) }); + } + + None + } + /// General matrix multiplication (internal implementation). /// /// Computes C = A ⊗ B where ⊗ is the semiring multiplication @@ -354,6 +475,23 @@ fn faer_gemm_f32(a: &[f32], m: usize, k: usize, b: &[f32], n: usize) -> Vec c } +fn faer_gemm_f32_layout(a: MatrixLayout<'_, f32>, b: MatrixLayout<'_, f32>) -> Vec { + use faer::{linalg::matmul::matmul, Accum, Mat, Par}; + + let a_mat = faer_mat_ref(a); + let b_mat = faer_mat_ref(b); + let mut c_mat = Mat::::zeros(a.rows, b.cols); + matmul(c_mat.as_mut(), Accum::Replace, a_mat, b_mat, 1.0f32, Par::Seq); + + let mut c = vec![0.0f32; a.rows * b.cols]; + for j in 0..b.cols { + for i in 0..a.rows { + c[j * a.rows + i] = c_mat[(i, j)]; + } + } + c +} + /// GEMM using faer for f64 (column-major layout). fn faer_gemm_f64(a: &[f64], m: usize, k: usize, b: &[f64], n: usize) -> Vec { use faer::Mat; @@ -372,6 +510,23 @@ fn faer_gemm_f64(a: &[f64], m: usize, k: usize, b: &[f64], n: usize) -> Vec c } +fn faer_gemm_f64_layout(a: MatrixLayout<'_, f64>, b: MatrixLayout<'_, f64>) -> Vec { + use faer::{linalg::matmul::matmul, Accum, Mat, Par}; + + let a_mat = faer_mat_ref(a); + let b_mat = faer_mat_ref(b); + let mut c_mat = Mat::::zeros(a.rows, b.cols); + matmul(c_mat.as_mut(), Accum::Replace, a_mat, b_mat, 1.0f64, Par::Seq); + + let mut c = vec![0.0f64; a.rows * b.cols]; + for j in 0..b.cols { + for i in 0..a.rows { + c[j * a.rows + i] = c_mat[(i, j)]; + } + } + c +} + /// Generic GEMM using semiring operations (column-major layout). fn generic_gemm( a: &[A::Scalar], @@ -662,6 +817,20 @@ mod tests { assert_eq!(c, vec![7.0, 10.0, 15.0, 22.0]); } + #[test] + fn test_faer_layout_gemm_accepts_rhs_transpose_view() { + let a = vec![1.0f32, 2.0, 3.0, 4.0]; + let b = vec![1.0f32, 2.0, 3.0, 4.0]; + + let c = faer_gemm_f32_layout( + MatrixLayout::column_major(&a, 2, 2), + MatrixLayout::column_major_transposed(&b, 2, 2), + ); + + let expected = faer_gemm_f32(&a, 2, 2, &[1.0, 3.0, 2.0, 4.0], 2); + assert_eq!(c, expected); + } + #[cfg(feature = "tropical")] #[test] fn test_cpu_gemm_maxplus() { diff --git a/tests/suites/backend_contract.rs b/tests/suites/backend_contract.rs index 31ddbd0..73ec8cb 100644 --- a/tests/suites/backend_contract.rs +++ b/tests/suites/backend_contract.rs @@ -198,6 +198,44 @@ fn test_cpu_contract_both_strided() { assert_eq!(c, vec![59.0, 78.0, 83.0, 110.0]); } +#[test] +fn test_cpu_contract_rhs_transpose_view_matches_contiguous_rhs() { + let cpu = Cpu; + + let a = vec![1.0f64, 2.0, 3.0, 4.0]; + let b = vec![1.0f64, 2.0, 3.0, 4.0]; + let b_contiguous_transpose = vec![1.0f64, 3.0, 2.0, 4.0]; + + let strided = cpu.contract::>( + &a, + &[2, 2], + &[1, 2], + &[0, 1], + &b, + &[2, 2], + &[2, 1], + &[1, 2], + &[2, 2], + &[0, 2], + ); + + let contiguous = cpu.contract::>( + &a, + &[2, 2], + &[1, 2], + &[0, 1], + &b_contiguous_transpose, + &[2, 2], + &[1, 2], + &[1, 2], + &[2, 2], + &[0, 2], + ); + + assert_eq!(strided, contiguous); + assert_eq!(strided, vec![10.0, 14.0, 14.0, 20.0]); +} + // ============================================================================ // Tests for output permutation // ============================================================================ From 627199b995dbf9cf358b319068994f3ff6aa9365 Mon Sep 17 00:00:00 2001 From: GiggleLiu Date: Sun, 12 Apr 2026 15:43:32 +0800 Subject: [PATCH 3/6] perf(cpu): extend standard fast path to batched contractions --- src/backend/cpu/contract.rs | 177 ++++++++++++++++++++----------- src/backend/cpu/mod.rs | 106 ++++++++++++++++++ tests/suites/backend_contract.rs | 21 ++-- 3 files changed, 229 insertions(+), 75 deletions(-) diff --git a/src/backend/cpu/contract.rs b/src/backend/cpu/contract.rs index e7220fa..2edbb20 100644 --- a/src/backend/cpu/contract.rs +++ b/src/backend/cpu/contract.rs @@ -375,17 +375,41 @@ fn matrix_layout_from_operand<'a, T>( } } +struct MaterializedMatrixOperand { + data: Vec, + shape: Vec, + strides: Vec, + modes: Vec, +} + fn materialize_matrix_operand( data: &[T], shape: &[usize], strides: &[usize], modes: &[i32], + batch_modes: &[i32], row_modes: &[i32], col_modes: &[i32], -) -> Vec { +) -> MaterializedMatrixOperand { let contiguous = ensure_contiguous(data, shape, strides); - let perm = compute_permutation(modes, row_modes, col_modes, &[]); - permute_data(&contiguous, shape, &perm) + let perm = compute_permutation(modes, batch_modes, row_modes, col_modes); + let target_modes: Vec = batch_modes + .iter() + .chain(row_modes.iter()) + .chain(col_modes.iter()) + .copied() + .collect(); + let target_shape: Vec = target_modes + .iter() + .map(|&mode| shape[mode_position(modes, mode)]) + .collect(); + + MaterializedMatrixOperand { + data: permute_data(&contiguous, shape, &perm), + strides: compute_contiguous_strides(&target_shape), + shape: target_shape, + modes: target_modes, + } } fn finalize_contraction_output( @@ -451,65 +475,76 @@ where if left_trace.is_empty() && right_trace.is_empty() { let plan = analyze_contraction_layout(shape_a, strides_a, modes_a, shape_b, strides_b, modes_b, modes_c); - if plan.batch_modes.is_empty() { - let left_nocopy = matches!(plan.left_materialization, MaterializationPlan::NoCopy); - let right_nocopy = matches!(plan.right_materialization, MaterializationPlan::NoCopy); - if left_nocopy || right_nocopy { - let left_materialized; - let a_layout = if left_nocopy { - matrix_layout_from_operand( - a, - shape_a, - strides_a, - modes_a, - &plan.left_modes, - &plan.contracted_modes, - ) - } else { - left_materialized = materialize_matrix_operand( - a, - shape_a, - strides_a, - modes_a, - &plan.left_modes, - &plan.contracted_modes, - ); - MatrixLayout::column_major( - &left_materialized, - plan.left_size, - plan.contract_size, - ) - }; - - let right_materialized; - let b_layout = if right_nocopy { - matrix_layout_from_operand( - b, - shape_b, - strides_b, - modes_b, - &plan.contracted_modes, - &plan.right_modes, - ) - } else { - right_materialized = materialize_matrix_operand( - b, - shape_b, - strides_b, - modes_b, - &plan.contracted_modes, - &plan.right_modes, - ); - MatrixLayout::column_major( - &right_materialized, - plan.contract_size, - plan.right_size, - ) - }; - - if let Some(c_data) = cpu.gemm_standard_layout_internal::(a_layout, b_layout) { - return finalize_contraction_output(c_data, &plan, shape_c, modes_c); - } + let left_nocopy = matches!(plan.left_materialization, MaterializationPlan::NoCopy); + let right_nocopy = matches!(plan.right_materialization, MaterializationPlan::NoCopy); + if left_nocopy || right_nocopy { + let left_materialized; + let a_layout = if left_nocopy { + matrix_layout_from_operand( + a, + shape_a, + strides_a, + modes_a, + &plan.left_modes, + &plan.contracted_modes, + ) + } else { + left_materialized = materialize_matrix_operand( + a, + shape_a, + strides_a, + modes_a, + &plan.batch_modes, + &plan.left_modes, + &plan.contracted_modes, + ); + matrix_layout_from_operand( + &left_materialized.data, + &left_materialized.shape, + &left_materialized.strides, + &left_materialized.modes, + &plan.left_modes, + &plan.contracted_modes, + ) + }; + + let right_materialized; + let b_layout = if right_nocopy { + matrix_layout_from_operand( + b, + shape_b, + strides_b, + modes_b, + &plan.contracted_modes, + &plan.right_modes, + ) + } else { + right_materialized = materialize_matrix_operand( + b, + shape_b, + strides_b, + modes_b, + &plan.batch_modes, + &plan.contracted_modes, + &plan.right_modes, + ); + matrix_layout_from_operand( + &right_materialized.data, + &right_materialized.shape, + &right_materialized.strides, + &right_materialized.modes, + &plan.contracted_modes, + &plan.right_modes, + ) + }; + + let c_data = if plan.batch_modes.is_empty() { + cpu.gemm_standard_layout_internal::(a_layout, b_layout) + } else { + cpu.gemm_batched_standard_layout_internal::(plan.batch_size, a_layout, b_layout) + }; + if let Some(c_data) = c_data { + return finalize_contraction_output(c_data, &plan, shape_c, modes_c); } } } @@ -844,4 +879,22 @@ mod tests { MaterializationPlan::Permute { .. } )); } + + #[test] + fn test_analyze_contraction_layout_preserves_batch_tail() { + let plan = analyze_contraction_layout( + &[2, 2, 2], + &[1, 2, 4], + &[0, 1, 2], + &[2, 2, 2], + &[1, 2, 4], + &[0, 2, 3], + &[0, 1, 3], + ); + + assert_eq!(plan.batch_modes, vec![0]); + assert_eq!(plan.batch_size, 2); + assert!(matches!(plan.left_materialization, MaterializationPlan::NoCopy)); + assert!(matches!(plan.right_materialization, MaterializationPlan::NoCopy)); + } } diff --git a/src/backend/cpu/mod.rs b/src/backend/cpu/mod.rs index d3f4a8e..28610dc 100644 --- a/src/backend/cpu/mod.rs +++ b/src/backend/cpu/mod.rs @@ -20,6 +20,7 @@ pub(crate) struct MatrixLayout<'a, T> { } impl<'a, T> MatrixLayout<'a, T> { + #[cfg(test)] pub(crate) fn column_major(data: &'a [T], rows: usize, cols: usize) -> Self { Self { data, @@ -86,6 +87,14 @@ fn faer_mat_ref<'a, T>(layout: MatrixLayout<'a, T>) -> faer::MatRef<'a, T> { } } +fn matrix_layout_batch_view<'a, T>(layout: MatrixLayout<'a, T>, batch: usize) -> MatrixLayout<'a, T> { + assert!(batch < layout.data.len(), "batch offset must be in bounds"); + MatrixLayout { + data: &layout.data[batch..], + ..layout + } +} + impl Cpu { pub(crate) fn gemm_standard_layout_internal( &self, @@ -132,6 +141,74 @@ impl Cpu { None } + pub(crate) fn gemm_batched_standard_layout_internal( + &self, + batch_size: usize, + a: MatrixLayout<'_, A::Scalar>, + b: MatrixLayout<'_, A::Scalar>, + ) -> Option> { + let c_batch_stride = a.rows * b.cols; + + if TypeId::of::() == TypeId::of::>() { + let a_f32 = MatrixLayout { + data: unsafe { std::mem::transmute::<&[A::Scalar], &[f32]>(a.data) }, + rows: a.rows, + cols: a.cols, + row_stride: a.row_stride, + col_stride: a.col_stride, + }; + let b_f32 = MatrixLayout { + data: unsafe { std::mem::transmute::<&[A::Scalar], &[f32]>(b.data) }, + rows: b.rows, + cols: b.cols, + row_stride: b.row_stride, + col_stride: b.col_stride, + }; + let mut result = vec![0.0f32; batch_size * c_batch_stride]; + + for batch in 0..batch_size { + let c_offset = batch * c_batch_stride; + let c_batch = faer_gemm_f32_layout( + matrix_layout_batch_view(a_f32, batch), + matrix_layout_batch_view(b_f32, batch), + ); + result[c_offset..c_offset + c_batch_stride].copy_from_slice(&c_batch); + } + + return Some(unsafe { std::mem::transmute::, Vec>(result) }); + } + if TypeId::of::() == TypeId::of::>() { + let a_f64 = MatrixLayout { + data: unsafe { std::mem::transmute::<&[A::Scalar], &[f64]>(a.data) }, + rows: a.rows, + cols: a.cols, + row_stride: a.row_stride, + col_stride: a.col_stride, + }; + let b_f64 = MatrixLayout { + data: unsafe { std::mem::transmute::<&[A::Scalar], &[f64]>(b.data) }, + rows: b.rows, + cols: b.cols, + row_stride: b.row_stride, + col_stride: b.col_stride, + }; + let mut result = vec![0.0f64; batch_size * c_batch_stride]; + + for batch in 0..batch_size { + let c_offset = batch * c_batch_stride; + let c_batch = faer_gemm_f64_layout( + matrix_layout_batch_view(a_f64, batch), + matrix_layout_batch_view(b_f64, batch), + ); + result[c_offset..c_offset + c_batch_stride].copy_from_slice(&c_batch); + } + + return Some(unsafe { std::mem::transmute::, Vec>(result) }); + } + + None + } + /// General matrix multiplication (internal implementation). /// /// Computes C = A ⊗ B where ⊗ is the semiring multiplication @@ -831,6 +908,35 @@ mod tests { assert_eq!(c, expected); } + #[test] + fn test_gemm_batched_standard_layout_internal_accepts_batch_major_views() { + let cpu = Cpu; + let a = vec![1.0f32, 5.0, 2.0, 6.0, 3.0, 7.0, 4.0, 8.0]; + let b = vec![1.0f32, 2.0, 0.0, 0.0, 0.0, 0.0, 1.0, 2.0]; + + let c = cpu + .gemm_batched_standard_layout_internal::>( + 2, + MatrixLayout { + data: &a, + rows: 2, + cols: 2, + row_stride: 2, + col_stride: 4, + }, + MatrixLayout { + data: &b, + rows: 2, + cols: 2, + row_stride: 2, + col_stride: 4, + }, + ) + .expect("standard layout helper should handle batch-major inputs"); + + assert_eq!(c, vec![1.0, 2.0, 3.0, 4.0, 10.0, 12.0, 14.0, 16.0]); + } + #[cfg(feature = "tropical")] #[test] fn test_cpu_gemm_maxplus() { diff --git a/tests/suites/backend_contract.rs b/tests/suites/backend_contract.rs index 73ec8cb..cadfde8 100644 --- a/tests/suites/backend_contract.rs +++ b/tests/suites/backend_contract.rs @@ -67,9 +67,10 @@ fn test_cpu_contract_batched() { // bij,bjk->bik (batched matmul) let cpu = Cpu; - // 2 batches of 2x2 matrices - let a = vec![1.0f64, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0]; - let b = vec![1.0, 0.0, 0.0, 1.0, 2.0, 0.0, 0.0, 2.0]; + // Batch-major storage is column-major with the batch axis first, so the + // per-batch matrices are interleaved in memory. + let a = vec![1.0f64, 5.0, 2.0, 6.0, 3.0, 7.0, 4.0, 8.0]; + let b = vec![1.0, 2.0, 0.0, 0.0, 0.0, 0.0, 1.0, 2.0]; let c = cpu.contract::>( &a, @@ -84,9 +85,7 @@ fn test_cpu_contract_batched() { &[0, 1, 3], ); - // Batch 0: identity @ [[1,2],[3,4]] = [[1,2],[3,4]] - // Batch 1: 2*identity @ [[5,6],[7,8]] = [[10,12],[14,16]] - assert_eq!(c.len(), 8); + assert_eq!(c, vec![1.0, 10.0, 2.0, 12.0, 3.0, 14.0, 4.0, 16.0]); } #[cfg(feature = "tropical")] @@ -273,11 +272,8 @@ fn test_cpu_contract_batched_output_permuted() { // bij,bjk->kib (complex permutation) let cpu = Cpu; - // Simple 2x2x2 tensors for easy manual verification - // A: batch=2, i=2, j=2 - let a = vec![1.0f64, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0]; - // B: batch=2, j=2, k=2 (identity matrices) - let b = vec![1.0, 0.0, 0.0, 1.0, 1.0, 0.0, 0.0, 1.0]; + let a = vec![1.0f64, 5.0, 2.0, 6.0, 3.0, 7.0, 4.0, 8.0]; + let b = vec![1.0, 1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 1.0]; // Result should be A with axes permuted to [k, i, b] let c = cpu.contract::>( @@ -293,8 +289,7 @@ fn test_cpu_contract_batched_output_permuted() { &[3, 1, 0], // k, i, b ); - assert_eq!(c.len(), 8); - // Since B is identity, result is A with permuted axes + assert_eq!(c, vec![1.0, 3.0, 2.0, 4.0, 5.0, 7.0, 6.0, 8.0]); } // ============================================================================ From f9631fdeb26941cb4f0a06433c84d9479304ff65 Mon Sep 17 00:00:00 2001 From: GiggleLiu Date: Sun, 12 Apr 2026 15:49:27 +0800 Subject: [PATCH 4/6] perf(cpu): reuse scratch buffers for contraction materialization --- src/backend/cpu/buffer_pool.rs | 78 ++++++++++++++++ src/backend/cpu/contract.rs | 158 +++++++++++++++++++++------------ src/backend/cpu/mod.rs | 1 + 3 files changed, 178 insertions(+), 59 deletions(-) create mode 100644 src/backend/cpu/buffer_pool.rs diff --git a/src/backend/cpu/buffer_pool.rs b/src/backend/cpu/buffer_pool.rs new file mode 100644 index 0000000..51340aa --- /dev/null +++ b/src/backend/cpu/buffer_pool.rs @@ -0,0 +1,78 @@ +#[derive(Default)] +pub(crate) struct ScratchPool { + free: Vec>, +} + +pub(crate) struct ScratchBuffer<'a, T> { + buf: Vec, + pool: &'a mut ScratchPool, +} + +impl ScratchPool { + pub(crate) fn acquire(&mut self, len: usize) -> ScratchBuffer<'_, T> { + let reuse_index = self + .free + .iter() + .enumerate() + .filter(|(_, buf)| buf.capacity() >= len) + .min_by_key(|(_, buf)| buf.capacity()) + .map(|(index, _)| index); + let mut buf = reuse_index + .map(|index| self.free.swap_remove(index)) + .unwrap_or_else(|| Vec::with_capacity(len)); + buf.clear(); + buf.resize_with(len, T::default); + + ScratchBuffer { buf, pool: self } + } +} + +impl ScratchBuffer<'_, T> { + #[cfg(test)] + pub(crate) fn as_mut_slice(&mut self) -> &mut [T] { + self.buf.as_mut_slice() + } + + pub(crate) fn as_mut_vec(&mut self) -> &mut Vec { + &mut self.buf + } + + #[cfg(test)] + pub(crate) fn capacity(&self) -> usize { + self.buf.capacity() + } +} + +impl Drop for ScratchBuffer<'_, T> { + fn drop(&mut self) { + let mut buf = std::mem::take(&mut self.buf); + buf.clear(); + self.pool.free.push(buf); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_scratch_pool_reuses_released_capacity() { + let mut pool = ScratchPool::::default(); + let mut first = pool.acquire(32); + first.as_mut_slice().fill(1.0); + drop(first); + + let second = pool.acquire(16); + assert!(second.capacity() >= 32); + } + + #[test] + fn test_scratch_pool_grows_when_requested_capacity_is_larger() { + let mut pool = ScratchPool::::default(); + let small = pool.acquire(8); + drop(small); + + let large = pool.acquire(128); + assert!(large.capacity() >= 128); + } +} diff --git a/src/backend/cpu/contract.rs b/src/backend/cpu/contract.rs index 2edbb20..98e9a57 100644 --- a/src/backend/cpu/contract.rs +++ b/src/backend/cpu/contract.rs @@ -286,6 +286,7 @@ fn analyze_contraction_layout( } } +use super::buffer_pool::ScratchPool; use super::MatrixLayout; use crate::algebra::Algebra; use crate::backend::Cpu; @@ -391,7 +392,6 @@ fn materialize_matrix_operand( row_modes: &[i32], col_modes: &[i32], ) -> MaterializedMatrixOperand { - let contiguous = ensure_contiguous(data, shape, strides); let perm = compute_permutation(modes, batch_modes, row_modes, col_modes); let target_modes: Vec = batch_modes .iter() @@ -403,9 +403,12 @@ fn materialize_matrix_operand( .iter() .map(|&mode| shape[mode_position(modes, mode)]) .collect(); + let mut scratch = Vec::new(); MaterializedMatrixOperand { - data: permute_data(&contiguous, shape, &perm), + data: materialize_with_permutation_into(data, shape, strides, &perm, &mut scratch) + .as_slice() + .to_vec(), strides: compute_contiguous_strides(&target_shape), shape: target_shape, modes: target_modes, @@ -436,6 +439,19 @@ fn finalize_contraction_output( } } +enum MaterializedSlice<'a, T> { + Borrowed(&'a [T]), + Scratch(&'a [T]), +} + +impl<'a, T> MaterializedSlice<'a, T> { + fn as_slice(&self) -> &'a [T] { + match self { + Self::Borrowed(data) | Self::Scratch(data) => data, + } + } +} + /// Execute tensor contraction on CPU via reshape→GEMM→reshape. #[allow(clippy::too_many_arguments)] pub(super) fn contract( @@ -549,30 +565,39 @@ where } } - // 3. Make inputs contiguous if needed for the generic materialized path. - let a_contig = ensure_contiguous(a, shape_a, strides_a); - let b_contig = ensure_contiguous(b, shape_b, strides_b); + let mut a_pool = ScratchPool::::default(); + let mut b_pool = ScratchPool::::default(); - // 4. Reduce trace modes before generic GEMM. - let (a_data, a_shape, a_modes) = if !left_trace.is_empty() { - reduce_trace_modes::(&a_contig, shape_a, modes_a, &left_trace) + // 3. Reduce trace modes before the generic GEMM fallback. + let (a_reduced, a_shape, a_modes, a_strides) = if !left_trace.is_empty() { + let (a_data, a_shape, a_modes) = { + let mut a_contig = a_pool.acquire(shape_a.iter().product::().max(1)); + let a_contig = ensure_contiguous_into(a, shape_a, strides_a, a_contig.as_mut_vec()); + reduce_trace_modes::(a_contig.as_slice(), shape_a, modes_a, &left_trace) + }; + let a_strides = compute_contiguous_strides(&a_shape); + (Some(a_data), a_shape, a_modes, a_strides) } else { - (a_contig, shape_a.to_vec(), modes_a.to_vec()) + (None, shape_a.to_vec(), modes_a.to_vec(), strides_a.to_vec()) }; - let (b_data, b_shape, b_modes) = if !right_trace.is_empty() { - reduce_trace_modes::(&b_contig, shape_b, modes_b, &right_trace) + let (b_reduced, b_shape, b_modes, b_strides) = if !right_trace.is_empty() { + let (b_data, b_shape, b_modes) = { + let mut b_contig = b_pool.acquire(shape_b.iter().product::().max(1)); + let b_contig = ensure_contiguous_into(b, shape_b, strides_b, b_contig.as_mut_vec()); + reduce_trace_modes::(b_contig.as_slice(), shape_b, modes_b, &right_trace) + }; + let b_strides = compute_contiguous_strides(&b_shape); + (Some(b_data), b_shape, b_modes, b_strides) } else { - (b_contig, shape_b.to_vec(), modes_b.to_vec()) + (None, shape_b.to_vec(), modes_b.to_vec(), strides_b.to_vec()) }; - let a_layout_strides = compute_contiguous_strides(&a_shape); - let b_layout_strides = compute_contiguous_strides(&b_shape); let plan = analyze_contraction_layout( &a_shape, - &a_layout_strides, + &a_strides, &a_modes, &b_shape, - &b_layout_strides, + &b_strides, &b_modes, modes_c, ); @@ -584,7 +609,12 @@ where &plan.contracted_modes, &plan.batch_modes, ); - let a_permuted = permute_data(&a_data, &a_shape, &a_perm); + let mut a_permuted_scratch = a_pool.acquire(a_shape.iter().product::().max(1)); + let a_permuted = if let Some(ref a_data) = a_reduced { + permute_data_into(a_data, &a_shape, &a_perm, a_permuted_scratch.as_mut_vec()) + } else { + materialize_with_permutation_into(a, &a_shape, &a_strides, &a_perm, a_permuted_scratch.as_mut_vec()) + }; // 6. Permute B to [contracted, right_free, batch] - batch LAST let b_perm = compute_permutation( @@ -593,24 +623,29 @@ where &plan.right_modes, &plan.batch_modes, ); - let b_permuted = permute_data(&b_data, &b_shape, &b_perm); + let mut b_permuted_scratch = b_pool.acquire(b_shape.iter().product::().max(1)); + let b_permuted = if let Some(ref b_data) = b_reduced { + permute_data_into(b_data, &b_shape, &b_perm, b_permuted_scratch.as_mut_vec()) + } else { + materialize_with_permutation_into(b, &b_shape, &b_strides, &b_perm, b_permuted_scratch.as_mut_vec()) + }; // 7. Call GEMM let c_data = if plan.batch_modes.is_empty() { cpu.gemm_internal::( - &a_permuted, + a_permuted.as_slice(), plan.left_size, plan.contract_size, - &b_permuted, + b_permuted.as_slice(), plan.right_size, ) } else { cpu.gemm_batched_internal::( - &a_permuted, + a_permuted.as_slice(), plan.batch_size, plan.left_size, plan.contract_size, - &b_permuted, + b_permuted.as_slice(), plan.right_size, ) }; @@ -620,53 +655,58 @@ where /// Ensure data is contiguous (copy if strided). fn ensure_contiguous(data: &[T], shape: &[usize], strides: &[usize]) -> Vec { - let expected_strides = compute_contiguous_strides(shape); - if strides == expected_strides { - data.to_vec() - } else { - // Copy with stride handling - let numel: usize = shape.iter().product(); - let mut result = vec![T::default(); numel]; - copy_strided_to_contiguous(data, &mut result, shape, strides); - result - } + let mut scratch = Vec::new(); + ensure_contiguous_into(data, shape, strides, &mut scratch) + .as_slice() + .to_vec() } -/// Copy strided data to contiguous buffer. -fn copy_strided_to_contiguous( - src: &[T], - dst: &mut [T], +fn ensure_contiguous_into<'a, T: Copy + Default>( + data: &'a [T], shape: &[usize], strides: &[usize], -) { - let numel: usize = shape.iter().product(); - - for (i, dst_elem) in dst.iter_mut().enumerate().take(numel) { - // Convert linear index to multi-index - let mut remaining = i; - let mut src_offset = 0; - for dim in 0..shape.len() { - let coord = remaining % shape[dim]; - remaining /= shape[dim]; - src_offset += coord * strides[dim]; - } - *dst_elem = src[src_offset]; - } + scratch: &'a mut Vec, +) -> MaterializedSlice<'a, T> { + let identity: Vec = (0..shape.len()).collect(); + materialize_with_permutation_into(data, shape, strides, &identity, scratch) } /// Permute data according to axis permutation. fn permute_data(data: &[T], shape: &[usize], perm: &[usize]) -> Vec { - if perm.iter().enumerate().all(|(i, &p)| i == p) { - return data.to_vec(); // Already in correct order + let mut scratch = Vec::new(); + permute_data_into(data, shape, perm, &mut scratch) + .as_slice() + .to_vec() +} + +fn permute_data_into<'a, T: Copy + Default>( + data: &'a [T], + shape: &[usize], + perm: &[usize], + scratch: &'a mut Vec, +) -> MaterializedSlice<'a, T> { + let strides = compute_contiguous_strides(shape); + materialize_with_permutation_into(data, shape, &strides, perm, scratch) +} + +fn materialize_with_permutation_into<'a, T: Copy + Default>( + data: &'a [T], + shape: &[usize], + strides: &[usize], + perm: &[usize], + scratch: &'a mut Vec, +) -> MaterializedSlice<'a, T> { + let expected_strides = compute_contiguous_strides(shape); + if strides == expected_strides && perm.iter().enumerate().all(|(i, &p)| i == p) { + return MaterializedSlice::Borrowed(data); } - let new_shape: Vec = perm.iter().map(|&p| shape[p]).collect(); + scratch.clear(); let numel: usize = shape.iter().product(); - let mut result = vec![T::default(); numel]; - - let old_strides = compute_contiguous_strides(shape); + scratch.resize(numel, T::default()); + let new_shape: Vec = perm.iter().map(|&p| shape[p]).collect(); - for (new_idx, result_elem) in result.iter_mut().enumerate().take(numel) { + for (new_idx, result_elem) in scratch.iter_mut().enumerate().take(numel) { // Convert new linear index to new multi-index let mut remaining = new_idx; let mut new_coords = vec![0; shape.len()]; @@ -678,13 +718,13 @@ fn permute_data(data: &[T], shape: &[usize], perm: &[usize]) // Map to old coordinates via inverse permutation let mut old_idx = 0; for (new_dim, &old_dim) in perm.iter().enumerate() { - old_idx += new_coords[new_dim] * old_strides[old_dim]; + old_idx += new_coords[new_dim] * strides[old_dim]; } *result_elem = data[old_idx]; } - result + MaterializedSlice::Scratch(scratch.as_slice()) } /// Execute tensor contraction with argmax tracking. diff --git a/src/backend/cpu/mod.rs b/src/backend/cpu/mod.rs index 28610dc..727df7e 100644 --- a/src/backend/cpu/mod.rs +++ b/src/backend/cpu/mod.rs @@ -1,5 +1,6 @@ //! CPU backend implementation. +mod buffer_pool; mod contract; use super::traits::{Backend, BackendScalar, Storage}; From c280b9875e171a683882a13c539a01bbd1594fdb Mon Sep 17 00:00:00 2001 From: GiggleLiu Date: Sun, 12 Apr 2026 15:53:15 +0800 Subject: [PATCH 5/6] perf(einsum): contract optimized roots in final order --- src/einsum/engine.rs | 122 ++++++++++++++++++++++++++++++++++++------- src/tensor/mod.rs | 1 + src/tensor/ops.rs | 90 ++++++++++++++++++++++++++++++- 3 files changed, 192 insertions(+), 21 deletions(-) diff --git a/src/einsum/engine.rs b/src/einsum/engine.rs index ef7a289..78854fc 100644 --- a/src/einsum/engine.rs +++ b/src/einsum/engine.rs @@ -6,7 +6,7 @@ use omeco::{optimize_code, EinCode, GreedyMethod, Label, NestedEinsum, TreeSA}; use crate::algebra::{Algebra, Scalar}; use crate::backend::{Backend, BackendScalar}; -use crate::tensor::Tensor; +use crate::tensor::{BinaryContractOptions, Tensor}; /// Einsum specification and execution engine. /// @@ -142,8 +142,18 @@ impl Einsum { &self.size_dict, ) } else { - let result = self.execute_tree::(tree, tensors); - finalize_optimized_result::(result, tree, &self.iy, &self.size_dict) + let emit_final_root_output = + can_emit_final_root_output(optimized_tree_output(tree), &self.iy); + let result = self.execute_tree::( + tree, + tensors, + emit_final_root_output.then_some(self.iy.as_slice()), + ); + if emit_final_root_output { + result + } else { + finalize_optimized_result::(result, tree, &self.iy, &self.size_dict) + } } } None => self.execute_pairwise::(tensors), @@ -195,15 +205,25 @@ impl Einsum { ) } } else { - let result = - self.execute_tree_with_argmax::(tree, tensors, &mut argmax_cache); - finalize_optimized_result_with_argmax::( - result, + let emit_final_root_output = + can_emit_final_root_output(optimized_tree_output(tree), &self.iy); + let result = self.execute_tree_with_argmax::( tree, - &self.iy, - &self.size_dict, + tensors, &mut argmax_cache, - ) + emit_final_root_output.then_some(self.iy.as_slice()), + ); + if emit_final_root_output { + result + } else { + finalize_optimized_result_with_argmax::( + result, + tree, + &self.iy, + &self.size_dict, + &mut argmax_cache, + ) + } } } None => self.execute_pairwise_with_argmax::(tensors, &mut argmax_cache), @@ -219,6 +239,7 @@ impl Einsum { tree: &NestedEinsum, tensors: &[&Tensor], argmax_cache: &mut Vec>, + preferred_output_indices: Option<&[usize]>, ) -> Tensor where A: Algebra, @@ -230,10 +251,18 @@ impl Einsum { NestedEinsum::Node { args, eins } => { assert_eq!(args.len(), 2, "Expected binary contraction tree"); - let left = - self.execute_tree_with_argmax::(&args[0], tensors, argmax_cache); - let right = - self.execute_tree_with_argmax::(&args[1], tensors, argmax_cache); + let left = self.execute_tree_with_argmax::( + &args[0], + tensors, + argmax_cache, + None, + ); + let right = self.execute_tree_with_argmax::( + &args[1], + tensors, + argmax_cache, + None, + ); let ia = &eins.ixs[0]; let ib = &eins.ixs[1]; @@ -249,12 +278,29 @@ impl Einsum { ); if A::needs_argmax() { - let (result, argmax) = - left.contract_binary_with_argmax::(&right, &ia, &ib, iy); + let (result, argmax) = if let Some(preferred_output_indices) = + preferred_output_indices + { + let options = BinaryContractOptions { + preferred_output_indices: Some(preferred_output_indices.to_vec()), + }; + left.contract_binary_with_argmax_with_options::( + &right, &ia, &ib, iy, &options, + ) + } else { + left.contract_binary_with_argmax::(&right, &ia, &ib, iy) + }; argmax_cache.push(argmax); result } else { - left.contract_binary::(&right, &ia, &ib, iy) + if let Some(preferred_output_indices) = preferred_output_indices { + let options = BinaryContractOptions { + preferred_output_indices: Some(preferred_output_indices.to_vec()), + }; + left.contract_binary_with_options::(&right, &ia, &ib, iy, &options) + } else { + left.contract_binary::(&right, &ia, &ib, iy) + } } } } @@ -353,6 +399,7 @@ impl Einsum { &self, tree: &NestedEinsum, tensors: &[&Tensor], + preferred_output_indices: Option<&[usize]>, ) -> Tensor where A: Algebra, @@ -364,8 +411,8 @@ impl Einsum { NestedEinsum::Node { args, eins } => { assert_eq!(args.len(), 2, "Expected binary contraction tree"); - let left = self.execute_tree::(&args[0], tensors); - let right = self.execute_tree::(&args[1], tensors); + let left = self.execute_tree::(&args[0], tensors, None); + let right = self.execute_tree::(&args[1], tensors, None); let ia = &eins.ixs[0]; let ib = &eins.ixs[1]; @@ -380,7 +427,14 @@ impl Einsum { &self.size_dict, ); - left.contract_binary::(&right, &ia, &ib, iy) + if let Some(preferred_output_indices) = preferred_output_indices { + let options = BinaryContractOptions { + preferred_output_indices: Some(preferred_output_indices.to_vec()), + }; + left.contract_binary_with_options::(&right, &ia, &ib, iy, &options) + } else { + left.contract_binary::(&right, &ia, &ib, iy) + } } } } @@ -526,6 +580,18 @@ fn optimized_tree_output(tree: &NestedEinsum) -> &[usize] { } } +fn can_emit_final_root_output(tree_output: &[usize], final_output: &[usize]) -> bool { + if tree_output.len() != final_output.len() { + return false; + } + + let tree_set: HashSet = tree_output.iter().copied().collect(); + let final_set: HashSet = final_output.iter().copied().collect(); + tree_set.len() == tree_output.len() + && final_set.len() == final_output.len() + && tree_set == final_set +} + fn finalize_optimized_result( result: Tensor, tree: &NestedEinsum, @@ -1001,6 +1067,22 @@ mod tests { // Tests for helper functions + #[test] + fn test_root_output_plan_uses_final_order_for_pure_permutation() { + let tree_output = vec![2, 0, 1]; + let final_output = vec![0, 1, 2]; + + assert!(can_emit_final_root_output(&tree_output, &final_output)); + } + + #[test] + fn test_root_output_plan_rejects_non_permutation_finalize_cases() { + let tree_output = vec![0, 1, 2]; + let final_output = vec![0, 0, 2]; + + assert!(!can_emit_final_root_output(&tree_output, &final_output)); + } + #[test] fn test_linear_to_multi_empty_shape() { // Empty shape should return empty multi-index diff --git a/src/tensor/mod.rs b/src/tensor/mod.rs index 73474ed..fa3d1e0 100644 --- a/src/tensor/mod.rs +++ b/src/tensor/mod.rs @@ -14,6 +14,7 @@ use crate::algebra::{Algebra, Scalar}; use crate::backend::{Backend, Storage}; pub use view::TensorView; +pub(crate) use ops::BinaryContractOptions; /// A multi-dimensional tensor with stride-based layout. /// diff --git a/src/tensor/ops.rs b/src/tensor/ops.rs index 3c9b252..2bf9087 100644 --- a/src/tensor/ops.rs +++ b/src/tensor/ops.rs @@ -4,6 +4,11 @@ use super::Tensor; use crate::algebra::{Algebra, Scalar}; use crate::backend::{Backend, BackendScalar}; +#[derive(Default)] +pub(crate) struct BinaryContractOptions { + pub preferred_output_indices: Option>, +} + /// Compute output shape from input shapes and modes. fn compute_output_shape( shape_a: &[usize], @@ -73,6 +78,52 @@ impl Tensor { (result, argmax.expect("argmax requested but not returned")) } + pub(crate) fn contract_binary_with_options>( + &self, + other: &Self, + ia: &[usize], + ib: &[usize], + iy: &[usize], + options: &BinaryContractOptions, + ) -> Self + where + T: BackendScalar, + { + let (result, _) = self.contract_binary_impl_with_options::( + other, + ia, + ib, + iy, + false, + options, + ); + result + } + + pub(crate) fn contract_binary_with_argmax_with_options< + A: Algebra, + >( + &self, + other: &Self, + ia: &[usize], + ib: &[usize], + iy: &[usize], + options: &BinaryContractOptions, + ) -> (Self, Tensor) + where + T: BackendScalar, + { + let (result, argmax) = self.contract_binary_impl_with_options::( + other, + ia, + ib, + iy, + true, + options, + ); + (result, argmax.expect("argmax requested but not returned")) + } + fn contract_binary_impl>( &self, other: &Self, @@ -81,16 +132,53 @@ impl Tensor { iy: &[usize], track_argmax: bool, ) -> (Self, Option>) + where + T: BackendScalar, + { + self.contract_binary_impl_with_options::( + other, + ia, + ib, + iy, + track_argmax, + &BinaryContractOptions::default(), + ) + } + + pub(crate) fn contract_binary_impl_with_options>( + &self, + other: &Self, + ia: &[usize], + ib: &[usize], + iy: &[usize], + track_argmax: bool, + options: &BinaryContractOptions, + ) -> (Self, Option>) where T: BackendScalar, { assert_eq!(ia.len(), self.ndim(), "ia length must match self.ndim()"); assert_eq!(ib.len(), other.ndim(), "ib length must match other.ndim()"); + let output_indices = options + .preferred_output_indices + .as_deref() + .unwrap_or(iy); + if let Some(preferred_output_indices) = &options.preferred_output_indices { + let mut preferred_sorted = preferred_output_indices.clone(); + preferred_sorted.sort_unstable(); + let mut output_sorted = iy.to_vec(); + output_sorted.sort_unstable(); + debug_assert_eq!( + preferred_sorted, output_sorted, + "preferred output indices must be a permutation of iy" + ); + } + // Convert usize indices to i32 modes let modes_a: Vec = ia.iter().map(|&i| i as i32).collect(); let modes_b: Vec = ib.iter().map(|&i| i as i32).collect(); - let modes_c: Vec = iy.iter().map(|&i| i as i32).collect(); + let modes_c: Vec = output_indices.iter().map(|&i| i as i32).collect(); // Compute output shape let shape_c = From 143b94d4f41534b599bba9b1e3643d458b9b7baf Mon Sep 17 00:00:00 2001 From: GiggleLiu Date: Sun, 12 Apr 2026 15:55:44 +0800 Subject: [PATCH 6/6] perf(cpu): reduce contraction materialization overhead --- src/backend/cpu/contract.rs | 70 +++++++++++++++++++++++++++---------- src/backend/cpu/mod.rs | 28 ++++++++++++--- src/einsum/engine.rs | 58 ++++++++++++++---------------- src/tensor/mod.rs | 2 +- src/tensor/ops.rs | 29 ++++----------- 5 files changed, 108 insertions(+), 79 deletions(-) diff --git a/src/backend/cpu/contract.rs b/src/backend/cpu/contract.rs index 98e9a57..1712020 100644 --- a/src/backend/cpu/contract.rs +++ b/src/backend/cpu/contract.rs @@ -116,7 +116,11 @@ fn is_flattenable_group( return true; } - let mut expected = if require_unit_base { 1 } else { strides[axes[0]] }; + let mut expected = if require_unit_base { + 1 + } else { + strides[axes[0]] + }; if require_unit_base && strides[axes[0]] != 1 { return false; } @@ -154,8 +158,14 @@ fn analyze_operand_materialization( .iter() .map(|&mode| mode_position(modes, mode)) .collect(); - let row_axes: Vec = row_modes.iter().map(|&mode| mode_position(modes, mode)).collect(); - let col_axes: Vec = col_modes.iter().map(|&mode| mode_position(modes, mode)).collect(); + let row_axes: Vec = row_modes + .iter() + .map(|&mode| mode_position(modes, mode)) + .collect(); + let col_axes: Vec = col_modes + .iter() + .map(|&mode| mode_position(modes, mode)) + .collect(); let physical_order = physical_axis_order(strides); let mut no_copy_orders = vec![batch_axes @@ -489,8 +499,9 @@ where .collect(); if left_trace.is_empty() && right_trace.is_empty() { - let plan = - analyze_contraction_layout(shape_a, strides_a, modes_a, shape_b, strides_b, modes_b, modes_c); + let plan = analyze_contraction_layout( + shape_a, strides_a, modes_a, shape_b, strides_b, modes_b, modes_c, + ); let left_nocopy = matches!(plan.left_materialization, MaterializationPlan::NoCopy); let right_nocopy = matches!(plan.right_materialization, MaterializationPlan::NoCopy); if left_nocopy || right_nocopy { @@ -593,13 +604,7 @@ where }; let plan = analyze_contraction_layout( - &a_shape, - &a_strides, - &a_modes, - &b_shape, - &b_strides, - &b_modes, - modes_c, + &a_shape, &a_strides, &a_modes, &b_shape, &b_strides, &b_modes, modes_c, ); // 5. Permute A to [left_free, contracted, batch] - batch LAST for current GEMM layout @@ -613,7 +618,13 @@ where let a_permuted = if let Some(ref a_data) = a_reduced { permute_data_into(a_data, &a_shape, &a_perm, a_permuted_scratch.as_mut_vec()) } else { - materialize_with_permutation_into(a, &a_shape, &a_strides, &a_perm, a_permuted_scratch.as_mut_vec()) + materialize_with_permutation_into( + a, + &a_shape, + &a_strides, + &a_perm, + a_permuted_scratch.as_mut_vec(), + ) }; // 6. Permute B to [contracted, right_free, batch] - batch LAST @@ -627,7 +638,13 @@ where let b_permuted = if let Some(ref b_data) = b_reduced { permute_data_into(b_data, &b_shape, &b_perm, b_permuted_scratch.as_mut_vec()) } else { - materialize_with_permutation_into(b, &b_shape, &b_strides, &b_perm, b_permuted_scratch.as_mut_vec()) + materialize_with_permutation_into( + b, + &b_shape, + &b_strides, + &b_perm, + b_permuted_scratch.as_mut_vec(), + ) }; // 7. Call GEMM @@ -897,8 +914,14 @@ mod tests { assert_eq!(plan.left_modes, vec![0]); assert_eq!(plan.right_modes, vec![2]); assert_eq!(plan.contracted_modes, vec![1]); - assert!(matches!(plan.left_materialization, MaterializationPlan::NoCopy)); - assert!(matches!(plan.right_materialization, MaterializationPlan::NoCopy)); + assert!(matches!( + plan.left_materialization, + MaterializationPlan::NoCopy + )); + assert!(matches!( + plan.right_materialization, + MaterializationPlan::NoCopy + )); } #[test] @@ -913,7 +936,10 @@ mod tests { &[0, 1, 3], ); - assert!(matches!(plan.left_materialization, MaterializationPlan::NoCopy)); + assert!(matches!( + plan.left_materialization, + MaterializationPlan::NoCopy + )); assert!(matches!( plan.right_materialization, MaterializationPlan::Permute { .. } @@ -934,7 +960,13 @@ mod tests { assert_eq!(plan.batch_modes, vec![0]); assert_eq!(plan.batch_size, 2); - assert!(matches!(plan.left_materialization, MaterializationPlan::NoCopy)); - assert!(matches!(plan.right_materialization, MaterializationPlan::NoCopy)); + assert!(matches!( + plan.left_materialization, + MaterializationPlan::NoCopy + )); + assert!(matches!( + plan.right_materialization, + MaterializationPlan::NoCopy + )); } } diff --git a/src/backend/cpu/mod.rs b/src/backend/cpu/mod.rs index 727df7e..cf1615b 100644 --- a/src/backend/cpu/mod.rs +++ b/src/backend/cpu/mod.rs @@ -65,7 +65,10 @@ fn layout_offset_bounds(layout: &MatrixLayout<'_, T>) -> (isize, isize) { fn faer_mat_ref<'a, T>(layout: MatrixLayout<'a, T>) -> faer::MatRef<'a, T> { let (min_offset, max_offset) = layout_offset_bounds(&layout); if layout.rows > 0 && layout.cols > 0 { - assert!(!layout.data.is_empty(), "matrix layout requires backing storage"); + assert!( + !layout.data.is_empty(), + "matrix layout requires backing storage" + ); assert!( max_offset >= min_offset, "matrix layout offsets must be ordered" @@ -88,7 +91,10 @@ fn faer_mat_ref<'a, T>(layout: MatrixLayout<'a, T>) -> faer::MatRef<'a, T> { } } -fn matrix_layout_batch_view<'a, T>(layout: MatrixLayout<'a, T>, batch: usize) -> MatrixLayout<'a, T> { +fn matrix_layout_batch_view<'a, T>( + layout: MatrixLayout<'a, T>, + batch: usize, +) -> MatrixLayout<'a, T> { assert!(batch < layout.data.len(), "batch offset must be in bounds"); MatrixLayout { data: &layout.data[batch..], @@ -559,7 +565,14 @@ fn faer_gemm_f32_layout(a: MatrixLayout<'_, f32>, b: MatrixLayout<'_, f32>) -> V let a_mat = faer_mat_ref(a); let b_mat = faer_mat_ref(b); let mut c_mat = Mat::::zeros(a.rows, b.cols); - matmul(c_mat.as_mut(), Accum::Replace, a_mat, b_mat, 1.0f32, Par::Seq); + matmul( + c_mat.as_mut(), + Accum::Replace, + a_mat, + b_mat, + 1.0f32, + Par::Seq, + ); let mut c = vec![0.0f32; a.rows * b.cols]; for j in 0..b.cols { @@ -594,7 +607,14 @@ fn faer_gemm_f64_layout(a: MatrixLayout<'_, f64>, b: MatrixLayout<'_, f64>) -> V let a_mat = faer_mat_ref(a); let b_mat = faer_mat_ref(b); let mut c_mat = Mat::::zeros(a.rows, b.cols); - matmul(c_mat.as_mut(), Accum::Replace, a_mat, b_mat, 1.0f64, Par::Seq); + matmul( + c_mat.as_mut(), + Accum::Replace, + a_mat, + b_mat, + 1.0f64, + Par::Seq, + ); let mut c = vec![0.0f64; a.rows * b.cols]; for j in 0..b.cols { diff --git a/src/einsum/engine.rs b/src/einsum/engine.rs index 78854fc..7006dc7 100644 --- a/src/einsum/engine.rs +++ b/src/einsum/engine.rs @@ -152,7 +152,12 @@ impl Einsum { if emit_final_root_output { result } else { - finalize_optimized_result::(result, tree, &self.iy, &self.size_dict) + finalize_optimized_result::( + result, + tree, + &self.iy, + &self.size_dict, + ) } } } @@ -251,18 +256,10 @@ impl Einsum { NestedEinsum::Node { args, eins } => { assert_eq!(args.len(), 2, "Expected binary contraction tree"); - let left = self.execute_tree_with_argmax::( - &args[0], - tensors, - argmax_cache, - None, - ); - let right = self.execute_tree_with_argmax::( - &args[1], - tensors, - argmax_cache, - None, - ); + let left = + self.execute_tree_with_argmax::(&args[0], tensors, argmax_cache, None); + let right = + self.execute_tree_with_argmax::(&args[1], tensors, argmax_cache, None); let ia = &eins.ixs[0]; let ib = &eins.ixs[1]; @@ -278,29 +275,26 @@ impl Einsum { ); if A::needs_argmax() { - let (result, argmax) = if let Some(preferred_output_indices) = - preferred_output_indices - { - let options = BinaryContractOptions { - preferred_output_indices: Some(preferred_output_indices.to_vec()), + let (result, argmax) = + if let Some(preferred_output_indices) = preferred_output_indices { + let options = BinaryContractOptions { + preferred_output_indices: Some(preferred_output_indices.to_vec()), + }; + left.contract_binary_with_argmax_with_options::( + &right, &ia, &ib, iy, &options, + ) + } else { + left.contract_binary_with_argmax::(&right, &ia, &ib, iy) }; - left.contract_binary_with_argmax_with_options::( - &right, &ia, &ib, iy, &options, - ) - } else { - left.contract_binary_with_argmax::(&right, &ia, &ib, iy) - }; argmax_cache.push(argmax); result + } else if let Some(preferred_output_indices) = preferred_output_indices { + let options = BinaryContractOptions { + preferred_output_indices: Some(preferred_output_indices.to_vec()), + }; + left.contract_binary_with_options::(&right, &ia, &ib, iy, &options) } else { - if let Some(preferred_output_indices) = preferred_output_indices { - let options = BinaryContractOptions { - preferred_output_indices: Some(preferred_output_indices.to_vec()), - }; - left.contract_binary_with_options::(&right, &ia, &ib, iy, &options) - } else { - left.contract_binary::(&right, &ia, &ib, iy) - } + left.contract_binary::(&right, &ia, &ib, iy) } } } diff --git a/src/tensor/mod.rs b/src/tensor/mod.rs index fa3d1e0..5eee1a5 100644 --- a/src/tensor/mod.rs +++ b/src/tensor/mod.rs @@ -13,8 +13,8 @@ use std::sync::Arc; use crate::algebra::{Algebra, Scalar}; use crate::backend::{Backend, Storage}; -pub use view::TensorView; pub(crate) use ops::BinaryContractOptions; +pub use view::TensorView; /// A multi-dimensional tensor with stride-based layout. /// diff --git a/src/tensor/ops.rs b/src/tensor/ops.rs index 2bf9087..41c1391 100644 --- a/src/tensor/ops.rs +++ b/src/tensor/ops.rs @@ -89,20 +89,12 @@ impl Tensor { where T: BackendScalar, { - let (result, _) = self.contract_binary_impl_with_options::( - other, - ia, - ib, - iy, - false, - options, - ); + let (result, _) = + self.contract_binary_impl_with_options::(other, ia, ib, iy, false, options); result } - pub(crate) fn contract_binary_with_argmax_with_options< - A: Algebra, - >( + pub(crate) fn contract_binary_with_argmax_with_options>( &self, other: &Self, ia: &[usize], @@ -113,14 +105,8 @@ impl Tensor { where T: BackendScalar, { - let (result, argmax) = self.contract_binary_impl_with_options::( - other, - ia, - ib, - iy, - true, - options, - ); + let (result, argmax) = + self.contract_binary_impl_with_options::(other, ia, ib, iy, true, options); (result, argmax.expect("argmax requested but not returned")) } @@ -160,10 +146,7 @@ impl Tensor { assert_eq!(ia.len(), self.ndim(), "ia length must match self.ndim()"); assert_eq!(ib.len(), other.ndim(), "ib length must match other.ndim()"); - let output_indices = options - .preferred_output_indices - .as_deref() - .unwrap_or(iy); + let output_indices = options.preferred_output_indices.as_deref().unwrap_or(iy); if let Some(preferred_output_indices) = &options.preferred_output_indices { let mut preferred_sorted = preferred_output_indices.clone(); preferred_sorted.sort_unstable();