saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-13 18:08:13 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

[js/webgpu] add kernel Not and Equal (#17306)

Browse source 
### 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
<!-- - Why is this change required? What problem does it solve?
- If it fixes an open issue, please link to the issue here. -->
This commit is contained in:
Yulong Wang 2023-08-27 19:50:17 -07:00 committed by GitHub
parent 4eedd3bb46
commit bb1871332f
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
10 changed files with 47 additions and 40 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
js/web/docs/webgpu-operators.md
View file

@ -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+) | |

25
js/web/lib/wasm/jsep/webgpu/gpu-data-manager.ts
View file

@ -57,10 +57,6 @@ interface StorageCacheValue {
originalSize: number;
}
interface DownloadCacheValue {
data: Promise<ArrayBufferLike>;
}
/**
* 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<GpuDataId, StorageCacheValue>;
// GPU Data ID => GPU Data ( read buffer )
downloadCache: Map<GpuDataId, DownloadCacheValue>;
// 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<ArrayBufferLike> {
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<ArrayBuffer>((resolve) => {
return new Promise<ArrayBuffer>((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();
}
}

2
js/web/lib/wasm/jsep/webgpu/op-resolve-rules.ts
View file

@ -52,6 +52,7 @@ export const WEBGPU_OP_RESOLVE_RULES: Map<string, OperatorImplementation> = 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<string, OperatorImplementation> = new
['MaxPool', [pool.maxPool, pool.parseMaxPoolAttributes]],
['Mul', [binaryOps.mul]],
['Neg', [unaryOps.neg]],
['Not', [unaryOps.not]],
['Pow', [binaryOps.pow]],
['Reciprocal', [unaryOps.reciprocal]],
['ReduceMin', [reduceMin, parseReduceAttributes]],

20
js/web/lib/wasm/jsep/webgpu/ops/binary-op.ts
View file

@ -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<u32>(${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<u32>(0), vec4<u32>(1), ${a}>${b})`
}),
context.inputs, 'Greater', ({scalar: (a, b) => `u32(${a}>${b})`, vector: (a, b) => `vec4<u32>(${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<u32>(0), vec4<u32>(1), ${a}<${b})`
}),
context.inputs, 'Less', ({scalar: (a, b) => `u32(${a}<${b})`, vector: (a, b) => `vec4<u32>(${a}<${b})`}),
undefined, undefined, DataType.bool));
};

5
js/web/lib/wasm/jsep/webgpu/ops/common.ts
View file

@ -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.
*

4
js/web/lib/wasm/jsep/webgpu/ops/unary-op.ts
View file

@ -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}`));
};

10
js/web/test/suite-test-list.jsonc
View file

@ -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",

10
onnxruntime/core/providers/js/js_execution_provider.cc
View file

@ -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<KernelRegistry> 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<KernelRegistry> 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),

6
onnxruntime/core/providers/js/operators/binary.cc
View file

@ -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);

3
onnxruntime/core/providers/js/operators/unary.cc
View file

@ -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)