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

[js/webgpu] Remove enableShapesUniforms (#19279)

Browse source
This commit is contained in:
Xu Xing 2024-01-30 09:49:06 +08:00 committed by GitHub
parent 00d048121b
commit 624b4e2063
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
9 changed files with 68 additions and 134 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

12
js/web/lib/wasm/jsep/webgpu/ops/3rd-party/matmul_packed_webgpu.ts vendored
View file

@ -443,9 +443,9 @@ export const createMatmulProgramInfo =
const components = isVec4 ? 4 : 1; const components = isVec4 ? 4 : 1;
const aShapeTemp = [...outerDimsA, dimAOuter, dimInner / components]; const aShapeTemp = [...outerDimsA, dimAOuter, dimInner / components];
const aShapeOrRank = aShapeTemp.length; const aRank = aShapeTemp.length;
const bShapeTemp = [...outerDimsB, dimInner, dimBOuter / components]; const bShapeTemp = [...outerDimsB, dimInner, dimBOuter / components];
const bShapeOrRank = bShapeTemp.length; const bRank = bShapeTemp.length;
const outputShapeTemp = [batchSize, dimAOuter, dimBOuter / components]; const outputShapeTemp = [batchSize, dimAOuter, dimBOuter / components];
const programUniforms: ProgramUniform[] = const programUniforms: ProgramUniform[] =
[{type: 'int32', data: dimAOuter}, {type: 'int32', data: dimBOuter}, {type: 'int32', data: dimInner}]; [{type: 'int32', data: dimAOuter}, {type: 'int32', data: dimBOuter}, {type: 'int32', data: dimInner}];
@ -467,12 +467,12 @@ export const createMatmulProgramInfo =
programUniforms.push(...createTensorShapeVariables(outputShapeTemp)); programUniforms.push(...createTensorShapeVariables(outputShapeTemp));
const getShaderSource = (shaderHelper: ShaderHelper) => { const getShaderSource = (shaderHelper: ShaderHelper) => {
const batchShapeOrRank = outerDims.length; const batchRank = outerDims.length;
const batchDims = internalVariable('batchDims', inputs[0].dataType, batchShapeOrRank, 1); const batchDims = internalVariable('batchDims', inputs[0].dataType, batchRank, 1);
const dataType = tensorTypeToWsglStorageType(inputs[0].dataType); const dataType = tensorTypeToWsglStorageType(inputs[0].dataType);
const A = inputVariable('a', inputs[0].dataType, aShapeOrRank, components); const A = inputVariable('a', inputs[0].dataType, aRank, components);
const B = inputVariable('b', inputs[1].dataType, bShapeOrRank, components); const B = inputVariable('b', inputs[1].dataType, bRank, components);
const output = outputVariable('result', inputs[0].dataType, outputShapeTemp.length, components); const output = outputVariable('result', inputs[0].dataType, outputShapeTemp.length, components);
const inputVariables = [A, B]; const inputVariables = [A, B];
if (hasBias) { if (hasBias) {

4
js/web/lib/wasm/jsep/webgpu/ops/batch-norm.ts
View file

@ -8,7 +8,7 @@ import {ShapeUtil} from '../../util';
import {AttributeWithCacheKey, createAttributeWithCacheKey} from '../attribute-with-cache-key'; import {AttributeWithCacheKey, createAttributeWithCacheKey} from '../attribute-with-cache-key';
import {ComputeContext, ProgramInfo} from '../types'; import {ComputeContext, ProgramInfo} from '../types';
import {createTensorShapeVariables, enableShapesUniforms, getMaxComponents, inputVariable, outputVariable, ShaderHelper} from './common'; import {createTensorShapeVariables, getMaxComponents, inputVariable, outputVariable, ShaderHelper} from './common';
export interface BatchNormAttributes extends AttributeWithCacheKey { export interface BatchNormAttributes extends AttributeWithCacheKey {
readonly epsilon: number; readonly epsilon: number;
@ -61,7 +61,7 @@ const createBatchNormInferenceProgramInfo =
const cComponents = format === 'NHWC' && yShape.length > 1 ? components : 1; const cComponents = format === 'NHWC' && yShape.length > 1 ? components : 1;
const outputSize = ShapeUtil.size(yShape) / components; const outputSize = ShapeUtil.size(yShape) / components;
// Only support uniforms for opset version >= 9 (spatial = true). // Only support uniforms for opset version >= 9 (spatial = true).
const useShapesUniforms = enableShapesUniforms(yShape.length) && spatial; const useShapesUniforms = spatial;
const shapeOrRank = useShapesUniforms ? yShape.length : yShape; const shapeOrRank = useShapesUniforms ? yShape.length : yShape;
const x = inputVariable('x', inputs[0].dataType, inputs[0].dims, components); const x = inputVariable('x', inputs[0].dataType, inputs[0].dims, components);
const scale = inputVariable('scale', inputs[1].dataType, inputs[1].dims, cComponents); const scale = inputVariable('scale', inputs[1].dataType, inputs[1].dims, cComponents);

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

@ -6,7 +6,7 @@ import {TensorView} from '../../tensor-view';
import {BroadcastUtil, ShapeUtil} from '../../util'; import {BroadcastUtil, ShapeUtil} from '../../util';
import {ComputeContext, ProgramInfo} from '../types'; import {ComputeContext, ProgramInfo} from '../types';
import {createTensorShapeVariables, enableShapesUniforms, inputVariable, outputVariable, ShaderHelper} from './common'; import {createTensorShapeVariables, inputVariable, outputVariable, ShaderHelper} from './common';
type BuiltinFunctionName = string; type BuiltinFunctionName = string;
type BinaryCustomExpression = (expressionA: string, expressionB: string) => string; type BinaryCustomExpression = (expressionA: string, expressionB: string) => string;
@ -18,8 +18,7 @@ type BinaryFunctionCall = BuiltinFunctionName|BinaryCustomExpression|{
const createBinaryOpProgramShader = const createBinaryOpProgramShader =
(shaderHelper: ShaderHelper, dimsA: readonly number[], dimsB: readonly number[], dimsOutput: readonly number[], (shaderHelper: ShaderHelper, dimsA: readonly number[], dimsB: readonly number[], dimsOutput: readonly number[],
vectorize: boolean, doBroadcast: boolean, sharedDimensionDivisibleBy4: boolean, funcCall: BinaryFunctionCall, vectorize: boolean, doBroadcast: boolean, sharedDimensionDivisibleBy4: boolean, funcCall: BinaryFunctionCall,
typeA: number, typeB: number, typeOutput: number, useShapesUniforms: boolean, typeA: number, typeB: number, typeOutput: number, additionalImplementation?: string) => {
additionalImplementation?: string) => {
let expressionScalar: BinaryCustomExpression; let expressionScalar: BinaryCustomExpression;
let expressionVector: BinaryCustomExpression; let expressionVector: BinaryCustomExpression;
if (typeof funcCall === 'string') { if (typeof funcCall === 'string') {
@ -31,12 +30,9 @@ const createBinaryOpProgramShader =
expressionVector = funcCall.vector; expressionVector = funcCall.vector;
} }
const inputAShapeOrRank = useShapesUniforms ? dimsA.length : dimsA; const output = outputVariable('outputData', typeOutput, dimsOutput.length, 4);
const inputBShapeOrRank = useShapesUniforms ? dimsB.length : dimsB; const a = inputVariable('aData', typeA, dimsA.length, 4);
const outputShapeOrRank = useShapesUniforms ? dimsOutput.length : dimsOutput; const b = inputVariable('bData', typeB, dimsB.length, 4);
const output = outputVariable('outputData', typeOutput, outputShapeOrRank, 4);
const a = inputVariable('aData', typeA, inputAShapeOrRank, 4);
const b = inputVariable('bData', typeB, inputBShapeOrRank, 4);
let assignment: string; let assignment: string;
if (vectorize) { if (vectorize) {
@ -169,30 +165,25 @@ const createBinaryOpProgramInfo =
vectorize = true; vectorize = true;
} }
cacheKeyAux.push(vectorize); cacheKeyAux.push(vectorize);
const useShapesUniforms = enableShapesUniforms(a.dims.length) && enableShapesUniforms(b.dims.length) &&
enableShapesUniforms(outputShape.length);
return { return {
name, name,
shaderCache: { shaderCache: {
hint: cacheKey + cacheKeyAux.map((x) => x.toString()).join('_'), hint: cacheKey + cacheKeyAux.map((x) => x.toString()).join('_'),
inputDependencies: useShapesUniforms ? ['rank', 'rank'] : ['dims', 'dims'], inputDependencies: ['rank', 'rank'],
}, },
getShaderSource: (shaderHelper) => createBinaryOpProgramShader( getShaderSource: (shaderHelper) => createBinaryOpProgramShader(
shaderHelper, a.dims, b.dims, outputShape, vectorize, isBroadcast, sharedDimensionDivisibleBy4, funcCall, shaderHelper, a.dims, b.dims, outputShape, vectorize, isBroadcast, sharedDimensionDivisibleBy4, funcCall,
a.dataType, b.dataType, outputDataType, useShapesUniforms, additionalImplementation), a.dataType, b.dataType, outputDataType, additionalImplementation),
getRunData: () => ({ getRunData: () => ({
outputs: [{dims: outputShape, dataType: outputDataType}], outputs: [{dims: outputShape, dataType: outputDataType}],
dispatchGroup: {x: Math.ceil(outputSize / 64 /* workgroup size */ / 4 /* component size */)}, dispatchGroup: {x: Math.ceil(outputSize / 64 /* workgroup size */ / 4 /* component size */)},
programUniforms: useShapesUniforms ? programUniforms: [
[ {type: 'uint32', data: Math.ceil(ShapeUtil.size(outputShape) / 4)},
{type: 'uint32', data: Math.ceil(ShapeUtil.size(outputShape) / 4)}, ...createTensorShapeVariables(a.dims),
...createTensorShapeVariables(a.dims), ...createTensorShapeVariables(b.dims),
...createTensorShapeVariables(b.dims), ...createTensorShapeVariables(outputShape),
...createTensorShapeVariables(outputShape), ],
] :
[
{type: 'uint32', data: Math.ceil(ShapeUtil.size(outputShape) / 4)},
],
}), }),
}; };
}; };

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

@ -922,6 +922,3 @@ export const getBroadcastDims = (inShape: readonly number[], outShape: readonly
} }
return dims; return dims;
}; };
// TODO: remove this when all related uses have been removed.
export const enableShapesUniforms = (_rank: number): boolean => true;

26
js/web/lib/wasm/jsep/webgpu/ops/concat.ts
View file

@ -6,7 +6,7 @@ import {ShapeUtil} from '../../util';
import {AttributeWithCacheKey, createAttributeWithCacheKey} from '../attribute-with-cache-key'; import {AttributeWithCacheKey, createAttributeWithCacheKey} from '../attribute-with-cache-key';
import {ComputeContext, ProgramInfo, ProgramInputTensorInfoDependency, ProgramUniform} from '../types'; import {ComputeContext, ProgramInfo, ProgramInputTensorInfoDependency, ProgramUniform} from '../types';
import {createTensorShapeVariables, enableShapesUniforms, IndicesHelper, inputVariable, outputVariable, ShaderHelper} from './common'; import {createTensorShapeVariables, IndicesHelper, inputVariable, outputVariable, ShaderHelper} from './common';
export interface ConcatAttributes extends AttributeWithCacheKey { export interface ConcatAttributes extends AttributeWithCacheKey {
readonly axis: number; readonly axis: number;
@ -94,32 +94,22 @@ const createConcatProgramInfo = (inputs: readonly TensorView[], axis: number): P
let previousSum = 0; let previousSum = 0;
const inputDependencies: ProgramInputTensorInfoDependency[] = []; const inputDependencies: ProgramInputTensorInfoDependency[] = [];
const inputShapeOrRanks = []; const inputRanks = [];
const enableInputShapesUniforms = [];
const programUniforms: ProgramUniform[] = [{type: 'uint32', data: outputSize}]; const programUniforms: ProgramUniform[] = [{type: 'uint32', data: outputSize}];
for (let i = 0; i < inputs.length; ++i) { for (let i = 0; i < inputs.length; ++i) {
previousSum += inputs[i].dims[adjustedAxis]; previousSum += inputs[i].dims[adjustedAxis];
sizeInConcatAxis[i] = previousSum; sizeInConcatAxis[i] = previousSum;
enableInputShapesUniforms.push(enableShapesUniforms(inputs[i].dims.length)); inputRanks.push(inputs[i].dims.length);
inputShapeOrRanks.push(enableInputShapesUniforms[i] ? inputs[i].dims.length : inputs[i].dims); inputVars[i] = inputVariable(`input${i}`, dataType, inputRanks[i]);
inputVars[i] = inputVariable(`input${i}`, dataType, inputShapeOrRanks[i]); inputDependencies.push('rank');
inputDependencies.push(enableInputShapesUniforms[i] ? 'rank' : 'dims');
programUniforms.push({type: 'uint32', data: sizeInConcatAxis[i]}); programUniforms.push({type: 'uint32', data: sizeInConcatAxis[i]});
} }
for (let i = 0; i < inputs.length; ++i) { for (let i = 0; i < inputs.length; ++i) {
if (enableInputShapesUniforms[i]) { programUniforms.push(...createTensorShapeVariables(inputs[i].dims));
programUniforms.push(...createTensorShapeVariables(inputs[i].dims));
}
} }
programUniforms.push(...createTensorShapeVariables(outputShape));
const enableOutputShapesUniforms = enableShapesUniforms(outputShape.length); const output = outputVariable('output', dataType, outputShape.length);
if (enableOutputShapesUniforms) {
programUniforms.push(...createTensorShapeVariables(outputShape));
}
const outputShapeOrRank = enableOutputShapesUniforms ? outputShape.length : outputShape;
const output = outputVariable('output', dataType, outputShapeOrRank);
const indicesAxis = output.indicesGet('indices', adjustedAxis); const indicesAxis = output.indicesGet('indices', adjustedAxis);
const sizeInConcatAxisStr = const sizeInConcatAxisStr =
Array.from(Array(sizeInConcatAxis.length).keys()).map(i => `uniforms.sizeInConcatAxis${i}`).join(','); Array.from(Array(sizeInConcatAxis.length).keys()).map(i => `uniforms.sizeInConcatAxis${i}`).join(',');

31
js/web/lib/wasm/jsep/webgpu/ops/einsum.ts
View file

@ -6,8 +6,7 @@ import {ShapeUtil} from '../../util';
import {AttributeWithCacheKey, createAttributeWithCacheKey} from '../attribute-with-cache-key'; import {AttributeWithCacheKey, createAttributeWithCacheKey} from '../attribute-with-cache-key';
import {ComputeContext, ProgramInfo, ProgramUniform} from '../types'; import {ComputeContext, ProgramInfo, ProgramUniform} from '../types';
import {createTensorShapeVariables, enableShapesUniforms, inputVariable, outputVariable, ShaderHelper} from './common'; import {createTensorShapeVariables, inputVariable, outputVariable, ShaderHelper} from './common';
export interface EinsumAttributes extends AttributeWithCacheKey { export interface EinsumAttributes extends AttributeWithCacheKey {
readonly equation: string; readonly equation: string;
@ -181,14 +180,12 @@ class EinsumEquation {
const appendMax = (name: string): string => name + '_max'; const appendMax = (name: string): string => name + '_max';
const createEinsumProgramInfo = const createEinsumProgramInfo =
(enableInputShapesUniforms: readonly boolean[], inputShapes: Array<readonly number[]>, dataType: number, (inputShapes: Array<readonly number[]>, dataType: number, einsumEquation: EinsumEquation,
einsumEquation: EinsumEquation, outputShape: readonly number[]): ProgramInfo => { outputShape: readonly number[]): ProgramInfo => {
const shapeOrRanks = inputShapes.map((dims, index) => enableInputShapesUniforms[index] ? dims.length : dims); const ranks = inputShapes.map((dims) => dims.length);
const inputVars = shapeOrRanks.map((shapeOrRank, index) => inputVariable(`input${index}`, dataType, shapeOrRank)); const inputVars = ranks.map((rank, index) => inputVariable(`input${index}`, dataType, rank));
const outputSize = ShapeUtil.size(outputShape); const outputSize = ShapeUtil.size(outputShape);
const enableOutputShapesUniforms = enableShapesUniforms(outputShape.length); const output = outputVariable('output', dataType, outputShape.length);
const outputShapeOrRank = enableOutputShapesUniforms ? outputShape.length : outputShape;
const output = outputVariable('output', dataType, outputShapeOrRank);
const uniformsSymbols = const uniformsSymbols =
[...einsumEquation.symbolToInfo.keys()].filter((symbol) => !einsumEquation.rhs.symbolToIndices.has(symbol)); [...einsumEquation.symbolToInfo.keys()].filter((symbol) => !einsumEquation.rhs.symbolToIndices.has(symbol));
const getShaderSource = (shaderHelper: ShaderHelper) => { const getShaderSource = (shaderHelper: ShaderHelper) => {
@ -269,10 +266,7 @@ const createEinsumProgramInfo =
}; };
return { return {
name: 'Einsum', name: 'Einsum',
shaderCache: { shaderCache: {hint: einsumEquation.equation, inputDependencies: inputShapes.map(() => 'rank')},
hint: einsumEquation.equation,
inputDependencies: enableInputShapesUniforms.map((enableShapeUniform) => enableShapeUniform ? 'rank' : 'dims')
},
getRunData: () => { getRunData: () => {
// The symbols from uniformSymbols array are guaranteed to exist in einsumEquations.symbolToInfo map. The // The symbols from uniformSymbols array are guaranteed to exist in einsumEquations.symbolToInfo map. The
// filter is added to make sure that dimValue is never 0. // filter is added to make sure that dimValue is never 0.
@ -281,12 +275,9 @@ const createEinsumProgramInfo =
.map((symbol) => ({type: 'uint32', data: einsumEquation.symbolToInfo.get(symbol)?.dimValue || 0})); .map((symbol) => ({type: 'uint32', data: einsumEquation.symbolToInfo.get(symbol)?.dimValue || 0}));
programUniformsInit.push({type: 'uint32', data: outputSize}); programUniformsInit.push({type: 'uint32', data: outputSize});
const programUniforms: ProgramUniform[] = const programUniforms: ProgramUniform[] =
inputShapes.filter((_, index) => enableInputShapesUniforms[index]) inputShapes.map((dims, _) => [...createTensorShapeVariables(dims)])
.map((dims, _) => [...createTensorShapeVariables(dims)])
.reduce((acc, inputProgramUniforms) => acc.concat(inputProgramUniforms), programUniformsInit); .reduce((acc, inputProgramUniforms) => acc.concat(inputProgramUniforms), programUniformsInit);
if (enableOutputShapesUniforms) { programUniforms.push(...createTensorShapeVariables(outputShape));
programUniforms.push(...createTensorShapeVariables(outputShape));
}
return ({ return ({
outputs: [{dims: outputShape, dataType}], outputs: [{dims: outputShape, dataType}],
dispatchGroup: {x: Math.ceil(outputSize / 64 /* workgroup size */)}, dispatchGroup: {x: Math.ceil(outputSize / 64 /* workgroup size */)},
@ -299,11 +290,9 @@ const createEinsumProgramInfo =
export const einsum = (context: ComputeContext, attributes: EinsumAttributes): void => { export const einsum = (context: ComputeContext, attributes: EinsumAttributes): void => {
const einsumEquation = new EinsumEquation(context.inputs, attributes.equation); const einsumEquation = new EinsumEquation(context.inputs, attributes.equation);
const enableInputShapesUniforms = context.inputs.map((input, _) => enableShapesUniforms(input.dims.length));
const outputShape = einsumEquation.outputDims; const outputShape = einsumEquation.outputDims;
const inputShapes = context.inputs.map((input, _) => input.dims); const inputShapes = context.inputs.map((input, _) => input.dims);
context.compute(createEinsumProgramInfo( context.compute(createEinsumProgramInfo(inputShapes, context.inputs[0].dataType, einsumEquation, outputShape));
enableInputShapesUniforms, inputShapes, context.inputs[0].dataType, einsumEquation, outputShape));
}; };
export const parseEinsumAttributes = (attributes: Record<string, unknown>): EinsumAttributes => { export const parseEinsumAttributes = (attributes: Record<string, unknown>): EinsumAttributes => {

25
js/web/lib/wasm/jsep/webgpu/ops/expand.ts
View file

@ -6,7 +6,7 @@ import {TensorView} from '../../tensor-view';
import {ShapeUtil} from '../../util'; import {ShapeUtil} from '../../util';
import {ComputeContext, ProgramInfo, ProgramUniform} from '../types'; import {ComputeContext, ProgramInfo, ProgramUniform} from '../types';
import {createTensorShapeVariables, enableShapesUniforms, inputVariable, outputVariable, ShaderHelper} from './common'; import {createTensorShapeVariables, inputVariable, outputVariable, ShaderHelper} from './common';
const validateInputs = (inputs: readonly TensorView[]): void => { const validateInputs = (inputs: readonly TensorView[]): void => {
if (!inputs || inputs.length !== 2) { if (!inputs || inputs.length !== 2) {
@ -49,15 +49,9 @@ const createExpandProgramInfo = (inputs: readonly TensorView[]): ProgramInfo =>
const components = dataType === DataType.bool ? 4 : 1; const components = dataType === DataType.bool ? 4 : 1;
const outputSize = Math.ceil(ShapeUtil.size(outputShape) / components); const outputSize = Math.ceil(ShapeUtil.size(outputShape) / components);
const enableInputShapeUniform = enableShapesUniforms(inputShape.length);
const enableOutputShapeUniform = enableShapesUniforms(outputShape.length);
const getShaderSource = (shaderHelper: ShaderHelper) => { const getShaderSource = (shaderHelper: ShaderHelper) => {
const inputShapeOrRank = enableInputShapeUniform ? inputShape.length : inputShape; const input = inputVariable('input', dataType, inputShape.length, components);
const outputShapeOrRank = enableOutputShapeUniform ? outputShape.length : outputShape; const output = outputVariable('output', dataType, outputShape.length, components);
const input = inputVariable('input', dataType, inputShapeOrRank, components);
const output = outputVariable('output', dataType, outputShapeOrRank, components);
let assignment: string; let assignment: string;
if (dataType === DataType.bool) { if (dataType === DataType.bool) {
const singleAssignment = (resStr: string, x: number, typeCast = '') => ` const singleAssignment = (resStr: string, x: number, typeCast = '') => `
@ -90,16 +84,13 @@ const createExpandProgramInfo = (inputs: readonly TensorView[]): ProgramInfo =>
${assignment}`; ${assignment}`;
}; };
const programUniforms: ProgramUniform[] = [{type: 'uint32', data: outputSize}]; const programUniforms: ProgramUniform[] = [
if (enableInputShapeUniform) { {type: 'uint32', data: outputSize}, ...createTensorShapeVariables(inputShape),
programUniforms.push(...createTensorShapeVariables(inputShape)); ...createTensorShapeVariables(outputShape)
} ];
if (enableOutputShapeUniform) {
programUniforms.push(...createTensorShapeVariables(outputShape));
}
return { return {
name: 'Expand', name: 'Expand',
shaderCache: {hint: `${outputShape.length}`, inputDependencies: [enableInputShapeUniform ? 'rank' : 'dims']}, shaderCache: {hint: `${outputShape.length}`, inputDependencies: ['rank']},
getShaderSource, getShaderSource,
getRunData: () => ({ getRunData: () => ({
outputs: [{dims: outputShape, dataType: inputs[0].dataType}], outputs: [{dims: outputShape, dataType: inputs[0].dataType}],

39
js/web/lib/wasm/jsep/webgpu/ops/gather.ts
View file

@ -5,9 +5,9 @@ import {DataType} from '../../../wasm-common';
import {TensorView} from '../../tensor-view'; import {TensorView} from '../../tensor-view';
import {ShapeUtil} from '../../util'; import {ShapeUtil} from '../../util';
import {AttributeWithCacheKey, createAttributeWithCacheKey} from '../attribute-with-cache-key'; import {AttributeWithCacheKey, createAttributeWithCacheKey} from '../attribute-with-cache-key';
import {ComputeContext, ProgramInfo, ProgramInputTensorInfoDependency, ProgramUniform} from '../types'; import {ComputeContext, ProgramInfo, ProgramUniform} from '../types';
import {createTensorShapeVariables, enableShapesUniforms, inputVariable, outputVariable, ShaderHelper} from './common'; import {createTensorShapeVariables, inputVariable, outputVariable, ShaderHelper} from './common';
export interface GatherAttributes extends AttributeWithCacheKey { export interface GatherAttributes extends AttributeWithCacheKey {
axis: number; axis: number;
@ -33,33 +33,16 @@ const createGatherProgramInfo = (inputs: readonly TensorView[], attributes: Gath
const components = inputs[0].dataType === DataType.bool ? 4 : 1; const components = inputs[0].dataType === DataType.bool ? 4 : 1;
const outputSize = Math.ceil(ShapeUtil.size(outputShape) / components); const outputSize = Math.ceil(ShapeUtil.size(outputShape) / components);
const enableInputShapesUniforms = enableShapesUniforms(inputs[0].dims.length); const programUniforms: ProgramUniform[] = [
const inputShapeOrRank = enableInputShapesUniforms ? inputs[0].dims.length : inputs[0].dims; {type: 'uint32', data: outputSize}, {type: 'int32', data: axisDimLimit}, {type: 'uint32', data: axis},
const enableIndicesShapesUniforms = enableShapesUniforms(inputs[1].dims.length); ...createTensorShapeVariables(inputs[0].dims), ...createTensorShapeVariables(inputs[1].dims),
const indicesShapeOrRank = enableIndicesShapesUniforms ? inputs[1].dims.length : inputs[1].dims; ...createTensorShapeVariables(outputShape)
const enableOutputShapesUniforms = enableShapesUniforms(outputShape.length); ];
const outputShapeOrRank = enableOutputShapesUniforms ? outputShape.length : outputShape;
const programUniforms: ProgramUniform[] =
[{type: 'uint32', data: outputSize}, {type: 'int32', data: axisDimLimit}, {type: 'uint32', data: axis}];
if (enableInputShapesUniforms) {
programUniforms.push(...createTensorShapeVariables(inputs[0].dims));
}
if (enableIndicesShapesUniforms) {
programUniforms.push(...createTensorShapeVariables(inputs[1].dims));
}
if (enableOutputShapesUniforms) {
programUniforms.push(...createTensorShapeVariables(outputShape));
}
const inputDependencies: ProgramInputTensorInfoDependency[] = [];
inputDependencies.push(enableInputShapesUniforms ? 'rank' : 'dims');
inputDependencies.push(enableIndicesShapesUniforms ? 'rank' : 'dims');
const getShaderSource = (shaderHelper: ShaderHelper) => { const getShaderSource = (shaderHelper: ShaderHelper) => {
const data = inputVariable('data', inputs[0].dataType, inputShapeOrRank, components); const data = inputVariable('data', inputs[0].dataType, inputs[0].dims.length, components);
const indices = inputVariable('inputIndices', inputs[1].dataType, indicesShapeOrRank); const indices = inputVariable('inputIndices', inputs[1].dataType, inputs[1].dims.length);
const output = outputVariable('output', inputs[0].dataType, outputShapeOrRank, components); const output = outputVariable('output', inputs[0].dataType, outputShape.length, components);
const calcDataIndices = (x: number|string): string => { const calcDataIndices = (x: number|string): string => {
const indicesRank = indicesShape.length; const indicesRank = indicesShape.length;
@ -127,7 +110,7 @@ const createGatherProgramInfo = (inputs: readonly TensorView[], attributes: Gath
}; };
return { return {
name: 'Gather', name: 'Gather',
shaderCache: {hint: attributes.cacheKey, inputDependencies}, shaderCache: {hint: attributes.cacheKey, inputDependencies: ['rank', 'rank']},
getRunData: () => ({ getRunData: () => ({
outputs: [ outputs: [
{dims: outputShape, dataType: inputs[0].dataType}, {dims: outputShape, dataType: inputs[0].dataType},

25
js/web/lib/wasm/jsep/webgpu/ops/transpose.ts
View file

@ -6,7 +6,7 @@ import {ShapeUtil} from '../../util';
import {AttributeWithCacheKey, createAttributeWithCacheKey} from '../attribute-with-cache-key'; import {AttributeWithCacheKey, createAttributeWithCacheKey} from '../attribute-with-cache-key';
import {ComputeContext, ProgramInfo} from '../types'; import {ComputeContext, ProgramInfo} from '../types';
import {createTensorShapeVariables, enableShapesUniforms, IndicesHelper, inputVariable, outputVariable, ShaderHelper} from './common'; import {createTensorShapeVariables, IndicesHelper, inputVariable, outputVariable, ShaderHelper} from './common';
export interface TransposeAttributes extends AttributeWithCacheKey { export interface TransposeAttributes extends AttributeWithCacheKey {
readonly perm: number[]; readonly perm: number[];
@ -39,12 +39,9 @@ export const createTransposeProgramInfo = (inputTensor: TensorView, permAttr: nu
const inputDataType = inputTensor.dataType; const inputDataType = inputTensor.dataType;
const inputRank = inputTensor.dims.length; const inputRank = inputTensor.dims.length;
const perm = getAdjustedPerm(inputRank, permAttr); const perm = getAdjustedPerm(inputRank, permAttr);
const useShapesUniforms = enableShapesUniforms(inputRank);
const outputShape = getOutputShape(inputTensor.dims, perm); const outputShape = getOutputShape(inputTensor.dims, perm);
const outShapeOrRank = useShapesUniforms ? outputShape.length : outputShape; const output = outputVariable('output', inputDataType, outputShape.length);
const inShapeOrRank = useShapesUniforms ? inputRank : inputTensor.dims; const input = inputVariable('a', inputDataType, inputRank);
const output = outputVariable('output', inputDataType, outShapeOrRank);
const input = inputVariable('a', inputDataType, inShapeOrRank);
const getShaderSource = (shaderHelper: ShaderHelper) => ` const getShaderSource = (shaderHelper: ShaderHelper) => `
${shaderHelper.registerUniform('output_size', 'u32').declareVariables(input, output)} ${shaderHelper.registerUniform('output_size', 'u32').declareVariables(input, output)}
@ -61,21 +58,17 @@ export const createTransposeProgramInfo = (inputTensor: TensorView, permAttr: nu
}`; }`;
return { return {
name: 'Transpose', name: 'Transpose',
shaderCache: {hint: `${permAttr}`, inputDependencies: useShapesUniforms ? ['rank'] : ['dims']}, shaderCache: {hint: `${permAttr}`, inputDependencies: ['rank']},
getRunData: (inputs) => { getRunData: (inputs) => {
const outputSize = ShapeUtil.size(outputShape); const outputSize = ShapeUtil.size(outputShape);
return { return {
outputs: [{dims: outputShape, dataType: inputs[0].dataType}], outputs: [{dims: outputShape, dataType: inputs[0].dataType}],
dispatchGroup: {x: Math.ceil(outputSize / 64 /* workgroup size */)}, dispatchGroup: {x: Math.ceil(outputSize / 64 /* workgroup size */)},
programUniforms: useShapesUniforms ? programUniforms: [
[ {type: 'uint32', data: outputSize},
{type: 'uint32', data: outputSize}, ...createTensorShapeVariables(inputs[0].dims),
...createTensorShapeVariables(inputs[0].dims), ...createTensorShapeVariables(outputShape),
...createTensorShapeVariables(outputShape), ],
] :
[
{type: 'uint32', data: outputSize},
],
}; };
}, },
getShaderSource, getShaderSource,