368 lines
14 KiB
TypeScript
368 lines
14 KiB
TypeScript
// Copyright (c) Microsoft Corporation. All rights reserved.
|
|
// Licensed under the MIT License.
|
|
|
|
import {AttributeWithCacheKey, createAttributeWithCacheKey} from '../../../attribute-with-cache-key';
|
|
import {Graph} from '../../../graph';
|
|
import {OperatorImplementation, OperatorInitialization} from '../../../operators';
|
|
import {Tensor} from '../../../tensor';
|
|
import {PoolConvUtil, ShapeUtil} from '../../../util';
|
|
import {WebGLInferenceHandler} from '../inference-handler';
|
|
import {ProgramInfo, ProgramMetadata, TextureType} from '../types';
|
|
|
|
export interface AveragePoolAttributes extends AttributeWithCacheKey {
|
|
readonly autoPad: string;
|
|
readonly ceilMode: number;
|
|
readonly countIncludePad: boolean;
|
|
readonly kernelShape: readonly number[];
|
|
readonly strides: readonly number[];
|
|
readonly pads: readonly number[];
|
|
}
|
|
|
|
export const averagePool: OperatorImplementation<AveragePoolAttributes> =
|
|
(inferenceHandler: WebGLInferenceHandler, inputs: Tensor[], attributes: AveragePoolAttributes): Tensor[] => {
|
|
validateInputs(inputs);
|
|
const metadata =
|
|
{name: 'AveragePool', inputNames: ['X'], inputTypes: [TextureType.unpacked], cacheHint: attributes.cacheKey};
|
|
const output = inferenceHandler.run(
|
|
{...metadata, get: () => createAveragePoolProgramInfo(inputs, metadata, false, attributes)}, inputs);
|
|
return [output];
|
|
};
|
|
|
|
export const parseAveragePoolAttributes: OperatorInitialization<AveragePoolAttributes> =
|
|
(node: Graph.Node): AveragePoolAttributes => {
|
|
const autoPad = node.attributes.getString('auto_pad', 'NOTSET');
|
|
const ceilMode = node.attributes.getInt('ceil_mode', 0);
|
|
const countIncludePad = (node.attributes.getInt('count_include_pad', 0) === 0 ? false : true);
|
|
const kernelShape = node.attributes.getInts('kernel_shape');
|
|
const strides = node.attributes.getInts('strides', []);
|
|
const pads = node.attributes.getInts('pads', []);
|
|
|
|
// TODO: support attribute 'ceil_mode'
|
|
if (ceilMode !== 0) {
|
|
throw new Error('using ceil() in shape computation is not yet supported for AveragePool');
|
|
}
|
|
|
|
return createAttributeWithCacheKey({autoPad, ceilMode, countIncludePad, kernelShape, strides, pads});
|
|
};
|
|
|
|
const createAveragePoolProgramInfo =
|
|
(inputs: Tensor[], metadata: ProgramMetadata, isGlobalOperator: boolean, attributes: AveragePoolAttributes):
|
|
ProgramInfo => {
|
|
const [adjustedAttributes, outputShape] =
|
|
getAdjustedPoolAttributesAndOutputShape(inputs, attributes, isGlobalOperator);
|
|
const kernelSize = ShapeUtil.size(adjustedAttributes.kernelShape);
|
|
const op1 = 'value += _X(x);';
|
|
let op2 = '';
|
|
if (adjustedAttributes.countIncludePad) {
|
|
op2 += `value /= float(${kernelSize});`;
|
|
} else {
|
|
op2 += `value /= float(${kernelSize} - pad);`;
|
|
}
|
|
const poolingCode = generatePoolingCode(inputs[0].dims, adjustedAttributes, op1, op2, '0.0');
|
|
const shaderSource = `
|
|
${poolingCode}
|
|
`;
|
|
return {
|
|
...metadata,
|
|
output: {dims: outputShape, type: inputs[0].type, textureType: TextureType.unpacked},
|
|
shaderSource
|
|
};
|
|
};
|
|
|
|
export const globalAveragePool: OperatorImplementation<AveragePoolAttributes> =
|
|
(inferenceHandler: WebGLInferenceHandler, inputs: Tensor[], attributes: AveragePoolAttributes): Tensor[] => {
|
|
validateInputs(inputs);
|
|
const metadata = {
|
|
name: 'GlobalAveragePool',
|
|
inputNames: ['X'],
|
|
inputTypes: [TextureType.unpacked],
|
|
cacheHint: `${attributes.countIncludePad}`
|
|
};
|
|
const output = inferenceHandler.run(
|
|
{...metadata, get: () => createAveragePoolProgramInfo(inputs, metadata, true, attributes)}, inputs);
|
|
return [output];
|
|
};
|
|
|
|
export const parseGlobalAveragePoolAttributes: OperatorInitialization<AveragePoolAttributes> =
|
|
(node: Graph.Node): AveragePoolAttributes => {
|
|
const countIncludePad = (node.attributes.getInt('count_include_pad', 0) === 0 ? false : true);
|
|
return createAttributeWithCacheKey(
|
|
{autoPad: '', ceilMode: 0, countIncludePad, kernelShape: [], strides: [], pads: []});
|
|
};
|
|
|
|
export interface MaxPoolAttributes extends AveragePoolAttributes {
|
|
readonly storageOrder: number;
|
|
readonly dilations: number[];
|
|
}
|
|
|
|
export const maxPool: OperatorImplementation<MaxPoolAttributes> =
|
|
(inferenceHandler: WebGLInferenceHandler, inputs: Tensor[], attributes: MaxPoolAttributes): Tensor[] => {
|
|
validateInputs(inputs);
|
|
const metadata =
|
|
{name: 'MaxPool', inputNames: ['X'], inputTypes: [TextureType.unpacked], cacheHint: attributes.cacheKey};
|
|
const output = inferenceHandler.run(
|
|
{...metadata, get: () => createMaxPoolProgramInfo(inputs, metadata, false, attributes)}, inputs);
|
|
return [output];
|
|
};
|
|
|
|
export const parseMaxPoolAttributes: OperatorInitialization<MaxPoolAttributes> =
|
|
(node: Graph.Node): MaxPoolAttributes => {
|
|
const autoPad = node.attributes.getString('auto_pad', 'NOTSET');
|
|
const ceilMode = node.attributes.getInt('ceil_mode', 0);
|
|
const kernelShape = node.attributes.getInts('kernel_shape');
|
|
const strides = node.attributes.getInts('strides', []);
|
|
const pads = node.attributes.getInts('pads', []);
|
|
const storageOrder = node.attributes.getInt('storage_order', 0);
|
|
const dilations = node.attributes.getInts('dilations', []);
|
|
|
|
// TODO: support attribute 'ceil_mode' and 'storage_order'
|
|
if (storageOrder !== 0) {
|
|
throw new Error('column major storage order is not yet supported for MaxPool');
|
|
}
|
|
if (ceilMode !== 0) {
|
|
throw new Error('using ceil() in shape computation is not yet supported for MaxPool');
|
|
}
|
|
|
|
return createAttributeWithCacheKey(
|
|
{autoPad, ceilMode, countIncludePad: false, kernelShape, strides, pads, storageOrder, dilations});
|
|
};
|
|
|
|
const createMaxPoolProgramInfo =
|
|
(inputs: Tensor[], metadata: ProgramMetadata, isGlobalOperator: boolean, attributes: MaxPoolAttributes):
|
|
ProgramInfo => {
|
|
const [adjustedAttributes, outputShape] =
|
|
getAdjustedPoolAttributesAndOutputShape(inputs, attributes, isGlobalOperator);
|
|
const op1 = `
|
|
value = max(_X(x), value);
|
|
`;
|
|
const op2 = '';
|
|
const poolingCode = generatePoolingCode(inputs[0].dims, adjustedAttributes, op1, op2, '-1e5');
|
|
const shaderSource = `
|
|
${poolingCode}
|
|
`;
|
|
return {
|
|
...metadata,
|
|
output: {dims: outputShape, type: inputs[0].type, textureType: TextureType.unpacked},
|
|
shaderSource
|
|
};
|
|
};
|
|
|
|
const getAdjustedPoolAttributesAndOutputShape =
|
|
(inputs: Tensor[], attributes: AveragePoolAttributes|MaxPoolAttributes, isGlobalOperator: boolean):
|
|
[AveragePoolAttributes|MaxPoolAttributes, number[]] => {
|
|
const inputShape = inputs[0].dims.slice();
|
|
const hasDilations = Object.hasOwnProperty.call(attributes, 'dilations');
|
|
const kernelShape = attributes.kernelShape.slice();
|
|
const strides = attributes.strides.slice();
|
|
const dilations: number[] = hasDilations ? (attributes as MaxPoolAttributes).dilations.slice() : [];
|
|
const pads = attributes.pads.slice();
|
|
PoolConvUtil.adjustPoolAttributes(isGlobalOperator, inputShape, kernelShape, strides, dilations, pads);
|
|
|
|
const outputShape = PoolConvUtil.computePoolOutputShape(
|
|
isGlobalOperator, inputShape, strides, dilations, kernelShape, pads, attributes.autoPad);
|
|
|
|
const newAttributes = Object.assign({}, attributes);
|
|
if (hasDilations) {
|
|
Object.assign(newAttributes, {kernelShape, strides, pads, dilations, cacheKey: attributes.cacheKey});
|
|
} else {
|
|
Object.assign(newAttributes, {kernelShape, strides, pads, cacheKey: attributes.cacheKey});
|
|
}
|
|
return [newAttributes, outputShape];
|
|
};
|
|
|
|
const globalMaxPoolAttributes = {
|
|
autoPad: '',
|
|
ceilMode: 0,
|
|
countIncludePad: false,
|
|
kernelShape: [],
|
|
strides: [],
|
|
pads: [],
|
|
storageOrder: 0,
|
|
dilations: [],
|
|
cacheKey: ''
|
|
};
|
|
|
|
const globalMaxPoolMetadata = {
|
|
name: 'GlobalMaxPool',
|
|
inputNames: ['X'],
|
|
inputTypes: [TextureType.unpacked]
|
|
};
|
|
|
|
export const globalMaxPool = (inferenceHandler: WebGLInferenceHandler, inputs: Tensor[]): Tensor[] => {
|
|
validateInputs(inputs);
|
|
const output = inferenceHandler.run(
|
|
{
|
|
...globalMaxPoolMetadata,
|
|
get: () => createMaxPoolProgramInfo(inputs, globalMaxPoolMetadata, true, globalMaxPoolAttributes)
|
|
},
|
|
inputs);
|
|
return [output];
|
|
};
|
|
|
|
const validateInputs = (inputs: Tensor[]): void => {
|
|
if (!inputs || inputs.length !== 1) {
|
|
throw new Error('Pool ops requires 1 input.');
|
|
}
|
|
if (inputs[0].type !== 'float32' && inputs[0].type !== 'float64') {
|
|
throw new Error('Invalid input type.');
|
|
}
|
|
};
|
|
|
|
const generatePoolingCode =
|
|
(inputDims: readonly number[], attributes: AveragePoolAttributes, op1: string, op2: string, start: string):
|
|
string => {
|
|
const rank = inputDims.length;
|
|
if (attributes.kernelShape.length <= 2) {
|
|
const kw = attributes.kernelShape[attributes.kernelShape.length - 1];
|
|
const sw = attributes.strides[attributes.strides.length - 1];
|
|
const pwStart = attributes.pads[attributes.pads.length / 2 - 1];
|
|
const pwEnd = attributes.pads[attributes.pads.length - 1];
|
|
const dimW = inputDims[rank - 1];
|
|
let codeW = '';
|
|
let codeH = '';
|
|
let codeHEnd = '';
|
|
if (pwStart + pwEnd !== 0) {
|
|
codeW = `
|
|
for (int i = 0; i < ${kw}; i++) {
|
|
x[${rank} - 1] = indices[${rank} - 1] * ${sw} - ${pwStart} + i;
|
|
if (x[${rank} - 1] < 0 || x[${rank} - 1] >= ${dimW}) {
|
|
pad++;
|
|
continue;
|
|
}
|
|
${op1}
|
|
}`;
|
|
} else {
|
|
codeW = `
|
|
for (int i = 0; i < ${kw}; i++) {
|
|
x[${rank} - 1] = indices[${rank} - 1] * ${sw} - ${pwStart} + i;
|
|
${op1}
|
|
}`;
|
|
}
|
|
|
|
if (attributes.kernelShape.length === 2) {
|
|
const kh = attributes.kernelShape[attributes.kernelShape.length - 2];
|
|
const sh = attributes.strides[attributes.strides.length - 2];
|
|
const phStart = attributes.pads[attributes.pads.length / 2 - 2];
|
|
const phEnd = attributes.pads[attributes.pads.length - 2];
|
|
const dimH = inputDims[rank - 2];
|
|
if (phStart + phEnd !== 0) {
|
|
codeH = `
|
|
for (int j = 0; j < ${kh}; j++) {
|
|
x[${rank} - 2] = indices[${rank} - 2] * ${sh} - ${phStart} + j;
|
|
if (x[${rank} - 2] < 0 || x[${rank} - 2] >= ${dimH}) {
|
|
pad+= ${kw};
|
|
continue;
|
|
}
|
|
`;
|
|
} else {
|
|
codeH = `
|
|
for (int j = 0; j < ${kh}; j++) {
|
|
x[${rank} - 2] = indices[${rank} - 2] * ${sh} - ${phStart} + j;
|
|
`;
|
|
}
|
|
codeHEnd = `
|
|
}
|
|
`;
|
|
}
|
|
|
|
const poolingCode = `
|
|
float process(int indices[${rank}]) {
|
|
int x[${rank}];
|
|
copyVec(indices, x);
|
|
|
|
float value = ${start};
|
|
int pad = 0;
|
|
${codeH}
|
|
${codeW}
|
|
${codeHEnd}
|
|
${op2}
|
|
return value;
|
|
}
|
|
`;
|
|
return poolingCode;
|
|
} else {
|
|
const kernelSize = ShapeUtil.size(attributes.kernelShape);
|
|
const kernelStrides = ShapeUtil.computeStrides(attributes.kernelShape);
|
|
const stridesRank = kernelStrides.length;
|
|
const padsRank = attributes.pads.length;
|
|
const offsetToIndicesFunction = offsetToIndices(stridesRank);
|
|
const copyInputDims = copyArray(inputDims, 'inputDims');
|
|
const copyPads = copyArray(attributes.pads, 'pads');
|
|
const copyKernelStrides = copyArray(kernelStrides, 'kernelStrides');
|
|
const copyStrides = copyArray(attributes.strides, 'strides');
|
|
const hasPads = attributes.pads.reduce((sum, cur) => sum + cur);
|
|
let padCode = '';
|
|
if (hasPads) {
|
|
padCode = `
|
|
if (x[j] >= inputDims[j] || x[j] < 0) {
|
|
pad++;
|
|
isPad = true;
|
|
break;
|
|
}
|
|
}
|
|
if (!isPad) {
|
|
${op1}
|
|
}`;
|
|
} else {
|
|
padCode = `
|
|
}
|
|
${op1}
|
|
`;
|
|
}
|
|
const poolingCode = `
|
|
${offsetToIndicesFunction}
|
|
float process(int indices[${rank}]) {
|
|
int x[${rank}];
|
|
copyVec(indices, x);
|
|
int offset[${stridesRank}];
|
|
int pads[${padsRank}];
|
|
int inputDims[${rank}];
|
|
int kernelStrides[${stridesRank}];
|
|
int strides[${stridesRank}];
|
|
${copyPads}
|
|
${copyInputDims}
|
|
${copyStrides}
|
|
${copyKernelStrides}
|
|
|
|
float value = ${start};
|
|
int pad = 0;
|
|
bool isPad = false;
|
|
for (int i = 0; i < ${kernelSize}; i++) {
|
|
offsetToIndices(i, kernelStrides, offset);
|
|
isPad = false;
|
|
for (int j = ${rank} - ${stridesRank}; j < ${rank}; j++) {
|
|
x[j] = indices[j] * strides[j - ${rank} + ${stridesRank}]
|
|
+ offset[j - ${rank} + ${stridesRank}] - pads[j - 2];
|
|
${padCode}
|
|
}
|
|
${op2}
|
|
|
|
return value;
|
|
}
|
|
`;
|
|
return poolingCode;
|
|
}
|
|
};
|
|
|
|
const copyArray = (array: readonly number[], arrayName: string): string => {
|
|
let block = '';
|
|
for (let i = 0; i < array.length; i++) {
|
|
block += `
|
|
${arrayName}[${i}] = ${array[i]};
|
|
`;
|
|
}
|
|
return block;
|
|
};
|
|
|
|
const offsetToIndices = (rank: number): string => `
|
|
void offsetToIndices(int offset, int[${rank}] strides, out int[${rank}] indices) {
|
|
if (${rank} == 0) {
|
|
return;
|
|
}
|
|
for (int i = 0; i < ${rank} - 1; ++i) {
|
|
indices[i] = offset / strides[i];
|
|
offset -= indices[i] * strides[i];
|
|
}
|
|
indices[${rank} - 1] = offset;
|
|
}`;
|