mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-29 20:14:01 +00:00
[js/webgpu] Support where (#17544)
Supported type: float. int32_t, uint32_t, bool. Case where_broadcast.jsonc is not enabled due to https://github.com/microsoft/onnxruntime/issues/17405. ### Description <!-- Describe your changes. --> ### 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. --> --------- Co-authored-by: Yulong Wang <7679871+fs-eire@users.noreply.github.com>
This commit is contained in:
Xu Xing
GitHub
parent
f8a8452a6b
commit
992f3e4609
9 changed files with 455 additions and 1 deletions
|
|
@ -96,3 +96,4 @@ Do not modify directly.*
|
|||
| Tile | ai.onnx(6-12,13+) | |
|
||||
| Transpose | ai.onnx(1-12,13+) | need perf optimization |
|
||||
| Unsqueeze | ai.onnx(1-10,11-12,13+) | |
|
||||
| Where | ai.onnx(9-15,16+) | |
|
||||
|
|
|
|||
|
|
@ -28,6 +28,7 @@ import {parseSplitAttributes, split} from './ops/split';
|
|||
import {tile} from './ops/tile';
|
||||
import {parseTransposeAttributes, transpose} from './ops/transpose';
|
||||
import * as unaryOps from './ops/unary-op';
|
||||
import {where} from './ops/where';
|
||||
import {ComputeContext} from './types';
|
||||
|
||||
export type RunFunction = (context: ComputeContext, attribute?: unknown) => void;
|
||||
|
|
@ -116,4 +117,5 @@ export const WEBGPU_OP_RESOLVE_RULES: Map<string, OperatorImplementation> = new
|
|||
['ThresholdedRelu', [unaryOps.thresholdedRelu, unaryOps.parseAlphaAttributes]],
|
||||
['Tile', [tile]],
|
||||
['Transpose', [transpose, parseTransposeAttributes]],
|
||||
['Where', [where]],
|
||||
]);
|
||||
|
|
|
|||
|
|
@ -102,6 +102,16 @@ export interface IndicesHelper {
|
|||
*/
|
||||
readonly indicesToOffset: (varIndices: string) => string;
|
||||
|
||||
/**
|
||||
* WGSL code of an `u32` expression for getting original offset from broadcasted indices.
|
||||
*
|
||||
* @param varIndices - a `type.indices` expression representing the output indices.
|
||||
* @param output - output IndicesHelper.
|
||||
*
|
||||
* @returns an `u32` expression
|
||||
*/
|
||||
readonly broadcastedIndicesToOffset: (varIndices: string, output: IndicesHelper) => string;
|
||||
|
||||
/**
|
||||
* WGSL code of generating an indices literal
|
||||
*
|
||||
|
|
@ -262,6 +272,7 @@ const createIndicesHelper =
|
|||
const implementationUsed = {
|
||||
offsetToIndices: false,
|
||||
indicesToOffset: false,
|
||||
broadcastedIndicesToOffset: false,
|
||||
set: false,
|
||||
setByIndices: false,
|
||||
get: false,
|
||||
|
|
@ -310,6 +321,26 @@ const createIndicesHelper =
|
|||
return rank < 2 ? varIndices : `i2o_${name}(${varIndices})`;
|
||||
};
|
||||
|
||||
const broadcastedIndicesToOffsetImplementation: {[key: string]: string} = {};
|
||||
const broadcastedIndicesToOffset = (varIndices: string, output: IndicesHelper) => {
|
||||
implementationUsed.broadcastedIndicesToOffset = true;
|
||||
const implKey = `${output.name}broadcastedIndicesTo${name}Offset`;
|
||||
if (implKey in broadcastedIndicesToOffsetImplementation) {
|
||||
return `${implKey}(${varIndices})`;
|
||||
}
|
||||
const offsets = [];
|
||||
for (let i = shape.length - 1; i >= 0; i--) {
|
||||
const idx = output.indicesGet('outputIndices', i + output.shape.length - shape.length);
|
||||
offsets.push(`${strides[i]}u * (${idx} % ${shape[i]}u)`);
|
||||
}
|
||||
broadcastedIndicesToOffsetImplementation[implKey] =
|
||||
`fn ${implKey}(outputIndices: ${output.type.indices}) -> u32 {
|
||||
return ${offsets.length > 0 ? offsets.join('+') : '0u'};
|
||||
}`;
|
||||
|
||||
return `${implKey}(${varIndices})`;
|
||||
};
|
||||
|
||||
const indices = (...init: ReadonlyArray<number|string>) =>
|
||||
rank === 0 ? '0u' : `${type.indices}(${init.map(normalizeDim).join(',')})`;
|
||||
|
||||
|
|
@ -462,6 +493,9 @@ const createIndicesHelper =
|
|||
if (implementationUsed.indicesToOffset) {
|
||||
impls.push(indicesToOffsetImplementation);
|
||||
}
|
||||
if (implementationUsed.broadcastedIndicesToOffset) {
|
||||
Object.values(broadcastedIndicesToOffsetImplementation).forEach(impl => impls.push(impl));
|
||||
}
|
||||
if (implementationUsed.set) {
|
||||
impls.push(setImplementation);
|
||||
}
|
||||
|
|
@ -482,6 +516,7 @@ const createIndicesHelper =
|
|||
type,
|
||||
offsetToIndices,
|
||||
indicesToOffset,
|
||||
broadcastedIndicesToOffset,
|
||||
indices,
|
||||
indicesGet,
|
||||
indicesSet,
|
||||
|
|
|
|||
110
js/web/lib/wasm/jsep/webgpu/ops/where.ts
Normal file
110
js/web/lib/wasm/jsep/webgpu/ops/where.ts
Normal file
|
|
@ -0,0 +1,110 @@
|
|||
// Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
// Licensed under the MIT License.
|
||||
|
||||
import {DataType} from '../../../wasm-common';
|
||||
import {TensorView} from '../../tensor-view';
|
||||
import {BroadcastUtil, ShapeUtil} from '../../util';
|
||||
import {ComputeContext, GpuDataType, ProgramInfo, ProgramInfoLoader, ProgramMetadata} from '../types';
|
||||
|
||||
import {inputVariable, outputVariable, ShaderHelper} from './common';
|
||||
|
||||
const createWhereOpProgramShader =
|
||||
(shaderHelper: ShaderHelper, inputs: readonly TensorView[], dimsOutput: readonly number[], isBroadcast: boolean,
|
||||
typeOutput: number) => {
|
||||
const outputSize = ShapeUtil.size(dimsOutput);
|
||||
const vecSize = Math.ceil(outputSize / 4);
|
||||
|
||||
const output = outputVariable('outputData', typeOutput, dimsOutput, 4);
|
||||
const a = inputVariable('aData', inputs[1].dataType, inputs[1].dims, 4);
|
||||
const b = inputVariable('bData', inputs[2].dataType, inputs[2].dims, 4);
|
||||
const c = inputVariable('cData', inputs[0].dataType, inputs[0].dims, 4);
|
||||
|
||||
let assignment: string;
|
||||
const expression = (a: string, b: string, c: string) => `select(${b}, ${a}, ${c})`;
|
||||
if (!isBroadcast) {
|
||||
assignment = output.setByOffset(
|
||||
'global_idx',
|
||||
expression(a.getByOffset('global_idx'), b.getByOffset('global_idx'), c.getByOffset('global_idx')));
|
||||
} else {
|
||||
const singleAssignment = (resStr: string, x: number, typeCast = '') => {
|
||||
const expressionA = `aData[indexA${x}][componentA${x}]`;
|
||||
const expressionB = `bData[indexB${x}][componentB${x}]`;
|
||||
// eslint-disable-next-line no-bitwise
|
||||
const expressionC = `bool(cData[indexC${x}] & ${0xff000000 >>> ((3 - x) * 8)}u)`;
|
||||
return `
|
||||
let outputIndices${x} = ${output.offsetToIndices(`global_idx * 4u + ${x}u`)};
|
||||
let offsetA${x} = ${a.broadcastedIndicesToOffset(`outputIndices${x}`, output)};
|
||||
let offsetB${x} = ${b.broadcastedIndicesToOffset(`outputIndices${x}`, output)};
|
||||
let offsetC${x} = ${c.broadcastedIndicesToOffset(`outputIndices${x}`, output)};
|
||||
let indexA${x} = offsetA${x} / 4u;
|
||||
let indexB${x} = offsetB${x} / 4u;
|
||||
let indexC${x} = offsetC${x} / 4u;
|
||||
let componentA${x} = offsetA${x} % 4u;
|
||||
let componentB${x} = offsetB${x} % 4u;
|
||||
${resStr}[${x}] = ${typeCast}(${expression(expressionA, expressionB, expressionC)});
|
||||
`;
|
||||
};
|
||||
if (typeOutput === DataType.bool) {
|
||||
assignment = `
|
||||
var data = vec4<u32>(0);
|
||||
${singleAssignment('data', 0, 'u32')}
|
||||
${singleAssignment('data', 1, 'u32')}
|
||||
${singleAssignment('data', 2, 'u32')}
|
||||
${singleAssignment('data', 3, 'u32')}
|
||||
outputData[global_idx] = dot(vec4<u32>(0x1, 0x100, 0x10000, 0x1000000), vec4<u32>(data));`;
|
||||
} else {
|
||||
assignment = `
|
||||
${singleAssignment('outputData[global_idx]', 0)}
|
||||
${singleAssignment('outputData[global_idx]', 1)}
|
||||
${singleAssignment('outputData[global_idx]', 2)}
|
||||
${singleAssignment('outputData[global_idx]', 3)}
|
||||
`;
|
||||
}
|
||||
}
|
||||
|
||||
return `
|
||||
${shaderHelper.declareVariables(c, a, b, output)}
|
||||
${shaderHelper.mainStart()}
|
||||
${shaderHelper.guardAgainstOutOfBoundsWorkgroupSizes(vecSize)}
|
||||
${assignment}
|
||||
}`;
|
||||
};
|
||||
|
||||
const createWhereOpProgramInfo = (metadata: ProgramMetadata, inputs: readonly TensorView[]): ProgramInfo => {
|
||||
const dimsA = inputs[1].dims;
|
||||
const dimsB = inputs[2].dims;
|
||||
const dimsC = inputs[0].dims;
|
||||
const outputDataType = inputs[1].dataType;
|
||||
|
||||
const isBroadcast = !(ShapeUtil.areEqual(dimsA, dimsB) && ShapeUtil.areEqual(dimsB, dimsC));
|
||||
let outputShape = dimsA;
|
||||
let outputSize = ShapeUtil.size(dimsA);
|
||||
// TODO: deal with zero-sized tensors (eg. dims=[1,0])
|
||||
|
||||
if (isBroadcast) {
|
||||
const calculatedShape = BroadcastUtil.calcShape(BroadcastUtil.calcShape(dimsA, dimsB, false)!, dimsC, false);
|
||||
if (!calculatedShape) {
|
||||
throw new Error('Can\'t perform where op on the given tensors');
|
||||
}
|
||||
outputShape = calculatedShape;
|
||||
outputSize = ShapeUtil.size(outputShape);
|
||||
}
|
||||
|
||||
return {
|
||||
...metadata,
|
||||
getShaderSource: (shaderHelper) =>
|
||||
createWhereOpProgramShader(shaderHelper, inputs, outputShape, isBroadcast, outputDataType),
|
||||
outputs: [{dims: outputShape, dataType: outputDataType, gpuDataType: GpuDataType.default}],
|
||||
dispatchGroup: () => ({x: Math.ceil(outputSize / 64 /* workgroup size */ / 4 /* vec size */)})
|
||||
};
|
||||
};
|
||||
|
||||
const createWhereOpProgramInfoLoader = (inputs: readonly TensorView[], name: string): ProgramInfoLoader => {
|
||||
const inputTypes = [GpuDataType.default, GpuDataType.default, GpuDataType.default];
|
||||
const metadata: ProgramMetadata = {name, inputTypes};
|
||||
return {...metadata, get: () => createWhereOpProgramInfo(metadata, inputs)};
|
||||
};
|
||||
|
||||
export const where = (context: ComputeContext): void => {
|
||||
context.compute(createWhereOpProgramInfoLoader(context.inputs, 'Where'));
|
||||
};
|
||||
172
js/web/test/data/ops/where.jsonc
Normal file
172
js/web/test/data/ops/where.jsonc
Normal file
|
|
@ -0,0 +1,172 @@
|
|||
[
|
||||
{
|
||||
"name": "Where with no attributes",
|
||||
"operator": "Where",
|
||||
"attributes": [],
|
||||
"cases": [
|
||||
{
|
||||
"name": "T[3] T[3] T[3] float32 T[3] ",
|
||||
"inputs": [
|
||||
{
|
||||
"data": [true, false, true, false, true, false, true, false],
|
||||
"dims": [8],
|
||||
"type": "bool"
|
||||
},
|
||||
{
|
||||
"data": [4.0, 8.0, 7.0, 2.0, 4.0, 8.0, 7.0, 1.0],
|
||||
"dims": [8],
|
||||
"type": "float32"
|
||||
},
|
||||
{
|
||||
"data": [1.0, 3.0, 9.0, 6.0, 1.0, 3.0, 9.0, 2.0],
|
||||
"dims": [8],
|
||||
"type": "float32"
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"data": [4.0, 3.0, 7.0, 6.0, 4.0, 3.0, 7.0, 2.0],
|
||||
"dims": [8],
|
||||
"type": "float32"
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"name": "Where with no attributes",
|
||||
"operator": "Where",
|
||||
"attributes": [],
|
||||
"cases": [
|
||||
{
|
||||
"name": "T[3] T[3] T[3] int32 T[3] ",
|
||||
"inputs": [
|
||||
{
|
||||
"data": [true, false, true, false, true, false, true, false],
|
||||
"dims": [8],
|
||||
"type": "bool"
|
||||
},
|
||||
{
|
||||
"data": [4, 8, 7, 2, 4, 8, 7, 1],
|
||||
"dims": [8],
|
||||
"type": "int32"
|
||||
},
|
||||
{
|
||||
"data": [1, 3, 9, 6, 1, 3, 9, 2],
|
||||
"dims": [8],
|
||||
"type": "int32"
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"data": [4, 3, 7, 6, 4, 3, 7, 2],
|
||||
"dims": [8],
|
||||
"type": "int32"
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"name": "Where with no attributes",
|
||||
"operator": "Where",
|
||||
"attributes": [],
|
||||
"cases": [
|
||||
{
|
||||
"name": "T[3] T[3] T[3] uint32 T[3] ",
|
||||
"inputs": [
|
||||
{
|
||||
"data": [true, false, true, false, true, false, true, false],
|
||||
"dims": [8],
|
||||
"type": "bool"
|
||||
},
|
||||
{
|
||||
"data": [4, 8, 7, 2, 4, 8, 7, 1],
|
||||
"dims": [8],
|
||||
"type": "uint32"
|
||||
},
|
||||
{
|
||||
"data": [1, 4294967295, 9, 6, 1, 3, 9, 2],
|
||||
"dims": [8],
|
||||
"type": "uint32"
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"data": [4, 4294967295, 7, 6, 4, 3, 7, 2],
|
||||
"dims": [8],
|
||||
"type": "uint32"
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"name": "Where with no attributes",
|
||||
"operator": "Where",
|
||||
"attributes": [],
|
||||
"cases": [
|
||||
{
|
||||
"name": "T[3] T[3] T[3] bool T[3] ",
|
||||
"inputs": [
|
||||
{
|
||||
"data": [true, false, true, false, true, false, true, false],
|
||||
"dims": [8],
|
||||
"type": "bool"
|
||||
},
|
||||
{
|
||||
"data": [true, true, true, true, true, true, true, true],
|
||||
"dims": [8],
|
||||
"type": "float32"
|
||||
},
|
||||
{
|
||||
"data": [true, false, true, false, true, false, true, false],
|
||||
"dims": [8],
|
||||
"type": "float32"
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"data": [true, false, true, false, true, false, true, false],
|
||||
"dims": [8],
|
||||
"type": "float32"
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"name": "Where with no attributes",
|
||||
"operator": "Where",
|
||||
"attributes": [],
|
||||
"cases": [
|
||||
{
|
||||
"name": "T[3 3] T[3 3] T[1] float32 broadcast",
|
||||
"inputs": [
|
||||
{
|
||||
"data": [true, true, true, true, true, false, false, false, false],
|
||||
"dims": [3, 3],
|
||||
"type": "bool"
|
||||
},
|
||||
{
|
||||
"data": [0, 1, 2, 3, 4, 5, 6, 7, 8],
|
||||
"dims": [3, 3],
|
||||
"type": "float32"
|
||||
},
|
||||
{
|
||||
"data": [-1.0],
|
||||
"dims": [1],
|
||||
"type": "float32"
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"data": [0, 1, 2, 3, 4, -1, -1, -1, -1],
|
||||
"dims": [3, 3],
|
||||
"type": "float32"
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
84
js/web/test/data/ops/where_broadcast.jsonc
Normal file
84
js/web/test/data/ops/where_broadcast.jsonc
Normal file
|
|
@ -0,0 +1,84 @@
|
|||
[
|
||||
{
|
||||
"name": "Where with no attributes",
|
||||
"operator": "Where",
|
||||
"attributes": [],
|
||||
"cases": [
|
||||
{
|
||||
// This failed due to: https://github.com/microsoft/onnxruntime/issues/17405.
|
||||
"name": "T[3 6] T[3 6] T[1] float32 broadcast",
|
||||
"inputs": [
|
||||
{
|
||||
"data": [
|
||||
true,
|
||||
true,
|
||||
true,
|
||||
true,
|
||||
true,
|
||||
false,
|
||||
false,
|
||||
false,
|
||||
false,
|
||||
false,
|
||||
false,
|
||||
true,
|
||||
true,
|
||||
true,
|
||||
true,
|
||||
true,
|
||||
true,
|
||||
true
|
||||
],
|
||||
"dims": [3, 6],
|
||||
"type": "bool"
|
||||
},
|
||||
{
|
||||
"data": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17],
|
||||
"dims": [3, 6],
|
||||
"type": "float32"
|
||||
},
|
||||
{
|
||||
"data": [-1.0],
|
||||
"dims": [1],
|
||||
"type": "float32"
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"data": [0, 1, 2, 3, 4, -1, -1, -1, -1, -1, -1, 11, 12, 13, 14, 15, 16, 17],
|
||||
"dims": [3, 6],
|
||||
"type": "float32"
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
// This failed due to: https://github.com/microsoft/onnxruntime/issues/17405.
|
||||
"name": "T[3 1] T[3 6] T[1] float32 broadcast",
|
||||
"inputs": [
|
||||
{
|
||||
"data": [true, false, true],
|
||||
"dims": [3, 1],
|
||||
"type": "bool"
|
||||
},
|
||||
{
|
||||
"data": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17],
|
||||
"dims": [3, 6],
|
||||
"type": "float32"
|
||||
},
|
||||
{
|
||||
"data": [-1.0],
|
||||
"dims": [1],
|
||||
"type": "float32"
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"data": [0, 1, 2, 3, 4, 5, -1, -1, -1, -1, -1, -1, 12, 13, 14, 15, 16, 17],
|
||||
"dims": [3, 6],
|
||||
"type": "float32"
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
|
|
@ -1388,7 +1388,10 @@
|
|||
"tan.jsonc",
|
||||
"tile.jsonc",
|
||||
"transpose.jsonc",
|
||||
"transpose_int32_uint32.jsonc"
|
||||
"transpose_int32_uint32.jsonc",
|
||||
"where.jsonc"
|
||||
// Turn on this when https://github.com/microsoft/onnxruntime/issues/17405 is fixed.
|
||||
//"where_broadcast.jsonc",
|
||||
//"xor.jsonc"
|
||||
]
|
||||
},
|
||||
|
|
|
|||
|
|
@ -229,6 +229,9 @@ class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomai
|
|||
class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 11, 12, Unsqueeze);
|
||||
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 13, Unsqueeze);
|
||||
|
||||
class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 9, 15, Where);
|
||||
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 16, Where);
|
||||
|
||||
class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 1, 12, Transpose);
|
||||
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 13, Transpose);
|
||||
|
||||
|
|
@ -496,6 +499,9 @@ std::unique_ptr<KernelRegistry> RegisterKernels() {
|
|||
BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 13, 17, ReduceLogSumExp)>,
|
||||
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 18, ReduceLogSumExp)>,
|
||||
|
||||
KERNEL_CREATE_INFO_VERSIONED(9, 15, Where),
|
||||
KERNEL_CREATE_INFO(16, Where),
|
||||
|
||||
BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 1, 12, Transpose)>,
|
||||
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 13, Transpose)>,
|
||||
|
||||
|
|
|
|||
41
onnxruntime/core/providers/js/operators/where.cc
Normal file
41
onnxruntime/core/providers/js/operators/where.cc
Normal file
|
|
@ -0,0 +1,41 @@
|
|||
// Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
// Licensed under the MIT License.
|
||||
|
||||
#include "core/providers/js/js_kernel.h"
|
||||
|
||||
namespace onnxruntime {
|
||||
namespace js {
|
||||
|
||||
#define REG_ELEMENTWISE_KERNEL(OP_TYPE, VERSION, KERNEL_CLASS) \
|
||||
ONNX_OPERATOR_KERNEL_EX( \
|
||||
OP_TYPE, \
|
||||
kOnnxDomain, \
|
||||
VERSION, \
|
||||
kJsExecutionProvider, \
|
||||
KernelDefBuilder() \
|
||||
.TypeConstraint("T", \
|
||||
{DataTypeImpl::GetTensorType<float>(), \
|
||||
DataTypeImpl::GetTensorType<int32_t>(), \
|
||||
DataTypeImpl::GetTensorType<uint32_t>(), \
|
||||
DataTypeImpl::GetTensorType<bool>()}), \
|
||||
KERNEL_CLASS);
|
||||
|
||||
#define REG_ELEMENTWISE_VERSIONED_KERNEL(OP_TYPE, VERSION_FROM, VERSION_TO, KERNEL_CLASS) \
|
||||
ONNX_OPERATOR_VERSIONED_KERNEL_EX( \
|
||||
OP_TYPE, \
|
||||
kOnnxDomain, \
|
||||
VERSION_FROM, VERSION_TO, \
|
||||
kJsExecutionProvider, \
|
||||
KernelDefBuilder() \
|
||||
.TypeConstraint("T", \
|
||||
{DataTypeImpl::GetTensorType<float>(), \
|
||||
DataTypeImpl::GetTensorType<int32_t>(), \
|
||||
DataTypeImpl::GetTensorType<uint32_t>(), \
|
||||
DataTypeImpl::GetTensorType<bool>()}), \
|
||||
KERNEL_CLASS);
|
||||
|
||||
JSEP_KERNEL_IMPL(Where, Where)
|
||||
REG_ELEMENTWISE_VERSIONED_KERNEL(Where, 9, 15, Where);
|
||||
REG_ELEMENTWISE_KERNEL(Where, 16, Where);
|
||||
} // namespace js
|
||||
} // namespace onnxruntime
|
||||
Loading…
Reference in a new issue