125 lines
5.4 KiB
JavaScript
125 lines
5.4 KiB
JavaScript
"use strict";
|
|
// Copyright (c) Microsoft Corporation. All rights reserved.
|
|
// Licensed under the MIT License.
|
|
Object.defineProperty(exports, "__esModule", { value: true });
|
|
exports.createPackedConcatProgramInfoLoader = void 0;
|
|
const glsl_source_1 = require("../glsl-source");
|
|
const types_1 = require("../types");
|
|
const utils_1 = require("../utils");
|
|
const packing_utils_1 = require("./packing-utils");
|
|
const createPackedConcatProgramMetadata = (inputCount, cacheHint) => ({
|
|
name: 'Concat (packed)',
|
|
inputNames: Array.from({ length: inputCount }, (v, i) => `X${i}`),
|
|
inputTypes: Array(inputCount).fill(types_1.TextureType.packed),
|
|
cacheHint
|
|
});
|
|
const createPackedConcatProgramInfo = (handler, metadata, inputs, axis) => {
|
|
const inputShape = inputs[0].dims.slice();
|
|
if (axis >= inputShape.length || axis < (-1 * inputShape.length)) {
|
|
throw new Error('axis specified for concat doesn\'t match input dimensionality');
|
|
}
|
|
if (axis < 0) {
|
|
axis = inputShape.length + axis;
|
|
}
|
|
// ensure all of the non-concatenated axes match each other
|
|
// calculate the shape of the output tensor while we do that
|
|
const outputShape = inputShape.slice(0);
|
|
for (let i = 1; i < inputs.length; i++) {
|
|
const dataNShape = inputs[i].dims.slice();
|
|
for (let axisIndex = 0; axisIndex < inputShape.length; axisIndex++) {
|
|
// add to the placeholder for computing output shape
|
|
if (axisIndex === axis) {
|
|
outputShape[axis] += dataNShape[axisIndex];
|
|
}
|
|
// ensure all non-cancatenated axes match each other
|
|
else if (inputShape[axisIndex] !== dataNShape[axisIndex]) {
|
|
throw new Error('non concat dimensions must match');
|
|
}
|
|
}
|
|
}
|
|
const rank = outputShape.length;
|
|
const coords = (0, packing_utils_1.getChannels)('coords', rank);
|
|
const dtype = (0, utils_1.getCoordsDataType)(rank);
|
|
const unpackChannel = (0, packing_utils_1.unpackFromChannel)();
|
|
const shapes = inputs.map(i => i.dims);
|
|
const channels = (0, utils_1.getGlChannels)(rank);
|
|
const offsets = new Array(shapes.length - 1);
|
|
offsets[0] = shapes[0][axis];
|
|
for (let i = 1; i < offsets.length; i++) {
|
|
offsets[i] = offsets[i - 1] + shapes[i][axis];
|
|
}
|
|
const channel = channels[axis];
|
|
const lastChannels = channels.slice(-2);
|
|
const allChannels = channels.join();
|
|
let getValueSnippet = `if (${channel} < ${offsets[0]}) {
|
|
return getChannel(
|
|
getX0(${allChannels}), vec2(${lastChannels.join()}));
|
|
}`;
|
|
for (let i = 1; i < offsets.length; i++) {
|
|
const shift = offsets[i - 1];
|
|
getValueSnippet += `
|
|
if (${channel} < ${offsets[i]} && ${channel} >= ${offsets[i - 1]}) {
|
|
return getChannel(
|
|
getX${i}(${getShiftedChannelsSnippet(channels, channel, shift)}),
|
|
vec2(${getShiftedChannelsSnippet(lastChannels, channel, shift)}));
|
|
}`;
|
|
}
|
|
const lastIndex = offsets.length;
|
|
const shift = offsets[offsets.length - 1];
|
|
getValueSnippet += `
|
|
return getChannel(
|
|
getX${lastIndex}(${getShiftedChannelsSnippet(channels, channel, shift)}),
|
|
vec2(${getShiftedChannelsSnippet(lastChannels, channel, shift)}));`;
|
|
const glsl = (0, glsl_source_1.getGlsl)(handler.session.backend.glContext.version);
|
|
const shaderSource = `
|
|
${unpackChannel}
|
|
float getValue(${channels.map(x => 'int ' + x)}) {
|
|
${getValueSnippet}
|
|
}
|
|
|
|
void main() {
|
|
${dtype} coords = getOutputCoords();
|
|
int lastDim = coords.${channels[rank - 1]};
|
|
coords.${channels[rank - 1]} = coords.${channels[rank - 2]};
|
|
coords.${channels[rank - 2]} = lastDim;
|
|
|
|
vec4 result = vec4(getValue(${coords}), 0., 0., 0.);
|
|
|
|
${coords[rank - 1]} = ${coords[rank - 1]} + 1;
|
|
if (${coords[rank - 1]} < ${outputShape[rank - 1]}) {
|
|
result.g = getValue(${coords});
|
|
}
|
|
|
|
${coords[rank - 2]} = ${coords[rank - 2]} + 1;
|
|
if (${coords[rank - 2]} < ${outputShape[rank - 2]}) {
|
|
result.a = getValue(${coords});
|
|
}
|
|
|
|
${coords[rank - 1]} = ${coords[rank - 1]} - 1;
|
|
if (${coords[rank - 2]} < ${outputShape[rank - 2]} &&
|
|
${coords[rank - 1]} < ${outputShape[rank - 1]}) {
|
|
result.b = getValue(${coords});
|
|
}
|
|
${glsl.output} = result;
|
|
}
|
|
`;
|
|
return Object.assign(Object.assign({}, metadata), { output: { dims: outputShape, type: inputs[0].type, textureType: types_1.TextureType.packed }, shaderSource, hasMain: true });
|
|
};
|
|
const createPackedConcatProgramInfoLoader = (handler, inputs, attributes) => {
|
|
const metadata = createPackedConcatProgramMetadata(inputs.length, attributes.cacheKey);
|
|
return Object.assign(Object.assign({}, metadata), { get: () => createPackedConcatProgramInfo(handler, metadata, inputs, attributes.axis) });
|
|
};
|
|
exports.createPackedConcatProgramInfoLoader = createPackedConcatProgramInfoLoader;
|
|
const getShiftedChannelsSnippet = (channels, channel, shift) => {
|
|
const channelIdx = channels.indexOf(channel);
|
|
const res = channels.map((c, idx) => {
|
|
if (idx === channelIdx) {
|
|
return `${c} - ${shift}`;
|
|
}
|
|
else {
|
|
return c;
|
|
}
|
|
});
|
|
return res.join();
|
|
};
|
|
//# sourceMappingURL=concat-packed.js.map
|