Skip to content

Commit

Permalink
combine rowIndexMap and colIndexMap for glsl texture coordinate trans…
Browse files Browse the repository at this point in the history 
…formations to reduce number of textures created
  • Loading branch information
@transcranial
transcranial committed Nov 24, 2017
1 parent 4872098 commit 53b07ec
Show file tree
Hide file tree
Showing 23 changed files with 249 additions and 400 deletions.
62 changes: 17 additions & 45 deletions src/layers/convolutional/Conv2D.js
View file
Original file line number Diff line number Diff line change
Expand Up @@ -330,23 +330,21 @@ export default class Conv2D extends Layer {
* @param {Object} indicesForReshaped
*/
_createIndexMap(indicesForReshaped) {
if (this.rowIndexMap && this.colIndexMap) {
if (this.indexMap) {
return
}

let [inputRows, inputCols, inputChannels] = this.inputShape

const indicesRow = new Tensor(indicesForReshaped.row.data, indicesForReshaped.row.shape, { type: Int32Array })
const indicesCol = new Tensor(indicesForReshaped.col.data, indicesForReshaped.col.shape, { type: Int32Array })
const indices = new Tensor(indicesForReshaped.data, indicesForReshaped.shape, { type: Int32Array })

// padding for border mode 'same'
if (this.padding === 'same') {
const [paddingRowBefore, paddingRowAfter, paddingColBefore, paddingColAfter] = this.inputPadding
inputRows = inputRows + paddingRowBefore + paddingRowAfter
inputCols = inputCols + paddingColBefore + paddingColAfter
const padValue = -1
this._padInput(indicesRow, padValue)
this._padInput(indicesCol, padValue)
this._padInput(indices, padValue)
}

const nbRow = this.kernelShape[1]
Expand All @@ -360,36 +358,25 @@ export default class Conv2D extends Layer {
const nbRowDilated = nbRow + (nbRow - 1) * (this.dilationRate[0] - 1)
const nbColDilated = nbCol + (nbCol - 1) * (this.dilationRate[1] - 1)

this.rowIndexMap = new Tensor([], [nbPatches, patchLen], { type: Int32Array })
this.colIndexMap = new Tensor([], [nbPatches, patchLen], { type: Int32Array })
this.indexMap = new Tensor([], [nbPatches, patchLen], { type: Int32Array })

const indicesRowPatch = new Tensor([], [nbRow, nbCol, inputChannels])
const indicesColPatch = new Tensor([], [nbRow, nbCol, inputChannels])
const indicesPatch = new Tensor([], [nbRow, nbCol, inputChannels])
let offset = 0
for (let i = 0, limit = inputRows - nbRowDilated; i <= limit; i += this.strides[0]) {
for (let j = 0, limit = inputCols - nbColDilated; j <= limit; j += this.strides[1]) {
ops.assign(
indicesRowPatch.tensor,
indicesRow.tensor
indicesPatch.tensor,
indices.tensor
.hi(i + nbRowDilated, j + nbColDilated, inputChannels)
.lo(i, j, 0)
.step(this.dilationRate[0], this.dilationRate[1], 1)
)
ops.assign(
indicesColPatch.tensor,
indicesCol.tensor
.hi(i + nbRowDilated, j + nbColDilated, inputChannels)
.lo(i, j, 0)
.step(this.dilationRate[0], this.dilationRate[1], 1)
)
this.rowIndexMap.tensor.data.set(indicesRowPatch.tensor.data, offset)
this.colIndexMap.tensor.data.set(indicesColPatch.tensor.data, offset)
this.indexMap.tensor.data.set(indicesPatch.tensor.data, offset)
offset += patchLen
}
}

this.rowIndexMap.createGLTexture({ type: '2d', format: 'int', supportsTextureFragments: true })
this.colIndexMap.createGLTexture({ type: '2d', format: 'int', supportsTextureFragments: true })
this.indexMap.createGLTexture({ type: '2d', format: 'int', supportsTextureFragments: true })
}

/**
Expand All @@ -413,35 +400,20 @@ export default class Conv2D extends Layer {
if (x.is2DReshaped || x.is2DSquareReshaped) {
this._createIndexMap(x.indicesForReshaped)
if (!this.mappedInput) {
this.mappedInput = new Tensor([], this.rowIndexMap.glTextureShape)
this.mappedInput = new Tensor([], this.indexMap.glTextureShape)
this.mappedInput.createGLTexture({ type: '2d', format: 'float', supportsTextureFragments: true })
}

if (x.glTextureFragments) {
x.convert2DRowFragmentedGLTextureToColStack()
webgl2.runProgram({
program: this.mapInputFragmentsProgram,
output: this.mappedInput,
inputs: [
{ input: x, name: 'x' },
{ input: this.rowIndexMap, name: 'rowIndexMap' },
{ input: this.colIndexMap, name: 'colIndexMap' }
],
uniforms: [{ value: x.glTextureShape[1], type: 'int', name: 'fragmentCols' }],
supportsTextureFragments: true
})
} else {
webgl2.runProgram({
program: this.mapInputProgram,
output: this.mappedInput,
inputs: [
{ input: x, name: 'x' },
{ input: this.rowIndexMap, name: 'rowIndexMap' },
{ input: this.colIndexMap, name: 'colIndexMap' }
],
supportsTextureFragments: true
})
}
webgl2.runProgram({
program: x.glTextureFragments ? this.mapInputFragmentsProgram : this.mapInputProgram,
output: this.mappedInput,
inputs: [{ input: x, name: 'x' }, { input: this.indexMap, name: 'indexMap' }],
uniforms: [{ value: x.glTextureShape[1], type: 'int', name: 'inputCols' }],
supportsTextureFragments: true
})
}

const input = x.is2DReshaped || x.is2DSquareReshaped ? this.mappedInput : this.imColsMat
Expand Down
16 changes: 8 additions & 8 deletions src/layers/convolutional/Conv2DTranspose.fragments.glsl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,8 +4,7 @@ precision highp isampler2D;

in vec2 outTex;
uniform sampler2D matMulResult;
uniform isampler2D rowIndexMap;
uniform isampler2D colIndexMap;
uniform isampler2D indexMap;
uniform sampler2D bias;
uniform bool use_bias;
uniform int inputFragmentCols;
Expand All @@ -14,17 +13,18 @@ out vec4 outColor;

void main() {
int inputFragmentRows = textureSize(matMulResult, 0)[1];
int outputFragmentRows = textureSize(rowIndexMap, 0)[1];
int summationLength = textureSize(rowIndexMap, 0)[0];
int outputFragmentRows = textureSize(indexMap, 0)[1];
int summationLength = textureSize(indexMap, 0)[0];
int out_x = int(float(outputFragmentCols) * outTex.x);
int out_y = int(float(outputFragmentRows) * outTex.y);

float sum = 0.;
for (int n = 0; n < summationLength; ++n) {
int rowIndex = texelFetch(rowIndexMap, ivec2(n, out_y), 0).r;
int colIndex = texelFetch(colIndexMap, ivec2(n, out_y), 0).r;
int fragmentIndex = int(floor(float(rowIndex) / float(inputFragmentRows)));
if (rowIndex != -1 && colIndex != -1) {
int index = texelFetch(indexMap, ivec2(n, out_y), 0).r;
if (index != -1) {
int rowIndex = int(floor(float(index) / float(inputFragmentCols)));
int colIndex = int(mod(float(index), float(inputFragmentCols)));
int fragmentIndex = int(floor(float(rowIndex) / float(inputFragmentRows)));
rowIndex = int(mod(float(rowIndex), float(inputFragmentRows)));
colIndex += fragmentIndex * inputFragmentCols;
sum += texelFetch(matMulResult, ivec2(colIndex + out_x, rowIndex), 0).r;
Expand Down
21 changes: 11 additions & 10 deletions src/layers/convolutional/Conv2DTranspose.glsl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,24 +4,25 @@ precision highp isampler2D;

in vec2 outTex;
uniform sampler2D matMulResult;
uniform isampler2D rowIndexMap;
uniform isampler2D colIndexMap;
uniform isampler2D indexMap;
uniform sampler2D bias;
uniform bool use_bias;
uniform int cols;
uniform int inputCols;
uniform int outputCols;
out vec4 outColor;

void main() {
int rows = textureSize(rowIndexMap, 0)[1];
int summationLength = textureSize(rowIndexMap, 0)[0];
int out_x = int(float(cols) * outTex.x);
int out_y = int(float(rows) * outTex.y);
int outputRows = textureSize(indexMap, 0)[1];
int summationLength = textureSize(indexMap, 0)[0];
int out_x = int(float(outputCols) * outTex.x);
int out_y = int(float(outputRows) * outTex.y);

float sum = 0.;
for (int n = 0; n < summationLength; ++n) {
int rowIndex = texelFetch(rowIndexMap, ivec2(n, out_y), 0).r;
int colIndex = texelFetch(colIndexMap, ivec2(n, out_y), 0).r;
if (rowIndex != -1 && colIndex != -1) {
int index = texelFetch(indexMap, ivec2(n, out_y), 0).r;
if (index != -1) {
int rowIndex = int(floor(float(index) / float(inputCols)));
int colIndex = int(mod(float(index), float(inputCols)));
sum += texelFetch(matMulResult, ivec2(colIndex + out_x, rowIndex), 0).r;
}
}
Expand Down
32 changes: 18 additions & 14 deletions src/layers/convolutional/Conv2DTranspose.js
View file
Original file line number Diff line number Diff line change
Expand Up @@ -315,7 +315,7 @@ export default class Conv2DTranspose extends Layer {
* arrays.
*/
_createIndexMap() {
if (this.rowIndexMap && this.colIndexMap) {
if (this.indexMap) {
return
}

Expand Down Expand Up @@ -384,20 +384,25 @@ export default class Conv2DTranspose extends Layer {

// combine first two dimensions
const tiledIndicesMapShape = [this.outputShape[0] * this.outputShape[1], effectiveKernelSize]
this.rowIndexMap = new Tensor([], tiledIndicesMapShape, { type: Int32Array })
this.colIndexMap = new Tensor([], tiledIndicesMapShape, { type: Int32Array })
this.indexMap = new Tensor([], tiledIndicesMapShape, { type: Int32Array })
const channelData = new Tensor([], [effectiveKernelSize], { type: Int32Array })
for (let i = 0; i < this.outputShape[0]; i++) {
for (let j = 0; j < this.outputShape[1]; j++) {
ops.assign(channelData.tensor, outputRowIndicesMap.tensor.pick(i, j, null))
ops.assign(this.rowIndexMap.tensor.pick(i * this.outputShape[1] + j, null), channelData.tensor)
ops.assign(channelData.tensor, outputColIndicesMap.tensor.pick(i, j, null))
ops.assign(this.colIndexMap.tensor.pick(i * this.outputShape[1] + j, null), channelData.tensor)
for (let k = 0; k < effectiveKernelSize; k++) {
// i * cols + j
const rowIndex = outputRowIndicesMap.tensor.get(i, j, k)
const colIndex = outputColIndicesMap.tensor.get(i, j, k)
if (rowIndex !== -1 && colIndex !== -1) {
channelData.tensor.set(k, rowIndex * this.weights['kernel'].glTextureShape[1] + colIndex)
} else {
channelData.tensor.set(k, -1)
}
}
ops.assign(this.indexMap.tensor.pick(i * this.outputShape[1] + j, null), channelData.tensor)
}
}

this.rowIndexMap.createGLTexture({ type: '2d', format: 'int', supportsTextureFragments: true })
this.colIndexMap.createGLTexture({ type: '2d', format: 'int', supportsTextureFragments: true })
this.indexMap.createGLTexture({ type: '2d', format: 'int', supportsTextureFragments: true })
}

/**
Expand Down Expand Up @@ -459,8 +464,7 @@ export default class Conv2DTranspose extends Layer {
output: this.activation === 'linear' ? this.output : this.outputPreactiv,
inputs: [
{ input: this.matMulResult, name: 'matMulResult' },
{ input: this.rowIndexMap, name: 'rowIndexMap' },
{ input: this.colIndexMap, name: 'colIndexMap' },
{ input: this.indexMap, name: 'indexMap' },
...(this.use_bias ? [{ input: this.weights['bias'], name: 'bias' }] : [])
],
uniforms: [
Expand All @@ -476,13 +480,13 @@ export default class Conv2DTranspose extends Layer {
output: this.activation === 'linear' ? this.output : this.outputPreactiv,
inputs: [
{ input: this.matMulResult, name: 'matMulResult' },
{ input: this.rowIndexMap, name: 'rowIndexMap' },
{ input: this.colIndexMap, name: 'colIndexMap' },
{ input: this.indexMap, name: 'indexMap' },
...(this.use_bias ? [{ input: this.weights['bias'], name: 'bias' }] : [])
],
uniforms: [
{ value: this.use_bias ? 1 : 0, type: 'bool', name: 'use_bias' },
{ value: this.outputShape[2], type: 'int', name: 'cols' }
{ value: this.matMulResult.glTextureShape[1], type: 'int', name: 'inputCols' },
{ value: this.outputShape[2], type: 'int', name: 'outputCols' }
],
supportsTextureFragments: true
})
Expand Down
62 changes: 17 additions & 45 deletions src/layers/convolutional/Conv3D.js
View file
Original file line number Diff line number Diff line change
Expand Up @@ -375,14 +375,13 @@ export default class Conv3D extends Layer {
* @param {Object} indicesForReshaped
*/
_createIndexMap(indicesForReshaped) {
if (this.rowIndexMap && this.colIndexMap) {
if (this.indexMap) {
return
}

let [inputDim1, inputDim2, inputDim3, inputChannels] = this.inputShape

const indicesRow = new Tensor(indicesForReshaped.row.data, indicesForReshaped.row.shape, { type: Int32Array })
const indicesCol = new Tensor(indicesForReshaped.col.data, indicesForReshaped.col.shape, { type: Int32Array })
const indices = new Tensor(indicesForReshaped.data, indicesForReshaped.shape, { type: Int32Array })

// padding for border mode 'same'
if (this.padding === 'same') {
Expand All @@ -398,8 +397,7 @@ export default class Conv3D extends Layer {
inputDim2 = inputDim2 + paddingDim2Before + paddingDim2After
inputDim3 = inputDim3 + paddingDim3Before + paddingDim3After
const padValue = -1
this._padInput(indicesRow, padValue)
this._padInput(indicesCol, padValue)
this._padInput(indices, padValue)
}

const kernelDim1 = this.kernelShape[1]
Expand All @@ -416,38 +414,27 @@ export default class Conv3D extends Layer {
const kernelDim2Dilated = kernelDim2 + (kernelDim2 - 1) * (this.dilationRate[1] - 1)
const kernelDim3Dilated = kernelDim3 + (kernelDim3 - 1) * (this.dilationRate[2] - 1)

this.rowIndexMap = new Tensor([], [nbPatches, patchLen], { type: Int32Array })
this.colIndexMap = new Tensor([], [nbPatches, patchLen], { type: Int32Array })
this.indexMap = new Tensor([], [nbPatches, patchLen], { type: Int32Array })

const indicesRowPatch = new Tensor([], [kernelDim1, kernelDim2, kernelDim3, inputChannels])
const indicesColPatch = new Tensor([], [kernelDim1, kernelDim2, kernelDim3, inputChannels])
const indicesPatch = new Tensor([], [kernelDim1, kernelDim2, kernelDim3, inputChannels])
let offset = 0
for (let i = 0, limit = inputDim1 - kernelDim1Dilated; i <= limit; i += this.strides[0]) {
for (let j = 0, limit = inputDim2 - kernelDim2Dilated; j <= limit; j += this.strides[1]) {
for (let k = 0, limit = inputDim3 - kernelDim3Dilated; k <= limit; k += this.strides[2]) {
ops.assign(
indicesRowPatch.tensor,
indicesRow.tensor
indicesPatch.tensor,
indices.tensor
.hi(i + kernelDim1Dilated, j + kernelDim2Dilated, k + kernelDim3Dilated, inputChannels)
.lo(i, j, k, 0)
.step(this.dilationRate[0], this.dilationRate[1], this.dilationRate[2], 1)
)
ops.assign(
indicesColPatch.tensor,
indicesCol.tensor
.hi(i + kernelDim1Dilated, j + kernelDim2Dilated, k + kernelDim3Dilated, inputChannels)
.lo(i, j, k, 0)
.step(this.dilationRate[0], this.dilationRate[1], this.dilationRate[2], 1)
)
this.rowIndexMap.tensor.data.set(indicesRowPatch.tensor.data, offset)
this.colIndexMap.tensor.data.set(indicesColPatch.tensor.data, offset)
this.indexMap.tensor.data.set(indicesPatch.tensor.data, offset)
offset += patchLen
}
}
}

this.rowIndexMap.createGLTexture({ type: '2d', format: 'int', supportsTextureFragments: true })
this.colIndexMap.createGLTexture({ type: '2d', format: 'int', supportsTextureFragments: true })
this.indexMap.createGLTexture({ type: '2d', format: 'int', supportsTextureFragments: true })
}

/**
Expand All @@ -471,35 +458,20 @@ export default class Conv3D extends Layer {
if (x.is2DReshaped || x.is2DSquareReshaped) {
this._createIndexMap(x.indicesForReshaped)
if (!this.mappedInput) {
this.mappedInput = new Tensor([], this.rowIndexMap.glTextureShape)
this.mappedInput = new Tensor([], this.indexMap.glTextureShape)
this.mappedInput.createGLTexture({ type: '2d', format: 'float', supportsTextureFragments: true })
}

if (x.glTextureFragments) {
x.convert2DRowFragmentedGLTextureToColStack()
webgl2.runProgram({
program: this.mapInputFragmentsProgram,
output: this.mappedInput,
inputs: [
{ input: x, name: 'x' },
{ input: this.rowIndexMap, name: 'rowIndexMap' },
{ input: this.colIndexMap, name: 'colIndexMap' }
],
uniforms: [{ value: x.glTextureShape[1], type: 'int', name: 'fragmentCols' }],
supportsTextureFragments: true
})
} else {
webgl2.runProgram({
program: this.mapInputProgram,
output: this.mappedInput,
inputs: [
{ input: x, name: 'x' },
{ input: this.rowIndexMap, name: 'rowIndexMap' },
{ input: this.colIndexMap, name: 'colIndexMap' }
],
supportsTextureFragments: true
})
}
webgl2.runProgram({
program: x.glTextureFragments ? this.mapInputFragmentsProgram : this.mapInputProgram,
output: this.mappedInput,
inputs: [{ input: x, name: 'x' }, { input: this.indexMap, name: 'indexMap' }],
uniforms: [{ value: x.glTextureShape[1], type: 'int', name: 'inputCols' }],
supportsTextureFragments: true
})
}

const input = x.is2DReshaped || x.is2DSquareReshaped ? this.mappedInput : this.volColsMat
Expand Down
Loading

0 comments on commit 53b07ec

Please sign in to comment.