Can't unpack recv buffer after MPI Waitall

[romanlee]

Hit this issue while debugging Oceananigans+Reactant+MPI halo exchange.

The data from send_buf is properly received into recv_buf, but then when we try to "unpack the data", ie copy it from recv_buf into data_buf, data_buf does not get modified.

Interestingly, it all works fine when we wait for the request with Wait instead of Waitall.

using MPI
using Reactant

Reactant.set_default_backend("cpu")

const USE_WAITALL = parse(Bool, get(ENV, "USE_WAITALL", "true"))

function exchange!(send_buf, recv_buf, data_buf)
    # send/recv
    if rank == 0
        dest = 1
        req = MPI.Isend(send_buf, dest, tag, comm)
    else
        src = 0
        req = MPI.Irecv!(recv_buf, src, tag, comm)
    end

    # waitall/wait
    if USE_WAITALL
        MPI.Waitall(vcat([req]...))
    else
        MPI.Wait(req)
    end

    # unpack recv_buf into data_buf
    if rank == 1
        # properly updated with wait, not modified with waitall
        data_buf .= recv_buf
    end

    return nothing
end

# initialize
MPI.Init()

comm = MPI.COMM_WORLD
rank = MPI.Comm_rank(comm)
tag = 43

# allocate/init buffers
Nx=1
FT=Float64

send_buf = Reactant.to_rarray(zeros(FT, Nx))
send_buf .= reshape(FT.(collect(1:Nx)), Nx)

recv_buf = Reactant.to_rarray(zeros(FT, Nx))

data_buf = Reactant.to_rarray(zeros(FT, Nx))

# run
@jit exchange!(send_buf, recv_buf, data_buf)

# verify
# on rank 1 should have: recv_buf and data_buf == send_buf
if rank==1
    println("send_buf = $(send_buf)")
    println("recv_buf = $(recv_buf), ==send_buf:$(recv_buf==send_buf)")
    println("data_buf = $(data_buf), ==send_buf:$(data_buf==send_buf)")
end

MPI.Finalize()

[romanlee]

It seems like the problem might be that the Waitall and copy get reordered downstream somewhere, so that the copy happens on stale data (ie, before the recv has completed).

Consider the code_hlo outputs for rank 1 with USE_WAITALL=true below. The reordering hasn't happened yet, but from what I can understand it would be legal for it to happen downstream.

@mofeing @wsmoses Thoughts on this hypothesis / Is there a way we can enforce that the Waitall happen before the copy?

code_hlo optimize=false:
module @"reactant_exchange!" attributes {mhlo.num_partitions = 1 : i64, mhlo.num_replicas = 1 : i64} {
  func.func private @identity_broadcast_scalar(%arg0: tensor<f64>) -> tensor<f64> {
    return %arg0 : tensor<f64>
  }
  func.func @main(%arg0: tensor<1xf64> {tf.aliasing_output = 0 : i32}, %arg1: tensor<1xf64> {tf.aliasing_output = 1 : i32}, %arg2: tensor<1xf64> {tf.aliasing_output = 2 : i32}) -> (tensor<1xf64>, tensor<1xf64>, tensor<1xf64>) {
    %0 = stablehlo.transpose %arg0, dims = [0] : (tensor<1xf64>) -> tensor<1xf64>
    %1 = stablehlo.transpose %arg1, dims = [0] : (tensor<1xf64>) -> tensor<1xf64>
    %2 = stablehlo.transpose %arg2, dims = [0] : (tensor<1xf64>) -> tensor<1xf64>
    %c = stablehlo.constant dense<0> : tensor<i32>
    %c_0 = stablehlo.constant dense<43> : tensor<i32>
    %c_1 = stablehlo.constant dense<1> : tensor<i32>
    %outbuf, %request = enzymexla.mpi.irecv(%1, %c_1, %c, %c_0) {datatype = #enzymexla.datatype<MPI_DOUBLE>} : (tensor<1xf64>, tensor<i32>, tensor<i32>, tensor<i32>) -> (tensor<1xf64>, tensor<i32>)
    %3 = stablehlo.broadcast_in_dim %request, dims = [] : (tensor<i32>) -> tensor<1xi32>
    %4 = stablehlo.concatenate %3, dim = 0 : (tensor<1xi32>) -> tensor<1xi32>
    %c_2 = stablehlo.constant dense<1> : tensor<i32>
    enzymexla.mpi.waitall(%c_2, %4) : tensor<i32>, tensor<1xi32>
    %5 = stablehlo.broadcast_in_dim %outbuf, dims = [0] : (tensor<1xf64>) -> tensor<1xf64>
    %6 = stablehlo.broadcast_in_dim %5, dims = [0] : (tensor<1xf64>) -> tensor<1xf64>
    %7 = enzyme.batch @identity_broadcast_scalar(%6) {batch_shape = array<i64: 1>} : (tensor<1xf64>) -> tensor<1xf64>
    %8 = stablehlo.transpose %0, dims = [0] : (tensor<1xf64>) -> tensor<1xf64>
    %9 = stablehlo.transpose %outbuf, dims = [0] : (tensor<1xf64>) -> tensor<1xf64>
    %10 = stablehlo.transpose %7, dims = [0] : (tensor<1xf64>) -> tensor<1xf64>
    return %8, %9, %10 : tensor<1xf64>, tensor<1xf64>, tensor<1xf64>
  }
}

code_hlo optimize="lower-enzymexla-mpi{backend=cpu}":
module @"reactant_exchange!" attributes {mhlo.num_partitions = 1 : i64, mhlo.num_replicas = 1 : i64} {
  llvm.mlir.global external constant @MPI_DOUBLE() {addr_space = 0 : i32} : !llvm.ptr
  llvm.mlir.global external constant @MPI_COMM_WORLD() {addr_space = 0 : i32} : !llvm.ptr
  llvm.func @MPI_Irecv(!llvm.ptr, i32, !llvm.ptr, i32, i32, !llvm.ptr, !llvm.ptr) -> i32
  llvm.func @enzymexla_wrapper_MPI_Irecv_MPI_DOUBLE(%arg0: !llvm.ptr {enzymexla.memory_effects = ["read", "write", "allocate", "free"]}, %arg1: !llvm.ptr {enzymexla.memory_effects = ["read", "write", "allocate", "free"]}, %arg2: !llvm.ptr {enzymexla.memory_effects = ["read", "write", "allocate", "free"]}, %arg3: !llvm.ptr {enzymexla.memory_effects = ["read", "write", "allocate", "free"]}, %arg4: !llvm.ptr {enzymexla.memory_effects = ["read", "write", "allocate", "free"]}) attributes {enzymexla.memory_effects = ["read", "write", "allocate", "free"]} {
    %0 = llvm.mlir.addressof @MPI_COMM_WORLD : !llvm.ptr
    %1 = llvm.mlir.addressof @MPI_DOUBLE : !llvm.ptr
    %2 = llvm.load %arg1 : !llvm.ptr -> i32
    %3 = llvm.load %arg2 : !llvm.ptr -> i32
    %4 = llvm.load %arg3 : !llvm.ptr -> i32
    %5 = llvm.call @MPI_Irecv(%arg0, %2, %1, %3, %4, %0, %arg4) : (!llvm.ptr, i32, !llvm.ptr, i32, i32, !llvm.ptr, !llvm.ptr) -> i32
    llvm.return
  }
  llvm.func @MPI_Waitall(i32, !llvm.ptr, !llvm.ptr) -> i32
  llvm.func @enzymexla_wrapper_MPI_Waitall(%arg0: !llvm.ptr {enzymexla.memory_effects = ["read", "write", "allocate", "free"]}, %arg1: !llvm.ptr {enzymexla.memory_effects = ["read", "write", "allocate", "free"]}) attributes {enzymexla.memory_effects = ["read", "write", "allocate", "free"]} {
    %0 = llvm.load %arg0 : !llvm.ptr -> i32
    %1 = llvm.alloca %0 x !llvm.array<6 x i32> : (i32) -> !llvm.ptr
    %2 = llvm.call @MPI_Waitall(%0, %arg1, %1) : (i32, !llvm.ptr, !llvm.ptr) -> i32
    llvm.return
  }
  func.func private @identity_broadcast_scalar(%arg0: tensor<f64> {enzymexla.memory_effects = []}) -> tensor<f64> attributes {enzymexla.memory_effects = []} {
    return %arg0 : tensor<f64>
  }
  func.func @main(%arg0: tensor<1xf64> {enzymexla.memory_effects = [], tf.aliasing_output = 0 : i32}, %arg1: tensor<1xf64> {enzymexla.memory_effects = ["read", "write", "allocate", "free"], tf.aliasing_output = 1 : i32}, %arg2: tensor<1xf64> {enzymexla.memory_effects = [], tf.aliasing_output = 2 : i32}) -> (tensor<1xf64>, tensor<1xf64>, tensor<1xf64>) attributes {enzymexla.memory_effects = ["read", "write", "allocate", "free"]} {
    %c = stablehlo.constant dense<-1> : tensor<i32>
    %c_0 = stablehlo.constant dense<1> : tensor<i32>
    %c_1 = stablehlo.constant dense<43> : tensor<i32>
    %c_2 = stablehlo.constant dense<0> : tensor<i32>
    %0 = stablehlo.transpose %arg0, dims = [0] : (tensor<1xf64>) -> tensor<1xf64>
    %1 = stablehlo.transpose %arg1, dims = [0] : (tensor<1xf64>) -> tensor<1xf64>
    %2:2 = enzymexla.jit_call @enzymexla_wrapper_MPI_Irecv_MPI_DOUBLE (%1, %c_0, %c_2, %c_1, %c) {output_operand_aliases = [#stablehlo.output_operand_alias<output_tuple_indices = [0], operand_index = 0, operand_tuple_indices = []>, #stablehlo.output_operand_alias<output_tuple_indices = [1], operand_index = 4, operand_tuple_indices = []>]} : (tensor<1xf64>, tensor<i32>, tensor<i32>, tensor<i32>, tensor<i32>) -> (tensor<1xf64>, tensor<i32>)
    %3 = stablehlo.broadcast_in_dim %2#1, dims = [] : (tensor<i32>) -> tensor<1xi32>
    %4 = stablehlo.concatenate %3, dim = 0 : (tensor<1xi32>) -> tensor<1xi32>
    enzymexla.jit_call @enzymexla_wrapper_MPI_Waitall (%c_0, %4) : (tensor<i32>, tensor<1xi32>) -> ()
    %5 = stablehlo.broadcast_in_dim %2#0, dims = [0] : (tensor<1xf64>) -> tensor<1xf64>
    %6 = stablehlo.broadcast_in_dim %5, dims = [0] : (tensor<1xf64>) -> tensor<1xf64>
    %7 = enzyme.batch @identity_broadcast_scalar(%6) {batch_shape = array<i64: 1>} : (tensor<1xf64>) -> tensor<1xf64>
    %8 = stablehlo.transpose %0, dims = [0] : (tensor<1xf64>) -> tensor<1xf64>
    %9 = stablehlo.transpose %2#0, dims = [0] : (tensor<1xf64>) -> tensor<1xf64>
    %10 = stablehlo.transpose %7, dims = [0] : (tensor<1xf64>) -> tensor<1xf64>
    return %8, %9, %10 : tensor<1xf64>, tensor<1xf64>, tensor<1xf64>
  }
}

code_hlo:
module @"reactant_exchange!" attributes {mhlo.num_partitions = 1 : i64, mhlo.num_replicas = 1 : i64} {
  llvm.mlir.global external constant @MPI_DOUBLE() {addr_space = 0 : i32} : !llvm.ptr
  llvm.mlir.global external constant @MPI_COMM_WORLD() {addr_space = 0 : i32} : !llvm.ptr
  llvm.func @MPI_Irecv(!llvm.ptr, i32, !llvm.ptr, i32, i32, !llvm.ptr, !llvm.ptr) -> i32
  llvm.func @MPI_Waitall(i32, !llvm.ptr, !llvm.ptr) -> i32
  func.func @main(%arg0: tensor<1xf64> {enzymexla.memory_effects = []}, %arg1: tensor<1xf64> {enzymexla.memory_effects = ["read", "write", "allocate", "free"], tf.aliasing_output = 0 : i32}, %arg2: tensor<1xf64> {enzymexla.memory_effects = [], tf.aliasing_output = 1 : i32}) -> (tensor<1xf64>, tensor<1xf64>) attributes {enzymexla.memory_effects = ["read", "write", "allocate", "free"]} {
    %c = stablehlo.constant dense<1> : tensor<i32>
    %c_0 = stablehlo.constant dense<43> : tensor<i32>
    %c_1 = stablehlo.constant dense<0> : tensor<i32>
    %c_2 = stablehlo.constant dense<-1> : tensor<i32>
    %0:2 = stablehlo.custom_call @enzymexla_compile_cpu(%arg1, %c, %c_1, %c_0, %c_2) {api_version = 3 : i32, backend_config = "\00`b\05\D7v\00\00\00\00\00\00\00\00\00\00", has_side_effect = true, output_operand_aliases = [#stablehlo.output_operand_alias<output_tuple_indices = [0], operand_index = 0, operand_tuple_indices = []>, #stablehlo.output_operand_alias<output_tuple_indices = [1], operand_index = 4, operand_tuple_indices = []>]} : (tensor<1xf64>, tensor<i32>, tensor<i32>, tensor<i32>, tensor<i32>) -> (tensor<1xf64>, tensor<i32>)
    %1 = stablehlo.reshape %0#1 : (tensor<i32>) -> tensor<1xi32>
    stablehlo.custom_call @enzymexla_compile_cpu(%c, %1) {api_version = 3 : i32, backend_config = "\00P\C3\04\D7v\00\00\00\00\00\00\00\00\00\00", has_side_effect = true} : (tensor<i32>, tensor<1xi32>) -> ()
    return %0#0, %0#0 : tensor<1xf64>, tensor<1xf64>
  }
}

Aside, as for why Wait still works: it seems like in principle Wait could suffer from the same reordering issue, but perhaps the differences in the way the request is constructed between Wait/Waitall lead to differences in code motion.

[wsmoses]

So in this case I think the ideal case would be to fuse the wait and send/recv into a single JIT call, which is something we could do as an MLIR side transformation

[romanlee]

What would that look like exactly? Happy to talk through it tomorrow in the meeting if that's easier

[wsmoses]

we could look for jit calls whose operands are all constants or other jit calls then fuse it into a bigger jit call. It'll also be faster too

[romanlee]

Ok cool. So this would be implemented in enzyme jax as a lowering and fed into reactant alongside these https://github.com/EnzymeAD/Reactant.jl/blob/main/src/Compiler.jl#L2030?