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 0fb1766..1712020 100644 --- a/src/backend/cpu/contract.rs +++ b/src/backend/cpu/contract.rs @@ -78,6 +78,226 @@ 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 super::buffer_pool::ScratchPool; +use super::MatrixLayout; use crate::algebra::Algebra; use crate::backend::Cpu; use crate::tensor::compute_contiguous_strides; @@ -142,6 +362,106 @@ 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), + } +} + +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], +) -> MaterializedMatrixOperand { + 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(); + let mut scratch = Vec::new(); + + MaterializedMatrixOperand { + 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, + } +} + +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 + } +} + +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( @@ -160,14 +480,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); + // 1. Classify modes + let (_batch, left, right, _contracted) = classify_modes(modes_a, modes_b, modes_c); - // 2. 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(); @@ -182,132 +498,232 @@ where .copied() .collect(); - let (a_data, a_shape, a_modes) = if !left_trace.is_empty() { - reduce_trace_modes::(&a_contig, shape_a, modes_a, &left_trace) + 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 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); + } + } + } + + let mut a_pool = ScratchPool::::default(); + let mut b_pool = ScratchPool::::default(); + + // 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()) }; - // 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_permuted = permute_data(&a_data, &a_shape, &a_perm); + let plan = analyze_contraction_layout( + &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 + let a_perm = compute_permutation( + &a_modes, + &plan.left_modes, + &plan.contracted_modes, + &plan.batch_modes, + ); + 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(&b_modes, &contracted, &right_free, &batch); - let b_permuted = permute_data(&b_data, &b_shape, &b_perm); + let b_perm = compute_permutation( + &b_modes, + &plan.contracted_modes, + &plan.right_modes, + &plan.batch_modes, + ); + 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 batch.is_empty() { + let c_data = if plan.batch_modes.is_empty() { cpu.gemm_internal::( - &a_permuted, - left_size, - contract_size, - &b_permuted, - right_size, + a_permuted.as_slice(), + plan.left_size, + plan.contract_size, + b_permuted.as_slice(), + plan.right_size, ) } else { cpu.gemm_batched_internal::( - &a_permuted, - batch_size, - left_size, - contract_size, - &b_permuted, - right_size, + a_permuted.as_slice(), + plan.batch_size, + plan.left_size, + plan.contract_size, + b_permuted.as_slice(), + plan.right_size, ) }; - // 8. Permute result to output order - // Result is in [left_free, right_free, batch] order - let current_order: Vec = left_free - .iter() - .chain(right_free.iter()) - .chain(batch.iter()) - .copied() - .collect(); - - if current_order == modes_c { - 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) - } + finalize_contraction_output(c_data, &plan, shape_c, modes_c) } /// 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()]; @@ -319,13 +735,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. @@ -481,4 +897,76 @@ 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 { .. } + )); + } + + #[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 2d60ba5..cf1615b 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}; @@ -10,7 +11,211 @@ 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> { + #[cfg(test)] + 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, + ) + } +} + +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, + 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 + } + + 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 @@ -354,6 +559,30 @@ 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 +601,30 @@ 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 +915,49 @@ 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); + } + + #[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/src/einsum/engine.rs b/src/einsum/engine.rs index ef7a289..7006dc7 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,23 @@ 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 +210,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 +244,7 @@ impl Einsum { tree: &NestedEinsum, tensors: &[&Tensor], argmax_cache: &mut Vec>, + preferred_output_indices: Option<&[usize]>, ) -> Tensor where A: Algebra, @@ -231,9 +257,9 @@ impl Einsum { assert_eq!(args.len(), 2, "Expected binary contraction tree"); let left = - self.execute_tree_with_argmax::(&args[0], tensors, argmax_cache); + self.execute_tree_with_argmax::(&args[0], tensors, argmax_cache, None); let right = - self.execute_tree_with_argmax::(&args[1], tensors, argmax_cache); + self.execute_tree_with_argmax::(&args[1], tensors, argmax_cache, None); let ia = &eins.ixs[0]; let ib = &eins.ixs[1]; @@ -250,9 +276,23 @@ impl Einsum { if A::needs_argmax() { let (result, argmax) = - left.contract_binary_with_argmax::(&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_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 { left.contract_binary::(&right, &ia, &ib, iy) } @@ -353,6 +393,7 @@ impl Einsum { &self, tree: &NestedEinsum, tensors: &[&Tensor], + preferred_output_indices: Option<&[usize]>, ) -> Tensor where A: Algebra, @@ -364,8 +405,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 +421,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 +574,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 +1061,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..5eee1a5 100644 --- a/src/tensor/mod.rs +++ b/src/tensor/mod.rs @@ -13,6 +13,7 @@ use std::sync::Arc; use crate::algebra::{Algebra, Scalar}; use crate::backend::{Backend, Storage}; +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 3c9b252..41c1391 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,38 @@ 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>( + &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 +118,50 @@ 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 = diff --git a/tests/suites/backend_contract.rs b/tests/suites/backend_contract.rs index 31ddbd0..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")] @@ -198,6 +197,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 // ============================================================================ @@ -235,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::>( @@ -255,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]); } // ============================================================================