NDReducibles.jl
NDReducibles.ndreducible
— Method.ndreducible(A => (:i, :j, ...), B => (...), ..., C => (...))
Given pairs of an array (or a Broadcasted
) and indices, ndreducible
"couples" the arrays with shared index and plan the best nested for
-loop along all the indices. ndreducible
returns a fold-able object which can be used with foldl
and foreach
defined in Transducers.jl. Use referenceable
from Referenceables.jl to mutate arrays.
See also: plan
.
Examples
\[C_{ij} = ∑_k A_{ik} B_{kj} + C_{ij}\]
julia> using NDReducibles
using Referenceables: referenceable
julia> A = rand(-10:10, 5, 4)
B = rand(-10:10, 4, 3)
C = zeros(Int, 5, 3);
julia> foreach(
ndreducible(
referenceable(C) => (:i, :j),
A => (:i, :k),
B => (:k, :j),
),
simd = :ivdep, # optional
) do (c, a, b)
c[] += a * b
return # not required but useful for performance
end;
julia> C == A * B
true
\[∑_i A_i B_i C_i\]
julia> using NDReducibles
using Transducers: MapSplat
julia> A = rand(-10:10, 5)
B = rand(-10:10, 5)
C = rand(-10:10, 5);
julia> foldl(
+,
MapSplat(*),
ndreducible(
A => (:i,),
B => (:i,),
C => (:i,),
),
simd = true, # optional
) == sum(A .* B .* C)
true
\[∑ \{ A_i B_i C_i \;|\; A_i > 0 \}\]
julia> using NDReducibles
using Transducers: MapSplat, Filter
julia> A = rand(-10:10, 5)
B = rand(-10:10, 5)
C = rand(-10:10, 5);
julia> foldl(
+,
Filter(((a, b, c),) -> a > 0) |> MapSplat(*),
ndreducible(
A => (:i,),
B => (:i,),
C => (:i,),
),
simd = true, # optional
) == sum(A .* (A .> 0) .* B .* C)
true