From bb1871332f5e37ebaa6a508fed460ab836fb23c5 Mon Sep 17 00:00:00 2001 From: Yulong Wang <7679871+fs-eire@users.noreply.github.com> Date: Sun, 27 Aug 2023 19:50:17 -0700 Subject: [PATCH] [js/webgpu] add kernel Not and Equal (#17306) ### Description This PR adds kernel implementation for operator "Not" and "Equal". Also removed download cache in gpu data manager. **Why removing download cache** The following test case failed. ("Or" is on CPU, "Greater" and "Equal" are on JSEP) ![image](https://github.com/microsoft/onnxruntime/assets/7679871/8d9798ad-2703-4fb9-907e-ff716c67d0b2) after debugging, I found that both "Equal" and "Greater" are using the same output GPU Data ID. This is because when ORT executes the graph, it first run "Equal", allowing its shader to write into GPU Data ID 2; then a Gpu2Cpu copy for it is issued (because currently "Or" is on CPU EP); at this point, ORT thinks GPU Data ID=2 is free to use; so it reuse it as output for "Greater". This means there is no allocation for output of "Greater" kernel, and both kernel writes to GPU Data ID=2. For gpu data manager, there will be 2 downloads from the same GPU buffer. Previously I think this is a waste of resource so I cached the data. But now it shoes that we need to perform 2 downloads because the GPU data is already different. The download data cache should be removed. ### Motivation and Context --- js/web/docs/webgpu-operators.md | 2 ++ .../lib/wasm/jsep/webgpu/gpu-data-manager.ts | 25 +------------------ .../lib/wasm/jsep/webgpu/op-resolve-rules.ts | 2 ++ js/web/lib/wasm/jsep/webgpu/ops/binary-op.ts | 20 +++++++-------- js/web/lib/wasm/jsep/webgpu/ops/common.ts | 5 ++++ js/web/lib/wasm/jsep/webgpu/ops/unary-op.ts | 4 +++ js/web/test/suite-test-list.jsonc | 10 ++++---- .../providers/js/js_execution_provider.cc | 10 ++++++++ .../core/providers/js/operators/binary.cc | 6 +++++ .../core/providers/js/operators/unary.cc | 3 +++ 10 files changed, 47 insertions(+), 40 deletions(-) diff --git a/js/web/docs/webgpu-operators.md b/js/web/docs/webgpu-operators.md index a210071bc1..c56bf4c6ff 100644 --- a/js/web/docs/webgpu-operators.md +++ b/js/web/docs/webgpu-operators.md @@ -31,6 +31,7 @@ Do not modify directly.* | Cosh | ai.onnx(9+) | | | Div | ai.onnx(7-12,13,14+) | | | Elu | ai.onnx(6+) | | +| Equal | ai.onnx(7-10,11-12,13-18,19+) | | | Erf | ai.onnx(9-12,13+) | | | Exp | ai.onnx(6-12,13+) | | | Expand | ai.onnx(8-12,13+) | | @@ -53,6 +54,7 @@ Do not modify directly.* | MemcpyToHost | ai.onnx(1+) | | | Mul | ai.onnx(7-12,13,14+) | | | Neg | ai.onnx(6-12,13+) | | +| Not | ai.onnx(1+) | | | Pow | ai.onnx(7-11,12,13-14,15+) | | | Reciprocal | ai.onnx(6-12,13+) | | | ReduceL1 | ai.onnx(1-10,11-12,13-17,18+) | | diff --git a/js/web/lib/wasm/jsep/webgpu/gpu-data-manager.ts b/js/web/lib/wasm/jsep/webgpu/gpu-data-manager.ts index 2e5bf19f5e..92fdd5abc3 100644 --- a/js/web/lib/wasm/jsep/webgpu/gpu-data-manager.ts +++ b/js/web/lib/wasm/jsep/webgpu/gpu-data-manager.ts @@ -57,10 +57,6 @@ interface StorageCacheValue { originalSize: number; } -interface DownloadCacheValue { - data: Promise; -} - /** * normalize the buffer size so that it fits the 128-bits (16 bytes) alignment. */ @@ -73,9 +69,6 @@ class GpuDataManagerImpl implements GpuDataManager { // GPU Data ID => GPU Data ( storage buffer ) storageCache: Map; - // GPU Data ID => GPU Data ( read buffer ) - downloadCache: Map; - // pending buffers for uploading ( data is unmapped ) private buffersForUploadingPending: GPUBuffer[]; // pending buffers for computing @@ -86,7 +79,6 @@ class GpuDataManagerImpl implements GpuDataManager { constructor(private backend: WebGpuBackend) { this.storageCache = new Map(); - this.downloadCache = new Map(); this.freeBuffers = new Map(); this.buffersForUploadingPending = []; this.buffersPending = []; @@ -198,20 +190,10 @@ class GpuDataManagerImpl implements GpuDataManager { this.buffersPending.push(cachedData.gpuData.buffer); // cachedData.gpuData.buffer.destroy(); - const downloadingData = this.downloadCache.get(id); - if (downloadingData) { - this.downloadCache.delete(id); - } - return cachedData.originalSize; } async download(id: GpuDataId): Promise { - const downloadData = this.downloadCache.get(id); - if (downloadData) { - return downloadData.data; - } - const cachedData = this.storageCache.get(id); if (!cachedData) { throw new Error('data does not exist'); @@ -229,17 +211,13 @@ class GpuDataManagerImpl implements GpuDataManager { ); this.backend.flush(); - const readDataPromise = new Promise((resolve) => { + return new Promise((resolve) => { gpuReadBuffer.mapAsync(GPUMapMode.READ).then(() => { const data = gpuReadBuffer.getMappedRange().slice(0); gpuReadBuffer.destroy(); resolve(data); }); }); - - this.downloadCache.set(id, {data: readDataPromise}); - - return readDataPromise; } refreshPendingBuffers(): void { @@ -272,7 +250,6 @@ class GpuDataManagerImpl implements GpuDataManager { }); this.storageCache = new Map(); - this.downloadCache = new Map(); this.freeBuffers = new Map(); } } diff --git a/js/web/lib/wasm/jsep/webgpu/op-resolve-rules.ts b/js/web/lib/wasm/jsep/webgpu/op-resolve-rules.ts index 11e54545c4..ae4b754f76 100644 --- a/js/web/lib/wasm/jsep/webgpu/op-resolve-rules.ts +++ b/js/web/lib/wasm/jsep/webgpu/op-resolve-rules.ts @@ -52,6 +52,7 @@ export const WEBGPU_OP_RESOLVE_RULES: Map = new ['Cosh', [unaryOps.cosh]], ['Div', [binaryOps.div]], ['Elu', [unaryOps.elu, unaryOps.parseAlphaAttributes]], + ['Equal', [binaryOps.equal]], ['Erf', [unaryOps.erf]], ['Exp', [unaryOps.exp]], ['Expand', [expand]], @@ -72,6 +73,7 @@ export const WEBGPU_OP_RESOLVE_RULES: Map = new ['MaxPool', [pool.maxPool, pool.parseMaxPoolAttributes]], ['Mul', [binaryOps.mul]], ['Neg', [unaryOps.neg]], + ['Not', [unaryOps.not]], ['Pow', [binaryOps.pow]], ['Reciprocal', [unaryOps.reciprocal]], ['ReduceMin', [reduceMin, parseReduceAttributes]], diff --git a/js/web/lib/wasm/jsep/webgpu/ops/binary-op.ts b/js/web/lib/wasm/jsep/webgpu/ops/binary-op.ts index a16aed7ae4..28284554f9 100644 --- a/js/web/lib/wasm/jsep/webgpu/ops/binary-op.ts +++ b/js/web/lib/wasm/jsep/webgpu/ops/binary-op.ts @@ -42,9 +42,7 @@ const createBinaryOpProgramShader = const strides = ShapeUtil.computeStrides(dims); const offsets: string[] = []; for (let i = dims.length - 1; i >= 0; i--) { - const idx = dimsOutput.length === 0 ? '0u' : - (dimsOutput.length === 1) ? 'outputIndices' : - `outputIndices[${i + dimsOutput.length - dims.length}]`; + const idx = output.indicesGet('outputIndices', i + dimsOutput.length - dims.length); offsets.push(`${strides[i]}u * (${idx} % ${dims[i]}u)`); } return offsets.length > 0 ? offsets.join('+') : '0u'; @@ -194,6 +192,12 @@ export const div = (context: ComputeContext): void => { context.compute(createBinaryOpProgramInfoLoader(context.inputs, 'Div', (a, b) => `${a}/${b}`)); }; +export const equal = (context: ComputeContext): void => { + context.compute(createBinaryOpProgramInfoLoader( + context.inputs, 'Equal', ({scalar: (a, b) => `u32(${a}==${b})`, vector: (a, b) => `vec4(${a}==${b})`}), + undefined, undefined, DataType.bool)); +}; + export const mul = (context: ComputeContext): void => { context.compute(createBinaryOpProgramInfoLoader(context.inputs, 'Mul', (a, b) => `${a}*${b}`)); }; @@ -227,18 +231,12 @@ export const sub = (context: ComputeContext): void => { export const greater = (context: ComputeContext): void => { context.compute(createBinaryOpProgramInfoLoader( - context.inputs, 'Greater', ({ - scalar: (a, b) => `select(0, 1, ${a}>${b})`, - vector: (a, b) => `select(vec4(0), vec4(1), ${a}>${b})` - }), + context.inputs, 'Greater', ({scalar: (a, b) => `u32(${a}>${b})`, vector: (a, b) => `vec4(${a}>${b})`}), undefined, undefined, DataType.bool)); }; export const less = (context: ComputeContext): void => { context.compute(createBinaryOpProgramInfoLoader( - context.inputs, 'Less', ({ - scalar: (a, b) => `select(0, 1, ${a}<${b})`, - vector: (a, b) => `select(vec4(0), vec4(1), ${a}<${b})` - }), + context.inputs, 'Less', ({scalar: (a, b) => `u32(${a}<${b})`, vector: (a, b) => `vec4(${a}<${b})`}), undefined, undefined, DataType.bool)); }; diff --git a/js/web/lib/wasm/jsep/webgpu/ops/common.ts b/js/web/lib/wasm/jsep/webgpu/ops/common.ts index 7da57bcb9c..75c37b3ed0 100644 --- a/js/web/lib/wasm/jsep/webgpu/ops/common.ts +++ b/js/web/lib/wasm/jsep/webgpu/ops/common.ts @@ -229,6 +229,11 @@ export const tensorTypeToWsglStorageType = (type: DataType, components: 1|2|3|4 return typeof mappedType === 'string' ? mappedType : mappedType[0]; }; +export const tensorTypeToWsglValueType = (type: DataType, components: 1|2|3|4 = 1) => { + const mappedType = getWgslMappedType(type, components); + return typeof mappedType === 'string' ? mappedType : mappedType[1]; +}; + /** * A helper function to get a IndicesHelper for a given input or output. * diff --git a/js/web/lib/wasm/jsep/webgpu/ops/unary-op.ts b/js/web/lib/wasm/jsep/webgpu/ops/unary-op.ts index 7e7ad5f4e6..ef63d11777 100644 --- a/js/web/lib/wasm/jsep/webgpu/ops/unary-op.ts +++ b/js/web/lib/wasm/jsep/webgpu/ops/unary-op.ts @@ -219,6 +219,10 @@ export const leakyRelu = (context: ComputeContext, attributes: AlphaAttributes): `const leaky_relu_alpha_: f32 = f32(${attributes.alpha});`, attributes.cacheKey)); }; +export const not = (context: ComputeContext): void => { + context.compute(createElementwiseProgramInfoLoader(context.inputs[0], 'Not', a => `!${a}`)); +}; + export const neg = (context: ComputeContext): void => { context.compute(createElementwiseProgramInfoLoader(context.inputs[0], 'Neg', a => `-${a}`)); }; diff --git a/js/web/test/suite-test-list.jsonc b/js/web/test/suite-test-list.jsonc index a7964d9ca1..e0b0207c9f 100644 --- a/js/web/test/suite-test-list.jsonc +++ b/js/web/test/suite-test-list.jsonc @@ -839,9 +839,9 @@ // "test_nonmaxsuppression_two_batches", // "test_nonmaxsuppression_two_classes", // "test_nonzero_example", - // "test_not_2d", - // "test_not_3d", - // "test_not_4d", + "test_not_2d", + "test_not_3d", + "test_not_4d", // // "test_onehot_negative_indices", // // "test_onehot_with_axis", // // "test_onehot_with_negative_axis", @@ -1335,7 +1335,7 @@ "div.jsonc", "div_int32.jsonc", //"depth-to-space.jsonc", - //"equal.jsonc", + "equal.jsonc", "exp.jsonc", "expand.jsonc", "floor.jsonc", @@ -1348,7 +1348,7 @@ "mul.jsonc", "mul_int32.jsonc", //"neg.jsonc", - //"not.jsonc", + "not.jsonc", //"or.jsonc", "layer-norm.jsonc", "leaky-relu.jsonc", diff --git a/onnxruntime/core/providers/js/js_execution_provider.cc b/onnxruntime/core/providers/js/js_execution_provider.cc index 2146d9a0c5..2732eb0c3d 100644 --- a/onnxruntime/core/providers/js/js_execution_provider.cc +++ b/onnxruntime/core/providers/js/js_execution_provider.cc @@ -111,6 +111,7 @@ class ONNX_OPERATOR_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 9, Acos class ONNX_OPERATOR_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 9, Atanh); class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 6, 12, Tanh); class ONNX_OPERATOR_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 13, Tanh); +class ONNX_OPERATOR_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 1, Not); class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 6, 8, Cast); class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 9, 12, Cast); @@ -197,6 +198,10 @@ class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomai class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 12, 12, Pow); class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 13, 14, Pow); class ONNX_OPERATOR_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 15, Pow); +class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 7, 10, Equal); +class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 11, 12, Equal); +class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 13, 18, Equal); +class ONNX_OPERATOR_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 19, Equal); class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 7, 8, Greater); class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 9, 12, Greater); class ONNX_OPERATOR_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 13, Greater); @@ -351,6 +356,7 @@ std::unique_ptr RegisterKernels() { KERNEL_CREATE_INFO(9, Atanh), KERNEL_CREATE_INFO_VERSIONED(6, 12, Tanh), KERNEL_CREATE_INFO(13, Tanh), + KERNEL_CREATE_INFO(1, Not), KERNEL_CREATE_INFO_VERSIONED(6, 8, Cast), KERNEL_CREATE_INFO_VERSIONED(9, 12, Cast), @@ -387,6 +393,10 @@ std::unique_ptr RegisterKernels() { KERNEL_CREATE_INFO_VERSIONED(12, 12, Pow), KERNEL_CREATE_INFO_VERSIONED(13, 14, Pow), KERNEL_CREATE_INFO(15, Pow), + KERNEL_CREATE_INFO_VERSIONED(7, 10, Equal), + KERNEL_CREATE_INFO_VERSIONED(11, 12, Equal), + KERNEL_CREATE_INFO_VERSIONED(13, 18, Equal), + KERNEL_CREATE_INFO(19, Equal), KERNEL_CREATE_INFO_VERSIONED(7, 8, Greater), KERNEL_CREATE_INFO_VERSIONED(9, 12, Greater), KERNEL_CREATE_INFO(13, Greater), diff --git a/onnxruntime/core/providers/js/operators/binary.cc b/onnxruntime/core/providers/js/operators/binary.cc index e26bb0e49f..2a96619c2c 100644 --- a/onnxruntime/core/providers/js/operators/binary.cc +++ b/onnxruntime/core/providers/js/operators/binary.cc @@ -52,6 +52,12 @@ REG_ELEMENTWISE_VERSIONED_KERNEL(Pow, 12, 12, Pow); REG_ELEMENTWISE_VERSIONED_KERNEL(Pow, 13, 14, Pow); REG_ELEMENTWISE_KERNEL(Pow, 15, Pow); +JSEP_KERNEL_IMPL(Equal, Equal) +REG_ELEMENTWISE_VERSIONED_KERNEL(Equal, 7, 10, Equal); +REG_ELEMENTWISE_VERSIONED_KERNEL(Equal, 11, 12, Equal); +REG_ELEMENTWISE_VERSIONED_KERNEL(Equal, 13, 18, Equal); +REG_ELEMENTWISE_KERNEL(Equal, 19, Equal); + JSEP_KERNEL_IMPL(Greater, Greater) REG_ELEMENTWISE_VERSIONED_KERNEL(Greater, 7, 8, Greater); REG_ELEMENTWISE_VERSIONED_KERNEL(Greater, 9, 12, Greater); diff --git a/onnxruntime/core/providers/js/operators/unary.cc b/onnxruntime/core/providers/js/operators/unary.cc index 0947221bbe..869d78f351 100644 --- a/onnxruntime/core/providers/js/operators/unary.cc +++ b/onnxruntime/core/providers/js/operators/unary.cc @@ -97,6 +97,9 @@ JSEP_KERNEL_IMPL(Tanh, Tanh) JSEP_ELEMENTWISE_VERSIONED_KERNEL(Tanh, 6, 12, float, Tanh) JSEP_ELEMENTWISE_KERNEL(Tanh, 13, float, Tanh) +JSEP_KERNEL_IMPL(Not, Not) +JSEP_ELEMENTWISE_KERNEL(Not, 1, bool, Not) + // activation JSEP_CLASS_IMPL_ATTRIBUTE_FLOAT_2_DEFAULT(ClipV10, ClipV10, min, 3.402823e+38f, max, -3.402823e+38f)