Home

NDReducibles.jl

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
source