jipolanco / PencilArrays.jl

Distributed Julia arrays using the MPI protocol
https://jipolanco.github.io/PencilArrays.jl/dev/
MIT License
60 stars 8 forks source link

cartesian indexing of a PencilArray with extra_indices #79

Closed fbignonnet closed 1 year ago

fbignonnet commented 1 year ago

Hello @jipolanco , thanks for your nice library.

There seems to be an incompatibility in how permutations are handled in the case of PencilArray with extra_indices.

In file PencilsArrays/src/arrays.jl, functions Base.getindex and Base.setindex! call function _genperm which assumes that the permutation of x::PencilArray{T,N} with extra_dims is a M-dims permutation where M = ndims_space(x) and thus splits indices. But it then calls function permutation which already appends extra_dims to the permutation, which results in applying a N-dims permutation to a M-dims index.

@inline function _genperm(x::PencilArray{T,N}, I::NTuple{N,Int}) where {T,N}
    # Split "spatial" and "extra" indices.
    M = ndims_space(x)
    E = ndims_extra(x)
    @assert M + E === N
    J = ntuple(n -> @inbounds(I[n]), Val(M))
    K = ntuple(n -> @inbounds(I[M + n]), Val(E))
    perm = permutation(x)
    ((perm * J)..., K...)
end

A quick fix might be just to redefine _genperm as:

@inline function _genperm(x::PencilArray{T,N}, I::NTuple{N,Int}) where {T,N}
    perm = permutation(x)
    perm * I
end

Another related minor bug is in PencilFFTs when displaying a PencilFFTPlan with extra_dims, a two line breaks occur after printing Extra dimensions instead of one before and one after.

Example:

pen = Pencil((4,8), comm)
plan = PencilFFTPlan(pen, Transforms.FFT!(), extra_dims = (3,))

displays plan as:

Transforms: (FFT!, FFT!)
Input type: ComplexF64
Global dimensions: (4, 8) -> (4, 8)Extra dimensions: (3,)

MPI topology: 1D decomposition (1 processes)

instead of

Transforms: (FFT!, FFT!)
Input type: ComplexF64
Global dimensions: (4, 8) -> (4, 8)
Extra dimensions: (3,)
MPI topology: 1D decomposition (1 process)
jipolanco commented 1 year ago

Hi @fbignonnet, thanks for catching these issues! To be honest extra_indices hasn't received a lot of attention and it's not sufficiently tested. I'll push a fix.

fbignonnet commented 1 year ago

Wow that was fast ! Thanks