@aislamov/onnxruntime-web64
Version:
A Javascript library for running ONNX models on browsers
159 lines • 7.4 kB
JavaScript
"use strict";
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
Object.defineProperty(exports, "__esModule", { value: true });
exports.init = void 0;
const wasm_common_1 = require("../wasm-common");
const backend_webgpu_1 = require("./backend-webgpu");
const log_1 = require("./log");
const util_1 = require("./util");
/* eslint-disable no-bitwise */
class TensorViewImpl {
constructor(module, dataType, data, dims) {
this.module = module;
this.dataType = dataType;
this.data = data;
this.dims = dims;
}
getFloat32Array() {
if (this.dataType !== 1 /* DataType.float */) {
throw new Error('Invalid data type');
}
const elementCount = util_1.ShapeUtil.size(this.dims);
return elementCount === 0 ? new Float32Array() :
new Float32Array(this.module.HEAP8.buffer, this.data, elementCount);
}
getBigInt64Array() {
if (this.dataType !== 7 /* DataType.int64 */) {
throw new Error('Invalid data type');
}
const elementCount = util_1.ShapeUtil.size(this.dims);
return elementCount === 0 ? new BigInt64Array() :
new BigInt64Array(this.module.HEAP8.buffer, this.data, elementCount);
}
getInt32Array() {
if (this.dataType !== 6 /* DataType.int32 */) {
throw new Error('Invalid data type');
}
const elementCount = util_1.ShapeUtil.size(this.dims);
return elementCount === 0 ? new Int32Array() : new Int32Array(this.module.HEAP8.buffer, this.data, elementCount);
}
reshape(newDims) {
if (util_1.ShapeUtil.size(newDims) !== util_1.ShapeUtil.size(this.dims)) {
throw new Error('Invalid new shape');
}
return new TensorViewImpl(this.module, this.dataType, this.data, newDims);
}
}
class ComputeContextImpl {
get kernelCustomData() {
return this.backend.currentKernelCustomData;
}
get customDataBuffer() {
return this.module.HEAPU8.subarray(this.customDataOffset, this.customDataOffset + this.customDataSize);
}
constructor(module, backend, contextDataOffset) {
this.module = module;
this.backend = backend;
this.customDataOffset = 0;
this.customDataSize = 0;
const heap = module.PTR_SIZE === 4 ? module.HEAPU32 : module.HEAPU64;
// extract context data
let dataIndex = module.PTR_SIZE === 8 ? (contextDataOffset / 2 ** 3) : (contextDataOffset >> 2);
this.opKernelContext = Number(heap[dataIndex++]);
const inputCount = Number(heap[dataIndex++]);
this.outputCount = Number(heap[dataIndex++]);
this.customDataOffset = Number(heap[dataIndex++]);
this.customDataSize = Number(heap[dataIndex++]);
const inputs = [];
for (let i = 0; i < inputCount; i++) {
const dataType = Number(heap[dataIndex++]);
const data = Number(heap[dataIndex++]);
const dim = Number(heap[dataIndex++]);
const dims = [];
for (let d = 0; d < dim; d++) {
dims.push(Number(heap[dataIndex++]));
}
inputs.push(new TensorViewImpl(module, dataType, data, dims));
}
this.inputs = inputs;
}
compute(program, inputsOutputsMapping) {
// prepare inputs. inputs should always be valid data.
const mappedInputs = inputsOutputsMapping?.inputs?.map(i => typeof i === 'number' ? this.inputs[i] : i) ?? this.inputs;
// prepare outputs.
const outputIndices = inputsOutputsMapping?.outputs ?? [];
const createKernelOutput = (index, dataType, dims) => new TensorViewImpl(this.module, dataType, this.output(index, dims), dims);
const createTemporaryOutput = (dataType, dims) => {
const elementSize = (0, wasm_common_1.getTensorElementSize)(dataType);
if (!elementSize) {
throw new Error(`Unsupported data type: ${dataType}`);
}
const bufferSize = elementSize * util_1.ShapeUtil.size(dims);
return new TensorViewImpl(this.module, dataType, this.backend.gpuDataManager.create(bufferSize).id, dims);
};
return this.backend.run(program, mappedInputs, outputIndices, createKernelOutput, createTemporaryOutput);
}
output(index, dims) {
const stack = this.module.stackSave();
try {
const ptrSize = this.module.PTR_SIZE;
const data = this.module.stackAlloc((1 + dims.length) * ptrSize /* sizeof(size_t) */);
this.module.setValue(data, dims.length, '*');
for (let i = 0; i < dims.length; i++) {
this.module.setValue(data + ptrSize * (i + 1), dims[i], '*');
}
return this.module._JsepOutput(this.opKernelContext, index, data);
}
finally {
this.module.stackRestore(stack);
}
}
}
const init = async (module, env) => {
const init = module.jsepInit;
if (init && navigator.gpu) {
if (!env.wasm.simd) {
throw new Error('Not supported for WebGPU=ON and SIMD=OFF. Please set `env.wasm.simd` to true when using WebGPU EP');
}
const backend = new backend_webgpu_1.WebGpuBackend();
await backend.initialize(env);
init(
// backend
{ backend },
// jsepAlloc()
(size) => backend.alloc(Number(size)),
// jsepFree()
(ptr) => backend.free(Number(ptr)),
// jsepCopy(src, dst, size, isSourceGpu)
(src, dst, size, isSourceGpu = false) => {
if (isSourceGpu) {
(0, log_1.LOG_DEBUG)('verbose', () => `[WebGPU] jsepCopyGpuToGpu: src=${src}, dst=${dst}, size=${size}`);
backend.memcpy(Number(src), Number(dst));
}
else {
(0, log_1.LOG_DEBUG)('verbose', () => `[WebGPU] jsepCopyCpuToGpu: dataOffset=${src}, gpuDataId=${dst}, size=${size}`);
const data = module.HEAPU8.subarray(Number(src), Number(src) + Number(size));
backend.upload(Number(dst), data);
}
},
// jsepCopyAsync(src, dst, size)
async (gpuDataId, dataOffset, size) => {
(0, log_1.LOG_DEBUG)('verbose', () => `[WebGPU] jsepCopyGpuToCpu: gpuDataId=${gpuDataId}, dataOffset=${dataOffset}, size=${size}`);
await backend.download(Number(gpuDataId), () => module.HEAPU8.subarray(Number(dataOffset), Number(dataOffset) + Number(size)));
},
// jsepCreateKernel
(name, kernel, attribute) => backend.createKernel(name, kernel, attribute, env.debug || env.webgpu.profilingMode === 'default' ? module.UTF8ToString(module._JsepGetNodeName(kernel)) :
`${kernel}`),
// jsepReleaseKernel
(kernel) => backend.releaseKernel(Number(kernel)),
// jsepRun
(kernel, contextDataOffset, sessionState) => {
(0, log_1.LOG_DEBUG)('verbose', () => `[WebGPU] jsepRun: sessionId=${sessionState.sessionId}, kernel=${kernel}, contextDataOffset=${contextDataOffset}`);
const context = new ComputeContextImpl(module, backend, Number(contextDataOffset));
return backend.computeKernel(kernel, context, sessionState.errors);
});
}
};
exports.init = init;
//# sourceMappingURL=init.js.map