/**
|
* @license
|
* Copyright 2018 Google LLC. All Rights Reserved.
|
* Licensed 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.
|
* =============================================================================
|
*/
|
|
#include <node_api.h>
|
#include "tfjs_backend.h"
|
#include "utils.h"
|
|
namespace tfnodejs {
|
|
/**
|
* Get the TFJSBackend associated with this environment instance.
|
* While this does throw a JS error if the instance data is not found,
|
* the caller should still check if the return value is a `nullptr`.
|
*/
|
static inline TFJSBackend *GetTFJSBackend(napi_env env)
|
{
|
TFJSBackend *backend = nullptr;
|
napi_status nstatus = napi_get_instance_data(env, reinterpret_cast<void **>(&backend));
|
ENSURE_NAPI_OK_RETVAL(env, nstatus, nullptr);
|
return backend;
|
}
|
|
static void AssignIntProperty(napi_env env, napi_value exports,
|
const char *name, int32_t value) {
|
|
napi_value js_value;
|
napi_status nstatus = napi_create_int32(env, value, &js_value);
|
ENSURE_NAPI_OK(env, nstatus);
|
|
napi_property_descriptor property = {name, nullptr, nullptr,
|
nullptr, nullptr, js_value,
|
napi_default, nullptr};
|
nstatus = napi_define_properties(env, exports, 1, &property);
|
ENSURE_NAPI_OK(env, nstatus);
|
}
|
|
static napi_value CreateTensor(napi_env env, napi_callback_info info) {
|
napi_status nstatus;
|
|
// Create tensor takes 3 params: shape, dtype, typed-array/array:
|
size_t argc = 3;
|
napi_value args[3];
|
napi_value js_this;
|
nstatus = napi_get_cb_info(env, info, &argc, args, &js_this, nullptr);
|
ENSURE_NAPI_OK_RETVAL(env, nstatus, nullptr);
|
|
if (argc < 3) {
|
NAPI_THROW_ERROR(env,
|
"Invalid number of args passed to createTensor(). "
|
"Expecting 3 args but got %d.",
|
argc);
|
return nullptr;
|
}
|
|
ENSURE_VALUE_IS_ARRAY_RETVAL(env, args[0], nullptr);
|
ENSURE_VALUE_IS_NUMBER_RETVAL(env, args[1], nullptr);
|
|
// The third array can either be a typed array or an array:
|
bool is_typed_array;
|
nstatus = napi_is_typedarray(env, args[2], &is_typed_array);
|
ENSURE_NAPI_OK_RETVAL(env, nstatus, nullptr);
|
if (!is_typed_array) {
|
ENSURE_VALUE_IS_ARRAY_RETVAL(env, args[2], nullptr);
|
}
|
|
TFJSBackend *const backend = GetTFJSBackend(env);
|
if (!backend) return nullptr;
|
|
return backend->CreateTensor(env, args[0], args[1], args[2]);
|
}
|
|
static napi_value DeleteTensor(napi_env env, napi_callback_info info) {
|
napi_status nstatus;
|
|
// Delete tensor takes 1 param: tensor ID;
|
size_t argc = 1;
|
napi_value args[1];
|
napi_value js_this;
|
nstatus = napi_get_cb_info(env, info, &argc, args, &js_this, nullptr);
|
ENSURE_NAPI_OK_RETVAL(env, nstatus, js_this);
|
|
if (argc < 1) {
|
NAPI_THROW_ERROR(env,
|
"Invalid number of args passed to deleteTensor(). "
|
"Expecting 1 arg but got %d.",
|
argc);
|
return js_this;
|
}
|
|
ENSURE_VALUE_IS_NUMBER_RETVAL(env, args[0], js_this);
|
|
TFJSBackend *const backend = GetTFJSBackend(env);
|
if (!backend) return nullptr;
|
|
backend->DeleteTensor(env, args[0]);
|
return js_this;
|
}
|
|
static napi_value TensorDataSync(napi_env env, napi_callback_info info) {
|
napi_status nstatus;
|
|
// Tensor data-sync takes 1 param: tensor ID;
|
size_t argc = 1;
|
napi_value args[1];
|
napi_value js_this;
|
nstatus = napi_get_cb_info(env, info, &argc, args, &js_this, nullptr);
|
ENSURE_NAPI_OK_RETVAL(env, nstatus, js_this);
|
|
if (argc < 1) {
|
NAPI_THROW_ERROR(env,
|
"Invalid number of args passed to tensorDataSync(). "
|
"Expecting 1 arg but got %d.",
|
argc);
|
return nullptr;
|
}
|
|
ENSURE_VALUE_IS_NUMBER_RETVAL(env, args[0], js_this);
|
|
TFJSBackend *const backend = GetTFJSBackend(env);
|
if (!backend) return nullptr;
|
|
return backend->GetTensorData(env, args[0]);
|
}
|
|
static napi_value ExecuteOp(napi_env env, napi_callback_info info) {
|
napi_status nstatus;
|
|
// Create tensor takes 4 params: op-name, op-attrs, input-tensor-ids,
|
// num-outputs:
|
size_t argc = 4;
|
napi_value args[4];
|
napi_value js_this;
|
nstatus = napi_get_cb_info(env, info, &argc, args, &js_this, nullptr);
|
ENSURE_NAPI_OK_RETVAL(env, nstatus, nullptr);
|
|
if (argc < 4) {
|
NAPI_THROW_ERROR(env,
|
"Invalid number of args passed to executeOp(). Expecting "
|
"4 args but got %d.",
|
argc);
|
return nullptr;
|
}
|
|
ENSURE_VALUE_IS_STRING_RETVAL(env, args[0], nullptr);
|
ENSURE_VALUE_IS_ARRAY_RETVAL(env, args[1], nullptr);
|
ENSURE_VALUE_IS_ARRAY_RETVAL(env, args[2], nullptr);
|
ENSURE_VALUE_IS_NUMBER_RETVAL(env, args[3], nullptr);
|
|
TFJSBackend *const backend = GetTFJSBackend(env);
|
if (!backend) return nullptr;
|
|
return backend->ExecuteOp(env, args[0], args[1], args[2], args[3]);
|
}
|
|
static napi_value IsUsingGPUDevice(napi_env env, napi_callback_info info) {
|
napi_value result;
|
|
TFJSBackend *const backend = GetTFJSBackend(env);
|
if (!backend) return nullptr;
|
|
napi_status nstatus;
|
nstatus = napi_get_boolean(env, backend->is_gpu_device, &result);
|
ENSURE_NAPI_OK_RETVAL(env, nstatus, nullptr);
|
|
return result;
|
}
|
|
static napi_value LoadSavedModel(napi_env env, napi_callback_info info) {
|
napi_status nstatus;
|
|
// Load saved model takes 2 params: export_dir, tags:
|
size_t argc = 2;
|
napi_value args[2];
|
napi_value js_this;
|
nstatus = napi_get_cb_info(env, info, &argc, args, &js_this, nullptr);
|
ENSURE_NAPI_OK_RETVAL(env, nstatus, nullptr);
|
|
if (argc < 2) {
|
NAPI_THROW_ERROR(env,
|
"Invalid number of args passed to LoadSavedModel(). "
|
"Expecting 2 args but got %d.",
|
argc);
|
return nullptr;
|
}
|
|
ENSURE_VALUE_IS_STRING_RETVAL(env, args[0], nullptr);
|
ENSURE_VALUE_IS_STRING_RETVAL(env, args[1], nullptr);
|
|
TFJSBackend *const backend = GetTFJSBackend(env);
|
if (!backend) return nullptr;
|
|
return backend->LoadSavedModel(env, args[0], args[1]);
|
}
|
|
static napi_value DeleteSavedModel(napi_env env, napi_callback_info info) {
|
napi_status nstatus;
|
|
// Delete SavedModel takes 1 param: savedModel ID;
|
size_t argc = 1;
|
napi_value args[1];
|
napi_value js_this;
|
nstatus = napi_get_cb_info(env, info, &argc, args, &js_this, nullptr);
|
ENSURE_NAPI_OK_RETVAL(env, nstatus, js_this);
|
|
if (argc < 1) {
|
NAPI_THROW_ERROR(env,
|
"Invalid number of args passed to deleteSavedModel(). "
|
"Expecting 1 arg but got %d.",
|
argc);
|
return js_this;
|
}
|
|
ENSURE_VALUE_IS_NUMBER_RETVAL(env, args[0], js_this);
|
|
TFJSBackend *const backend = GetTFJSBackend(env);
|
if (!backend) return nullptr;
|
|
backend->DeleteSavedModel(env, args[0]);
|
return js_this;
|
}
|
|
static napi_value RunSavedModel(napi_env env, napi_callback_info info) {
|
napi_status nstatus;
|
|
// Run SavedModel takes 4 params: session_id, input_tensor_ids,
|
// input_op_names, output_op_names.
|
size_t argc = 4;
|
napi_value args[4];
|
napi_value js_this;
|
nstatus = napi_get_cb_info(env, info, &argc, args, &js_this, nullptr);
|
ENSURE_NAPI_OK_RETVAL(env, nstatus, nullptr);
|
|
if (argc < 4) {
|
NAPI_THROW_ERROR(env, "Invalid number of args passed to RunSavedModel()");
|
return nullptr;
|
}
|
|
ENSURE_VALUE_IS_NUMBER_RETVAL(env, args[0], nullptr);
|
ENSURE_VALUE_IS_ARRAY_RETVAL(env, args[1], nullptr);
|
ENSURE_VALUE_IS_STRING_RETVAL(env, args[2], nullptr);
|
ENSURE_VALUE_IS_STRING_RETVAL(env, args[3], nullptr);
|
|
TFJSBackend *const backend = GetTFJSBackend(env);
|
if (!backend) return nullptr;
|
|
return backend->RunSavedModel(env, args[0], args[1], args[2], args[3]);
|
}
|
|
static napi_value GetNumOfTensors(napi_env env, napi_callback_info info) {
|
TFJSBackend *const backend = GetTFJSBackend(env);
|
if (!backend) return nullptr;
|
// Delete SavedModel takes 0 param;
|
return backend->GetNumOfTensors(env);
|
}
|
|
static napi_value GetNumOfSavedModels(napi_env env, napi_callback_info info) {
|
TFJSBackend *const backend = GetTFJSBackend(env);
|
if (!backend) return nullptr;
|
// Delete SavedModel takes 0 param;
|
return backend->GetNumOfSavedModels(env);
|
}
|
|
/**
|
* Called by Node to cleanup our instance data, which is
|
* the TFJSBackend allocated in `InitTFNodeJSBinding`.
|
*/
|
static void FinalizeTFNodeJSBinding(napi_env env, void *finalize_data, void *finalize_hint)
|
{
|
delete reinterpret_cast<TFJSBackend *>(finalize_data);
|
}
|
|
static napi_value InitTFNodeJSBinding(napi_env env, napi_value exports) {
|
|
napi_status nstatus;
|
|
TFJSBackend *const backend = TFJSBackend::Create(env);
|
ENSURE_VALUE_IS_NOT_NULL_RETVAL(env, backend, nullptr);
|
|
// Store the backend in node's instance data for this addon
|
nstatus = napi_set_instance_data(env, backend, &FinalizeTFNodeJSBinding, nullptr);
|
ENSURE_NAPI_OK_RETVAL(env, nstatus, exports);
|
|
// TF version
|
napi_value tf_version;
|
nstatus = napi_create_string_latin1(env, TF_Version(), -1, &tf_version);
|
ENSURE_NAPI_OK_RETVAL(env, nstatus, exports);
|
|
// Set all export values list here.
|
napi_property_descriptor exports_properties[] = {
|
{"createTensor", nullptr, CreateTensor, nullptr, nullptr, nullptr,
|
napi_default, nullptr},
|
{"deleteTensor", nullptr, DeleteTensor, nullptr, nullptr, nullptr,
|
napi_default, nullptr},
|
{"tensorDataSync", nullptr, TensorDataSync, nullptr, nullptr, nullptr,
|
napi_default, nullptr},
|
{"executeOp", nullptr, ExecuteOp, nullptr, nullptr, nullptr, napi_default,
|
nullptr},
|
{"loadSavedModel", nullptr, LoadSavedModel, nullptr, nullptr, nullptr,
|
napi_default, nullptr},
|
{"deleteSavedModel", nullptr, DeleteSavedModel, nullptr, nullptr, nullptr,
|
napi_default, nullptr},
|
{"runSavedModel", nullptr, RunSavedModel, nullptr, nullptr, nullptr,
|
napi_default, nullptr},
|
{"TF_Version", nullptr, nullptr, nullptr, nullptr, tf_version,
|
napi_default, nullptr},
|
{"isUsingGpuDevice", nullptr, IsUsingGPUDevice, nullptr, nullptr, nullptr,
|
napi_default, nullptr},
|
{"getNumOfSavedModels", nullptr, GetNumOfSavedModels, nullptr, nullptr,
|
nullptr, napi_default, nullptr},
|
{"getNumOfTensors", nullptr, GetNumOfTensors, nullptr, nullptr,
|
nullptr, napi_default, nullptr},
|
|
};
|
nstatus = napi_define_properties(env, exports, ARRAY_SIZE(exports_properties),
|
exports_properties);
|
ENSURE_NAPI_OK_RETVAL(env, nstatus, exports);
|
|
// Export TF property types to JS
|
#define EXPORT_INT_PROPERTY(v) AssignIntProperty(env, exports, #v, v)
|
// Types
|
EXPORT_INT_PROPERTY(TF_FLOAT);
|
EXPORT_INT_PROPERTY(TF_INT32);
|
EXPORT_INT_PROPERTY(TF_INT64);
|
EXPORT_INT_PROPERTY(TF_BOOL);
|
EXPORT_INT_PROPERTY(TF_COMPLEX64);
|
EXPORT_INT_PROPERTY(TF_STRING);
|
EXPORT_INT_PROPERTY(TF_RESOURCE);
|
EXPORT_INT_PROPERTY(TF_UINT8);
|
|
// Op AttrType
|
EXPORT_INT_PROPERTY(TF_ATTR_STRING);
|
EXPORT_INT_PROPERTY(TF_ATTR_INT);
|
EXPORT_INT_PROPERTY(TF_ATTR_FLOAT);
|
EXPORT_INT_PROPERTY(TF_ATTR_BOOL);
|
EXPORT_INT_PROPERTY(TF_ATTR_TYPE);
|
EXPORT_INT_PROPERTY(TF_ATTR_SHAPE);
|
#undef EXPORT_INT_PROPERTY
|
|
return exports;
|
}
|
|
NAPI_MODULE(tfe_binding, InitTFNodeJSBinding)
|
|
} // namespace tfnodejs
|
