Mapslices and @threads

General Usage
,
1

I am fooling around with 1.30-rc1 and would like to understand how to make a call on mapslices threaded.

Since mapslices can emulate a loop over an array dimension, it would nice to make that “loop” Threads.@threads. Does anyone have a hint?

2

Writing a loop over eachcol(x) is usually faster than mapslices even before you add threading. Or you can use my SliceMap, which has tmapcols. Or use JuliennedArrays for more complicated slicing.

4 Likes
3

My use case us that I want to write code that can take an AbstractArray as input and then do operations along one dimension. In some calls the array will be a Matrix, but in other cases there will be more dimensions.

This rules out eachcol and probably SliceMap, or? Not even eachslice seems to work well since it only allows a single dimension in dims.

4

If I understand right, then yes, eachslice allows only one “outer” dimension, meaning that your function will get N-1-dimensional arrays when acting on an N-array. Slightly confusingly mapslices takes the “inner” dimensions as dims instead, as do Slices and slicemap. But these should be more general:

julia> map(ndims, eachslice(ones(1,2,3), dims=3))
3-element Array{Int64,1}:
 2
 2
 2

julia> mapslices(ndims, ones(1,2,3), dims=(1,2))
1×1×3 Array{Int64,3}:
[:, :, 1] =
 2

[:, :, 2] =
 2

[:, :, 3] =
 2

julia> using JuliennedArrays

julia> map(ndims, Slices(ones(1,2,3), 1,2))
3-element Array{Int64,1}:
 2
 2
 2

julia> using SliceMap

julia> slicemap(ndims, ones(1,2,3), dims=(1,2))
3-element Array{Int64,1}:
 2
 2
 2

Here of course the function gives a scalar; had it made an array then the map examples would be arrays of arrays, needing reduce(hcat,... or Align etc, which is built into mapslices / slicemap.

1 Like
5

Here is my solution. Hope it helps.

import Base.Threads
import Base: Slice, concatenate_setindex!


"""
    par_mapslices(f, A::AbstractArray{<:Real,N}, args...; dims=N, kw...)

# Arguments
- `dims`: the dimension apply f

@seealso `mapslices`

# Example
```julia
A = rand(10, 10, 30, 4)
par_mapslices(mean, A)
```
"""
function par_mapslices(f, A::AbstractArray{<:Real,N}, args...; dims=N, kw...) where {N}
  idx1 = ntuple(d -> d in dims ? (:) : firstindex(A, d), ndims(A))
  Aslice = A[idx1...]
  r1 = f(Aslice)

  _dims = size(A) |> collect
  _dims[dims] = length(r1)

  dim_mask = ntuple(d -> d in dims, ndims(A))

  itershape = ntuple(d -> d in dims ? Base.OneTo(1) : axes(A, d), ndims(A))
  indices = CartesianIndices(itershape)

  R = zeros(eltype(r1), _dims...)

  slice_A = Slice.(axes(A))
  slice_R = Slice.(axes(R))

  @inbounds Threads.@threads for I in indices
    idx = ifelse.(dim_mask, slice_A, Tuple(I))
    ridx = ifelse.(dim_mask, slice_R, Tuple(I))

    # _unsafe_getindex!(Aslice, A, idx...)
    Aslice = @view A[idx...]
    r = f(Aslice, args...; kw...)
    concatenate_setindex!(R, r, ridx...)
  end
  R |> squeeze
end
```