mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-19 19:00:47 +00:00
d9b9c5a537
### Description support using uniform buffer. This PR allows to use uniform buffer in shader program, so that some runtime information (eg. input/output shape) is no longer need to be hardcoded into shader code. There are 2 commits in this PR: - [667f31c](667f31c83d): framework changes to support uniform buffer, as well as updates in program manager, gpu data manager and indices helper. - [09e1d2a](09e1d2ad1d): an example change for operator `Transpose` to use input's rank-only instead of dims as shader key. With this change, model mobilenetv2-12 shader compile times dropped from 71 to 52.
92 lines
3.8 KiB
TypeScript
92 lines
3.8 KiB
TypeScript
// Copyright (c) Microsoft Corporation. All rights reserved.
|
|
// Licensed under the MIT License.
|
|
|
|
// TODO: this is the same naive implementation we use for reduce that has
|
|
// performance limitations when the reduced axis is long. Need to add
|
|
// a optimized codepath for this.
|
|
|
|
import {DataType} from '../../../wasm-common';
|
|
import {TensorView} from '../../tensor-view';
|
|
import {AttributeWithCacheKey, createAttributeWithCacheKey} from '../attribute-with-cache-key';
|
|
import {ComputeContext} from '../types';
|
|
|
|
import {createReduceProgramInfo, ReduceOp} from './reduce';
|
|
|
|
const validateInputs = (inputs: readonly TensorView[]): void => {
|
|
if (!inputs || inputs.length === 0 || inputs.length > 2) {
|
|
throw new Error('ArgMinMaxOp op requires 1 or 2 inputs.');
|
|
}
|
|
if (inputs[0].dataType !== DataType.float) {
|
|
throw new Error('Invalid input type.');
|
|
}
|
|
};
|
|
|
|
export interface ArgMinMaxAttributes extends AttributeWithCacheKey {
|
|
keepDims: boolean;
|
|
axis: number;
|
|
selectLastIndex: number;
|
|
}
|
|
|
|
const createArgMinMaxAttributesFromInputs =
|
|
(inputs: readonly TensorView[], attributes: ArgMinMaxAttributes): ArgMinMaxAttributes =>
|
|
createAttributeWithCacheKey(
|
|
{axis: attributes.axis, keepDims: attributes.keepDims, selectLastIndex: attributes.selectLastIndex});
|
|
|
|
export const argMin = (context: ComputeContext, attributes: ArgMinMaxAttributes): void => {
|
|
validateInputs(context.inputs);
|
|
const argMinMaxOp: ReduceOp = (input, output, axes) => {
|
|
const idxZero = [];
|
|
for (let k = 0; k < input.rank; k++) {
|
|
if (axes.indexOf(k) >= 0 || axes.length === 0) {
|
|
idxZero.push(`inputIndices[${k}] = 0;`); // first element
|
|
}
|
|
}
|
|
return [
|
|
`${idxZero.join('\n')}`, `var value = ${input.getByOffset('inputOffset')};\nvar bestIndex : i32 = 0;`,
|
|
`if (${input.getByOffset('inputOffset')} ${attributes.selectLastIndex > 0 ? '<=' : '<'} value) {
|
|
value = ${input.getByOffset('inputOffset')};
|
|
bestIndex = i32(lastIndex);
|
|
}`,
|
|
'', output.setByOffset('global_idx', 'bestIndex')
|
|
];
|
|
};
|
|
|
|
const updatedAttributes: ArgMinMaxAttributes =
|
|
context.inputs.length === 1 ? attributes : createArgMinMaxAttributesFromInputs(context.inputs, attributes);
|
|
context.compute(
|
|
createReduceProgramInfo(
|
|
'ArgMin', {hint: updatedAttributes.cacheKey}, [context.inputs[0]], argMinMaxOp, [updatedAttributes.axis],
|
|
DataType.int64, updatedAttributes.keepDims),
|
|
{inputs: [0]});
|
|
};
|
|
|
|
export const argMax = (context: ComputeContext, attributes: ArgMinMaxAttributes): void => {
|
|
validateInputs(context.inputs);
|
|
const argMinMaxOp: ReduceOp = (input, output, axes) => {
|
|
const idxZero = [];
|
|
for (let k = 0; k < input.rank; k++) {
|
|
if (axes.indexOf(k) >= 0 || axes.length === 0) {
|
|
idxZero.push(`inputIndices[${k}] = 0;`); // first element
|
|
}
|
|
}
|
|
return [
|
|
`${idxZero.join('\n')}`, `var value = ${input.getByOffset('inputOffset')};\nvar bestIndex : i32 = 0;`,
|
|
`if (${input.getByOffset('inputOffset')} ${attributes.selectLastIndex > 0 ? '>=' : '>'} value) {
|
|
value = ${input.getByOffset('inputOffset')};
|
|
bestIndex = i32(lastIndex);
|
|
}`,
|
|
'', output.setByOffset('global_idx', 'bestIndex')
|
|
];
|
|
};
|
|
|
|
const updatedAttributes: ArgMinMaxAttributes =
|
|
context.inputs.length === 1 ? attributes : createArgMinMaxAttributesFromInputs(context.inputs, attributes);
|
|
context.compute(
|
|
createReduceProgramInfo(
|
|
'argMax', {hint: updatedAttributes.cacheKey}, [context.inputs[0]], argMinMaxOp, [updatedAttributes.axis],
|
|
DataType.int64, updatedAttributes.keepDims),
|
|
{inputs: [0]});
|
|
};
|
|
|
|
export const parseArgMinMaxAttributes = (attributes: Record<string, unknown>): ArgMinMaxAttributes =>
|
|
createAttributeWithCacheKey(attributes as Omit<ArgMinMaxAttributes, keyof AttributeWithCacheKey>);
|