Tests fail on macOS (arm64 architecture)

[lamBOOO]

After playing around with GridapPETSc (see gridap/GridapPETSc.jl#102), I noticed that on macOS (Apple M4 chip), some custom MPI operations produce an error if custom reduction are used (see the error message below).

One can reproduce theses errors when running the PartionedArrays test suite either:

1. using a Github macos (arm64) runner as done in https://github.com/lamBOOO/PartitionedArrays.jl/actions/runs/16026309753/job/45215040606
2. or locally with Docker and the docker run --rm -it --platform linux/amd64 julia command followed by the usual ] test command within Julia.

The problem seems to be known, see JuliaParallel/MPI.jl#404 or the MPI.jl docs.

I came up with 2 ways to fix the tests:

1. Using @RegisterOps as recommend in the MPI.jl docs (also done here): Statically register MPI reduction ops for distances and fix global AN… lamBOOO/PartitionedArrays.jl#3
2. Using a try/catch that uses a fallback when running on macOS: add macos fix (not as performant) lamBOOO/PartitionedArrays.jl#1

I think the second option in not optimal in terms of runtime. Do you have any other idea to get this working?

Details

Error: 
│   exception =
│    User-defined reduction operators are currently not supported on non-Intel architectures.
│    See https://github.com/JuliaParallel/MPI.jl/issues/404 for more details.
│    
│    You may want to use `@RegisterOp` to statically register `f`.
│    
│    Stacktrace:
│      [1] error(s::String)
│        @ Base ./error.jl:35
│      [2] MPI.Op(f::Function, T::Type; iscommutative::Bool)
│        @ MPI ~/.julia/packages/MPI/TKXAj/src/operators.jl:109
│      [3] MPI.Op(f::Function, T::Type)
│        @ MPI ~/.julia/packages/MPI/TKXAj/src/operators.jl:102
│      [4] reduction_impl(op::PartitionedArrays.var"#and#68", a::PartitionedArrays.MPIArray{Bool, 1}, destination::Symbol; init::Bool)
│        @ PartitionedArrays ~/work/PartitionedArrays.jl/PartitionedArrays.jl/src/mpi_array.jl:481
│      [5] #reduction#49
│        @ ~/work/PartitionedArrays.jl/PartitionedArrays.jl/src/primitives.jl:682 [inlined]
│      [6] (::PartitionedArrays.var"#is_included#65")(snd_ids_a::PartitionedArrays.MPIArray{Vector{Int32}, 1}, snd_ids_b::PartitionedArrays.MPIArray{Vector{Int32}, 1})
│        @ PartitionedArrays ~/work/PartitionedArrays.jl/PartitionedArrays.jl/src/primitives.jl:811
│      [7] ExchangeGraph_impl_with_neighbors(snd_ids::PartitionedArrays.MPIArray{Vector{Int32}, 1}, neighbors::PartitionedArrays.ExchangeGraph{PartitionedArrays.MPIArray{Vector{Int32}, 1}})
│        @ PartitionedArrays ~/work/PartitionedArrays.jl/PartitionedArrays.jl/src/primitives.jl:817
│      [8] #ExchangeGraph#57
│        @ ~/work/PartitionedArrays.jl/PartitionedArrays.jl/src/primitives.jl:797 [inlined]
│      [9] ExchangeGraph
│        @ ~/work/PartitionedArrays.jl/PartitionedArrays.jl/src/primitives.jl:787 [inlined]
│     [10] #compute_assembly_neighbors#153
│        @ ~/work/PartitionedArrays.jl/PartitionedArrays.jl/src/p_range.jl:448 [inlined]
│     [11] compute_assembly_neighbors
│        @ ~/work/PartitionedArrays.jl/PartitionedArrays.jl/src/p_range.jl:436 [inlined]
│     [12] #assembly_neighbors#146
│        @ ~/work/PartitionedArrays.jl/PartitionedArrays.jl/src/p_range.jl:423 [inlined]
│     [13] assembly_neighbors
│        @ ~/work/PartitionedArrays.jl/PartitionedArrays.jl/src/p_range.jl:417 [inlined]
│     [14] psparse(f::PartitionedArrays.var"#f#410"{SparseArrays.SparseMatrixCSC{Float64, Int64}}, I::PartitionedArrays.MPIArray{Vector{Int64}, 1}, J::PartitionedArrays.MPIArray{Vector{Int64}, 1}, V::PartitionedArrays.MPIArray{Vector{Float64}, 1}, rows::PartitionedArrays.MPIArray{PartitionedArrays.LocalIndicesWithVariableBlockSize{1}, 1}, cols::PartitionedArrays.MPIArray{PartitionedArrays.LocalIndicesWithVariableBlockSize{1}, 1}; split_format::Val{true}, subassembled::Bool, assembled::Bool, assemble::Bool, indices::Symbol, restore_ids::Bool, assembly_neighbors_options_rows::@NamedTuple{neighbors::PartitionedArrays.ExchangeGraph{PartitionedArrays.MPIArray{Vector{Int32}, 1}}}, assembly_neighbors_options_cols::@NamedTuple{neighbors::PartitionedArrays.ExchangeGraph{PartitionedArrays.MPIArray{Vector{Int32}, 1}}}, assembled_rows::Nothing, reuse::Bool)
│        @ PartitionedArrays ~/work/PartitionedArrays.jl/PartitionedArrays.jl/src/p_sparse_matrix.jl:1192
│     [15] psparse
│        @ ~/work/PartitionedArrays.jl/PartitionedArrays.jl/src/p_sparse_matrix.jl:1150 [inlined]
│     [16] #psparse#409
│        @ ~/work/PartitionedArrays.jl/PartitionedArrays.jl/src/p_sparse_matrix.jl:1139 [inlined]
│     [17] psparse
│        @ ~/work/PartitionedArrays.jl/PartitionedArrays.jl/src/p_sparse_matrix.jl:1137 [inlined]
│     [18] #psparse#408
│        @ ~/work/PartitionedArrays.jl/PartitionedArrays.jl/src/p_sparse_matrix.jl:1134 [inlined]
│     [19] psystem(I::PartitionedArrays.MPIArray{Vector{Int64}, 1}, J::PartitionedArrays.MPIArray{Vector{Int64}, 1}, V::PartitionedArrays.MPIArray{Vector{Float64}, 1}, I2::PartitionedArrays.MPIArray{Vector{Int64}, 1}, V2::PartitionedArrays.MPIArray{Vector{Float64}, 1}, rows::PartitionedArrays.MPIArray{PartitionedArrays.LocalIndicesWithVariableBlockSize{1}, 1}, cols::PartitionedArrays.MPIArray{PartitionedArrays.LocalIndicesWithVariableBlockSize{1}, 1}; subassembled::Bool, assembled::Bool, assemble::Bool, indices::Symbol, restore_ids::Bool, assembly_neighbors_options_rows::@NamedTuple{neighbors::PartitionedArrays.ExchangeGraph{PartitionedArrays.MPIArray{Vector{Int32}, 1}}}, assembly_neighbors_options_cols::@NamedTuple{neighbors::PartitionedArrays.ExchangeGraph{PartitionedArrays.MPIArray{Vector{Int32}, 1}}}, assembled_rows::Nothing, reuse::Bool)
│        @ PartitionedArrays ~/work/PartitionedArrays.jl/PartitionedArrays.jl/src/p_sparse_matrix.jl:2818
│     [20] fem_example(distribute::PartitionedArrays.var"#105#106"{MPI.Comm, Bool})
│        @ Main.MPIArrayFEMExample ~/work/PartitionedArrays.jl/PartitionedArrays.jl/test/fem_example.jl:314
│     [21] with_mpi(f::typeof(Main.MPIArrayFEMExample.fem_example); comm::MPI.Comm, duplicate_comm::Bool)
│        @ PartitionedArrays ~/work/PartitionedArrays.jl/PartitionedArrays.jl/src/mpi_array.jl:73
│     [22] with_mpi(f::Function)
│        @ PartitionedArrays ~/work/PartitionedArrays.jl/PartitionedArrays.jl/src/mpi_array.jl:64
│     [23] top-level scope
│        @ ~/work/PartitionedArrays.jl/PartitionedArrays.jl/test/mpi_array/drivers/fem_example.jl:7
│     [24] include(mod::Module, _path::String)
│        @ Base ./Base.jl:557
│     [25] exec_options(opts::Base.JLOptions)
│        @ Base ./client.jl:323
│     [26] _start()
│        @ Base ./client.jl:531
└ @ PartitionedArrays ~/work/PartitionedArrays.jl/PartitionedArrays.jl/src/mpi_array.jl:75```

</details>

[fverdugo]

Hi @lamBOOO,

I was not aware of this issue (I am not using an arm64 computer).

Regarding fix 1: Are you sure this is not some sort of type piracy? What does it happen if other packages register these operators?

[lamBOOO]

I think it is not a type piracy since the helper functions (_reduce_euclidean, ..) are defined in the PartitionedArrays package, so simply registering them doesn't extend anyone else's methods.

I'm not expert, but as far as I understand, if another package would later call MPI.@RegisterOp(_reduce_euclidean, Int) again, MPI.jl sees the same (function, Type) and would reuse the existing registration.

Moving it into Distances.jl itself would introduce MPI as an dependency there (not optimal).

[lamBOOO]

Another option might be to do a package extension, that is only loaded when both MPI and PartionedArrays are loaded (and potentially checks if Sys.ARCH == :aarch64). But I have never done that.

Do you have another idea, or should I open a draft PR that implements Approach 1 for now?

[fverdugo]

Hi @lamBOOO

I am wondering why we are using this custom operations in the first place. I guess it should be possible to implement these distances with builtin operators such as + and max. Ideally, we should be able to get rid of the custom operators. This would be the cleanest solution in my opinion.

Can you look into this?

[lamBOOO]

Hi @fverdugo,

sorry for the late response. Yes, I like that idea.

I looked into Distances.jl and for the standard distances that return a Number, we can simply use + and max as you proposed. I tested that in lamBOOO#4 and it works on aarch64 (Apple M4).

Some distances however, return Tuples. And the RenyiDivergence, e.g., has some more complex arithmetics inside the eval_reduce as seen here. Do we need to support them?

[fverdugo]

Hi @lamBOOO,

I think your last comment goes in the right direction, but I would do this instead:

Implement a fallback function

function _combine_distance_partials(d, partials, init)
        reduce((i, j) -> Distances.eval_reduce(d, i, j), partials;init)
end

Then you can add a per-distance method to solve your issue. For instance,

function _combine_distance_partials(d::Distances.Chebyshev, partials, init)
        reduce(max, partials;init)
end

Then, we can add as many methods as needed. I would start for the ones you need. I would try to use builtin julia functions in the first argument of reduce as much as possible. If not, we can register MPI Operators, but only if really needed.

Please, also rename _combine_distance_partials -> reduce_for_distance.

[lamBOOO]

Good idea!

I did that in lamBOOO@bb2a12f. This covers the most basic metrics for now.

[fverdugo]

Looks good to me!

Open a PR please.

[lamBOOO]

Closed with #209