Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[js/webgpu] revise uniform support #17871

Merged
merged 3 commits into from
Oct 11, 2023
+77 −131
Merged

[js/webgpu] revise uniform support #17871

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
85d4891
following up for some comments
fs-eire Oct 10, 2023
b35ef8a
remove type cross check
fs-eire Oct 10, 2023
898f222
Merge remote-tracking branch 'origin/main' into fs-eire/uniform-followup
fs-eire Oct 11, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 4 additions & 8 deletions js/web/lib/wasm/jsep/backend-webgpu.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -259,10 +259,6 @@ export class WebGpuBackend {
run(program: ProgramInfo, inputTensorViews: readonly TensorView[], outputIndices: readonly number[],
createKernelOutput: (index: number, dataType: number, dims: readonly number[]) => TensorView,
createIntermediateOutput: (dataType: number, dims: readonly number[]) => TensorView): TensorView[] {
if (inputTensorViews.length !== program.inputTypes.length) {
throw new Error(`Input size must be equal to ${program.inputTypes.length}.`);
}

// create info for inputs
const inputDatas: GpuData[] = [];
for (let i = 0; i < inputTensorViews.length; ++i) {
Expand All @@ -277,7 +273,7 @@ export class WebGpuBackend {
const key = getProgramInfoUniqueKey(program, inputTensorViews);
let artifact = this.programManager.getArtifact(key);

const {outputs, dispatchGroup, variables} = program.getRunData(inputTensorViews);
const {outputs, dispatchGroup, programUniforms} = program.getRunData(inputTensorViews);

// check output indices
const validatedOutputIndices = outputIndices.length === 0 ? outputs.map((_, i) => i) : outputIndices;
Expand Down Expand Up @@ -328,12 +324,12 @@ export class WebGpuBackend {
// TODO: add cache for uniform (is it necessary?)
//
let uniformBufferBinding: GPUBindingResource|undefined;
if (variables) {
if (programUniforms) {
let currentOffset = 0;
let preLength = 0;
const offsets: number[] = [];
let maxAlignmentOfField = 1;
variables.forEach(v => {
programUniforms.forEach(v => {
const data = typeof v.data === 'number' ? [v.data] : v.data;
// https://www.w3.org/TR/WGSL/#alignof
let baseAlignment: number;
Expand Down Expand Up @@ -374,7 +370,7 @@ export class WebGpuBackend {

currentOffset = Math.ceil(currentOffset / maxAlignmentOfField) * maxAlignmentOfField;
const arrayBuffer = new ArrayBuffer(currentOffset);
variables.forEach((v, i) => {
programUniforms.forEach((v, i) => {
const offset = offsets[i];
const data = typeof v.data === 'number' ? [v.data] : v.data;
if (v.type === 'int32') {
Expand Down
6 changes: 2 additions & 4 deletions js/web/lib/wasm/jsep/webgpu/ops/3rd-party/conv2d_mm_webgpu.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@
import {LOG_DEBUG} from '../../../log';
import {TensorView} from '../../../tensor-view';
import {ShapeUtil} from '../../../util';
import {GpuDataType, ProgramInfo} from '../../types';
import {ProgramInfo} from '../../types';
import {tensorTypeToWsglStorageType} from '../common';
import {ConvAttributes} from '../conv';

Expand Down Expand Up @@ -213,11 +213,9 @@ export const createConv2DMatMulProgramInfo =

return {
name: 'Conv2DMatMul',
inputTypes: hasBias ? [GpuDataType.default, GpuDataType.default, GpuDataType.default] :
[GpuDataType.default, GpuDataType.default],
shaderCache: {hint: attributes.cacheKey},
getRunData: () => ({
outputs: [{dims: outputShape, dataType: inputs[0].dataType, gpuDataType: GpuDataType.default}],
outputs: [{dims: outputShape, dataType: inputs[0].dataType}],
dispatchGroup: {x: dispatch[0], y: dispatch[1], z: dispatch[2]},
}),
getShaderSource: () => `
Expand Down
6 changes: 2 additions & 4 deletions js/web/lib/wasm/jsep/webgpu/ops/3rd-party/conv_backprop_mm_webgpu.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@
import {LOG_DEBUG} from '../../../log';
import {TensorView} from '../../../tensor-view';
import {ShapeUtil} from '../../../util';
import {GpuDataType, ProgramInfo} from '../../types';
import {ProgramInfo} from '../../types';
import {ConvTransposeAttributes} from '../conv-transpose';

import {Activation, activationFnSnippet, biasActivationSnippet, typeSnippet} from './activation_util';
Expand Down Expand Up @@ -200,11 +200,9 @@ export const createConv2DTransposeMatMulProgramInfo =
}
return {
name: 'Conv2DTransposeMatMul',
inputTypes: hasBias ? [GpuDataType.default, GpuDataType.default, GpuDataType.default] :
[GpuDataType.default, GpuDataType.default],
shaderCache: {hint: attributes.cacheKey},
getRunData: () => ({
outputs: [{dims: outputShape, dataType: inputs[0].dataType, gpuDataType: GpuDataType.default}],
outputs: [{dims: outputShape, dataType: inputs[0].dataType}],
dispatchGroup: {x: dispatch[0], y: dispatch[1], z: dispatch[2]}
}),
getShaderSource: () => `
Expand Down
7 changes: 2 additions & 5 deletions js/web/lib/wasm/jsep/webgpu/ops/3rd-party/conv_backprop_webgpu.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@
import {LOG_DEBUG} from '../../../log';
import {TensorView} from '../../../tensor-view';
import {ShapeUtil} from '../../../util';
import {GpuDataType, ProgramInfo} from '../../types';
import {ProgramInfo} from '../../types';
import {inputVariable, outputVariable, ShaderHelper, tensorTypeToWsglStorageType} from '../common';
import {ConvTransposeAttributes} from '../conv-transpose';

Expand Down Expand Up @@ -260,15 +260,12 @@ export const createConvTranspose2DProgramInfo =
const dataType = tensorTypeToWsglStorageType(inputs[0].dataType);
return {
name: 'ConvTranspose2D',
inputTypes: hasBias ? [GpuDataType.default, GpuDataType.default, GpuDataType.default] :
[GpuDataType.default, GpuDataType.default],
shaderCache: {hint: attributes.cacheKey},
getRunData: () => ({
dispatchGroup: {x: dispatch[0], y: dispatch[1], z: dispatch[2]},
outputs: [{
dims: squeezeOutputShapeFunction ? squeezeOutputShapeFunction(outputShape) : outputShape,
dataType: inputs[0].dataType,
gpuDataType: GpuDataType.default
dataType: inputs[0].dataType
}]
}),
getShaderSource: (shaderHelper: ShaderHelper) => createConvTranspose2DOpProgramShaderSource(
Expand Down
6 changes: 2 additions & 4 deletions js/web/lib/wasm/jsep/webgpu/ops/3rd-party/matmul_packed_webgpu.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@

import {TensorView} from '../../../tensor-view';
import {ShapeUtil} from '../../../util';
import {GpuDataType, ProgramInfo} from '../../types';
import {ProgramInfo} from '../../types';
import {getBroadcastDims, IndicesHelper, inputVariable, outputVariable, ShaderHelper, tensorTypeToWsglStorageType} from '../common';
import {getActicationSnippet, InternalActivationAttributes} from '../fuse-utils';

Expand Down Expand Up @@ -481,11 +481,9 @@ export const createMatmulProgramInfo =
${batchDims.impl()}`;
return {
name: 'MatMul',
inputTypes: hasBias ? [GpuDataType.default, GpuDataType.default, GpuDataType.default] :
[GpuDataType.default, GpuDataType.default],
shaderCache: {hint: activationAttributes.activationCacheKey},
getRunData: () => ({
outputs: [{dims: outputShape, dataType: inputs[0].dataType, gpuDataType: GpuDataType.default}],
outputs: [{dims: outputShape, dataType: inputs[0].dataType}],
dispatchGroup: {x: dispatch[0], y: dispatch[1], z: dispatch[2]}
}),
getShaderSource,
Expand Down
5 changes: 2 additions & 3 deletions js/web/lib/wasm/jsep/webgpu/ops/bias-add.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@

import {TensorView} from '../../tensor-view';
import {ShapeUtil} from '../../util';
import {ComputeContext, GpuDataType, ProgramInfo} from '../types';
import {ComputeContext, ProgramInfo} from '../types';

import {inputVariable, outputVariable, ShaderHelper} from './common';

Expand Down Expand Up @@ -51,9 +51,8 @@ const createBiasAddProgramInfo = (inputs: readonly TensorView[]): ProgramInfo =>

return {
name: 'BiasAdd',
inputTypes: Array(inputs.length).fill(GpuDataType.default),
getRunData: () => ({
outputs: [{dims: outputShape, dataType: inputs[0].dataType, gpuDataType: GpuDataType.default}],
outputs: [{dims: outputShape, dataType: inputs[0].dataType}],
dispatchGroup: {x: Math.ceil(outputSize / 64 /* workgroup size */)}
}),
getShaderSource,
Expand Down
5 changes: 2 additions & 3 deletions js/web/lib/wasm/jsep/webgpu/ops/bias-split-gelu.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@

import {TensorView} from '../../tensor-view';
import {ShapeUtil} from '../../util';
import {ComputeContext, GpuDataType, ProgramInfo} from '../types';
import {ComputeContext, ProgramInfo} from '../types';

import {inputVariable, outputVariable, ShaderHelper} from './common';
import {erfImpl} from './unary-op';
Expand Down Expand Up @@ -58,9 +58,8 @@ const createBiasSplitGeluProgramInfo = (inputs: readonly TensorView[]): ProgramI

return {
name: 'BiasSplitGelu',
inputTypes: [GpuDataType.default, GpuDataType.default],
getRunData: () => ({
outputs: [{dims: outputShape, dataType: inputs[0].dataType, gpuDataType: GpuDataType.default}],
outputs: [{dims: outputShape, dataType: inputs[0].dataType}],
dispatchGroup: {x: Math.ceil(outputSize / 64 /* workgroup size */)}
}),
getShaderSource,
Expand Down
5 changes: 2 additions & 3 deletions js/web/lib/wasm/jsep/webgpu/ops/binary-op.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
import {DataType} from '../../../wasm-common';
import {TensorView} from '../../tensor-view';
import {BroadcastUtil, ShapeUtil} from '../../util';
import {ComputeContext, GpuDataType, ProgramInfo} from '../types';
import {ComputeContext, ProgramInfo} from '../types';

import {inputVariable, outputVariable, ShaderHelper} from './common';

Expand Down Expand Up @@ -175,13 +175,12 @@ const createBinaryOpProgramInfo =

return {
name,
inputTypes: [GpuDataType.default, GpuDataType.default],
shaderCache: {hint: cacheKey},
getShaderSource: (shaderHelper) => createBinaryOpProgramShader(
shaderHelper, a.dims, b.dims, outputShape, vectorize, isBroadcast, funcCall, a.dataType, b.dataType,
outputDataType, additionalImplementation),
getRunData: () => ({
outputs: [{dims: outputShape, dataType: outputDataType, gpuDataType: GpuDataType.default}],
outputs: [{dims: outputShape, dataType: outputDataType}],
dispatchGroup: {x: Math.ceil(outputSize / 64 /* workgroup size */ / 4 /* component size */)}
}),
};
Expand Down
4 changes: 2 additions & 2 deletions js/web/lib/wasm/jsep/webgpu/ops/common.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@

import {DataType} from '../../../wasm-common';
import {ShapeUtil} from '../../util';
import {ProgramVariable} from '../types';
import {ProgramUniform} from '../types';

/**
* constant value for a workgroup size.
Expand Down Expand Up @@ -259,7 +259,7 @@ export const tensorTypeToWsglValueType = (type: DataType, components: 1|2|3|4 =
};

export const createTensorShapeVariables = (dims: readonly number[]):
ProgramVariable[] => [{type: 'uint32', data: dims}, {type: 'uint32', data: ShapeUtil.computeStrides(dims)}];
ProgramUniform[] => [{type: 'uint32', data: dims}, {type: 'uint32', data: ShapeUtil.computeStrides(dims)}];

/**
* A helper function to get a IndicesHelper for a given input or output.
Expand Down
5 changes: 2 additions & 3 deletions js/web/lib/wasm/jsep/webgpu/ops/concat.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
import {TensorView} from '../../tensor-view';
import {ShapeUtil} from '../../util';
import {AttributeWithCacheKey, createAttributeWithCacheKey} from '../attribute-with-cache-key';
import {ComputeContext, GpuDataType, ProgramInfo} from '../types';
import {ComputeContext, ProgramInfo} from '../types';

import {IndicesHelper, inputVariable, outputVariable, ShaderHelper} from './common';

Expand Down Expand Up @@ -122,10 +122,9 @@ const createConcatProgramInfo = (inputs: readonly TensorView[], axis: number): P
}`;
return {
name: 'Concat',
inputTypes: Array(inputs.length).fill(GpuDataType.default),
shaderCache: {hint: `${axis}`},
getRunData: () => ({
outputs: [{dims: outputShape, dataType: inputs[0].dataType, gpuDataType: GpuDataType.default}],
outputs: [{dims: outputShape, dataType: inputs[0].dataType}],
dispatchGroup: {x: Math.ceil(outputSize / 64 /* workgroup size */)}
}),
getShaderSource,
Expand Down
7 changes: 2 additions & 5 deletions js/web/lib/wasm/jsep/webgpu/ops/conv-grouped.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@

import {TensorView} from '../../tensor-view';
import {ShapeUtil} from '../../util';
import {GpuDataType, ProgramInfo} from '../types';
import {ProgramInfo} from '../types';

import {inputVariable, outputVariable, ShaderHelper} from './common';
import {calculateOutputShape, ConvAttributes} from './conv';
Expand Down Expand Up @@ -85,14 +85,11 @@ export const createGroupedConvProgramInfo =
}`;
return {
name: 'GroupedConv',
inputTypes: hasBias ? [GpuDataType.default, GpuDataType.default, GpuDataType.default] :
[GpuDataType.default, GpuDataType.default],
shaderCache: {hint: attributes.cacheKey},
getRunData: () => ({
outputs: [{
dims: squeezeOutputShapeFunction ? squeezeOutputShapeFunction(outputShape) : outputShape,
dataType: inputs[0].dataType,
gpuDataType: GpuDataType.default
dataType: inputs[0].dataType
}],
dispatchGroup: {x: Math.ceil(outputSize / 64 /* workgroup size */)},
}),
Expand Down
5 changes: 2 additions & 3 deletions js/web/lib/wasm/jsep/webgpu/ops/einsum.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
import {TensorView} from '../../tensor-view';
import {ShapeUtil} from '../../util';
import {AttributeWithCacheKey, createAttributeWithCacheKey} from '../attribute-with-cache-key';
import {ComputeContext, GpuDataType, ProgramInfo} from '../types';
import {ComputeContext, ProgramInfo} from '../types';

import {IndicesHelper, inputVariable, outputVariable, ShaderHelper} from './common';

Expand Down Expand Up @@ -260,10 +260,9 @@ const createEinsumProgramInfo = (inputs: readonly TensorView[], einsumEquation:
}`;
return {
name: 'Einsum',
inputTypes: Array(inputs.length).fill(GpuDataType.default),
shaderCache: {hint: einsumEquation.equation},
getRunData: () => ({
outputs: [{dims: outputShape, dataType: inputs[0].dataType, gpuDataType: GpuDataType.default}],
outputs: [{dims: outputShape, dataType: inputs[0].dataType}],
dispatchGroup: {x: Math.ceil(outputSize / 64 /* workgroup size */)}
}),
getShaderSource,
Expand Down
5 changes: 2 additions & 3 deletions js/web/lib/wasm/jsep/webgpu/ops/expand.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@

import {TensorView} from '../../tensor-view';
import {ShapeUtil} from '../../util';
import {ComputeContext, GpuDataType, ProgramInfo} from '../types';
import {ComputeContext, ProgramInfo} from '../types';

import {inputVariable, outputVariable, ShaderHelper} from './common';

Expand Down Expand Up @@ -70,11 +70,10 @@ const createExpandProgramInfo = (inputs: readonly TensorView[]): ProgramInfo =>
}`;
return {
name: 'Expand',
inputTypes: [GpuDataType.default],
shaderCache: {hint: `${outputShape}`},
getShaderSource,
getRunData: () => ({
outputs: [{dims: outputShape, dataType: inputs[0].dataType, gpuDataType: GpuDataType.default}],
outputs: [{dims: outputShape, dataType: inputs[0].dataType}],
dispatchGroup: {x: Math.ceil(outputSize / 64 /* workgroup size */)}
})
};
Expand Down
5 changes: 2 additions & 3 deletions js/web/lib/wasm/jsep/webgpu/ops/gather-elements.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
import {TensorView} from '../../tensor-view';
import {ShapeUtil} from '../../util';
import {AttributeWithCacheKey, createAttributeWithCacheKey} from '../attribute-with-cache-key';
import {ComputeContext, GpuDataType, ProgramInfo} from '../types';
import {ComputeContext, ProgramInfo} from '../types';

import {inputVariable, outputVariable, ShaderHelper} from './common';

Expand Down Expand Up @@ -87,10 +87,9 @@ const createGatherElementsProgramInfo =

return {
name: 'GatherElements',
inputTypes: [GpuDataType.default, GpuDataType.default],
shaderCache: {hint: attributes.cacheKey},
getRunData: () => ({
outputs: [{dims: outputShape, dataType: inputs[0].dataType, gpuDataType: GpuDataType.default}],
outputs: [{dims: outputShape, dataType: inputs[0].dataType}],
dispatchGroup: {x: Math.ceil(outputSize / 64 /* workgroup size */)}
}),
getShaderSource,
Expand Down
5 changes: 2 additions & 3 deletions js/web/lib/wasm/jsep/webgpu/ops/gather.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
import {TensorView} from '../../tensor-view';
import {ShapeUtil} from '../../util';
import {AttributeWithCacheKey, createAttributeWithCacheKey} from '../attribute-with-cache-key';
import {ComputeContext, GpuDataType, ProgramInfo} from '../types';
import {ComputeContext, ProgramInfo} from '../types';

import {inputVariable, outputVariable, ShaderHelper} from './common';

Expand Down Expand Up @@ -72,11 +72,10 @@ const createGatherProgramInfo = (inputs: readonly TensorView[], attributes: Gath
}`;
return {
name: 'Gather',
inputTypes: [GpuDataType.default, GpuDataType.default],
shaderCache: {hint: attributes.cacheKey},
getRunData: () => ({
outputs: [
{dims: outputShape, dataType: inputs[0].dataType, gpuDataType: GpuDataType.default},
{dims: outputShape, dataType: inputs[0].dataType},
],
dispatchGroup: {x: Math.ceil(outputSize / 64 /* workgroup size */)}
}),
Expand Down
6 changes: 2 additions & 4 deletions js/web/lib/wasm/jsep/webgpu/ops/gemm.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
import {TensorView} from '../../tensor-view';
import {GemmUtil, ShapeUtil} from '../../util';
import {AttributeWithCacheKey, createAttributeWithCacheKey} from '../attribute-with-cache-key';
import {ComputeContext, GpuDataType, ProgramInfo} from '../types';
import {ComputeContext, ProgramInfo} from '../types';

import {ShaderHelper, tensorTypeToWsglStorageType} from './common';

Expand Down Expand Up @@ -112,11 +112,9 @@ const createGemmProgramInfo = (inputs: readonly TensorView[], attributes: GemmAt
}`;
return {
name: 'Gemm',
inputTypes: inputs.length === 3 ? [GpuDataType.default, GpuDataType.default, GpuDataType.default] :
[GpuDataType.default, GpuDataType.default],
shaderCache: {hint: attributes.cacheKey},
getRunData: () => ({
outputs: [{dims: outputShape, dataType: inputs[0].dataType, gpuDataType: GpuDataType.default}],
outputs: [{dims: outputShape, dataType: inputs[0].dataType}],
dispatchGroup: {x: Math.ceil(outputSize / 64 /* workgroup size */)}
}),
getShaderSource,
Expand Down
Loading
Loading