MPI support in WAMR (#731)

* wamr: barebones mpi implmenetation to run lammps

* wamr: mpi asynchronous api working

* wamr: more mpi functions working

* wamr: fix MPI_IN_PLACE detection to get MPI_Reduce to work properly

* wamr: implement more mpi calls, mpi_cart_create not working

* wamr: fix mpi_cart_create

* wamr: all mpi stubs for lammps in place

* wamr: mpi tests passing locally

* nits: self-review cleanup

* mpi: add logging

* wamr: unaligned write to make ubsan happy

* wamr: make ubsan happy with mpi cart creation

* wamr: more ubsan fixes

faasmcli/faasmcli/tasks/codegen.py

@@ -18,35 +18,6 @@
     join(FAASM_RUNTIME_ROOT, "lib", "fake", "libfakeLibB.so"),
 ]
 
-WAMR_ALLOWED_FUNCS = [
-    # Misc
-    ["demo", "chain"],
-    ["demo", "chain_named_a"],
-    ["demo", "chain_named_b"],
-    ["demo", "chain_named_c"],
-    # Environment
-    ["demo", "argc_argv_test"],
-    ["demo", "exit"],
-    ["demo", "getenv"],
-    ["errors", "ret_one"],
-    # Memory
-    ["demo", "brk"],
-    ["demo", "mmap"],
-    ["demo", "mmap_big"],
-    # Filesystem
-    ["demo", "fcntl"],
-    ["demo", "file"],
-    ["demo", "filedescriptor"],
-    ["demo", "fstat"],
-    ["demo", "fread"],
-    ["demo", "shared_file"],
-    # Input output
-    ["demo", "check_input"],
-    ["demo", "echo"],
-    ["demo", "stdout"],
-    ["demo", "stderr"],
-]
-
 SGX_ALLOWED_FUNCS = [
     ["demo", "hello"],
     ["demo", "chain_named_a"],
@@ -81,11 +52,11 @@ def codegen(ctx, user, function, clean=False):
 
 
 @task
-def user(ctx, user):
+def user(ctx, user, clean=False):
     """
     Generates machine for all the functions belonging to the given user
     """
-    _do_codegen_user(user)
+    _do_codegen_user(user, clean=clean)
 
 
 def _do_codegen_user(user, clean=False):
@@ -95,7 +66,6 @@ def _do_codegen_user(user, clean=False):
     env = copy(environ)
     env.update(
         {
-            "WASM_VM": "wavm",
             "LD_LIBRARY_PATH": "/usr/local/lib/",
         }
     )
@@ -158,8 +128,8 @@ def wamr(ctx, clean=False):
     """
     env = copy(environ)
     env.update({"WASM_VM": "wamr"})
-    for user, func in WAMR_ALLOWED_FUNCS:
-        codegen(ctx, user, func, clean)
+    _do_codegen_user("demo", clean)
+    _do_codegen_user("mpi", clean)
 
 
 @task

include/wamr/WAMRModuleMixin.h

@@ -36,6 +36,12 @@ struct WAMRModuleMixin
           moduleInstance, nativePtr, size);
     }
 
+    void* wasmOffsetToNativePointer(uint32_t wasmOffset)
+    {
+        auto moduleInstance = this->underlying().getModuleInstance();
+        return wasm_runtime_addr_app_to_native(moduleInstance, wasmOffset);
+    }
+
     // Convert a native pointer to the corresponding offset in the WASM linear
     // memory.
     uint32_t nativePointerToWasmOffset(void* nativePtr)

include/wamr/native.h

@@ -57,6 +57,8 @@ uint32_t getFaasmFunctionsApi(NativeSymbol** nativeSymbols);
 
 uint32_t getFaasmMemoryApi(NativeSymbol** nativeSymbols);
 
+uint32_t getFaasmMpiApi(NativeSymbol** nativeSymbols);
+
 uint32_t getFaasmProcessApi(NativeSymbol** nativeSymbols);
 
 uint32_t getFaasmPthreadApi(NativeSymbol** nativeSymbols);

src/wamr/CMakeLists.txt

@@ -81,6 +81,7 @@ faasm_private_lib(wamrmodule
     filesystem.cpp
     funcs.cpp
     memory.cpp
+    mpi.cpp
     native.cpp
     process.cpp
     pthread.cpp

src/wamr/filesystem.cpp

@@ -178,6 +178,13 @@ static int32_t wasi_fd_filestat_get(wasm_exec_env_t exec_env,
     return doFileStat(fd, "", statWasm);
 }
 
+static int32_t wasi_fd_filestat_set_size(wasm_exec_env_t execEnv,
+                                         int32_t a,
+                                         int64_t b)
+{
+    throw std::runtime_error("wasi_fd_filestat_set_size not implemented!");
+}
+
 static uint32_t wasi_fd_pread(wasm_exec_env_t exec_env,
                               __wasi_fd_t fd,
                               iovec_app_t* iovecWasm,
@@ -251,7 +258,7 @@ static int32_t wasi_fd_read(wasm_exec_env_t exec_env,
     storage::FileSystem& fileSystem = module->getFileSystem();
     std::string path = fileSystem.getPathForFd(fd);
 
-    SPDLOG_DEBUG("S - fd_read {} ({})", fd, path);
+    SPDLOG_TRACE("S - fd_read {} ({})", fd, path);
 
     storage::FileDescriptor fileDesc = fileSystem.getFileDescriptor(fd);
 
@@ -334,7 +341,7 @@ static int32_t wasi_fd_write(wasm_exec_env_t exec_env,
     storage::FileSystem& fileSystem = module->getFileSystem();
     std::string path = fileSystem.getPathForFd(fd);
 
-    SPDLOG_DEBUG("S - fd_write {} ({})", fd, path);
+    SPDLOG_TRACE("S - fd_write {} ({})", fd, path);
 
     // Check pointers
     module->validateNativePointer(reinterpret_cast<void*>(ioVecBuffWasm),
@@ -559,6 +566,7 @@ static NativeSymbol wasiNs[] = {
     REG_WASI_NATIVE_FUNC(fd_fdstat_set_flags, "(ii)i"),
     REG_WASI_NATIVE_FUNC(fd_fdstat_set_rights, "(iII)i"),
     REG_WASI_NATIVE_FUNC(fd_filestat_get, "(i*)i"),
+    REG_WASI_NATIVE_FUNC(fd_filestat_set_size, "(iI)i"),
     REG_WASI_NATIVE_FUNC(fd_pread, "(i*iI*)i"),
     REG_WASI_NATIVE_FUNC(fd_prestat_dir_name, "(i*~)i"),
     REG_WASI_NATIVE_FUNC(fd_prestat_get, "(i*)i"),

src/wamr/funcs.cpp

@@ -3,17 +3,78 @@
 #include <faabric/util/bytes.h>
 #include <faabric/util/logging.h>
 #include <faabric/util/macros.h>
-
+#include <wamr/WAMRWasmModule.h>
 #include <wamr/native.h>
 #include <wasm/WasmExecutionContext.h>
 #include <wasm/WasmModule.h>
 #include <wasm/chaining.h>
+
 #include <wasm_export.h>
 
 using namespace faabric::scheduler;
 
 namespace wasm {
 
+static std::shared_ptr<faabric::state::StateKeyValue> getStateKV(
+  int32_t* keyPtr,
+  size_t size)
+{
+    WAMRWasmModule* module = getExecutingWAMRModule();
+    module->validateNativePointer(keyPtr, sizeof(int32_t));
+
+    const faabric::Message* call = &ExecutorContext::get()->getMsg();
+    char* key = reinterpret_cast<char*>(keyPtr); // second
+
+    faabric::state::State& s = faabric::state::getGlobalState();
+    auto kv = s.getKV(call->user(), key, size);
+
+    return kv;
+}
+
+/**
+ * Await a chained function's completion
+ */
+static int32_t __faasm_await_call_wrapper(wasm_exec_env_t exec_env,
+                                          int32_t callId)
+{
+    SPDLOG_DEBUG("S - faasm_await_call {}", callId);
+
+    int32_t result = wasm::awaitChainedCall((uint32_t)callId);
+    return result;
+}
+
+/**
+ * Chain a function by function pointer
+ */
+static int32_t __faasm_chain_ptr_wrapper(wasm_exec_env_t exec_env,
+                                         int32_t wasmFuncPtr,
+                                         char* inBuff,
+                                         int32_t inLen)
+{
+    SPDLOG_DEBUG("S - faasm_chain_ptr {} {} {}", wasmFuncPtr, inBuff, inLen);
+
+    faabric::Message& call = ExecutorContext::get()->getMsg();
+    std::vector<uint8_t> inputData(BYTES(inBuff), BYTES(inBuff) + inLen);
+    return makeChainedCall(call.function(), wasmFuncPtr, nullptr, inputData);
+}
+
+static void __faasm_pull_state_wrapper(wasm_exec_env_t execEnv,
+                                       int32_t* keyPtr,
+                                       int32_t stateLen)
+{
+    auto kv = getStateKV(keyPtr, stateLen);
+    SPDLOG_DEBUG("S - pull_state - {} {}", kv->key, stateLen);
+
+    kv->pull();
+}
+
+static void __faasm_push_state_wrapper(wasm_exec_env_t execEnv, int32_t* keyPtr)
+{
+    auto kv = getStateKV(keyPtr, 0);
+    SPDLOG_DEBUG("S - push_state - {}", kv->key);
+    kv->pushFull();
+}
+
 /**
  * Read the function input
  */
@@ -51,38 +112,13 @@ static void __faasm_write_output_wrapper(wasm_exec_env_t exec_env,
     call.set_outputdata(outBuff, outLen);
 }
 
-/**
- * Chain a function by function pointer
- */
-static int32_t __faasm_chain_ptr_wrapper(wasm_exec_env_t exec_env,
-                                         int32_t wasmFuncPtr,
-                                         char* inBuff,
-                                         int32_t inLen)
-{
-    SPDLOG_DEBUG("S - faasm_chain_ptr {} {} {}", wasmFuncPtr, inBuff, inLen);
-
-    faabric::Message& call = ExecutorContext::get()->getMsg();
-    std::vector<uint8_t> inputData(BYTES(inBuff), BYTES(inBuff) + inLen);
-    return makeChainedCall(call.function(), wasmFuncPtr, nullptr, inputData);
-}
-
-/**
- * Await a chained function's completion
- */
-static int32_t __faasm_await_call_wrapper(wasm_exec_env_t exec_env,
-                                          int32_t callId)
-{
-    SPDLOG_DEBUG("S - faasm_await_call {}", callId);
-
-    int32_t result = wasm::awaitChainedCall((uint32_t)callId);
-    return result;
-}
-
 static NativeSymbol ns[] = {
-    REG_NATIVE_FUNC(__faasm_write_output, "($i)"),
-    REG_NATIVE_FUNC(__faasm_read_input, "($i)i"),
-    REG_NATIVE_FUNC(__faasm_chain_ptr, "(i$i)i"),
     REG_NATIVE_FUNC(__faasm_await_call, "(i)i"),
+    REG_NATIVE_FUNC(__faasm_chain_ptr, "(i$i)i"),
+    REG_NATIVE_FUNC(__faasm_pull_state, "(*i)"),
+    REG_NATIVE_FUNC(__faasm_push_state, "(*)"),
+    REG_NATIVE_FUNC(__faasm_read_input, "($i)i"),
+    REG_NATIVE_FUNC(__faasm_write_output, "($i)"),
 };
 
 uint32_t getFaasmFunctionsApi(NativeSymbol** nativeSymbols)

src/wamr/memory.cpp

@@ -33,7 +33,7 @@ static int32_t mmap_wrapper(wasm_exec_env_t exec_env,
                             int32_t fd,
                             int64_t offset)
 {
-    SPDLOG_DEBUG(
+    SPDLOG_TRACE(
       "S - mmap - {} {} {} {} {} {}", addr, length, prot, flags, fd, offset);
 
     if (offset != 0) {