@aislamov/onnxruntime-web64
Version:
A Javascript library for running ONNX models on browsers
359 lines • 16.6 kB
JavaScript
;
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
Object.defineProperty(exports, "__esModule", { value: true });
exports.extractTransferableBuffers = exports.endProfiling = exports.run = exports.releaseSession = exports.createSession = exports.createSessionFinalize = exports.createSessionAllocate = exports.initRuntime = void 0;
const run_options_1 = require("./run-options");
const session_options_1 = require("./session-options");
const wasm_common_1 = require("./wasm-common");
const wasm_factory_1 = require("./wasm-factory");
const wasm_utils_1 = require("./wasm-utils");
/**
* get the input/output count of the session.
* @param sessionHandle the handle representing the session. should be non-zero.
* @returns a tuple including 2 numbers, representing the input count and output count.
*/
const getSessionInputOutputCount = (sessionHandle) => {
const wasm = (0, wasm_factory_1.getInstance)();
const stack = wasm.stackSave();
try {
const ptrSize = wasm.PTR_SIZE;
const dataOffset = wasm.stackAlloc(2 * ptrSize);
const errorCode = wasm._OrtGetInputOutputCount(sessionHandle, dataOffset, dataOffset + ptrSize);
if (errorCode !== 0) {
(0, wasm_utils_1.checkLastError)('Can\'t get session input/output count.');
}
return [wasm.getValue(dataOffset, '*'), wasm.getValue(dataOffset + ptrSize, '*')];
}
finally {
wasm.stackRestore(stack);
}
};
/**
* initialize ORT environment.
* @param numThreads SetGlobalIntraOpNumThreads(numThreads)
* @param loggingLevel CreateEnv(static_cast<OrtLoggingLevel>(logging_level))
*/
const initOrt = (numThreads, loggingLevel) => {
const errorCode = (0, wasm_factory_1.getInstance)()._OrtInit(numThreads, loggingLevel);
if (errorCode !== 0) {
(0, wasm_utils_1.checkLastError)('Can\'t initialize onnxruntime.');
}
};
/**
* intialize runtime environment.
* @param env passed in the environment config object.
*/
const initRuntime = async (env) => {
// init ORT
initOrt(10, (0, wasm_common_1.logLevelStringToEnum)(env.logLevel));
if (!BUILD_DEFS.DISABLE_WEBGPU) {
// init JSEP if available
// eslint-disable-next-line @typescript-eslint/no-require-imports, @typescript-eslint/no-var-requires
const initJsep = require('./jsep/init').init;
await initJsep((0, wasm_factory_1.getInstance)(), env);
}
};
exports.initRuntime = initRuntime;
const activeSessions = new Map();
/**
* allocate the memory and memcpy the model bytes, preparing for creating an instance of InferenceSession.
* @returns a 3-elements tuple - the pointer, size of the allocated buffer, and optional weights.pb FS node
*/
const createSessionAllocate = (model, weights) => {
const wasm = (0, wasm_factory_1.getInstance)();
const modelDataOffset = wasm._malloc(model.byteLength);
wasm.HEAPU8.set(model, modelDataOffset);
let weightsFile;
if (weights) {
weightsFile = wasm.FS.create('/home/web_user/weights.pb');
weightsFile.contents = new Uint8Array(weights);
weightsFile.usedBytes = weights.byteLength;
wasm.FS.chdir('/home/web_user');
}
return [modelDataOffset, model.byteLength, weightsFile];
};
exports.createSessionAllocate = createSessionAllocate;
/**
* create an inference session using the prepared buffer containing the model data.
* @param modelData a 2-elements tuple containing the pointer and size of the model data buffer.
* @param options an optional session options object.
* @returns a 3-elements tuple containing [session handle, input names, output names]
*/
const createSessionFinalize = (modelData, options) => {
const wasm = (0, wasm_factory_1.getInstance)();
let sessionHandle = 0;
let sessionOptionsHandle = 0;
let allocs = [];
const inputNamesUTF8Encoded = [];
const outputNamesUTF8Encoded = [];
try {
[sessionOptionsHandle, allocs] = (0, session_options_1.setSessionOptions)(options);
sessionHandle = wasm._OrtCreateSession(modelData[0], modelData[1], sessionOptionsHandle);
if (sessionHandle === 0) {
(0, wasm_utils_1.checkLastError)('Can\'t create a session.');
}
const [inputCount, outputCount] = getSessionInputOutputCount(sessionHandle);
const inputNames = [];
const outputNames = [];
for (let i = 0; i < inputCount; i++) {
const name = wasm._OrtGetInputName(sessionHandle, i);
if (name === 0) {
(0, wasm_utils_1.checkLastError)('Can\'t get an input name.');
}
inputNamesUTF8Encoded.push(name);
inputNames.push(wasm.UTF8ToString(name));
}
for (let i = 0; i < outputCount; i++) {
const name = wasm._OrtGetOutputName(sessionHandle, i);
if (name === 0) {
(0, wasm_utils_1.checkLastError)('Can\'t get an output name.');
}
outputNamesUTF8Encoded.push(name);
outputNames.push(wasm.UTF8ToString(name));
}
activeSessions.set(sessionHandle, [sessionHandle, inputNamesUTF8Encoded, outputNamesUTF8Encoded]);
return [sessionHandle, inputNames, outputNames];
}
catch (e) {
inputNamesUTF8Encoded.forEach(buf => wasm._OrtFree(buf));
outputNamesUTF8Encoded.forEach(buf => wasm._OrtFree(buf));
if (sessionHandle !== 0) {
wasm._OrtReleaseSession(sessionHandle);
}
throw e;
}
finally {
wasm._free(modelData[0]);
if (sessionOptionsHandle !== 0) {
wasm._OrtReleaseSessionOptions(sessionOptionsHandle);
}
allocs.forEach(alloc => wasm._free(alloc));
if (modelData[2]) {
wasm.FS.unlink('/home/web_user/weights.pb');
}
}
};
exports.createSessionFinalize = createSessionFinalize;
/**
* create an instance of InferenceSession.
* @returns the metadata of InferenceSession. 0-value handle for failure.
*/
const createSession = async (model, options) => {
const modelData = (0, exports.createSessionAllocate)(model);
return (0, exports.createSessionFinalize)(modelData, options);
};
exports.createSession = createSession;
const releaseSession = (sessionId) => {
const wasm = (0, wasm_factory_1.getInstance)();
const session = activeSessions.get(sessionId);
if (!session) {
throw new Error(`cannot release session. invalid session id: ${sessionId}`);
}
const [sessionHandle, inputNamesUTF8Encoded, outputNamesUTF8Encoded] = session;
inputNamesUTF8Encoded.forEach(buf => wasm._OrtFree(buf));
outputNamesUTF8Encoded.forEach(buf => wasm._OrtFree(buf));
wasm._OrtReleaseSession(sessionHandle);
activeSessions.delete(sessionId);
};
exports.releaseSession = releaseSession;
/**
* perform inference run
*/
const run = async (sessionId, inputIndices, inputs, outputIndices, options) => {
const wasm = (0, wasm_factory_1.getInstance)();
const session = activeSessions.get(sessionId);
if (!session) {
throw new Error(`cannot run inference. invalid session id: ${sessionId}`);
}
const [sessionHandle, inputNamesUTF8Encoded, outputNamesUTF8Encoded] = session;
const inputCount = inputIndices.length;
const outputCount = outputIndices.length;
let runOptionsHandle = 0;
let runOptionsAllocs = [];
const inputValues = [];
const inputAllocs = [];
try {
[runOptionsHandle, runOptionsAllocs] = (0, run_options_1.setRunOptions)(options);
const ptrSize = wasm.PTR_SIZE;
// create input tensors
for (let i = 0; i < inputCount; i++) {
const dataType = inputs[i][0];
const dims = inputs[i][1];
const data = inputs[i][2];
let dataOffset;
let dataByteLength;
if (Array.isArray(data)) {
// string tensor
dataByteLength = 4 * data.length;
dataOffset = wasm._malloc(dataByteLength);
inputAllocs.push(dataOffset);
let dataIndex = dataOffset / 4;
for (let i = 0; i < data.length; i++) {
if (typeof data[i] !== 'string') {
throw new TypeError(`tensor data at index ${i} is not a string`);
}
wasm.HEAPU32[dataIndex++] = (0, wasm_utils_1.allocWasmString)(data[i], inputAllocs);
}
}
else {
dataByteLength = data.byteLength;
dataOffset = wasm._malloc(dataByteLength);
inputAllocs.push(dataOffset);
wasm.HEAPU8.set(new Uint8Array(data.buffer, data.byteOffset, dataByteLength), dataOffset);
}
const stack = wasm.stackSave();
const dimsOffset = wasm.stackAlloc(ptrSize * dims.length);
try {
dims.forEach((d, index) => wasm.setValue(dimsOffset + (index * ptrSize), d, '*'));
const tensor = wasm._OrtCreateTensor((0, wasm_common_1.tensorDataTypeStringToEnum)(dataType), dataOffset, dataByteLength, dimsOffset, dims.length);
if (tensor === 0) {
(0, wasm_utils_1.checkLastError)(`Can't create tensor for input[${i}].`);
}
inputValues.push(tensor);
}
finally {
wasm.stackRestore(stack);
}
}
const beforeRunStack = wasm.stackSave();
const inputValuesOffset = wasm.stackAlloc(inputCount * ptrSize);
const inputNamesOffset = wasm.stackAlloc(inputCount * ptrSize);
const outputValuesOffset = wasm.stackAlloc(outputCount * ptrSize);
const outputNamesOffset = wasm.stackAlloc(outputCount * ptrSize);
try {
let inputValuesIndex = inputValuesOffset / 8;
let inputNamesIndex = inputNamesOffset / 8;
let outputValuesIndex = outputValuesOffset / 8;
let outputNamesIndex = outputNamesOffset / 8;
for (let i = 0; i < inputCount; i++) {
wasm.HEAPU64[inputValuesIndex++] = BigInt(inputValues[i]);
wasm.HEAPU64[inputNamesIndex++] = BigInt(inputNamesUTF8Encoded[inputIndices[i]]);
}
for (let i = 0; i < outputCount; i++) {
wasm.HEAPU64[outputValuesIndex++] = BigInt(0);
wasm.HEAPU64[outputNamesIndex++] = BigInt(outputNamesUTF8Encoded[outputIndices[i]]);
}
// jsepOnRunStart is only available when JSEP is enabled.
wasm.jsepOnRunStart?.(sessionId);
// support RunOptions
let errorCode = await wasm._OrtRun(sessionHandle, inputNamesOffset, inputValuesOffset, inputCount, outputNamesOffset, outputCount, outputValuesOffset, runOptionsHandle);
const runPromise = wasm.jsepRunPromise;
if (runPromise) {
// jsepRunPromise is a Promise object. It is only available when JSEP is enabled.
//
// OrtRun() is a synchrnous call, but it internally calls async functions. Emscripten's ASYNCIFY allows it to
// work in this way. However, OrtRun() does not return a promise, so when code reaches here, it is earlier than
// the async functions are finished.
//
// To make it work, we created a Promise and resolve the promise when the C++ code actually reaches the end of
// OrtRun(). If the promise exists, we need to await for the promise to be resolved.
errorCode = await runPromise;
}
const jsepOnRunEnd = wasm.jsepOnRunEnd;
if (jsepOnRunEnd) {
// jsepOnRunEnd is only available when JSEP is enabled.
//
// it returns a promise, which is resolved or rejected when the following async functions are finished:
// - collecting GPU validation errors.
await jsepOnRunEnd(sessionId);
}
const output = [];
if (errorCode !== 0) {
(0, wasm_utils_1.checkLastError)('failed to call OrtRun().');
}
const ptrSize = 8;
for (let i = 0; i < outputCount; i++) {
const tensor = wasm.getValue(outputValuesOffset + i * ptrSize, '*');
const beforeGetTensorDataStack = wasm.stackSave();
// stack allocate 4 pointer value
const tensorDataOffset = wasm.stackAlloc(4 * 8);
let type, dataOffset = 0;
try {
errorCode = wasm._OrtGetTensorData(tensor, tensorDataOffset, tensorDataOffset + ptrSize, tensorDataOffset + ptrSize * 2, tensorDataOffset + ptrSize * 3);
if (errorCode !== 0) {
(0, wasm_utils_1.checkLastError)(`Can't access output tensor data on index ${i}.`);
}
const dataType = wasm.getValue(tensorDataOffset, '*');
dataOffset = wasm.getValue(tensorDataOffset + ptrSize, '*');
const dimsOffset = wasm.getValue(tensorDataOffset + ptrSize * 2, '*');
const dimsLength = wasm.getValue(tensorDataOffset + ptrSize * 3, '*');
const dims = [];
for (let i = 0; i < dimsLength; i++) {
dims.push(wasm.getValue(dimsOffset + i * ptrSize, '*'));
}
wasm._OrtFree(dimsOffset);
const size = dims.length === 0 ? 1 : dims.reduce((a, b) => a * b);
type = (0, wasm_common_1.tensorDataTypeEnumToString)(dataType);
if (type === 'string') {
const stringData = [];
let dataIndex = dataOffset / 4;
for (let i = 0; i < size; i++) {
const offset = wasm.HEAPU32[dataIndex++];
const maxBytesToRead = i === size - 1 ? undefined : wasm.HEAPU32[dataIndex] - offset;
stringData.push(wasm.UTF8ToString(offset, maxBytesToRead));
}
output.push([type, dims, stringData]);
}
else {
const typedArrayConstructor = (0, wasm_common_1.tensorTypeToTypedArrayConstructor)(type);
const data = new typedArrayConstructor(size);
new Uint8Array(data.buffer, data.byteOffset, data.byteLength)
.set(wasm.HEAPU8.subarray(dataOffset, dataOffset + data.byteLength));
output.push([type, dims, data]);
}
}
finally {
wasm.stackRestore(beforeGetTensorDataStack);
if (type === 'string' && dataOffset) {
wasm._free(dataOffset);
}
wasm._OrtReleaseTensor(tensor);
}
}
return output;
}
finally {
wasm.stackRestore(beforeRunStack);
}
}
finally {
inputValues.forEach(v => wasm._OrtReleaseTensor(v));
inputAllocs.forEach(p => wasm._free(p));
if (runOptionsHandle !== 0) {
wasm._OrtReleaseRunOptions(runOptionsHandle);
}
runOptionsAllocs.forEach(p => wasm._free(p));
}
};
exports.run = run;
/**
* end profiling
*/
const endProfiling = (sessionId) => {
const wasm = (0, wasm_factory_1.getInstance)();
const session = activeSessions.get(sessionId);
if (!session) {
throw new Error('invalid session id');
}
const sessionHandle = session[0];
// profile file name is not used yet, but it must be freed.
const profileFileName = wasm._OrtEndProfiling(sessionHandle);
if (profileFileName === 0) {
(0, wasm_utils_1.checkLastError)('Can\'t get an profile file name.');
}
wasm._OrtFree(profileFileName);
};
exports.endProfiling = endProfiling;
const extractTransferableBuffers = (tensors) => {
const buffers = [];
for (const tensor of tensors) {
const data = tensor[2];
if (!Array.isArray(data) && data.buffer) {
buffers.push(data.buffer);
}
}
return buffers;
};
exports.extractTransferableBuffers = extractTransferableBuffers;
//# sourceMappingURL=wasm-core-impl.js.map