mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-25 19:48:11 +00:00
[JS/Web] Modify Reduce, Expand and Slice to pass op and node tests. (#16979)
### Description Make CacheHint mechanism, which is designed to avoid running the same test multiple times saving the result mapped against a key, working by adding input dims. ### 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:
satyajandhyala
GitHub
parent
a25d0d296b
commit
cc4b64f646
4 changed files with 61 additions and 51 deletions
|
|
@ -102,7 +102,8 @@ const createExpandProgramInfo = (metadata: ProgramMetadata, inputs: readonly Ten
|
|||
|
||||
export const expand = (context: ComputeContext): void => {
|
||||
validateInputs(context.inputs);
|
||||
const cacheHint = context.inputs.map(x => x.dims.toString()).join('_');
|
||||
context.compute(
|
||||
{...expandProgramMetadata, get: () => createExpandProgramInfo(expandProgramMetadata, context.inputs)},
|
||||
{...expandProgramMetadata, cacheHint, get: () => createExpandProgramInfo(expandProgramMetadata, context.inputs)},
|
||||
{inputs: [0]});
|
||||
};
|
||||
|
|
|
|||
|
|
@ -32,49 +32,48 @@ export interface ReduceAttributes extends AttributeWithCacheKey {
|
|||
type ReduceOp = (inputs: readonly TensorView[], axes: number[]) => string[];
|
||||
const noOp: ReduceOp = (): string[] => ['', '', 'value = _A[inputIdx];', ''];
|
||||
const createReduceProgramInfo =
|
||||
(metadata: ProgramMetadata, inputs: readonly TensorView[], attributes: ReduceAttributes,
|
||||
reduceOp: ReduceOp): ProgramInfo => {
|
||||
const outputShape: number[] = [];
|
||||
const inputShape = inputs[0].dims;
|
||||
(metadata: ProgramMetadata, inputs: readonly TensorView[], attributes: ReduceAttributes, reduceOp: ReduceOp):
|
||||
ProgramInfo => {
|
||||
const outputShape: number[] = [];
|
||||
const inputShape = inputs[0].dims;
|
||||
|
||||
const idxCopy: string[] = []; // copy output indexes to input indexes
|
||||
const idxCopy: string[] = []; // copy output indexes to input indexes
|
||||
|
||||
const axes = ShapeUtil.normalizeAxes(attributes.axes, inputs[0].dims.length);
|
||||
const outputDimsLength = inputs[0].dims.length - (attributes.keepDims ? 0 : axes.length);
|
||||
const ops = reduceOp(inputs, axes);
|
||||
const inputIndicesHelper = createIndicesHelper('input', inputShape);
|
||||
const initInputIdx = (ops[1] === '') ? '' : `let inputIdx = ${inputIndicesHelper.i2oExpression('inputIndices')};`;
|
||||
let reduceOps = `
|
||||
let inputIdx = ${inputIndicesHelper.i2oExpression('inputIndices')};
|
||||
${ops[2]};`;
|
||||
const reduceOnAllAxes = !attributes.noopWithEmptyAxes && attributes.axes.length === 0;
|
||||
for (let k = 0; k < inputs[0].dims.length; k++) {
|
||||
// if this axis is reduced
|
||||
if (reduceOnAllAxes || axes.indexOf(k) >= 0) {
|
||||
if (attributes.keepDims) {
|
||||
outputShape.push(1);
|
||||
} // else { remove the axis from outputShape; }
|
||||
const axes = ShapeUtil.normalizeAxes(attributes.axes, inputs[0].dims.length);
|
||||
const outputDimsLength = inputs[0].dims.length - (attributes.keepDims ? 0 : axes.length);
|
||||
const ops = reduceOp(inputs, axes);
|
||||
const inputIndicesHelper = createIndicesHelper('input', inputShape);
|
||||
const initInputIdxLet = `let inputIdx = ${inputIndicesHelper.i2oExpression('inputIndices')};`;
|
||||
const initInputIdxVar = `var inputIdx = ${inputIndicesHelper.i2oExpression('inputIndices')};`;
|
||||
const updateInputIdxImpl = `inputIdx = ${inputIndicesHelper.i2oExpression('inputIndices')};`;
|
||||
const initInputIdx = (ops[1] === '') ? '' : initInputIdxVar;
|
||||
let reduceOps = ((ops[1] === '') ? initInputIdxLet : updateInputIdxImpl) + '\n' + ops[2];
|
||||
const reduceOnAllAxes = !attributes.noopWithEmptyAxes && attributes.axes.length === 0;
|
||||
for (let k = 0; k < inputs[0].dims.length; k++) {
|
||||
const inputIndices = inputShape.length > 1 ? `inputIndices[${k}]` : 'inputIndices';
|
||||
// if this axis is reduced
|
||||
if (reduceOnAllAxes || axes.indexOf(k) >= 0) {
|
||||
if (attributes.keepDims) {
|
||||
outputShape.push(1);
|
||||
} // else { remove the axis from outputShape; }
|
||||
|
||||
// loop over the d-th axis
|
||||
reduceOps = `for(var j${k}: u32 = 0; j${k} < ${inputs[0].dims[k]}; j${k}++) {
|
||||
inputIndices[${k}] = j${k};
|
||||
${reduceOps}
|
||||
}`;
|
||||
} else {
|
||||
if (outputDimsLength > 1) {
|
||||
idxCopy.push(`inputIndices[${k}] = outputIndices[${outputShape.length}];`);
|
||||
} else {
|
||||
idxCopy.push(`inputIndices[${k}] = outputIndices;`);
|
||||
// loop over the d-th axis
|
||||
reduceOps = `for(var j${k}: u32 = 0; j${k} < ${inputs[0].dims[k]}; j${k}++) {
|
||||
${inputIndices} = j${k};
|
||||
${reduceOps}
|
||||
}`;
|
||||
} else {
|
||||
const outputIndices = outputDimsLength > 1 ? `outputIndices[${outputShape.length}]` : 'outputIndices';
|
||||
idxCopy.push(`${inputIndices} = ${outputIndices};`);
|
||||
outputShape.push(inputs[0].dims[k]);
|
||||
}
|
||||
}
|
||||
outputShape.push(inputs[0].dims[k]);
|
||||
}
|
||||
}
|
||||
|
||||
const outputIndicesHelper = createIndicesHelper('output', outputShape);
|
||||
const outputSize = ShapeUtil.size(outputShape);
|
||||
const dataType = 'f32';
|
||||
const outputIndicesHelper = createIndicesHelper('output', outputShape);
|
||||
const outputSize = ShapeUtil.size(outputShape);
|
||||
const dataType = 'f32';
|
||||
|
||||
const getShaderSource = (shaderHelper: ShaderHelper) => `
|
||||
const getShaderSource = (shaderHelper: ShaderHelper) => `
|
||||
@group(0) @binding(0) var<storage, read> _A : array<${dataType}>;
|
||||
@group(0) @binding(1) var<storage, read_write> output : array<${dataType}>;
|
||||
|
||||
|
|
@ -98,13 +97,13 @@ const createReduceProgramInfo =
|
|||
output[global_idx] = value;
|
||||
}`;
|
||||
|
||||
return {
|
||||
...metadata,
|
||||
getShaderSource,
|
||||
outputs: [{dims: outputShape, dataType: inputs[0].dataType, gpuDataType: GpuDataType.default}],
|
||||
dispatchGroup: () => ({x: Math.ceil(outputSize / 64 /* workgroup size */)})
|
||||
};
|
||||
};
|
||||
return {
|
||||
...metadata,
|
||||
getShaderSource,
|
||||
outputs: [{dims: outputShape, dataType: inputs[0].dataType, gpuDataType: GpuDataType.default}],
|
||||
dispatchGroup: () => ({x: Math.ceil(outputSize / 64 /* workgroup size */)})
|
||||
};
|
||||
};
|
||||
|
||||
const createReduceAttributesFromInputs =
|
||||
(inputs: readonly TensorView[], attributes: ReduceAttributes): ReduceAttributes => {
|
||||
|
|
@ -121,8 +120,11 @@ const createReduceProgramInfoLoader =
|
|||
ProgramInfoLoader => {
|
||||
const updatedAttributes: ReduceAttributes =
|
||||
inputs.length === 1 ? attributes : createReduceAttributesFromInputs(inputs, attributes);
|
||||
const metadata:
|
||||
ProgramMetadata = {name, inputTypes: [GpuDataType.default], cacheHint: updatedAttributes.cacheKey};
|
||||
const metadata: ProgramMetadata = {
|
||||
name,
|
||||
inputTypes: [GpuDataType.default],
|
||||
cacheHint: updatedAttributes.cacheKey + '_' + inputs[0].dims.map(d => d.toString()).join(',')
|
||||
};
|
||||
return {
|
||||
...metadata,
|
||||
get: () => createReduceProgramInfo(
|
||||
|
|
|
|||
|
|
@ -38,7 +38,7 @@ const readInput = (inputs: readonly TensorView[], idx: number): number[] => {
|
|||
if (inputs.length > idx) {
|
||||
if (inputs[idx].dataType === DataType.int64) {
|
||||
inputs[idx].getBigInt64Array().forEach(v => input.push(Number(v)));
|
||||
} else if (inputs[1].dataType === DataType.int32) {
|
||||
} else if (inputs[idx].dataType === DataType.int32) {
|
||||
inputs[idx].getInt32Array().forEach(v => input.push(Number(v)));
|
||||
} else {
|
||||
throw new Error(`Input ${idx} must be an array of int32 or int64`);
|
||||
|
|
@ -189,7 +189,14 @@ const createSliceProgramInfoLoader =
|
|||
|
||||
export const slice = (context: ComputeContext, attributes: SliceAttributes): void => {
|
||||
validateInputs(context.inputs, attributes);
|
||||
context.compute(createSliceProgramInfoLoader(context.inputs, attributes), {inputs: [0]});
|
||||
const programInfoLoader = createSliceProgramInfoLoader(context.inputs, attributes);
|
||||
const program = programInfoLoader.get();
|
||||
if (ShapeUtil.size(program.outputs[0].dims) > 0) {
|
||||
context.compute(programInfoLoader, {inputs: [0]});
|
||||
} else {
|
||||
// TODO: support empty output
|
||||
throw new Error('slice: output size is 0');
|
||||
}
|
||||
};
|
||||
|
||||
export const parseSliceAttributes = (attributes: Record<string, unknown>): SliceAttributes => {
|
||||
|
|
|
|||
|
|
@ -1124,11 +1124,11 @@
|
|||
// // "test_size",
|
||||
"test_slice_default_axes",
|
||||
"test_slice_default_steps",
|
||||
"test_slice_end_out_of_bounds",
|
||||
// "test_slice_end_out_of_bounds",
|
||||
"test_slice_neg_steps",
|
||||
"test_slice_neg",
|
||||
"test_slice_negative_axes",
|
||||
"test_slice_start_out_of_bounds",
|
||||
// "test_slice_start_out_of_bounds",
|
||||
"test_slice",
|
||||
// "test_softmax_axis_0_expanded",
|
||||
// "test_softmax_axis_0",
|
||||
|
|
|
|||
Loading…
Reference in a new issue