Example code that looks like MPI code

Summary: Pull Request resolved: facebookincubator#157

Differential Revision: D18332319

Pulled By: pietern

fbshipit-source-id: 1c0c0cb273984cfd11af42051caf23a1c29638d0

gloo/CMakeLists.txt

@@ -194,3 +194,7 @@ endif()
 if(BUILD_BENCHMARK)
 add_subdirectory(benchmark)
 endif()
+
+if(BUILD_EXAMPLES)
+  add_subdirectory(examples)
+endif()

gloo/examples/CMakeLists.txt

@@ -1,2 +1,4 @@
 add_executable(example1 example1.cc)
 target_link_libraries(example1 gloo)
+add_executable(looks_like_mpi looks_like_mpi.cc)
+target_link_libraries(looks_like_mpi gloo)

gloo/examples/looks_like_mpi.cc

@@ -0,0 +1,210 @@
+#include <sys/types.h>
+#include <unistd.h>
+
+#include <iostream>
+#include <memory>
+#include <string>
+
+#include <gloo/allreduce.h>
+#include <gloo/barrier.h>
+#include <gloo/rendezvous/context.h>
+#include <gloo/rendezvous/file_store.h>
+#include <gloo/rendezvous/prefix_store.h>
+#include <gloo/transport/tcp/device.h>
+
+#define ASSERT(expr)                                            \
+  do {                                                          \
+    if (!(expr)) {                                              \
+      throw std::runtime_error("Assertion failed: " #expr);     \
+    }                                                           \
+  } while (0);
+
+// Global context
+std::shared_ptr<gloo::Context> kContext;
+
+// Make this example look like MPI code
+using MPI_Comm = int;
+const MPI_Comm MPI_COMM_WORLD = 0;
+
+enum MPI_Datatype {
+  MPI_INT,
+};
+
+enum MPI_Op {
+  MPI_SUM,
+};
+
+// Same prototype as MPI API.
+int MPI_Comm_rank(MPI_Comm comm, int* rank) {
+  ASSERT(comm == MPI_COMM_WORLD);
+  if (rank) {
+    *rank = kContext->rank;
+  }
+}
+
+// Same prototype as MPI API.
+int MPI_Comm_size(MPI_Comm comm, int* size) {
+  ASSERT(comm == MPI_COMM_WORLD);
+  if (size) {
+    *size = kContext->size;
+  }
+}
+
+// Same prototype as MPI API.
+int MPI_Barrier(MPI_Comm comm) {
+  ASSERT(comm == MPI_COMM_WORLD);
+  gloo::BarrierOptions opts(kContext);
+  gloo::barrier(opts);
+}
+
+// Same prototype
+int MPI_Allreduce(
+    const void *sendbuf,
+    void *recvbuf,
+    int count,
+    MPI_Datatype datatype,
+    MPI_Op op,
+    MPI_Comm comm)
+{
+  ASSERT(datatype == MPI_INT);
+  ASSERT(op == MPI_SUM);
+  ASSERT(comm == MPI_COMM_WORLD);
+  gloo::AllreduceOptions opts(kContext);
+  opts.setInput(const_cast<int*>((const int*) sendbuf), count);
+  opts.setOutput((int*) recvbuf, count);
+  opts.setReduceFunction(
+      static_cast<void(*)(void*, const void*, const void*, size_t)>(
+          &gloo::sum<int>));
+  gloo::allreduce(opts);
+}
+
+// Actual prototype:
+//
+// int MPI_Recv(
+//     void *buf,
+//     int count,
+//     MPI_Datatype datatype,
+//     int source,
+//     int tag,
+//     MPI_Comm comm,
+//     MPI_Status *status);
+//
+// Implementation below:
+//   - doesn't use MPI_Datatype
+//   - doesn't take MPI_Comm argument but uses global
+//   - doesn't return an MPI_Status object
+//
+int MPI_Recv(
+    void *buf,
+    ssize_t bytes,
+    int source,
+    int tag,
+    MPI_Comm comm)
+{
+  auto ubuf = kContext->createUnboundBuffer(buf, bytes);
+  ubuf->recv(source, tag);
+  ubuf->waitRecv();
+}
+
+// Actual prototype:
+//
+// int MPI_Send(
+//     const void *buf,
+//     int count,
+//     MPI_Datatype datatype,
+//     int dest,
+//     int tag,
+//     MPI_Comm comm);
+//
+// Implementation below:
+//   - doesn't use MPI_Datatype
+//
+int MPI_Send(
+    const void *cbuf,
+    ssize_t bytes,
+    int dest,
+    int tag,
+    MPI_Comm comm)
+{
+  // Argument is logically const if we're only sending.
+  auto ubuf = kContext->createUnboundBuffer(const_cast<void*>(cbuf), bytes);
+  ubuf->send(dest, tag);
+  ubuf->waitSend();
+}
+
+// Entrypoint of this example.
+int run() {
+  int rank;
+  int size;
+
+  MPI_Comm_rank(MPI_COMM_WORLD, &rank);
+  MPI_Comm_size(MPI_COMM_WORLD, &size);
+
+  // Send on rank 0
+  if (rank == 0) {
+    const int dst = 1;
+    const int tag = 1234;
+    int pid = getpid();
+    MPI_Send(&pid, sizeof(pid), dst, tag, MPI_COMM_WORLD);
+    std::cout << "Sent to rank " << dst << ": " << pid << std::endl;
+  }
+
+  // Recv on rank 1
+  if (rank == 1) {
+    const int src = 0;
+    const int tag = 1234;
+    int pid = -1;
+    MPI_Recv(&pid, sizeof(pid), src, tag, MPI_COMM_WORLD);
+    std::cout << "Received from rank " << src << ": " << pid << std::endl;
+  }
+
+  // Run allreduce on the number 1
+  {
+    const int input = 1;
+    int output = -1;
+    MPI_Allreduce(&input, &output, 1, MPI_INT, MPI_SUM, MPI_COMM_WORLD);
+    std::cout << "Result of allreduce on " <<  input << ": " << output << std::endl;
+  }
+
+  // Barrier before exit
+  MPI_Barrier(MPI_COMM_WORLD);
+}
+
+// See example1.cc in this directory for a walkthrough of initialization.
+void init(const std::string& path) {
+  if (getenv("PREFIX") == nullptr ||
+      getenv("SIZE") == nullptr ||
+      getenv("RANK") == nullptr) {
+    std::cerr
+      << "Please set environment variables PREFIX, SIZE, and RANK."
+      << std::endl;
+    exit(1);
+  }
+
+  const std::string prefix = getenv("PREFIX");
+  const int rank = atoi(getenv("RANK"));
+  const int size = atoi(getenv("SIZE"));
+
+  // Initialize store
+  auto fileStore = gloo::rendezvous::FileStore(path);
+  auto prefixStore = gloo::rendezvous::PrefixStore(prefix, fileStore);
+
+  // Initialize device
+  gloo::transport::tcp::attr attr;
+  attr.iface = "eth0";
+  auto dev = gloo::transport::tcp::CreateDevice("localhost");
+
+  // Initialize global context
+  auto context = std::make_shared<gloo::rendezvous::Context>(rank, size);
+  context->connectFullMesh(prefixStore, dev);
+  kContext = std::move(context);
+}
+
+int main(int argc, char** argv) {
+  if (argc != 2) {
+    std::cerr << "Usage: " << argv[0] << " PATH" << std::endl;
+    exit(1);
+  }
+  init(argv[1]);
+  return run();
+}