strided-kernel 0.4.0

Cache-optimized kernels for strided multidimensional array operations in Rust (ported from Julia Strided.jl/StridedViews.jl).
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121

#[cfg(feature = "parallel")]
use std::num::NonZeroUsize;

use strided_kernel::{
    compare_into, CompareOp, Identity, StridedArray, StridedError, StridedView, StridedViewMut,
};
#[cfg(feature = "parallel")]
use strided_kernel::{with_execution_policy, ExecutionPolicy};

macro_rules! assert_all_compare_ops {
    ($ty:ty, $lhs:expr, $rhs:expr) => {{
        let lhs = StridedArray::from_parts($lhs, &[4], &[1], 0).unwrap();
        let rhs = StridedArray::from_parts($rhs, &[4], &[1], 0).unwrap();
        let cases = [
            (CompareOp::Eq, [false, true, false, true]),
            (CompareOp::Lt, [true, false, false, false]),
            (CompareOp::Le, [true, true, false, true]),
            (CompareOp::Gt, [false, false, true, false]),
            (CompareOp::Ge, [false, true, true, true]),
        ];

        for (op, expected) in cases {
            let mut out = StridedArray::<bool>::col_major(&[4]);
            compare_into(&mut out.view_mut(), &lhs.view(), &rhs.view(), op).unwrap();
            assert_eq!(out.data(), &expected, "operation {op:?}");
        }
    }};
}

#[test]
fn compare_into_covers_ordered_kernel_dtypes_and_operations() {
    assert_all_compare_ops!(f32, vec![1.0, 2.0, 3.0, 2.0], vec![2.0, 2.0, 2.0, 2.0]);
    assert_all_compare_ops!(f64, vec![1.0, 2.0, 3.0, 2.0], vec![2.0, 2.0, 2.0, 2.0]);
    assert_all_compare_ops!(i32, vec![1, 2, 3, 2], vec![2, 2, 2, 2]);
    assert_all_compare_ops!(i64, vec![1, 2, 3, 2], vec![2, 2, 2, 2]);

    let lhs = StridedArray::from_parts(vec![false, false, true, true], &[4], &[1], 0).unwrap();
    let rhs = StridedArray::from_parts(vec![false, true, false, true], &[4], &[1], 0).unwrap();
    let cases = [
        (CompareOp::Eq, [true, false, false, true]),
        (CompareOp::Lt, [false, true, false, false]),
        (CompareOp::Le, [true, true, false, true]),
        (CompareOp::Gt, [false, false, true, false]),
        (CompareOp::Ge, [true, false, true, true]),
    ];
    for (op, expected) in cases {
        let mut out = StridedArray::<bool>::col_major(&[4]);
        compare_into(&mut out.view_mut(), &lhs.view(), &rhs.view(), op).unwrap();
        assert_eq!(out.data(), &expected, "operation {op:?}");
    }
}

#[test]
fn compare_into_preserves_nan_unordered_semantics() {
    let lhs = StridedArray::from_parts(vec![f32::NAN], &[1], &[1], 0).unwrap();
    let rhs = StridedArray::from_parts(vec![1.0_f32], &[1], &[1], 0).unwrap();

    for op in [
        CompareOp::Eq,
        CompareOp::Lt,
        CompareOp::Le,
        CompareOp::Gt,
        CompareOp::Ge,
    ] {
        let mut out = StridedArray::<bool>::col_major(&[1]);
        compare_into(&mut out.view_mut(), &lhs.view(), &rhs.view(), op).unwrap();
        assert_eq!(out.data(), &[false], "operation {op:?}");
    }
}

#[test]
fn compare_into_preserves_noncontiguous_view_semantics() {
    let dims = [5];
    let lhs_data = [1_i64, 2, 3, 4, 5];
    let rhs = StridedArray::from_parts(vec![0_i64, 3, 3, 3, 6], &dims, &[1], 0).unwrap();
    let lhs = StridedView::<i64, Identity>::new(&lhs_data, &dims, &[-1], 4).unwrap();
    let mut out = StridedArray::<bool>::col_major(&dims);

    compare_into(&mut out.view_mut(), &lhs, &rhs.view(), CompareOp::Gt).unwrap();

    assert_eq!(out.data(), &[true, true, false, false, false]);
}

#[test]
fn compare_into_rejects_noninjective_destination_before_mutation() {
    let lhs = StridedArray::from_parts(vec![1_i64; 4], &[4], &[1], 0).unwrap();
    let rhs = StridedArray::from_parts(vec![2_i64; 4], &[4], &[1], 0).unwrap();
    let mut output = [true];
    let mut dest = StridedViewMut::new(&mut output, &[4], &[0], 0).unwrap();

    let error = compare_into(&mut dest, &lhs.view(), &rhs.view(), CompareOp::Lt);

    assert!(matches!(error, Err(StridedError::NonInjectiveOutputLayout)));
    assert_eq!(output, [true]);
}

#[cfg(feature = "parallel")]
#[test]
fn compare_into_rejects_large_noninjective_destination_before_parallel_mutation() {
    const LEN: usize = 100_000;
    let lhs = StridedArray::from_parts(vec![1_i64; LEN], &[LEN], &[1], 0).unwrap();
    let rhs = StridedArray::from_parts(vec![2_i64; LEN], &[LEN], &[1], 0).unwrap();
    let mut output = [true];
    let mut dest = StridedViewMut::new(&mut output, &[LEN], &[0], 0).unwrap();
    let pool = rayon::ThreadPoolBuilder::new()
        .num_threads(2)
        .build()
        .unwrap();

    let error = pool.install(|| {
        with_execution_policy(
            ExecutionPolicy::Rayon {
                max_threads: NonZeroUsize::new(2).unwrap(),
            },
            || compare_into(&mut dest, &lhs.view(), &rhs.view(), CompareOp::Lt),
        )
    });

    assert!(matches!(error, Err(StridedError::NonInjectiveOutputLayout)));
    assert_eq!(output, [true]);
}