MPI Storage Cache Initialization Can Deadlock

[ZwFink]

There are two hard problems in computer science:

1. Cache invalidation
2. Lazy management of MPI-based Proteus caches
3. Off-by-one errors

MPI-based caches can cause deadlock if some rank does an MPI operation while other rank(s) enter Proteus JIT and try to initialize the sub-communicator (collective MPI_COMM_DUPE).

Here is (codex's) minimal reproducer. It will complete with StorageCache but hang with either MPI cache:

// Root cause: MPICommHandle's constructor (MPIHelpers.cpp:55) calls
// MPI_Comm_dup(MPI_COMM_WORLD) — a collective — from inside the
// DispatcherHost singleton constructor, which is triggered lazily on the
// first JitModule creation per rank. If ranks reach their first
// JitModule at different times, with an intervening MPI collective on
// MPI_COMM_WORLD, the mismatched collectives cause a deadlock.
//
// Call chain on first JitModule("host") per rank:
//   JitModule ctor
//     → Dispatcher::getDispatcher(HOST)
//       → DispatcherHost::instance()        [static-local singleton]
//         → Dispatcher("DispatcherHost")     [base ctor]
//           → ObjectCacheChain("DispatcherHost")
//             → MPILocalLookupCache()
//               → MPIStorageCache()
//                 → MPICommHandle()          [member init]
//                   → MPI_Comm_dup(MPI_COMM_WORLD)   ← COLLECTIVE
//
// In CleverLeaf at 128+ GPUs, AMR load imbalance causes ranks to
// enter their first RAJA forall (triggering this chain) while other
// ranks are still in a SAMRAI collective on MPI_COMM_WORLD.
//
// Deadlocks:
//   PROTEUS_OBJECT_CACHE_CHAIN=mpi-local-lookup  mpirun -np 2 ./mwe
//   PROTEUS_OBJECT_CACHE_CHAIN=mpi-remote-lookup mpirun -np 2 ./mwe
//
// Succeeds (no MPI_Comm_dup):
//   PROTEUS_OBJECT_CACHE_CHAIN=storage mpirun -np 2 ./mwe
//   mpirun -np 2 ./mwe       (storage is the default)

#include <cstdio>
#include <mpi.h>
#include <proteus/JitFrontend.h>

int main(int argc, char **argv) {
  int provided;
  MPI_Init_thread(&argc, &argv, MPI_THREAD_MULTIPLE, &provided);
  if (provided < MPI_THREAD_MULTIPLE) {
    fprintf(stderr, "MPI_THREAD_MULTIPLE not available\n");
    MPI_Abort(MPI_COMM_WORLD, 1);
  }

  int rank;
  MPI_Comm_rank(MPI_COMM_WORLD, &rank);

  // Simulate the work imbalance that occurs in CleverLeaf with AMR:
  // rank 0 reaches its first JIT call before other ranks, while the
  // remaining ranks are still in an application-level MPI collective.
  if (rank == 0) {
    printf("[rank 0] Creating JitModule → MPI_Comm_dup(MPI_COMM_WORLD)...\n");
    fflush(stdout);
    proteus::JitModule J("host");
    // ^^^ With MPI cache: blocks here forever.
    //     Rank 0 is inside MPI_Comm_dup(MPI_COMM_WORLD).
    //     Other ranks never call MPI_Comm_dup — deadlock.
    printf("[rank 0] JitModule created\n");
    fflush(stdout);
  }

  // Models a SAMRAI collective (Allreduce, halo sync, regrid barrier, etc.).
  // Rank 0 never reaches this: it is stuck in MPI_Comm_dup.
  // Other ranks enter MPI_Barrier — a *different* collective on the
  // same communicator (MPI_COMM_WORLD). Mismatched collectives → deadlock.
  printf("[rank %d] Entering MPI_Barrier(MPI_COMM_WORLD)\n", rank);
  fflush(stdout);
  MPI_Barrier(MPI_COMM_WORLD);

  if (rank != 0) {
    proteus::JitModule J("host");
  }

  printf("[rank %d] Done\n", rank);
  MPI_Finalize();
  return 0;
}

[ggeorgakoudis]

The design decision that haunts us is to have the cache service inside the library. We should have made it a separate program that runs alongside the app.

[ggeorgakoudis]

There are two hard problems in computer science:

1. Cache invalidation
2. Lazy management of MPI-based Proteus caches
3. Off-by-one errors

I love the off-by-one error