wehub-resource-sync
315 lines
12 KiB
C++
315 lines
12 KiB
C++
/*
|
|
* Licensed to the Apache Software Foundation (ASF) under one
|
|
* or more contributor license agreements. See the NOTICE file
|
|
* distributed with this work for additional information
|
|
* regarding copyright ownership. The ASF licenses this file
|
|
* to you under the Apache License, Version 2.0 (the
|
|
* "License"); you may not use this file except in compliance
|
|
* with the License. You may obtain a copy of the License at
|
|
*
|
|
* http://www.apache.org/licenses/LICENSE-2.0
|
|
*
|
|
* Unless required by applicable law or agreed to in writing,
|
|
* software distributed under the License is distributed on an
|
|
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
|
* KIND, either express or implied. See the License for the
|
|
* specific language governing permissions and limitations
|
|
* under the License.
|
|
*/
|
|
|
|
/*
|
|
* \file tvmjs_support.cc
|
|
* \brief Support functions to be linked with wasm_runtime to provide
|
|
* ffi::Function callbacks in tvmjs.
|
|
* We do not need to link this file in standalone wasm.
|
|
*/
|
|
|
|
// configurations for tvm logging
|
|
#define TVM_LOG_DEBUG 0
|
|
#define TVM_LOG_CUSTOMIZE 1
|
|
#define TVM_FFI_ALWAYS_LOG_BEFORE_THROW 1
|
|
|
|
#include <tvm/ffi/extra/module.h>
|
|
#include <tvm/ffi/function.h>
|
|
#include <tvm/ffi/reflection/registry.h>
|
|
#include <tvm/runtime/device_api.h>
|
|
|
|
#include <unordered_set>
|
|
|
|
#include "../../src/runtime/rpc/rpc_local_session.h"
|
|
|
|
extern "C" {
|
|
// --- Additional C API for the Wasm runtime ---
|
|
/*!
|
|
* \brief Allocate space aligned to 64 bit.
|
|
* \param size The size of the space.
|
|
* \return The allocated space.
|
|
*/
|
|
TVM_DLL void* TVMWasmAllocSpace(int size);
|
|
|
|
/*!
|
|
* \brief Free the space allocated by TVMWasmAllocSpace.
|
|
* \param data The data pointer.
|
|
*/
|
|
TVM_DLL void TVMWasmFreeSpace(void* data);
|
|
|
|
/*!
|
|
* \brief Create ffi::Function from a resource handle.
|
|
* \param resource_handle The handle to the resource.
|
|
* \param out The output ffi::Function.
|
|
* \sa TVMWasmPackedCFunc, TVMWasmPackedCFuncFinalizer
|
|
3A * \return 0 if success.
|
|
*/
|
|
TVM_DLL int TVMFFIWasmFunctionCreate(void* resource_handle, TVMFFIObjectHandle* 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 new ffi call function caller.
|
|
*
|
|
* \param self The pointer to the ffi::Function.
|
|
* \param args The arguments
|
|
* \param num_args Number of arguments.
|
|
* \param result The return value handle.
|
|
* \return 0 if success, -1 if failure happens, set error via TVMAPISetLastError.
|
|
*/
|
|
extern int TVMFFIWasmSafeCall(void* self, const TVMFFIAny* args, int32_t num_args,
|
|
TVMFFIAny* result);
|
|
/*!
|
|
* \brief Delete ffi::Function.
|
|
* \param self The pointer to the ffi::Function.
|
|
*/
|
|
extern void TVMFFIWasmFunctionDeleter(void* self);
|
|
|
|
} // extern "C"
|
|
|
|
void* TVMWasmAllocSpace(int size) {
|
|
int num_count = (size + 7) / 8;
|
|
return new int64_t[num_count];
|
|
}
|
|
|
|
void TVMWasmFreeSpace(void* arr) { delete[] static_cast<int64_t*>(arr); }
|
|
|
|
int TVMFFIWasmFunctionCreate(void* self, TVMFFIObjectHandle* 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 {
|
|
namespace runtime {
|
|
|
|
// A special local session that can interact with async
|
|
// functions in the JS runtime.
|
|
class AsyncLocalSession : public LocalSession {
|
|
public:
|
|
AsyncLocalSession() {}
|
|
|
|
PackedFuncHandle GetFunction(const std::string& name) final {
|
|
if (name == "runtime.RPCTimeEvaluator") {
|
|
return get_time_eval_placeholder_.get();
|
|
} else if (auto fp = tvm::ffi::Function::GetGlobal(name)) {
|
|
TVMFFIAny val = tvm::ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(tvm::ffi::Any(*fp));
|
|
return val.v_obj;
|
|
} else if (auto fp = tvm::ffi::Function::GetGlobal("__async." + name)) {
|
|
TVMFFIAny val = tvm::ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(tvm::ffi::Any(*fp));
|
|
async_func_set_.insert(val.v_obj);
|
|
return val.v_obj;
|
|
} else {
|
|
return nullptr;
|
|
}
|
|
}
|
|
|
|
void FreeHandle(void* handle) final {
|
|
{
|
|
auto it = async_func_set_.find(handle);
|
|
if (it != async_func_set_.end()) {
|
|
async_func_set_.erase(it);
|
|
}
|
|
}
|
|
if (handle != get_time_eval_placeholder_.get()) {
|
|
LocalSession::FreeHandle(handle);
|
|
}
|
|
}
|
|
|
|
void AsyncCallFunc(PackedFuncHandle func, ffi::PackedArgs args, FAsyncCallback callback) final {
|
|
auto it = async_func_set_.find(func);
|
|
if (it != async_func_set_.end()) {
|
|
ffi::Function packed_callback([callback, this](ffi::PackedArgs args, ffi::Any*) {
|
|
int code = args[0].cast<int>();
|
|
ffi::Any rv;
|
|
rv = args[1];
|
|
if (code == static_cast<int>(RPCCode::kReturn)) {
|
|
this->EncodeReturn(std::move(rv), [&](ffi::PackedArgs encoded_args) {
|
|
callback(RPCCode::kReturn, encoded_args);
|
|
});
|
|
} else {
|
|
// for exception, we can pass through as since this is just normal encoding.
|
|
TVM_FFI_ICHECK_EQ(code, static_cast<int>(RPCCode::kException));
|
|
callback(RPCCode::kException, args);
|
|
}
|
|
});
|
|
|
|
std::vector<AnyView> packed_args(args.data(), args.data() + args.size());
|
|
// pass the callback as the last argument.
|
|
packed_args.emplace_back(AnyView(packed_callback));
|
|
auto* pf = static_cast<ffi::FunctionObj*>(func);
|
|
Any temp;
|
|
pf->CallPacked(packed_args.data(), packed_args.size(), &temp);
|
|
} else if (func == get_time_eval_placeholder_.get()) {
|
|
// special handle time evaluator.
|
|
try {
|
|
ffi::Function retfunc = this->GetTimeEvaluator(
|
|
args[0].cast<ffi::Optional<ffi::Module>>(), args[1].cast<std::string>(),
|
|
args[2].cast<int>(), args[3].cast<int>(), args[4].cast<int>(), args[5].cast<int>(),
|
|
args[6].cast<int>(), args[7].cast<int>(), args[8].cast<int>(), args[9].cast<int>());
|
|
ffi::Any rv;
|
|
rv = retfunc;
|
|
this->EncodeReturn(std::move(rv), [&](ffi::PackedArgs encoded_args) {
|
|
const void* pf = encoded_args[0].as<ffi::FunctionObj>();
|
|
TVM_FFI_ICHECK(pf != nullptr);
|
|
// mark as async.
|
|
async_func_set_.insert(const_cast<void*>(pf));
|
|
callback(RPCCode::kReturn, encoded_args);
|
|
});
|
|
} catch (const std::runtime_error& e) {
|
|
this->SendException(callback, e.what());
|
|
}
|
|
} else {
|
|
LocalSession::AsyncCallFunc(func, args, callback);
|
|
}
|
|
}
|
|
|
|
void AsyncCopyToRemote(void* local_from_bytes, DLTensor* remote_to, uint64_t nbytes,
|
|
FAsyncCallback on_complete) final {
|
|
try {
|
|
DLTensor local_from;
|
|
local_from.data = local_from_bytes;
|
|
local_from.device = Device{kDLCPU, 0};
|
|
local_from.ndim = remote_to->ndim;
|
|
local_from.shape = remote_to->shape;
|
|
local_from.dtype = remote_to->dtype;
|
|
local_from.strides = nullptr;
|
|
local_from.byte_offset = 0;
|
|
this->GetDeviceAPI(remote_to->device)->CopyDataFromTo(&local_from, remote_to, nullptr);
|
|
this->AsyncStreamWait(remote_to->device, nullptr, on_complete);
|
|
} catch (const std::runtime_error& e) {
|
|
this->SendException(on_complete, e.what());
|
|
}
|
|
}
|
|
|
|
void AsyncCopyFromRemote(DLTensor* remote_from, void* local_to_bytes, uint64_t nbytes,
|
|
FAsyncCallback on_complete) final {
|
|
try {
|
|
DLTensor local_to;
|
|
local_to.data = local_to_bytes;
|
|
local_to.device = Device{kDLCPU, 0};
|
|
local_to.ndim = remote_from->ndim;
|
|
local_to.shape = remote_from->shape;
|
|
local_to.dtype = remote_from->dtype;
|
|
local_to.strides = nullptr;
|
|
local_to.byte_offset = 0;
|
|
this->GetDeviceAPI(remote_from->device)->CopyDataFromTo(remote_from, &local_to, nullptr);
|
|
this->AsyncStreamWait(remote_from->device, nullptr, on_complete);
|
|
} catch (const std::runtime_error& e) {
|
|
this->SendException(on_complete, e.what());
|
|
}
|
|
}
|
|
|
|
void AsyncStreamWait(Device dev, TVMStreamHandle stream, FAsyncCallback on_complete) final {
|
|
if (dev.device_type == kDLCPU) {
|
|
AnyView packed_args[1];
|
|
packed_args[0] = nullptr;
|
|
on_complete(RPCCode::kReturn, ffi::PackedArgs(packed_args, 1));
|
|
} else {
|
|
TVM_FFI_ICHECK(dev.device_type == static_cast<DLDeviceType>(kDLWebGPU));
|
|
if (!async_wait_.has_value()) {
|
|
async_wait_ = tvm::ffi::Function::GetGlobal("__async.wasm.WebGPUWaitForTasks");
|
|
}
|
|
TVM_FFI_ICHECK(async_wait_.has_value());
|
|
ffi::Function packed_callback([on_complete](ffi::PackedArgs args, ffi::Any*) {
|
|
int code = args[0].cast<int>();
|
|
on_complete(static_cast<RPCCode>(code), args.Slice(1));
|
|
});
|
|
(*async_wait_)(packed_callback);
|
|
}
|
|
}
|
|
|
|
bool IsAsync() const final { return true; }
|
|
|
|
private:
|
|
std::unordered_set<void*> async_func_set_;
|
|
std::unique_ptr<ffi::Function> get_time_eval_placeholder_ = std::make_unique<ffi::Function>();
|
|
std::optional<ffi::Function> async_wait_;
|
|
|
|
// time evaluator
|
|
ffi::Function GetTimeEvaluator(ffi::Optional<ffi::Module> opt_mod, std::string name,
|
|
int device_type, int device_id, int number, int repeat,
|
|
int min_repeat_ms, int limit_zero_time_iterations,
|
|
int cooldown_interval_ms, int repeats_to_cooldown) {
|
|
Device dev;
|
|
dev.device_type = static_cast<DLDeviceType>(device_type);
|
|
dev.device_id = device_id;
|
|
|
|
if (opt_mod.has_value()) {
|
|
ffi::Module m = opt_mod.value();
|
|
std::string tkey = m->kind();
|
|
return WrapWasmTimeEvaluator(m->GetFunction(name, false).value(), dev, number, repeat,
|
|
min_repeat_ms, limit_zero_time_iterations, cooldown_interval_ms,
|
|
repeats_to_cooldown);
|
|
} else {
|
|
auto pf = tvm::ffi::Function::GetGlobal(name);
|
|
TVM_FFI_ICHECK(pf.has_value()) << "Cannot find " << name << " in the global function";
|
|
return WrapWasmTimeEvaluator(*pf, dev, number, repeat, min_repeat_ms,
|
|
limit_zero_time_iterations, cooldown_interval_ms,
|
|
repeats_to_cooldown);
|
|
}
|
|
}
|
|
|
|
// time evaluator
|
|
ffi::Function WrapWasmTimeEvaluator(ffi::Function pf, Device dev, int number, int repeat,
|
|
int min_repeat_ms, int limit_zero_time_iterations,
|
|
int cooldown_interval_ms, int repeats_to_cooldown) {
|
|
auto ftimer = [pf, dev, number, repeat, min_repeat_ms, limit_zero_time_iterations,
|
|
cooldown_interval_ms, repeats_to_cooldown](ffi::PackedArgs args, ffi::Any* rv) {
|
|
// the function is a async function.
|
|
ffi::Function on_complete = args[args.size() - 1].cast<ffi::Function>();
|
|
|
|
std::vector<AnyView> packed_args(args.data(), args.data() + args.size() - 1);
|
|
auto finvoke = [pf, packed_args](int n) {
|
|
ffi::Any temp;
|
|
ffi::PackedArgs invoke_args(packed_args.data(), packed_args.size());
|
|
for (int i = 0; i < n; ++i) {
|
|
pf.CallPacked(invoke_args, &temp);
|
|
}
|
|
};
|
|
auto time_exec = tvm::ffi::Function::GetGlobal("__async.wasm.TimeExecution");
|
|
TVM_FFI_ICHECK(time_exec.has_value()) << "Cannot find wasm.GetTimer in the global function";
|
|
(*time_exec)(ffi::TypedFunction<void(int)>(finvoke), dev, number, repeat, min_repeat_ms,
|
|
limit_zero_time_iterations, cooldown_interval_ms, repeats_to_cooldown,
|
|
/*cache_flush_bytes=*/0, on_complete);
|
|
};
|
|
return ffi::Function(ftimer);
|
|
}
|
|
};
|
|
|
|
TVM_FFI_STATIC_INIT_BLOCK() {
|
|
namespace refl = tvm::ffi::reflection;
|
|
refl::GlobalDef().def("wasm.LocalSession", []() {
|
|
return CreateRPCSessionModule(std::make_shared<AsyncLocalSession>());
|
|
});
|
|
}
|
|
|
|
} // namespace runtime
|
|
} // namespace tvm
|