mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-05 04:17:53 +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.
65 lines
2.5 KiB
TypeScript
65 lines
2.5 KiB
TypeScript
// Copyright (c) Microsoft Corporation. All rights reserved.
|
|
// Licensed under the MIT License.
|
|
|
|
import {env} from 'onnxruntime-common';
|
|
|
|
import {DataType} from '../../../wasm-common';
|
|
import {ComputeContext, GpuDataType, ProgramInfo} from '../types';
|
|
|
|
import {outputVariable, ShaderHelper} from './common';
|
|
|
|
const validateInputsContent = (start: number, limit: number, delta: number): void => {
|
|
const sameStartLimit = start === limit;
|
|
const increasingRangeNegativeStep = start < limit && delta < 0;
|
|
const decreasingRangePositiveStep = start > limit && delta > 0;
|
|
|
|
if (sameStartLimit || increasingRangeNegativeStep || decreasingRangePositiveStep) {
|
|
throw new Error('Range these inputs\' contents are invalid.');
|
|
}
|
|
};
|
|
|
|
const createRangeProgramInfo = (start: number, limit: number, delta: number, dataType: DataType): ProgramInfo => {
|
|
const numElements = Math.abs(Math.ceil((limit - start) / delta));
|
|
const outputShape: number[] = [numElements];
|
|
const outputSize = numElements;
|
|
|
|
const output = outputVariable('output', dataType, outputShape);
|
|
const wgslType = output.type.storage;
|
|
|
|
const getShaderSource = (shaderHelper: ShaderHelper) => `
|
|
${shaderHelper.declareVariables(output)}
|
|
${shaderHelper.mainStart()}
|
|
${shaderHelper.guardAgainstOutOfBoundsWorkgroupSizes(outputSize)}
|
|
output[global_idx] = ${wgslType}(${start}) + ${wgslType}(global_idx) * ${wgslType}(${delta});
|
|
}`;
|
|
return {
|
|
name: 'Range',
|
|
inputTypes: [],
|
|
shaderCache: {hint: [start, limit, delta].map(x => x.toString()).join('_')},
|
|
getShaderSource,
|
|
getRunData: () => ({
|
|
outputs: [{dims: outputShape, dataType, gpuDataType: GpuDataType.default}],
|
|
dispatchGroup: {x: Math.ceil(outputSize / 64 /* workgroup size */)}
|
|
})
|
|
};
|
|
};
|
|
|
|
export const range = (context: ComputeContext): void => {
|
|
let start = 0;
|
|
let limit = 0;
|
|
let delta = 0;
|
|
if (context.inputs[0].dataType === DataType.int32) {
|
|
start = context.inputs[0].getInt32Array()[0];
|
|
limit = context.inputs[1].getInt32Array()[0];
|
|
delta = context.inputs[2].getInt32Array()[0];
|
|
} else if (context.inputs[0].dataType === DataType.float) {
|
|
start = context.inputs[0].getFloat32Array()[0];
|
|
limit = context.inputs[1].getFloat32Array()[0];
|
|
delta = context.inputs[2].getFloat32Array()[0];
|
|
}
|
|
if (env.webgpu.validateInputContent) {
|
|
validateInputsContent(start, limit, delta);
|
|
}
|
|
|
|
context.compute(createRangeProgramInfo(start, limit, delta, context.inputs[0].dataType), {inputs: []});
|
|
};
|