Skip to content

Commit

Permalink
[js/node] enable float16 support for Node.js binding (#20581)
Browse files Browse the repository at this point in the history
### Description
enable float16 support for Node.js binding.

data of float16 tensor uses `Uint16Array`.
  • Loading branch information
fs-eire authored Jul 28, 2024
1 parent a4d3a1c commit dbff0cd
Showing 1 changed file with 7 additions and 9 deletions.
16 changes: 7 additions & 9 deletions js/node/src/tensor_helper.cc
Original file line number Diff line number Diff line change
Expand Up @@ -38,13 +38,13 @@ constexpr size_t DATA_TYPE_ELEMENT_SIZE_MAP[] = {
2, // ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT16
2, // ONNX_TENSOR_ELEMENT_DATA_TYPE_INT16
4, // ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32
8, // ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64 INT64 not working in Javascript
8, // ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64
0, // ONNX_TENSOR_ELEMENT_DATA_TYPE_STRING N/A
1, // ONNX_TENSOR_ELEMENT_DATA_TYPE_BOOL
0, // ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16 FLOAT16 not working in Javascript
2, // ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16
8, // ONNX_TENSOR_ELEMENT_DATA_TYPE_DOUBLE
4, // ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT32
8, // ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT64 UINT64 not working in Javascript
8, // ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT64
0, // ONNX_TENSOR_ELEMENT_DATA_TYPE_COMPLEX64 not supported
0, // ONNX_TENSOR_ELEMENT_DATA_TYPE_COMPLEX128 not supported
0 // ONNX_TENSOR_ELEMENT_DATA_TYPE_BFLOAT16 not supported
Expand All @@ -60,13 +60,13 @@ constexpr napi_typedarray_type DATA_TYPE_TYPEDARRAY_MAP[] = {
napi_uint16_array, // ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT16
napi_int16_array, // ONNX_TENSOR_ELEMENT_DATA_TYPE_INT16
napi_int32_array, // ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32
napi_bigint64_array, // ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64 INT64 not working i
napi_bigint64_array, // ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64
(napi_typedarray_type)(-1), // ONNX_TENSOR_ELEMENT_DATA_TYPE_STRING not supported
napi_uint8_array, // ONNX_TENSOR_ELEMENT_DATA_TYPE_BOOL
(napi_typedarray_type)(-1), // ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16 FLOAT16 not working
napi_uint16_array, // ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16 FLOAT16 uses Uint16Array
napi_float64_array, // ONNX_TENSOR_ELEMENT_DATA_TYPE_DOUBLE
napi_uint32_array, // ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT32
napi_biguint64_array, // ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT64 UINT64 not working
napi_biguint64_array, // ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT64
(napi_typedarray_type)(-1), // ONNX_TENSOR_ELEMENT_DATA_TYPE_COMPLEX64 not supported
(napi_typedarray_type)(-1), // ONNX_TENSOR_ELEMENT_DATA_TYPE_COMPLEX128 not supported
(napi_typedarray_type)(-1) // ONNX_TENSOR_ELEMENT_DATA_TYPE_BFLOAT16 not supported
Expand Down Expand Up @@ -182,9 +182,7 @@ Ort::Value NapiValueToOrtValue(Napi::Env env, Napi::Value value, OrtMemoryInfo *

char *buffer = reinterpret_cast<char *>(tensorDataTypedArray.ArrayBuffer().Data());
size_t bufferByteOffset = tensorDataTypedArray.ByteOffset();
// there is a bug in TypedArray::ElementSize(): https://github.com/nodejs/node-addon-api/pull/705
// TODO: change to TypedArray::ByteLength() in next node-addon-api release.
size_t bufferByteLength = tensorDataTypedArray.ElementLength() * DATA_TYPE_ELEMENT_SIZE_MAP[elemType];
size_t bufferByteLength = tensorDataTypedArray.ByteLength();
return Ort::Value::CreateTensor(memory_info, buffer + bufferByteOffset, bufferByteLength,
dims.empty() ? nullptr : &dims[0], dims.size(), elemType);
}
Expand Down

0 comments on commit dbff0cd

Please sign in to comment.