burn_ndarray/
parallel.rs

1/// Macro for running a function in parallel.
2#[cfg(feature = "multi-threads")]
3#[macro_export(local_inner_macros)]
4macro_rules! run_par {
5    (
6        $func:expr
7    ) => {{
8        use burn_common::rayon::prelude::*;
9
10        #[allow(clippy::redundant_closure_call)]
11        burn_common::rayon::scope(|_| $func())
12    }};
13}
14
15/// Macro for running a function in parallel.
16#[cfg(not(feature = "multi-threads"))]
17#[macro_export(local_inner_macros)]
18macro_rules! run_par {
19    (
20        $func:expr
21    ) => {{ $func() }};
22}
23
24/// Macro for iterating in parallel.
25#[cfg(not(feature = "multi-threads"))]
26#[macro_export(local_inner_macros)]
27macro_rules! iter_par {
28    (
29        $iter:expr
30    ) => {{ $iter }};
31}
32
33/// Macro for iterating in parallel.
34#[cfg(feature = "multi-threads")]
35#[macro_export(local_inner_macros)]
36macro_rules! iter_par {
37    (
38        $iter:expr
39    ) => {{ $iter.into_par_iter() }};
40}
41
42/// Macro for iterating in parallel.
43#[cfg(feature = "multi-threads")]
44#[macro_export(local_inner_macros)]
45macro_rules! iter_slice_par {
46    (
47        $slice:expr
48    ) => {{ $slice.into_par_iter() }};
49}
50
51/// Macro for iterating in parallel.
52#[cfg(not(feature = "multi-threads"))]
53#[macro_export(local_inner_macros)]
54macro_rules! iter_slice_par {
55    (
56        $slice:expr
57    ) => {{ $slice.iter() }};
58}
59
60/// Macro for iterating over a range in parallel.
61#[cfg(feature = "multi-threads")]
62#[macro_export(local_inner_macros)]
63macro_rules! iter_range_par {
64    (
65        $start:expr, $end:expr
66    ) => {{ ($start..$end).into_par_iter() }};
67}
68
69/// Macro for iterating over a range in parallel.
70#[cfg(not(feature = "multi-threads"))]
71#[macro_export(local_inner_macros)]
72macro_rules! iter_range_par {
73    (
74        $start:expr, $end:expr
75    ) => {{ ($start..$end) }};
76}