use criterion::{Criterion, criterion_group, criterion_main};
use ndarray::{Array2, Array3, Axis, Zip};
use std::{f64::consts::PI, hint::black_box};
use tensorism::new_ndarray;
use utils::C64;
struct SumAccu {
value: C64,
}
impl SumAccu {
fn new() -> Self {
Self { value: C64::ZERO }
}
fn accumulate(&mut self, x: C64) {
self.value = self.value + x;
}
fn get(&self) -> C64 {
self.value
}
}
fn pure_ndarray(t_a: &Array3<C64>, t_b: &Array3<C64>) -> Array2<C64> {
let intermediate: Array3<C64> = Zip::from(t_a).and(t_b).map_collect(|&a_val, &b_val| {
if a_val > C64::ZERO {
a_val * b_val
} else {
(C64::ONE + b_val.exp()).ln()
}
});
intermediate.sum_axis(Axis(1)).t().to_owned()
}
fn ndarray_non_custom(t_a: &Array3<C64>, t_b: &Array3<C64>) -> Array2<C64> {
let intermediate: Array3<C64> = Zip::from(t_a).and(t_b).map_collect(|&a_val, &b_val| {
if a_val > C64::ZERO {
a_val * b_val
} else {
(C64::ONE + b_val.exp()).ln()
}
});
intermediate
.map_axis(Axis(1), |row| row.iter().map(|x| *x).sum())
.t()
.to_owned()
}
fn tensorism(t_a: &Array3<C64>, t_b: &Array3<C64>) -> Array2<C64> {
new_ndarray! {for k i => (for j =>
let elt_a = t_a[i, j, k];
let elt_b = t_b[i, j, k];
if elt_a > C64::ZERO {
elt_a * elt_b
} else {
(C64::ONE + elt_b.exp()).ln()
}).sum::<C64>()
}
}
fn flatten_factorized(t_a: &Array3<C64>, t_b: &Array3<C64>) -> Array2<C64> {
let dim_number_2 = ::ndarray::ArrayBase::<_, _>::dim(&t_a).1;
if dim_number_2 != ::ndarray::ArrayBase::<_, _>::dim(&t_b).1 {
panic!("Dimensions are not matching between t_a[_, j, _] and t_b[_, j, _]");
}
let dim_number_0 = ::ndarray::ArrayBase::<_, _>::dim(&t_a).2;
if dim_number_0 != ::ndarray::ArrayBase::<_, _>::dim(&t_b).2 {
panic!("Dimensions are not matching between t_a[_, _, k] and t_b[_, _, k]");
}
let dim_number_1 = ::ndarray::ArrayBase::<_, _>::dim(&t_a).0;
if dim_number_1 != ::ndarray::ArrayBase::<_, _>::dim(&t_b).0 {
panic!("Dimensions are not matching between t_a[i, _, _] and t_b[i, _, _]");
}
let strides_a = t_a.strides();
if strides_a != [20isize, 4isize, 1isize] {
panic!("t_a has unexpected memory layout {:?}", strides_a);
}
let strides_b = t_b.strides();
if strides_b != [20isize, 4isize, 1isize] {
panic!("t_b has unexpected memory layout {:?}", strides_b);
}
let stride_a0 = strides_a[0];
let stride_a1 = strides_a[1];
let stride_a2 = strides_a[2];
let stride_b0 = strides_b[0];
let stride_b1 = strides_b[1];
let stride_b2 = strides_b[2];
let p_a = t_a.as_ptr();
let p_b = t_b.as_ptr();
::ndarray::Array::<_, ::ndarray::Dim<[::ndarray::Ix; 2usize]>>::from_shape_fn(
(dim_number_0, dim_number_1),
|(k, i)| {
let pre_a = (i as isize) * stride_a0 + (k as isize) * stride_a2;
let pre_b = (i as isize) * stride_b0 + (k as isize) * stride_b2;
((0usize..dim_number_2).map(|j| {
let elt_a: C64 = unsafe { *p_a.offset(pre_a + (j as isize) * stride_a1) };
let elt_b: C64 = unsafe { *p_b.offset(pre_b + (j as isize) * stride_b1) };
if elt_a > C64::ZERO {
elt_a * elt_b
} else {
(C64::ONE + elt_b.exp()).ln()
}
}))
.sum::<C64>()
},
)
}
fn flatten_normal(t_a: &Array3<C64>, t_b: &Array3<C64>) -> Array2<C64> {
let dim_number_2 = ::ndarray::ArrayBase::<_, _>::dim(&t_a).1;
if dim_number_2 != ::ndarray::ArrayBase::<_, _>::dim(&t_b).1 {
panic!("Dimensions are not matching between t_a[_, j, _] and t_b[_, j, _]");
}
let dim_number_0 = ::ndarray::ArrayBase::<_, _>::dim(&t_a).2;
if dim_number_0 != ::ndarray::ArrayBase::<_, _>::dim(&t_b).2 {
panic!("Dimensions are not matching between t_a[_, _, k] and t_b[_, _, k]");
}
let dim_number_1 = ::ndarray::ArrayBase::<_, _>::dim(&t_a).0;
if dim_number_1 != ::ndarray::ArrayBase::<_, _>::dim(&t_b).0 {
panic!("Dimensions are not matching between t_a[i, _, _] and t_b[i, _, _]");
}
let strides_a = t_a.strides();
if strides_a != [20isize, 4isize, 1isize] {
panic!("t_a has unexpected memory layout {:?}", strides_a);
}
let strides_b = t_b.strides();
if strides_b != [20isize, 4isize, 1isize] {
panic!("t_b has unexpected memory layout {:?}", strides_b);
}
let stride_a0 = strides_a[0];
let stride_a1 = strides_a[1];
let stride_a2 = strides_a[2];
let stride_b0 = strides_b[0];
let stride_b1 = strides_b[1];
let stride_b2 = strides_b[2];
let p_a = t_a.as_ptr();
let p_b = t_b.as_ptr();
let mut result = ::ndarray::Array::<C64, ::ndarray::Dim<[::ndarray::Ix; 2usize]>>::uninit((
dim_number_0,
dim_number_1,
));
let mut ptr_on_result = result.as_mut_ptr() as *mut C64;
for k in 0..dim_number_0 {
for i in 0..dim_number_1 {
let value = ((0usize..dim_number_2).map(|j| {
let elt_a: C64 = unsafe {
*p_a.offset(
(i as isize) * stride_a0
+ (k as isize) * stride_a2
+ (j as isize) * stride_a1,
)
};
let elt_b: C64 = unsafe {
*p_b.offset(
(i as isize) * stride_b0
+ (k as isize) * stride_b2
+ (j as isize) * stride_b1,
)
};
if elt_a > C64::ZERO {
elt_a * elt_b
} else {
(C64::ONE + elt_b.exp()).ln()
}
}))
.sum::<C64>();
unsafe {
ptr_on_result.write(value);
ptr_on_result = ptr_on_result.add(1);
}
}
}
unsafe { result.assume_init() }
}
fn flatten_accumulator(t_a: &Array3<C64>, t_b: &Array3<C64>) -> Array2<C64> {
let dim_number_2 = ::ndarray::ArrayBase::<_, _>::dim(&t_a).1;
if dim_number_2 != ::ndarray::ArrayBase::<_, _>::dim(&t_b).1 {
panic!("Dimensions are not matching between t_a[_, j, _] and t_b[_, j, _]");
}
let dim_number_0 = ::ndarray::ArrayBase::<_, _>::dim(&t_a).2;
if dim_number_0 != ::ndarray::ArrayBase::<_, _>::dim(&t_b).2 {
panic!("Dimensions are not matching between t_a[_, _, k] and t_b[_, _, k]");
}
let dim_number_1 = ::ndarray::ArrayBase::<_, _>::dim(&t_a).0;
if dim_number_1 != ::ndarray::ArrayBase::<_, _>::dim(&t_b).0 {
panic!("Dimensions are not matching between t_a[i, _, _] and t_b[i, _, _]");
}
let strides_a = t_a.strides();
if strides_a != [20isize, 4isize, 1isize] {
panic!("t_a has unexpected memory layout {:?}", strides_a);
}
let strides_b = t_b.strides();
if strides_b != [20isize, 4isize, 1isize] {
panic!("t_b has unexpected memory layout {:?}", strides_b);
}
let stride_a0 = strides_a[0];
let stride_a1 = strides_a[1];
let stride_a2 = strides_a[2];
let stride_b0 = strides_b[0];
let stride_b1 = strides_b[1];
let stride_b2 = strides_b[2];
let p_a = t_a.as_ptr();
let p_b = t_b.as_ptr();
let mut result = ::ndarray::Array::<C64, ::ndarray::Dim<[::ndarray::Ix; 2usize]>>::uninit((
dim_number_0,
dim_number_1,
));
let mut ptr_on_result = result.as_mut_ptr() as *mut C64;
for k in 0..dim_number_0 {
for i in 0..dim_number_1 {
let mut accu = SumAccu::new();
for j in 0..dim_number_2 {
let elt_a: C64 = unsafe {
*p_a.offset(
(i as isize) * stride_a0
+ (k as isize) * stride_a2
+ (j as isize) * stride_a1,
)
};
let elt_b: C64 = unsafe {
*p_b.offset(
(i as isize) * stride_b0
+ (k as isize) * stride_b2
+ (j as isize) * stride_b1,
)
};
accu.accumulate(if elt_a > C64::ZERO {
elt_a * elt_b
} else {
(C64::ONE + elt_b.exp()).ln()
});
}
unsafe {
ptr_on_result.write(accu.get());
ptr_on_result = ptr_on_result.add(1);
}
}
}
unsafe { result.assume_init() }
}
fn flatten_accumulator_factorized(t_a: &Array3<C64>, t_b: &Array3<C64>) -> Array2<C64> {
let dim_number_2 = ::ndarray::ArrayBase::<_, _>::dim(&t_a).1;
if dim_number_2 != ::ndarray::ArrayBase::<_, _>::dim(&t_b).1 {
panic!("Dimensions are not matching between t_a[_, j, _] and t_b[_, j, _]");
}
let dim_number_0 = ::ndarray::ArrayBase::<_, _>::dim(&t_a).2;
if dim_number_0 != ::ndarray::ArrayBase::<_, _>::dim(&t_b).2 {
panic!("Dimensions are not matching between t_a[_, _, k] and t_b[_, _, k]");
}
let dim_number_1 = ::ndarray::ArrayBase::<_, _>::dim(&t_a).0;
if dim_number_1 != ::ndarray::ArrayBase::<_, _>::dim(&t_b).0 {
panic!("Dimensions are not matching between t_a[i, _, _] and t_b[i, _, _]");
}
let strides_a = t_a.strides();
if strides_a != [20isize, 4isize, 1isize] {
panic!("t_a has unexpected memory layout {:?}", strides_a);
}
let strides_b = t_b.strides();
if strides_b != [20isize, 4isize, 1isize] {
panic!("t_b has unexpected memory layout {:?}", strides_b);
}
let stride_a0 = strides_a[0];
let stride_a1 = strides_a[1];
let stride_a2 = strides_a[2];
let stride_b0 = strides_b[0];
let stride_b1 = strides_b[1];
let stride_b2 = strides_b[2];
let p_a = t_a.as_ptr();
let p_b = t_b.as_ptr();
let mut result = ::ndarray::Array::<C64, ::ndarray::Dim<[::ndarray::Ix; 2usize]>>::uninit((
dim_number_0,
dim_number_1,
));
let mut ptr_on_result = result.as_mut_ptr() as *mut C64;
for k in 0..dim_number_0 {
for i in 0..dim_number_1 {
let mut accu = SumAccu::new();
let pre_a = (i as isize) * stride_a0 + (k as isize) * stride_a2;
let pre_b = (i as isize) * stride_b0 + (k as isize) * stride_b2;
for j in 0..dim_number_2 {
let elt_a: C64 = unsafe { *p_a.offset(pre_a + (j as isize) * stride_a1) };
let elt_b: C64 = unsafe { *p_b.offset(pre_b + (j as isize) * stride_b1) };
accu.accumulate(if elt_a > C64::ZERO {
elt_a * elt_b
} else {
(C64::ONE + elt_b.exp()).ln()
});
}
unsafe {
ptr_on_result.write(accu.get());
ptr_on_result = ptr_on_result.add(1);
}
}
}
unsafe { result.assume_init() }
}
fn benchmark(c: &mut Criterion) {
let shape = (8, 5, 4);
let tensor1 = Array3::<C64>::from_shape_fn(shape, |(i, j, k)| {
let i = C64::from_usize(i) * PI;
let j = C64::from_usize(j) * PI + C64::I;
let k = C64::from_usize(k) * PI;
(i * k + 3.0).sin() * 0.9 - (j / 2.0 + 1.44).sin() * 1.5 + (k / 3.0 + 4.9 - j).cos() * 2.7
});
let tensor2 = Array3::<C64>::from_shape_fn(shape, |(i, j, k)| {
let i = C64::from_usize(i) * PI;
let j = C64::from_usize(j) * PI;
let k = C64::from_usize(k) * PI + C64::I;
(j * i + 2.9).sin() * 1.1 - (k / 1.9 + 1.54).sin() * 1.6 + (j / 3.1 + 4.8 - k).cos() * 2.6
});
let result_ndarray = pure_ndarray(&tensor1, &tensor2);
let result_ndarray_non_custom = ndarray_non_custom(&tensor1, &tensor2);
let result_tensorism = tensorism(&tensor1, &tensor2);
let result_flatten = flatten_factorized(&tensor1, &tensor2);
let result_flatten_var = flatten_normal(&tensor1, &tensor2);
assert_eq!(result_ndarray, result_tensorism);
assert_eq!(result_ndarray_non_custom, result_tensorism);
assert_eq!(result_flatten, result_tensorism);
assert_eq!(result_flatten_var, result_tensorism);
c.bench_function("Conditional value - Ndarray", |b| {
b.iter(|| pure_ndarray(black_box(&tensor1), black_box(&tensor2)))
});
c.bench_function("Conditional value - Ndarray Non Custom", |b| {
b.iter(|| ndarray_non_custom(black_box(&tensor1), black_box(&tensor2)))
});
c.bench_function("Conditional value - Flatten normal", |b| {
b.iter(|| flatten_normal(black_box(&tensor1), black_box(&tensor2)))
});
c.bench_function("Conditional value - Flatten factorized", |b| {
b.iter(|| flatten_factorized(black_box(&tensor1), black_box(&tensor2)))
});
c.bench_function("Conditional value - Tensorism", |b| {
b.iter(|| tensorism(black_box(&tensor1), black_box(&tensor2)))
});
c.bench_function("Conditional value - Flatten accumulator", |b| {
b.iter(|| flatten_accumulator(black_box(&tensor1), black_box(&tensor2)))
});
c.bench_function("Conditional value - Flatten accumulator factorized", |b| {
b.iter(|| flatten_accumulator_factorized(black_box(&tensor1), black_box(&tensor2)))
});
}
criterion_group!(benches, benchmark);
criterion_main!(benches);