[REFACTOR][FFI] Upgrade Web Runtime to new FFI

This PR refactors the web runtime to the new FFI protocol.
Tested through rpc tests and local tests.

Author: tqchen

ffi/include/tvm/ffi/c_api.h

@@ -579,8 +579,10 @@ TVM_FFI_DLL int TVMFFIDataTypeFromString(const TVMFFIByteArray* str, DLDataType*
  * \return 0 when success, nonzero when failure happens
  * \note out is a String object that needs to be freed by the caller via TVMFFIObjectFree.
          The content of string can be accessed via TVMFFIObjectGetByteArrayPtr.
+
+ * \note The input dtype is a pointer to the DLDataType to avoid ABI compatibility issues.
  */
-TVM_FFI_DLL int TVMFFIDataTypeToString(DLDataType dtype, TVMFFIObjectHandle* out);
+TVM_FFI_DLL int TVMFFIDataTypeToString(const DLDataType* dtype, TVMFFIObjectHandle* out);
 
 //------------------------------------------------------------
 // Section: Backend noexcept functions for internal use

ffi/include/tvm/ffi/dtype.h

@@ -121,7 +121,7 @@ inline DLDataType StringToDLDataType(const String& str) {
 
 inline String DLDataTypeToString(DLDataType dtype) {
   TVMFFIObjectHandle out;
-  TVM_FFI_CHECK_SAFE_CALL(TVMFFIDataTypeToString(dtype, &out));
+  TVM_FFI_CHECK_SAFE_CALL(TVMFFIDataTypeToString(&dtype, &out));
   return String(details::ObjectUnsafe::ObjectPtrFromOwned<Object>(static_cast<TVMFFIObject*>(out)));
 }
 

ffi/include/tvm/ffi/error.h

@@ -51,6 +51,10 @@
 #define TVM_FFI_BACKTRACE_ON_SEGFAULT 1
 #endif
 
+#ifndef TVM_FFI_ALWAYS_LOG_BEFORE_THROW
+#define TVM_FFI_ALWAYS_LOG_BEFORE_THROW 0
+#endif
+
 namespace tvm {
 namespace ffi {
 
@@ -212,8 +216,10 @@ class ErrorBuilder {
  *
  * \endcode
  */
-#define TVM_FFI_THROW(ErrorKind) \
-  ::tvm::ffi::details::ErrorBuilder(#ErrorKind, TVM_FFI_TRACEBACK_HERE, false).stream()
+#define TVM_FFI_THROW(ErrorKind)                                        \
+  ::tvm::ffi::details::ErrorBuilder(#ErrorKind, TVM_FFI_TRACEBACK_HERE, \
+                                    TVM_FFI_ALWAYS_LOG_BEFORE_THROW)    \
+      .stream()
 
 /*!
  * \brief Explicitly log error in stderr and then throw the error.

ffi/src/ffi/dtype.cc

@@ -320,9 +320,9 @@ int TVMFFIDataTypeFromString(const TVMFFIByteArray* str, DLDataType* out) {
   TVM_FFI_SAFE_CALL_END();
 }
 
-int TVMFFIDataTypeToString(DLDataType dtype, TVMFFIObjectHandle* out) {
+int TVMFFIDataTypeToString(const DLDataType* dtype, TVMFFIObjectHandle* out) {
   TVM_FFI_SAFE_CALL_BEGIN();
-  tvm::ffi::String out_str(tvm::ffi::DLDataTypeToString_(dtype));
+  tvm::ffi::String out_str(tvm::ffi::DLDataTypeToString_(*dtype));
   *out = tvm::ffi::details::ObjectUnsafe::MoveObjectRefToTVMFFIObjectPtr(std::move(out_str));
   TVM_FFI_SAFE_CALL_END();
 }

python/tvm/ffi/cython/base.pxi

@@ -150,7 +150,7 @@ cdef extern from "tvm/ffi/c_api.h":
     int TVMFFIEnvRegisterCAPI(TVMFFIByteArray* name, void* ptr) nogil
     int TVMFFITypeKeyToIndex(TVMFFIByteArray* type_key, int32_t* out_tindex) nogil
     int TVMFFIDataTypeFromString(TVMFFIByteArray* str, DLDataType* out) nogil
-    int TVMFFIDataTypeToString(DLDataType dtype, TVMFFIObjectHandle* out) nogil
+    int TVMFFIDataTypeToString(const DLDataType* dtype, TVMFFIObjectHandle* out) nogil
     const TVMFFIByteArray* TVMFFITraceback(const char* filename, int lineno, const char* func) nogil;
     int TVMFFINDArrayFromDLPack(DLManagedTensor* src, int32_t require_alignment,
                                 int32_t require_contiguous, TVMFFIObjectHandle* out) nogil

python/tvm/ffi/cython/dtype.pxi

@@ -94,7 +94,7 @@ cdef class DataType:
     def __str__(self):
         cdef TVMFFIObjectHandle dtype_str
         cdef TVMFFIByteArray* bytes
-        CHECK_CALL(TVMFFIDataTypeToString(self.cdtype, &dtype_str))
+        CHECK_CALL(TVMFFIDataTypeToString(&(self.cdtype), &dtype_str))
         bytes = TVMFFIBytesGetByteArrayPtr(dtype_str)
         res = py_str(PyBytes_FromStringAndSize(bytes.data, bytes.size))
         CHECK_CALL(TVMFFIObjectFree(dtype_str))

web/.eslintignore

@@ -1,2 +1,4 @@
 dist
 debug
+tvmjs_runtime_wasi.js
+lib

web/apps/node/example.js

@@ -31,7 +31,7 @@ const wasmSource = fs.readFileSync(path.join(wasmPath, "tvmjs_runtime.wasm"));
 tvmjs.instantiate(wasmSource, tvmjs.createPolyfillWASI())
 .then((tvm) => {
     tvm.beginScope();
-    const log_info = tvm.getGlobalFunc("testing.log_info_str");
+    const log_info = tvm.getGlobalFunc("tvmjs.testing.log_info_str");
     log_info("hello world");
     // List all the global functions from the runtime.
     console.log("Runtime functions using EmccWASI\n", tvm.listGlobalFuncNames());

web/emcc/tvmjs_support.cc

@@ -28,12 +28,11 @@
 #define TVM_LOG_STACK_TRACE 0
 #define TVM_LOG_DEBUG 0
 #define TVM_LOG_CUSTOMIZE 1
+#define TVM_FFI_ALWAYS_LOG_BEFORE_THROW 1
 #define DMLC_USE_LOGGING_LIBRARY <tvm/runtime/logging.h>
 
-#include <tvm/runtime/c_runtime_api.h>
+#include <tvm/ffi/function.h>
 #include <tvm/runtime/device_api.h>
-#include <tvm/runtime/packed_func.h>
-#include <tvm/runtime/registry.h>
 
 #include "../../src/runtime/rpc/rpc_local_session.h"
 
@@ -59,27 +58,33 @@ TVM_DLL void TVMWasmFreeSpace(void* data);
  * \sa TVMWasmPackedCFunc, TVMWasmPackedCFuncFinalizer
 3A * \return 0 if success.
  */
-TVM_DLL int TVMWasmFuncCreateFromCFunc(void* resource_handle, TVMFunctionHandle* out);
+TVM_DLL int TVMFFIWasmFunctionCreate(void* resource_handle, TVMFunctionHandle* out);
+
+/*!
+ * \brief Get the last error message.
+ * \return The last error message.
+ */
+TVM_DLL const char* TVMFFIWasmGetLastError();
 
 // --- APIs to be implemented by the frontend. ---
+
 /*!
- * \brief Wasm frontend packed function caller.
+ * \brief Wasm frontend new ffi call function caller.
  *
+ * \param self The pointer to the ffi::Function.
  * \param args The arguments
- * \param type_codes The type codes of the arguments
  * \param num_args Number of arguments.
- * \param ret The return value handle.
- * \param resource_handle The handle additional resource handle from front-end.
+ * \param result The return value handle.
  * \return 0 if success, -1 if failure happens, set error via TVMAPISetLastError.
  */
-extern int TVMWasmPackedCFunc(TVMValue* args, int* type_codes, int num_args, TVMRetValueHandle ret,
-                              void* resource_handle);
-
+extern int TVMFFIWasmSafeCall(void* self, const TVMFFIAny* args, int32_t num_args,
+                              TVMFFIAny* result);
 /*!
- * \brief Wasm frontend resource finalizer.
- * \param resource_handle The pointer to the external resource.
+ * \brief Delete ffi::Function.
+ * \param self The pointer to the ffi::Function.
  */
-extern void TVMWasmPackedCFuncFinalizer(void* resource_handle);
+extern void TVMFFIWasmFunctionDeleter(void* self);
+
 }  // extern "C"
 
 void* TVMWasmAllocSpace(int size) {
@@ -89,9 +94,14 @@ void* TVMWasmAllocSpace(int size) {
 
 void TVMWasmFreeSpace(void* arr) { delete[] static_cast<int64_t*>(arr); }
 
-int TVMWasmFuncCreateFromCFunc(void* resource_handle, TVMFunctionHandle* out) {
-  return TVMFuncCreateFromCFunc(TVMWasmPackedCFunc, resource_handle, TVMWasmPackedCFuncFinalizer,
-                                out);
+int TVMFFIWasmFunctionCreate(void* self, TVMFunctionHandle* out) {
+  return TVMFFIFunctionCreate(self, TVMFFIWasmSafeCall, TVMFFIWasmFunctionDeleter, out);
+}
+
+const char* TVMFFIWasmGetLastError() {
+  static thread_local std::string last_error;
+  last_error = ::tvm::ffi::details::MoveFromSafeCallRaised().what();
+  return last_error.c_str();
 }
 
 namespace tvm {
@@ -291,7 +301,7 @@ class AsyncLocalSession : public LocalSession {
   }
 };
 
-TVM_REGISTER_GLOBAL("wasm.LocalSession").set_body_typed([]() {
+TVM_FFI_REGISTER_GLOBAL("wasm.LocalSession").set_body_typed([]() {
   return CreateRPCSessionModule(std::make_shared<AsyncLocalSession>());
 });
 

web/emcc/wasm_runtime.cc

@@ -27,9 +27,9 @@
 #define TVM_LOG_DEBUG 0
 #define TVM_LOG_CUSTOMIZE 1
 #define TVM_FFI_USE_LIBBACKTRACE 0
+#define TVM_FFI_ALWAYS_LOG_BEFORE_THROW 1
 #define DMLC_USE_LOGGING_LIBRARY <tvm/runtime/logging.h>
 
-#include <tvm/runtime/c_runtime_api.h>
 #include <tvm/runtime/logging.h>
 
 #include "src/runtime/c_runtime_api.cc"
@@ -107,45 +107,24 @@ void LogMessageImpl(const std::string& file, int lineno, int level, const std::s
 
 }  // namespace detail
 
-TVM_REGISTER_GLOBAL("testing.echo").set_body_packed([](ffi::PackedArgs args, ffi::Any* ret) {
-  *ret = args[0];
-});
-
-TVM_REGISTER_GLOBAL("testing.call").set_body_packed([](ffi::PackedArgs args, ffi::Any* ret) {
-  (args[0].cast<ffi::Function>()).CallPacked(args.Slice(1), ret);
-});
-
-TVM_REGISTER_GLOBAL("testing.ret_string").set_body_packed([](ffi::PackedArgs args, ffi::Any* ret) {
-  *ret = args[0].cast<String>();
-});
-
-TVM_REGISTER_GLOBAL("testing.log_info_str")
+TVM_FFI_REGISTER_GLOBAL("tvmjs.testing.call")
     .set_body_packed([](ffi::PackedArgs args, ffi::Any* ret) {
-      LOG(INFO) << args[0].cast<String>();
+      (args[0].cast<ffi::Function>()).CallPacked(args.Slice(1), ret);
     });
 
-TVM_REGISTER_GLOBAL("testing.log_fatal_str")
+TVM_FFI_REGISTER_GLOBAL("tvmjs.testing.log_info_str")
     .set_body_packed([](ffi::PackedArgs args, ffi::Any* ret) {
-      LOG(FATAL) << args[0].cast<String>();
+      LOG(INFO) << args[0].cast<String>();
     });
 
-TVM_REGISTER_GLOBAL("testing.add_one").set_body_typed([](int x) { return x + 1; });
+TVM_FFI_REGISTER_GLOBAL("tvmjs.testing.add_one").set_body_typed([](int x) { return x + 1; });
 
-TVM_REGISTER_GLOBAL("testing.wrap_callback")
+TVM_FFI_REGISTER_GLOBAL("tvmjs.testing.wrap_callback")
     .set_body_packed([](ffi::PackedArgs args, ffi::Any* ret) {
       ffi::Function pf = args[0].cast<ffi::Function>();
       *ret = ffi::TypedFunction<void()>([pf]() { pf(); });
     });
 
-// internal function used for debug and testing purposes
-TVM_REGISTER_GLOBAL("testing.object_use_count")
-    .set_body_packed([](ffi::PackedArgs args, ffi::Any* ret) {
-      auto obj = args[0].cast<ffi::ObjectRef>();
-      // subtract the current one because we always copy
-      // and get another value.
-      *ret = (obj.use_count() - 1);
-    });
-
 void ArrayDecodeStorage(NDArray cpu_arr, std::string bytes, std::string format, std::string dtype) {
   if (format == "f32-to-bf16" && dtype == "float32") {
     std::vector<uint16_t> buffer(bytes.length() / 2);
@@ -167,10 +146,10 @@ void ArrayDecodeStorage(NDArray cpu_arr, std::string bytes, std::string format,
   }
 }
 
-TVM_REGISTER_GLOBAL("tvmjs.array.decode_storage").set_body_typed(ArrayDecodeStorage);
+TVM_FFI_REGISTER_GLOBAL("tvmjs.array.decode_storage").set_body_typed(ArrayDecodeStorage);
 
 // Concatenate n TVMArrays
-TVM_REGISTER_GLOBAL("tvmjs.runtime.ArrayConcat")
+TVM_FFI_REGISTER_GLOBAL("tvmjs.runtime.ArrayConcat")
     .set_body_packed([](ffi::PackedArgs args, ffi::Any* ret) {
       std::vector<Any> data;
       for (int i = 0; i < args.size(); ++i) {
@@ -220,7 +199,7 @@ NDArray ConcatEmbeddings(const std::vector<NDArray>& embeddings) {
 }
 
 // Concatenate n NDArrays
-TVM_REGISTER_GLOBAL("tvmjs.runtime.ConcatEmbeddings")
+TVM_FFI_REGISTER_GLOBAL("tvmjs.runtime.ConcatEmbeddings")
     .set_body_packed([](ffi::PackedArgs args, ffi::Any* ret) {
       std::vector<NDArray> embeddings;
       for (int i = 0; i < args.size(); ++i) {
@@ -230,5 +209,19 @@ TVM_REGISTER_GLOBAL("tvmjs.runtime.ConcatEmbeddings")
       *ret = result;
     });
 
+TVM_FFI_REGISTER_GLOBAL("tvmjs.runtime.NDArrayCopyFromBytes")
+    .set_body_typed([](NDArray nd, TVMFFIByteArray* bytes) {
+      nd.CopyFromBytes(bytes->data, bytes->size);
+    });
+
+TVM_FFI_REGISTER_GLOBAL("tvmjs.runtime.NDArrayCopyToBytes")
+    .set_body_typed([](NDArray nd) -> ffi::Bytes {
+      size_t size = GetDataSize(*(nd.operator->()));
+      std::string bytes;
+      bytes.resize(size);
+      nd.CopyToBytes(bytes.data(), size);
+      return ffi::Bytes(bytes);
+    });
+
 }  // namespace runtime
 }  // namespace tvm