mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-29 20:14:01 +00:00
[js/web] FP16 LayerNorm, InstanceNorm, SkipLayerNorm (#17630)
### Description This PR includes fixes for Norm operations to support FP16 and also some optimizations to use vec2/vec4 if possible
This commit is contained in:
Arthur Islamov
GitHub
parent
f9694c5b97
commit
22947109f2
9 changed files with 341 additions and 187 deletions
|
|
@ -261,6 +261,69 @@ export const tensorTypeToWsglValueType = (type: DataType, components: 1|2|3|4 =
|
|||
export const createTensorShapeVariables = (dims: readonly number[]):
|
||||
ProgramUniform[] => [{type: 'uint32', data: dims}, {type: 'uint32', data: ShapeUtil.computeStrides(dims)}];
|
||||
|
||||
/**
|
||||
* A helper function to get maximum vector size for specified data length
|
||||
* @param size
|
||||
*/
|
||||
export const getMaxComponents = (size: number) => {
|
||||
// we cannot use vec3 type since it has alignment of 16 bytes
|
||||
if (size % 4 === 0) {
|
||||
return 4;
|
||||
} else if (size % 2 === 0) {
|
||||
return 2;
|
||||
}
|
||||
|
||||
return 1;
|
||||
};
|
||||
|
||||
/**
|
||||
* A helper function that initializes variable as a scalar or vector. e.g. f32(0) or vec4f(0,0,0,0)
|
||||
* @param dataType
|
||||
* @param components
|
||||
* @param value
|
||||
*/
|
||||
export const fillVector = (dataType = 'f32', components?: number, value = '0') => {
|
||||
if (!components || components === 1) {
|
||||
return `${dataType}(${value})`;
|
||||
}
|
||||
|
||||
return `vec${components}<${dataType}>(${value})`;
|
||||
};
|
||||
|
||||
/**
|
||||
* A helper function that casts value or vector to f32
|
||||
* @param dataType
|
||||
* @param components
|
||||
* @param value
|
||||
*/
|
||||
export const castToF32 = (dataType: string, components: number, value: string) => {
|
||||
if (dataType === 'f32') {
|
||||
return value;
|
||||
}
|
||||
if (components === 1) {
|
||||
return `f32(${value})`;
|
||||
}
|
||||
|
||||
return `vec${components}f(${value})`;
|
||||
};
|
||||
|
||||
/**
|
||||
* A helper function that returns scalar or sums all components of a vector
|
||||
* @param name
|
||||
* @param components
|
||||
*/
|
||||
export const sumVector = (name: string, components: number) => {
|
||||
if (components === 4) {
|
||||
return `(${name}.x + ${name}.y + ${name}.z + ${name}.w)`;
|
||||
} else if (components === 2) {
|
||||
return `(${name}.x + ${name}.y)`;
|
||||
} else if (components === 3) {
|
||||
return `(${name}.x + ${name}.y + ${name}.z)`;
|
||||
}
|
||||
|
||||
return name;
|
||||
};
|
||||
|
||||
/**
|
||||
* A helper function to get a IndicesHelper for a given input or output.
|
||||
*
|
||||
|
|
|
|||
|
|
@ -1,12 +1,13 @@
|
|||
// Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
// Licensed under the MIT License.
|
||||
|
||||
import {DataType} from '../../../wasm-common';
|
||||
import {TensorView} from '../../tensor-view';
|
||||
import {ShapeUtil} from '../../util';
|
||||
import {AttributeWithCacheKey, createAttributeWithCacheKey} from '../attribute-with-cache-key';
|
||||
import {ComputeContext, ProgramInfo} from '../types';
|
||||
|
||||
import {inputVariable, outputVariable, ShaderHelper, tensorTypeToWsglStorageType} from './common';
|
||||
import {fillVector, getMaxComponents, inputVariable, outputVariable, ShaderHelper, tensorTypeToWsglStorageType} from './common';
|
||||
|
||||
export interface InstanceNormAttributes extends AttributeWithCacheKey {
|
||||
epsilon: number;
|
||||
|
|
@ -111,69 +112,159 @@ const createInstanceNormProgramInfo =
|
|||
};
|
||||
};
|
||||
|
||||
const computeMean =
|
||||
(context: ComputeContext, input: TensorView, scale: TensorView, bias: TensorView, n: number, h: number, c: number,
|
||||
epsilon: number) => {
|
||||
const components = getMaxComponents(c);
|
||||
const inputHelper = inputVariable('input', input.dataType, input.dims, components);
|
||||
const scaleHelper = inputVariable('scale', scale.dataType, scale.dims, components);
|
||||
const biasHelper = inputVariable('bias', bias.dataType, bias.dims, components);
|
||||
|
||||
const WG = 64;
|
||||
// we will store channel scale and channel shift in [2, components] matrix
|
||||
// or in vec2 when components == 1
|
||||
const outputType = components === 1 ? 'vec2f' : `mat2x${components}f`;
|
||||
const sumCastType = components === 1 ? 'f32' : `vec${components}f`;
|
||||
const setOutputValue = (var1: string, var2: string) => `${outputType}(${var1}, ${var2})`;
|
||||
const unitsOfWork = n * c / components;
|
||||
const wgSize = Math.ceil(h / WG);
|
||||
|
||||
const getMeanShaderSource = (shaderHelper: ShaderHelper) => `
|
||||
const H: u32 = ${h};
|
||||
const C: u32 = ${c / components};
|
||||
const imageSize: u32 = ${h * c / components};
|
||||
|
||||
${shaderHelper.declareVariables(inputHelper)}
|
||||
@group(0) @binding(1) var<storage, read_write> output : array<${outputType}>;
|
||||
|
||||
${shaderHelper.mainStart(WG)}
|
||||
let currentImageNumber = global_idx / ${WG} / C;
|
||||
let currentChannelNumber = (global_idx / ${WG}) % C;
|
||||
let wgId = global_idx % ${WG};
|
||||
let wgOffset = wgId * ${wgSize};
|
||||
if (wgOffset >= H) {
|
||||
return;
|
||||
}
|
||||
let wgMax = min(wgOffset + ${wgSize}, H);
|
||||
|
||||
let offset = currentImageNumber * imageSize + currentChannelNumber;
|
||||
var sum = ${fillVector('f32', components)};
|
||||
var squaredSum = ${fillVector('f32', components)};
|
||||
for (var i: u32 = wgOffset; i < wgMax; i++) {
|
||||
let value = ${sumCastType}(input[offset + i * C]);
|
||||
sum += value;
|
||||
squaredSum += value * value;
|
||||
}
|
||||
output[global_idx] = ${setOutputValue('sum', 'squaredSum')};
|
||||
}`;
|
||||
|
||||
const meanValues = context.compute(
|
||||
{
|
||||
name: 'InstanceNormComputeMean',
|
||||
shaderCache: {hint: JSON.stringify({components, n, h, c})},
|
||||
getRunData: () => ({
|
||||
outputs: [
|
||||
{dims: [n, c, WG, 2], dataType: DataType.float},
|
||||
],
|
||||
dispatchGroup: {x: n * c / components},
|
||||
}),
|
||||
getShaderSource: getMeanShaderSource,
|
||||
},
|
||||
{inputs: [input], outputs: [-1]})[0];
|
||||
const getShaderSource = (shaderHelper: ShaderHelper) => `
|
||||
const H: u32 = ${h};
|
||||
const C: u32 = ${c / components};
|
||||
const imageSize: u32 = ${WG * c / components};
|
||||
const epsilon: f32 = ${epsilon};
|
||||
|
||||
@group(0) @binding(0) var<storage, read> input : array<${outputType}>;
|
||||
@group(0) @binding(1) var<storage, read> scale : array<${scaleHelper.type.storage}>;
|
||||
@group(0) @binding(2) var<storage, read> bias : array<${biasHelper.type.storage}>;
|
||||
@group(0) @binding(3) var<storage, read_write> output : array<${outputType}>;
|
||||
|
||||
${shaderHelper.mainStart()}
|
||||
${shaderHelper.guardAgainstOutOfBoundsWorkgroupSizes(unitsOfWork)}
|
||||
let currentImageNumber = global_idx / C;
|
||||
let currentChannelNumber = global_idx % C;
|
||||
|
||||
let offset = currentImageNumber * imageSize;
|
||||
var sum = ${fillVector('f32', components)};
|
||||
var squaredSum = ${fillVector('f32', components)};
|
||||
for (var i: u32 = 0; i < ${WG}; i++) {
|
||||
let value = input[offset + i + currentChannelNumber * ${WG}];
|
||||
sum += value[0];
|
||||
squaredSum += value[1];
|
||||
}
|
||||
sum = sum / f32(H);
|
||||
squaredSum = squaredSum / f32(H);
|
||||
let invStdDev = 1 / sqrt(squaredSum - sum * sum + epsilon);
|
||||
let channelScale = invStdDev * ${sumCastType}(scale[currentChannelNumber]);
|
||||
let channelShift = ${sumCastType}(bias[currentChannelNumber]) - sum * channelScale;
|
||||
|
||||
output[global_idx] = ${setOutputValue('channelScale', 'channelShift')};
|
||||
}`;
|
||||
|
||||
return context.compute(
|
||||
{
|
||||
name: 'InstanceNormComputeChannelScaleShift',
|
||||
shaderCache: {hint: JSON.stringify({components, n, h, c, epsilon})},
|
||||
getRunData: () => ({
|
||||
outputs: [
|
||||
{dims: [n, c, 2], dataType: DataType.float},
|
||||
],
|
||||
dispatchGroup: {x: Math.ceil(unitsOfWork / 64 /* workgroup size */)},
|
||||
}),
|
||||
getShaderSource,
|
||||
},
|
||||
{inputs: [meanValues, scale, bias], outputs: [-1]})[0];
|
||||
};
|
||||
|
||||
const createInstanceNormNHWCProgramInfo =
|
||||
(inputs: readonly TensorView[], attributes: InstanceNormAttributes): ProgramInfo => {
|
||||
(context: ComputeContext, inputs: readonly TensorView[], attributes: InstanceNormAttributes) => {
|
||||
const xShape = inputs[0].dims;
|
||||
const outputShape = xShape;
|
||||
const outputSize = ShapeUtil.size(outputShape);
|
||||
const N = xShape[0];
|
||||
const C = xShape[xShape.length - 1];
|
||||
const H = ShapeUtil.sizeFromDimension(xShape, 1) / C;
|
||||
|
||||
const components = getMaxComponents(C);
|
||||
const outputSize = ShapeUtil.size(outputShape) / components;
|
||||
const inputHelper = inputVariable('input', inputs[0].dataType, inputs[0].dims, components);
|
||||
const outputHelper = outputVariable('output', inputs[0].dataType, outputShape, components);
|
||||
|
||||
const dataType = tensorTypeToWsglStorageType(inputs[0].dataType);
|
||||
const scaleType = components === 1 ? 'vec2f' : `mat2x${components}f`;
|
||||
const scaleCastType = components === 1 ? dataType : `vec${components}<${dataType}>`;
|
||||
// first compute mean
|
||||
const channelScaleShift = computeMean(context, inputs[0], inputs[1], inputs[2], N, H, C, attributes.epsilon);
|
||||
|
||||
const normCount = C * N;
|
||||
const getShaderSource = (shaderHelper: ShaderHelper) => `
|
||||
const N: u32 = ${N};
|
||||
const H: u32 = ${H};
|
||||
const C: u32 = ${C};
|
||||
const normSizeTyped: ${dataType} = ${H};
|
||||
const imageSize: u32 = ${H * C};
|
||||
const epsilon: f32 = ${attributes.epsilon};
|
||||
const C: u32 = ${C / components};
|
||||
|
||||
@group(0) @binding(0) var<storage, read> x : array<${dataType}>;
|
||||
@group(0) @binding(1) var<storage, read> scale : array<${dataType}>;
|
||||
@group(0) @binding(2) var<storage, read> bias : array<${dataType}>;
|
||||
@group(0) @binding(3) var<storage, read_write> output : array<${dataType}>;
|
||||
@group(0) @binding(0) var<storage, read> input : array<${inputHelper.type.storage}>;
|
||||
@group(0) @binding(1) var<storage, read> scaleInput : array<${scaleType}>;
|
||||
@group(0) @binding(2) var<storage, read_write> output : array<${outputHelper.type.storage}>;
|
||||
|
||||
${shaderHelper.mainStart()}
|
||||
let currentImageNumber = global_idx / C;
|
||||
let currentImageNumber = global_idx / (C * H);
|
||||
let currentChannelNumber = global_idx % C;
|
||||
|
||||
// offset is channel num * N
|
||||
let offset = currentImageNumber * imageSize;
|
||||
if (offset >= ${outputSize}) { return; }
|
||||
var mean: ${dataType} = 0;
|
||||
|
||||
for (var i: u32 = 0u; i < H; i++) {
|
||||
mean = mean + x[offset + i * C + currentChannelNumber];
|
||||
}
|
||||
mean = mean / normSizeTyped;
|
||||
|
||||
var squaredNorm: ${dataType} = 0;
|
||||
for (var i: u32 = 0u; i < H; i++) {
|
||||
let deviation: f32 = x[offset + i * C + currentChannelNumber] - mean;
|
||||
squaredNorm = squaredNorm + deviation * deviation;
|
||||
}
|
||||
let invStdDev = 1 / sqrt(squaredNorm / normSizeTyped + epsilon);
|
||||
let channelScale = invStdDev * scale[currentChannelNumber];
|
||||
let channelShift = bias[currentChannelNumber] - mean * channelScale;
|
||||
for (var i: u32 = 0u; i < H; i++) {
|
||||
let currentOffset = offset + i * C + currentChannelNumber;
|
||||
output[currentOffset] = x[currentOffset] * channelScale + channelShift;
|
||||
}
|
||||
let scaleOffset = currentImageNumber * C + currentChannelNumber;
|
||||
let scale = scaleInput[scaleOffset];
|
||||
output[global_idx] = fma(input[global_idx], ${scaleCastType}(scale[0]), ${scaleCastType}(scale[1]));
|
||||
}`;
|
||||
return {
|
||||
...metadata,
|
||||
shaderCache: {hint: attributes.cacheKey},
|
||||
getRunData: () => ({
|
||||
outputs: [
|
||||
{dims: outputShape, dataType: inputs[0].dataType},
|
||||
],
|
||||
dispatchGroup: {x: Math.ceil(normCount / 64 /* workgroup size */)}
|
||||
}),
|
||||
getShaderSource,
|
||||
};
|
||||
context.compute(
|
||||
{
|
||||
name: 'InstanceNormalization',
|
||||
shaderCache: {hint: `${attributes.cacheKey}`},
|
||||
getRunData: () => ({
|
||||
outputs: [{dims: outputShape, dataType: inputs[0].dataType}],
|
||||
dispatchGroup: {x: Math.ceil(outputSize / 64 /* workgroup size */)}
|
||||
}),
|
||||
getShaderSource,
|
||||
},
|
||||
{inputs: [inputs[0], channelScaleShift]});
|
||||
};
|
||||
|
||||
export const parseInstanceNormAttributes = (attributes: InstanceNormAttributes): InstanceNormAttributes =>
|
||||
|
|
@ -181,7 +272,7 @@ export const parseInstanceNormAttributes = (attributes: InstanceNormAttributes):
|
|||
|
||||
export const instanceNorm = (context: ComputeContext, attributes: InstanceNormAttributes): void => {
|
||||
if (attributes.format === 'NHWC') {
|
||||
context.compute(createInstanceNormNHWCProgramInfo(context.inputs, attributes));
|
||||
createInstanceNormNHWCProgramInfo(context, context.inputs, attributes);
|
||||
} else {
|
||||
context.compute(createInstanceNormProgramInfo(context.inputs, attributes));
|
||||
}
|
||||
|
|
|
|||
|
|
@ -7,7 +7,7 @@ import {ShapeUtil} from '../../util';
|
|||
import {AttributeWithCacheKey, createAttributeWithCacheKey} from '../attribute-with-cache-key';
|
||||
import {ComputeContext, ProgramInfo} from '../types';
|
||||
|
||||
import {ShaderHelper, tensorTypeToWsglStorageType} from './common';
|
||||
import {castToF32, fillVector, getMaxComponents, inputVariable, outputVariable, ShaderHelper, sumVector, tensorTypeToWsglStorageType,} from './common';
|
||||
|
||||
export interface LayerNormAttributes extends AttributeWithCacheKey {
|
||||
axis: number;
|
||||
|
|
@ -18,10 +18,6 @@ const validateInputs = (inputs: readonly TensorView[]): void => {
|
|||
if (!inputs || inputs.length < 2) {
|
||||
throw new Error('layerNorm requires at least 2 inputs.');
|
||||
}
|
||||
|
||||
if (inputs[0].dataType !== DataType.float || inputs[1].dataType !== DataType.float) {
|
||||
throw new Error('inputs should be float type');
|
||||
}
|
||||
};
|
||||
|
||||
const createLayerNormProgramInfo =
|
||||
|
|
@ -31,7 +27,6 @@ const createLayerNormProgramInfo =
|
|||
const bias = inputs[2];
|
||||
|
||||
const outputShape = xShape;
|
||||
const outputSize = ShapeUtil.size(outputShape);
|
||||
const axis = ShapeUtil.normalizeAxis(attributes.axis, xShape.length);
|
||||
const normCount = ShapeUtil.sizeToDimension(xShape, axis);
|
||||
const normSize = ShapeUtil.sizeFromDimension(xShape, axis);
|
||||
|
|
@ -53,44 +48,54 @@ const createLayerNormProgramInfo =
|
|||
}
|
||||
}
|
||||
|
||||
const components = getMaxComponents(normSize);
|
||||
const dataType = tensorTypeToWsglStorageType(inputs[0].dataType);
|
||||
const variables = [
|
||||
inputVariable('x', inputs[0].dataType, inputs[0].dims, components),
|
||||
inputVariable('scale', scale.dataType, scale.dims, components),
|
||||
];
|
||||
if (bias) {
|
||||
variables.push(inputVariable('bias', bias.dataType, bias.dims, components));
|
||||
}
|
||||
variables.push(outputVariable('output', inputs[0].dataType, outputShape, components));
|
||||
|
||||
const hasMeanDataOutput = outputCount > 1;
|
||||
const hasInvStdOutput = outputCount > 2;
|
||||
let bindingIndex = 0;
|
||||
|
||||
if (hasMeanDataOutput) {
|
||||
variables.push(outputVariable('meanDataOutput', DataType.float, meanInvStdDevDim));
|
||||
}
|
||||
if (hasInvStdOutput) {
|
||||
variables.push(outputVariable('invStdOutput', DataType.float, meanInvStdDevDim));
|
||||
}
|
||||
|
||||
const getShaderSource = (shaderHelper: ShaderHelper) => `
|
||||
const normSize: u32 = ${normSize};
|
||||
const normSizeTyped: ${dataType} = ${normSize};
|
||||
const normSize: f32 = ${normSize};
|
||||
const normSizeVectorized: u32 = ${normSize / components};
|
||||
const epsilon: f32 = ${attributes.epsilon};
|
||||
|
||||
@group(0) @binding(${bindingIndex++}) var<storage, read> x : array<${dataType}>;
|
||||
@group(0) @binding(${bindingIndex++}) var<storage, read> scale : array<${dataType}>;
|
||||
${bias ? `@group(0) @binding(${bindingIndex++}) var<storage, read> bias : array<${dataType}>;` : ''}
|
||||
@group(0) @binding(${bindingIndex++}) var<storage, read_write> output : array<${dataType}>;
|
||||
${
|
||||
hasMeanDataOutput ?
|
||||
`@group(0) @binding(${bindingIndex++}) var<storage, read_write> meanDataOutput : array<${dataType}>` :
|
||||
''};
|
||||
${
|
||||
hasInvStdOutput ?
|
||||
`@group(0) @binding(${bindingIndex++}) var<storage, read_write> invStdOutput : array<${dataType}>` :
|
||||
''};
|
||||
|
||||
${shaderHelper.declareVariables(...variables)}
|
||||
${shaderHelper.mainStart()}
|
||||
let offset = global_idx * normSize;
|
||||
if (offset >= ${outputSize}) { return; }
|
||||
var mean: ${dataType} = 0;
|
||||
var meanSquare: ${dataType} = 0;
|
||||
${shaderHelper.guardAgainstOutOfBoundsWorkgroupSizes(normCount)}
|
||||
let offset = global_idx * normSizeVectorized;
|
||||
var meanVector = ${fillVector('f32', components)};
|
||||
var meanSquareVector = ${fillVector('f32', components)};
|
||||
|
||||
for (var h: u32 = 0u; h < normSize; h++) {
|
||||
mean = mean + x[h + offset];
|
||||
meanSquare = meanSquare + x[h + offset] * x[h + offset];
|
||||
for (var h: u32 = 0u; h < normSizeVectorized; h++) {
|
||||
let value = ${castToF32(dataType, components, 'x[h + offset]')};
|
||||
meanVector += value;
|
||||
meanSquareVector += value * value;
|
||||
}
|
||||
mean = mean / normSizeTyped;
|
||||
meanSquare = sqrt(meanSquare / normSizeTyped - mean * mean + epsilon);
|
||||
let mean = ${sumVector('meanVector', components)} / normSize;
|
||||
let meanSquare = sqrt(${sumVector('meanSquareVector', components)}
|
||||
/ normSize - mean * mean + epsilon);
|
||||
|
||||
for (var j: u32 = 0; j < normSize; j++) {
|
||||
output[j + offset] = (x[j + offset] - mean) / meanSquare * scale[j] ${bias ? '+ bias[j]' : ''};
|
||||
for (var j: u32 = 0; j < normSizeVectorized; j++) {
|
||||
let f32input = ${castToF32(dataType, components, 'x[j + offset]')};
|
||||
let f32scale = ${castToF32(dataType, components, 'scale[j]')};
|
||||
output[j + offset] = ${variables[0].type.value}((f32input - mean) / meanSquare * f32scale
|
||||
${bias ? `+ ${castToF32(dataType, components, 'bias[j]')}` : ''}
|
||||
);
|
||||
}
|
||||
|
||||
${hasMeanDataOutput ? 'meanDataOutput[global_idx] = mean' : ''};
|
||||
|
|
@ -98,14 +103,10 @@ const createLayerNormProgramInfo =
|
|||
}`;
|
||||
const outputs = [{dims: outputShape, dataType: inputs[0].dataType}];
|
||||
if (hasMeanDataOutput) {
|
||||
outputs.push(
|
||||
{dims: meanInvStdDevDim, dataType: inputs[0].dataType},
|
||||
);
|
||||
outputs.push({dims: meanInvStdDevDim, dataType: DataType.float});
|
||||
}
|
||||
if (hasInvStdOutput) {
|
||||
outputs.push(
|
||||
{dims: meanInvStdDevDim, dataType: inputs[0].dataType},
|
||||
);
|
||||
outputs.push({dims: meanInvStdDevDim, dataType: DataType.float});
|
||||
}
|
||||
|
||||
return {
|
||||
|
|
|
|||
|
|
@ -7,7 +7,7 @@ import {ShapeUtil} from '../../util';
|
|||
import {AttributeWithCacheKey, createAttributeWithCacheKey} from '../attribute-with-cache-key';
|
||||
import {ComputeContext, ProgramInfo} from '../types';
|
||||
|
||||
import {ShaderHelper, tensorTypeToWsglStorageType} from './common';
|
||||
import {castToF32, fillVector, getMaxComponents, inputVariable, outputVariable, ShaderHelper, sumVector, tensorTypeToWsglStorageType,} from './common';
|
||||
|
||||
export interface SkipLayerNormAttributes extends AttributeWithCacheKey {
|
||||
epsilon: number;
|
||||
|
|
@ -18,9 +18,6 @@ const validateInputs = (inputs: readonly TensorView[]): void => {
|
|||
throw new Error('layerNorm requires at least 3 inputs.');
|
||||
}
|
||||
|
||||
if (inputs[0].dataType !== DataType.float || inputs[1].dataType !== DataType.float) {
|
||||
throw new Error('inputs should be float type');
|
||||
}
|
||||
const input: TensorView = inputs[0];
|
||||
const skip: TensorView = inputs[1];
|
||||
const gamma: TensorView = inputs[2];
|
||||
|
|
@ -74,85 +71,93 @@ const validateInputs = (inputs: readonly TensorView[]): void => {
|
|||
};
|
||||
|
||||
const createSkipLayerNormProgramInfo =
|
||||
(inputs: readonly TensorView[], attributes: SkipLayerNormAttributes, outputCount: number,
|
||||
isTraining: boolean): ProgramInfo => {
|
||||
const inputShape = inputs[0].dims;
|
||||
const inputSize = ShapeUtil.size(inputShape);
|
||||
const outputShape = inputShape;
|
||||
const outputSize = inputSize;
|
||||
const hiddenSize = inputShape.slice(-1)[0];
|
||||
const meanInvStdDevDim = isTraining ? inputShape.slice(0, -1).concat(1) : [];
|
||||
const hasBetaInput = inputs.length > 3;
|
||||
const hasBiasInput = inputs.length > 4;
|
||||
const dataType = tensorTypeToWsglStorageType(inputs[0].dataType);
|
||||
const hasMeanOutput = isTraining && outputCount > 1;
|
||||
const hasInvStdDevOutput = isTraining && outputCount > 2;
|
||||
const hasInputSkipBiasSumOutput = outputCount > 3;
|
||||
let bindingNumber = 0;
|
||||
const getShaderSource = (shaderHelper: ShaderHelper) => `
|
||||
const hiddenSize: u32 = ${hiddenSize};
|
||||
(inputs: readonly TensorView[], attributes: SkipLayerNormAttributes, outputCount: number, isTraining: boolean):
|
||||
ProgramInfo => {
|
||||
const inputShape = inputs[0].dims;
|
||||
const inputSize = ShapeUtil.size(inputShape);
|
||||
const outputShape = inputShape;
|
||||
const outputSize = inputSize;
|
||||
const hiddenSize = inputShape.slice(-1)[0];
|
||||
const meanInvStdDevDim = isTraining ? inputShape.slice(0, -1).concat(1) : [];
|
||||
const hasBetaInput = inputs.length > 3;
|
||||
const hasBiasInput = inputs.length > 4;
|
||||
const hasMeanOutput = isTraining && outputCount > 1;
|
||||
const hasInvStdDevOutput = isTraining && outputCount > 2;
|
||||
const hasInputSkipBiasSumOutput = outputCount > 3;
|
||||
|
||||
const components = getMaxComponents(hiddenSize);
|
||||
const variables = [
|
||||
inputVariable('x', inputs[0].dataType, inputs[0].dims, components),
|
||||
inputVariable('skip', inputs[1].dataType, inputs[1].dims, components),
|
||||
inputVariable('gamma', inputs[2].dataType, inputs[2].dims, components),
|
||||
];
|
||||
if (hasBetaInput) {
|
||||
variables.push(inputVariable('beta', inputs[3].dataType, inputs[3].dims, components));
|
||||
}
|
||||
if (hasBiasInput) {
|
||||
variables.push(inputVariable('bias', inputs[4].dataType, inputs[4].dims, components));
|
||||
}
|
||||
variables.push(outputVariable('output', inputs[0].dataType, outputShape, components));
|
||||
if (hasMeanOutput) {
|
||||
variables.push(outputVariable('meanOutput', DataType.float, meanInvStdDevDim));
|
||||
}
|
||||
if (hasInvStdDevOutput) {
|
||||
variables.push(outputVariable('invStdOutput', DataType.float, meanInvStdDevDim));
|
||||
}
|
||||
if (hasInputSkipBiasSumOutput) {
|
||||
variables.push(outputVariable('inputSkipBiasSum', inputs[0].dataType, outputShape, components));
|
||||
}
|
||||
const dataType = tensorTypeToWsglStorageType(inputs[0].dataType);
|
||||
const getShaderSource = (shaderHelper: ShaderHelper) => `
|
||||
const hiddenSize: f32 = ${hiddenSize};
|
||||
const hiddenSizeVectorized: u32 = ${hiddenSize / components};
|
||||
const epsilon: f32 = ${attributes.epsilon};
|
||||
|
||||
@group(0) @binding(${bindingNumber++}) var<storage, read> x : array<${dataType}>;
|
||||
@group(0) @binding(${bindingNumber++}) var<storage, read> skip : array<${dataType}>;
|
||||
@group(0) @binding(${bindingNumber++}) var<storage, read> gamma : array<${dataType}>;
|
||||
${hasBetaInput ? `@group(0) @binding(${bindingNumber++}) var<storage, read> beta : array<${dataType}>;` : ''}
|
||||
${hasBiasInput ? `@group(0) @binding(${bindingNumber++}) var<storage, read> bias : array<${dataType}>;` : ''}
|
||||
@group(0) @binding(${bindingNumber++}) var<storage, read_write> output : array<${dataType}>;
|
||||
${
|
||||
hasMeanOutput ?
|
||||
`@group(0) @binding(${bindingNumber++}) var<storage, read_write> meanOutput : array<${dataType}>;` :
|
||||
''}
|
||||
${
|
||||
hasInvStdDevOutput ?
|
||||
`@group(0) @binding(${bindingNumber++}) var<storage, read_write> invStdOutput : array<${dataType}>;` :
|
||||
''}
|
||||
${
|
||||
hasInputSkipBiasSumOutput ?
|
||||
`@group(0) @binding(${bindingNumber++}) var<storage, read_write> inputSkipBiasSum : array<${dataType}>;` :
|
||||
''}
|
||||
${shaderHelper.declareVariables(...variables)}
|
||||
|
||||
${shaderHelper.mainStart()}
|
||||
${shaderHelper.guardAgainstOutOfBoundsWorkgroupSizes(outputSize / hiddenSize)}
|
||||
let offset = global_idx * hiddenSize;
|
||||
var sum: f32 = 0.0;
|
||||
var squareSum: f32 = 0.0;
|
||||
for (var i: u32 = 0; i < hiddenSize; i++) {
|
||||
let offset = global_idx * hiddenSizeVectorized;
|
||||
var sum = ${fillVector('f32', components)};
|
||||
var squareSum = ${fillVector('f32', components)};
|
||||
for (var i: u32 = 0; i < hiddenSizeVectorized; i++) {
|
||||
let skipValue = skip[offset + i];
|
||||
let biasValue = ${hasBiasInput ? 'bias[i]' : '0.0'};
|
||||
let inputValue = x[offset + i];
|
||||
let value = inputValue + skipValue + biasValue;
|
||||
${hasInputSkipBiasSumOutput ? 'inputSkipBiasSum[offset + i] = value;' : ''}
|
||||
output[offset + i] = value;
|
||||
sum += value;
|
||||
squareSum += value * value;
|
||||
let f32Value = ${castToF32(dataType, components, 'value')};
|
||||
sum += f32Value;
|
||||
squareSum += f32Value * f32Value;
|
||||
}
|
||||
let mean: f32 = sum / f32(hiddenSize);
|
||||
let variance: f32 = sqrt(squareSum / f32(hiddenSize) - mean * mean + epsilon);
|
||||
let mean = ${sumVector('sum', components)} / hiddenSize;
|
||||
let variance = sqrt(${sumVector('squareSum', components)} / hiddenSize - mean * mean + epsilon);
|
||||
${hasMeanOutput ? 'meanOutput[global_idx] = mean;' : ''}
|
||||
${hasInvStdDevOutput ? 'invStdOutput[global_idx] = 1.0 / variance;' : ''}
|
||||
for (var i: u32 = 0; i < hiddenSize; i++) {
|
||||
output[offset + i] = (output[offset + i] - mean) / variance * gamma[i] + ${hasBetaInput ? 'beta[i]' : '0.0'};
|
||||
for (var i: u32 = 0; i < hiddenSizeVectorized; i++) {
|
||||
output[offset + i] = (output[offset + i] - ${dataType}(mean)) / ${dataType}(variance) * gamma[i]
|
||||
+ ${hasBetaInput ? 'beta[i]' : '0.0'};
|
||||
}
|
||||
}`;
|
||||
const outputs = [{dims: outputShape, dataType: inputs[0].dataType}];
|
||||
if (outputCount > 1) {
|
||||
outputs.push({dims: meanInvStdDevDim, dataType: inputs[0].dataType});
|
||||
}
|
||||
if (outputCount > 2) {
|
||||
outputs.push({dims: meanInvStdDevDim, dataType: inputs[0].dataType});
|
||||
}
|
||||
if (outputCount > 3) {
|
||||
outputs.push({dims: inputShape, dataType: inputs[0].dataType});
|
||||
}
|
||||
const outputs = [{dims: outputShape, dataType: inputs[0].dataType}];
|
||||
if (outputCount > 1) {
|
||||
outputs.push({dims: meanInvStdDevDim, dataType: DataType.float});
|
||||
}
|
||||
if (outputCount > 2) {
|
||||
outputs.push({dims: meanInvStdDevDim, dataType: DataType.float});
|
||||
}
|
||||
if (outputCount > 3) {
|
||||
outputs.push({dims: inputShape, dataType: inputs[0].dataType});
|
||||
}
|
||||
|
||||
return {
|
||||
name: 'SkipLayerNormalization',
|
||||
shaderCache: {hint: attributes.cacheKey},
|
||||
getShaderSource,
|
||||
getRunData: () => ({outputs, dispatchGroup: {x: Math.ceil(outputSize / hiddenSize / 64)}}),
|
||||
};
|
||||
};
|
||||
return {
|
||||
name: 'SkipLayerNormalization',
|
||||
shaderCache: {hint: attributes.cacheKey},
|
||||
getShaderSource,
|
||||
getRunData: () => ({outputs, dispatchGroup: {x: Math.ceil(outputSize / hiddenSize / 64)}}),
|
||||
};
|
||||
};
|
||||
|
||||
export const skipLayerNorm = (context: ComputeContext, attributes: SkipLayerNormAttributes): void => {
|
||||
// TODO: initialize isTraining from ComputeContext
|
||||
|
|
|
|||
|
|
@ -7,14 +7,16 @@ namespace onnxruntime {
|
|||
namespace contrib {
|
||||
namespace js {
|
||||
|
||||
using onnxruntime::js::JsepSupportedFloatTypes;
|
||||
|
||||
ONNX_OPERATOR_KERNEL_EX(
|
||||
SkipLayerNormalization,
|
||||
kMSDomain,
|
||||
1,
|
||||
kJsExecutionProvider,
|
||||
(*KernelDefBuilder::Create())
|
||||
.TypeConstraint("T", DataTypeImpl::GetTensorType<float>())
|
||||
.TypeConstraint("U", DataTypeImpl::GetTensorType<float>()),
|
||||
.TypeConstraint("T", JsepSupportedFloatTypes())
|
||||
.TypeConstraint("U", JsepSupportedFloatTypes()),
|
||||
SkipLayerNorm);
|
||||
|
||||
} // namespace js
|
||||
|
|
|
|||
|
|
@ -321,9 +321,9 @@ class ONNX_OPERATOR_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 13, Fla
|
|||
class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 6, 12, Tile);
|
||||
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 13, Tile);
|
||||
|
||||
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 17, float, LayerNormalization);
|
||||
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kJsExecutionProvider, kMSInternalNHWCDomain, 6, float, InstanceNormalization);
|
||||
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 6, float, InstanceNormalization);
|
||||
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 17, LayerNormalization);
|
||||
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kJsExecutionProvider, kMSInternalNHWCDomain, 6, InstanceNormalization);
|
||||
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 6, InstanceNormalization);
|
||||
|
||||
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 11, Range);
|
||||
|
||||
|
|
@ -593,9 +593,9 @@ std::unique_ptr<KernelRegistry> RegisterKernels() {
|
|||
BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 6, 12, Tile)>,
|
||||
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 13, Tile)>,
|
||||
|
||||
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 17, float, LayerNormalization)>,
|
||||
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kJsExecutionProvider, kMSInternalNHWCDomain, 6, float, InstanceNormalization)>,
|
||||
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 6, float, InstanceNormalization)>,
|
||||
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 17, LayerNormalization)>,
|
||||
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kJsExecutionProvider, kMSInternalNHWCDomain, 6, InstanceNormalization)>,
|
||||
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 6, InstanceNormalization)>,
|
||||
|
||||
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 11, Range)>,
|
||||
|
||||
|
|
|
|||
|
|
@ -6,18 +6,17 @@
|
|||
namespace onnxruntime {
|
||||
namespace js {
|
||||
|
||||
#define INSTANCE_NORM_KERNEL(op_name, domain, data_type, since_version, is_channels_last) \
|
||||
ONNX_OPERATOR_TYPED_KERNEL_EX( \
|
||||
op_name, \
|
||||
domain, \
|
||||
since_version, \
|
||||
data_type, \
|
||||
kJsExecutionProvider, \
|
||||
(*KernelDefBuilder::Create()).TypeConstraint("T", DataTypeImpl::GetTensorType<data_type>()), \
|
||||
#define INSTANCE_NORM_KERNEL(op_name, domain, since_version, is_channels_last) \
|
||||
ONNX_OPERATOR_KERNEL_EX( \
|
||||
op_name, \
|
||||
domain, \
|
||||
since_version, \
|
||||
kJsExecutionProvider, \
|
||||
(*KernelDefBuilder::Create()).TypeConstraint("T", JsepSupportedFloatTypes()), \
|
||||
InstanceNorm<is_channels_last>);
|
||||
|
||||
INSTANCE_NORM_KERNEL(InstanceNormalization, kOnnxDomain, float, 6, false)
|
||||
INSTANCE_NORM_KERNEL(InstanceNormalization, kMSInternalNHWCDomain, float, 6, true)
|
||||
INSTANCE_NORM_KERNEL(InstanceNormalization, kOnnxDomain, 6, false)
|
||||
INSTANCE_NORM_KERNEL(InstanceNormalization, kMSInternalNHWCDomain, 6, true)
|
||||
|
||||
} // namespace js
|
||||
} // namespace onnxruntime
|
||||
|
|
|
|||
|
|
@ -8,21 +8,15 @@
|
|||
namespace onnxruntime {
|
||||
namespace js {
|
||||
|
||||
#define REGISTER_KERNEL_TYPED(T) \
|
||||
ONNX_OPERATOR_TYPED_KERNEL_EX( \
|
||||
LayerNormalization, \
|
||||
kOnnxDomain, \
|
||||
17, \
|
||||
T, \
|
||||
kJsExecutionProvider, \
|
||||
(*KernelDefBuilder::Create()) \
|
||||
.TypeConstraint("T", DataTypeImpl::GetTensorType<T>()) \
|
||||
.TypeConstraint("U", DataTypeImpl::GetTensorType<float>()), \
|
||||
LayerNorm<T, float>);
|
||||
|
||||
REGISTER_KERNEL_TYPED(float)
|
||||
// REGISTER_KERNEL_TYPED(double)
|
||||
// REGISTER_KERNEL_TYPED(MLFloat16)
|
||||
ONNX_OPERATOR_KERNEL_EX(
|
||||
LayerNormalization,
|
||||
kOnnxDomain,
|
||||
17,
|
||||
kJsExecutionProvider,
|
||||
(*KernelDefBuilder::Create())
|
||||
.TypeConstraint("T", JsepSupportedFloatTypes())
|
||||
.TypeConstraint("U", JsepSupportedFloatTypes()),
|
||||
LayerNorm);
|
||||
|
||||
} // namespace js
|
||||
} // namespace onnxruntime
|
||||
|
|
|
|||
|
|
@ -8,7 +8,6 @@
|
|||
namespace onnxruntime {
|
||||
namespace js {
|
||||
|
||||
template <typename T, typename U>
|
||||
class LayerNorm : public JsKernel {
|
||||
public:
|
||||
LayerNorm(const OpKernelInfo& info) : JsKernel(info) {
|
||||
|
|
|
|||
Loading…
Reference in a new issue