Skip to content

Commit 64814ca

Browse files
author
Nikhil Thorat
authored
cache uniform locations (tensorflow#127)
1 parent 101bf24 commit 64814ca

14 files changed

+114
-44
lines changed

demos/benchmarks/benchmark-demo.html

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -14,16 +14,17 @@
1414
==============================================================================-->
1515
<!DOCTYPE html>
1616
<html>
17+
1718
<head>
1819
<meta charset="UTF-8">
1920
<script src="../../bower_components/webcomponentsjs/webcomponents-lite.min.js"></script>
2021
<link rel="import" href="math-benchmark.html">
2122
<link rel="import" href="../demo-header.html">
2223
<link rel="import" href="../demo-footer.html">
23-
<link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/dialog-polyfill/0.4.8/dialog-polyfill.min.css"/>
24+
<link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/dialog-polyfill/0.4.8/dialog-polyfill.min.css" />
2425
<script type="text/javascript" src="https://cdnjs.cloudflare.com/ajax/libs/dialog-polyfill/0.4.8/dialog-polyfill.min.js"></script>
2526
<script src="../homepage/assets/support.js"></script>
26-
<title>LearnJS Benchmarks</title>
27+
<title>deeplearn.js benchmark</title>
2728
<style>
2829
html {
2930
width: 100%;
@@ -36,15 +37,18 @@
3637
width: 100%;
3738
height: 100%;
3839
}
40+
3941
head {
4042
height: 0;
4143
width: 0;
4244
}
4345
</style>
4446
</head>
47+
4548
<body>
4649
<demo-header name="benchmarks"></demo-header>
4750
<math-benchmark></math-benchmark>
4851
<demo-footer></demo-footer>
4952
</body>
53+
5054
</html>

demos/imagenet/imagenet-demo.html

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -14,16 +14,17 @@
1414
==============================================================================-->
1515
<!DOCTYPE html>
1616
<html>
17+
1718
<head>
1819
<meta charset="UTF-8">
1920
<script src="../../bower_components/webcomponentsjs/webcomponents-lite.min.js"></script>
2021
<link rel="import" href="imagenet.html">
2122
<link rel="import" href="../demo-header.html">
2223
<link rel="import" href="../demo-footer.html">
23-
<link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/dialog-polyfill/0.4.8/dialog-polyfill.min.css"/>
24+
<link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/dialog-polyfill/0.4.8/dialog-polyfill.min.css" />
2425
<script type="text/javascript" src="https://cdnjs.cloudflare.com/ajax/libs/dialog-polyfill/0.4.8/dialog-polyfill.min.js"></script>
2526
<script src="../homepage/assets/support.js"></script>
26-
<title>LearnJS Imagenet Demo</title>
27+
<title>deeplearn.js imagenet webcam demo</title>
2728
<style>
2829
html {
2930
width: 100%;
@@ -36,15 +37,18 @@
3637
width: 100%;
3738
height: 100%;
3839
}
40+
3941
head {
4042
height: 0;
4143
width: 0;
4244
}
4345
</style>
4446
</head>
47+
4548
<body>
4649
<demo-header name="webcam classifier"></demo-header>
4750
<imagenet-demo></imagenet-demo>
4851
<demo-footer></demo-footer>
4952
</body>
53+
5054
</html>

demos/model-builder/model-builder-demo.html

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
==============================================================================-->
1515
<!DOCTYPE html>
1616
<html>
17+
1718
<head>
1819
<meta charset="UTF-8">
1920
<script src="../../bower_components/webcomponentsjs/webcomponents-lite.min.js"></script>
@@ -22,10 +23,10 @@
2223
<link rel="import" href="model-builder.html">
2324
<link rel="import" href="../../bower_components/iron-icons/image-icons.html">
2425
<link rel="import" href="../../bower_components/iron-icons/iron-icons.html">
25-
<link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/dialog-polyfill/0.4.8/dialog-polyfill.min.css"/>
26+
<link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/dialog-polyfill/0.4.8/dialog-polyfill.min.css" />
2627
<script type="text/javascript" src="https://cdnjs.cloudflare.com/ajax/libs/dialog-polyfill/0.4.8/dialog-polyfill.min.js"></script>
2728
<script src="../homepage/assets/support.js"></script>
28-
<title>LearnJS Demo Page</title>
29+
<title>deeplearn.js model builder demo</title>
2930
<style>
3031
html {
3132
width: 100%;
@@ -38,15 +39,18 @@
3839
width: 100%;
3940
height: 100%;
4041
}
42+
4143
head {
4244
height: 0;
4345
width: 0;
4446
}
4547
</style>
4648
</head>
49+
4750
<body>
4851
<demo-header name="model builder"></demo-header>
4952
<model-builder></model-builder>
5053
<demo-footer></demo-footer>
5154
</body>
55+
5256
</html>

demos/models/imagenet_util.ts

Lines changed: 22 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -64,7 +64,9 @@ export function preprocessInput(
6464
shapeRowCol: [number, number]) {
6565
gpgpu.setOutputMatrixTexture(resultTex, shapeRowCol[0], shapeRowCol[1]);
6666
gpgpu.setProgram(preprocessInputShader);
67-
gpgpu.setInputMatrixTexture(sourceTex, 'source', 0);
67+
const samplerLocation = webgl_util.getProgramUniformLocationOrThrow(
68+
gpgpu.gl, preprocessInputShader, 'source');
69+
gpgpu.setInputMatrixTexture(sourceTex, samplerLocation, 0);
6870
gpgpu.executeProgram();
6971
}
7072

@@ -142,20 +144,32 @@ export function renderGrayscaleChannelsCollage(
142144
imageSize: number, channels: number, textureSize: number, numRows: number) {
143145
webgl_util.bindCanvasToFramebuffer(gpgpu.gl);
144146
gpgpu.setProgram(unpackChannelsShader);
145-
gpgpu.setInputMatrixTexture(sourceTex, 'source', 0);
146-
gpgpu.setInputMatrixTexture(minValuesTex, 'minValues', 1);
147-
gpgpu.setInputMatrixTexture(maxValuesTex, 'maxValues', 2);
148147

149-
const imageSizeLoc = gpgpu.getUniformLocation('imageSize');
148+
const sourceSamplerLocation = webgl_util.getProgramUniformLocationOrThrow(
149+
gpgpu.gl, unpackChannelsShader, 'source');
150+
const minValuesSamplerLocation = webgl_util.getProgramUniformLocationOrThrow(
151+
gpgpu.gl, unpackChannelsShader, 'minValues');
152+
const maxValuesSamplerLocation = webgl_util.getProgramUniformLocationOrThrow(
153+
gpgpu.gl, unpackChannelsShader, 'maxValues');
154+
155+
gpgpu.setInputMatrixTexture(sourceTex, sourceSamplerLocation, 0);
156+
gpgpu.setInputMatrixTexture(minValuesTex, minValuesSamplerLocation, 1);
157+
gpgpu.setInputMatrixTexture(maxValuesTex, maxValuesSamplerLocation, 2);
158+
159+
const imageSizeLoc =
160+
gpgpu.getUniformLocation(unpackChannelsShader, 'imageSize');
150161
gpgpu.gl.uniform1f(imageSizeLoc, imageSize);
151162

152-
const channelsLoc = gpgpu.getUniformLocation('channels');
163+
const channelsLoc =
164+
gpgpu.getUniformLocation(unpackChannelsShader, 'channels');
153165
gpgpu.gl.uniform1f(channelsLoc, channels);
154166

155-
const imagesPerRowLoc = gpgpu.getUniformLocation('imagesPerRow');
167+
const imagesPerRowLoc =
168+
gpgpu.getUniformLocation(unpackChannelsShader, 'imagesPerRow');
156169
gpgpu.gl.uniform1f(imagesPerRowLoc, Math.floor(textureSize / imageSize));
157170

158-
const inputShapeCRLoc = gpgpu.getUniformLocation('inputShapeCR');
171+
const inputShapeCRLoc =
172+
gpgpu.getUniformLocation(unpackChannelsShader, 'inputShapeCR');
159173
gpgpu.gl.uniform2f(inputShapeCRLoc, inputShapeRC[1], inputShapeRC[0]);
160174

161175
gpgpu.executeProgram();

demos/nn-art/nn-art-demo.html

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
==============================================================================-->
1515
<!DOCTYPE html>
1616
<html>
17+
1718
<head>
1819
<meta charset="UTF-8">
1920
<!-- This script is stripped during vulcanization, available only during development. It provides
@@ -26,11 +27,11 @@
2627
<link rel="import" href="../../bower_components/iron-icons/image-icons.html">
2728
<link rel="import" href="../../bower_components/iron-icons/iron-icons.html">
2829

29-
<link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/dialog-polyfill/0.4.8/dialog-polyfill.min.css"/>
30+
<link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/dialog-polyfill/0.4.8/dialog-polyfill.min.css" />
3031
<script type="text/javascript" src="https://cdnjs.cloudflare.com/ajax/libs/dialog-polyfill/0.4.8/dialog-polyfill.min.js"></script>
3132
<script src="../homepage/assets/support.js"></script>
3233

33-
<title>Abstract CPPN Art</title>
34+
<title>deeplearn.js abstract CPPN art demo</title>
3435
<style>
3536
html {
3637
width: 100%;
@@ -43,15 +44,18 @@
4344
width: 100%;
4445
height: 100%;
4546
}
47+
4648
head {
4749
height: 0;
4850
width: 0;
4951
}
5052
</style>
5153
</head>
54+
5255
<body>
5356
<demo-header name="cppn"></demo-header>
5457
<nn-art></nn-art>
5558
<demo-footer></demo-footer>
5659
</body>
60+
5761
</html>

demos/nn-art/nn_art_util.ts

Lines changed: 12 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -66,8 +66,12 @@ export function addLatentVariables(
6666
z2: number) {
6767
gpgpu.setOutputMatrixTexture(resultTex, shapeRowCol[0], shapeRowCol[1]);
6868
gpgpu.setProgram(addZShader);
69-
gpgpu.setInputMatrixTexture(sourceTex, 'source', 0);
70-
const zLoc = gpgpu.getUniformLocation('z');
69+
70+
const sourceSamplerLocation = webgl_util.getProgramUniformLocationOrThrow(
71+
gpgpu.gl, addZShader, 'source');
72+
gpgpu.setInputMatrixTexture(sourceTex, sourceSamplerLocation, 0);
73+
74+
const zLoc = gpgpu.getUniformLocation(addZShader, 'z');
7175
gpgpu.gl.uniform2f(zLoc, z1, z2);
7276
gpgpu.executeProgram();
7377
}
@@ -149,11 +153,14 @@ export function render(
149153
outputNumDimensions: number, colorMode: number) {
150154
webgl_util.bindCanvasToFramebuffer(gpgpu.gl);
151155
gpgpu.setProgram(renderShader);
152-
gpgpu.setInputMatrixTexture(sourceTex, 'source', 0);
153-
const colorModeLoc = gpgpu.getUniformLocation('colorMode');
156+
157+
const sourceSamplerLocation = webgl_util.getProgramUniformLocationOrThrow(
158+
gpgpu.gl, renderShader, 'source');
159+
gpgpu.setInputMatrixTexture(sourceTex, sourceSamplerLocation, 0);
160+
const colorModeLoc = gpgpu.getUniformLocation(renderShader, 'colorMode');
154161
gpgpu.gl.uniform1i(colorModeLoc, colorMode);
155162
const outputNumDimensionsLoc =
156-
gpgpu.getUniformLocation('outputNumDimensions');
163+
gpgpu.getUniformLocation(renderShader, 'outputNumDimensions');
157164
gpgpu.gl.uniform1f(outputNumDimensionsLoc, outputNumDimensions);
158165
gpgpu.executeProgram();
159166
}

src/math/math_gpu.ts

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -128,7 +128,8 @@ export class NDArrayMathGPU extends NDArrayMath {
128128

129129
private compileAndRun<T extends NDArray, K extends NDArray>(
130130
program: GPGPUProgram, inputs: T[], output?: K,
131-
customSetup?: (gpgpu: GPGPUContext) => void): K {
131+
customSetup?: (gpgpu: GPGPUContext, webGLProgram: WebGLProgram) => void):
132+
K {
132133
if (output == null) {
133134
output = this.makeOutputArray<K>(program.outputShape);
134135
}

src/math/webgl/conv_gpu.ts

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,12 +19,16 @@ import {ConvInfo} from '../conv_util';
1919
import {GPGPUProgram} from './gpgpu_math';
2020

2121
export class Conv2DProgram implements GPGPUProgram {
22-
variableNames = ['x', 'W', 'bias'];
22+
variableNames = ['x', 'W'];
2323
params: Array<{}>;
2424
outputShape: number[];
2525
userCode: string;
2626

2727
constructor(convInfo: ConvInfo, hasBias: boolean) {
28+
if (hasBias) {
29+
this.variableNames.push('bias');
30+
}
31+
2832
this.outputShape = convInfo.outShape;
2933
const biasSnippet = hasBias ? 'dotProd += getBias(d2);' : '';
3034
const [xNumRows, xNumCols, inputDepth] = convInfo.inShape;

src/math/webgl/copy_gpu.ts

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -44,12 +44,14 @@ export class Copy2DProgram implements GPGPUProgram {
4444
getCustomSetupFunc(
4545
sourceStart: [number, number], destStart: [number, number],
4646
destSize: [number, number]) {
47-
return (gpgpu: GPGPUContext) => {
47+
return (gpgpu: GPGPUContext, webGLProgram: WebGLProgram) => {
4848
gpgpu.setOutputMatrixWriteRegion(
4949
destStart[0], destSize[0], destStart[1], destSize[1]);
50-
const sourceStartCRLoc = gpgpu.getUniformLocation('sourceStart');
50+
const sourceStartCRLoc =
51+
gpgpu.getUniformLocation(webGLProgram, 'sourceStart');
5152
gpgpu.gl.uniform2i(sourceStartCRLoc, sourceStart[0], sourceStart[1]);
52-
const destStartCRLoc = gpgpu.getUniformLocation('destStart');
53+
const destStartCRLoc =
54+
gpgpu.getUniformLocation(webGLProgram, 'destStart');
5355
gpgpu.gl.uniform2i(destStartCRLoc, destStart[0], destStart[1]);
5456
};
5557
}

src/math/webgl/gpgpu_context.ts

Lines changed: 9 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -52,8 +52,9 @@ export class GPGPUContext {
5252
webgl_util.getExtensionOrThrow(this.gl, 'EXT_color_buffer_float');
5353
}
5454

55-
this.loseContextExtension = webgl_util.getExtensionOrThrow(
56-
this.gl, 'WEBGL_lose_context') as WebGLLoseContextExtension;
55+
this.loseContextExtension =
56+
webgl_util.getExtensionOrThrow(this.gl, 'WEBGL_lose_context') as
57+
WebGLLoseContextExtension;
5758
this.vertexBuffer = gpgpu_util.createVertexBuffer(this.gl);
5859
this.indexBuffer = gpgpu_util.createIndexBuffer(this.gl);
5960
this.framebuffer = webgl_util.createFramebuffer(this.gl);
@@ -188,20 +189,21 @@ export class GPGPUContext {
188189
webgl_util.callAndCheck(this.gl, () => this.gl.useProgram(program));
189190
}
190191

191-
public getUniformLocation(uniformName: string): WebGLUniformLocation {
192+
public getUniformLocation(program: WebGLProgram, uniformName: string):
193+
WebGLUniformLocation {
192194
this.throwIfDisposed();
193-
this.throwIfNoProgram();
194195
return webgl_util.getProgramUniformLocationOrThrow(
195-
this.gl, this.program, uniformName);
196+
this.gl, program, uniformName);
196197
}
197198

198199
public setInputMatrixTexture(
199-
inputMatrixTexture: WebGLTexture, uniformName: string,
200+
inputMatrixTexture: WebGLTexture, uniformLocation: WebGLUniformLocation,
200201
textureUnit: number) {
201202
this.throwIfDisposed();
202203
this.throwIfNoProgram();
203204
webgl_util.bindTextureToProgramUniformSampler(
204-
this.gl, this.program, inputMatrixTexture, uniformName, textureUnit);
205+
this.gl, this.program, inputMatrixTexture, uniformLocation,
206+
textureUnit);
205207
}
206208

207209
public setOutputMatrixTexture(

src/math/webgl/gpgpu_math.ts

Lines changed: 22 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@ export interface GPGPUProgram {
3333
export interface GPGPUBinary {
3434
webGLProgram: WebGLProgram;
3535
program: GPGPUProgram;
36+
uniformLocations: {[name: string]: WebGLUniformLocation};
3637
gpgpu: GPGPUContext;
3738
source: string;
3839
inShapeInfos: ShapeInfo[];
@@ -58,10 +59,24 @@ export function compileProgram<T extends NDArray, K extends NDArray>(
5859
const source = shader_compiler.makeShader(
5960
inputInfos, outShapeInfo, userCode,
6061
program.supportsBroadcasting === true);
62+
63+
const webGLProgram = gpgpu.createProgram(source);
64+
65+
const uniformLocations: {[name: string]: WebGLUniformLocation} = {};
66+
for (let i = 0; i < program.variableNames.length; i++) {
67+
const uniformName = program.variableNames[i];
68+
uniformLocations[uniformName] =
69+
gpgpu.getUniformLocation(webGLProgram, uniformName);
70+
}
71+
6172
return {
6273
program,
6374
source,
64-
webGLProgram: gpgpu.createProgram(source), gpgpu, inShapeInfos, outShapeInfo
75+
webGLProgram,
76+
uniformLocations,
77+
gpgpu,
78+
inShapeInfos,
79+
outShapeInfo
6580
};
6681
}
6782

@@ -93,7 +108,8 @@ function validateBinaryAndProgram(shapeInfos: ShapeInfo[], inputs: NDArray[]) {
93108

94109
export function runProgram<T extends NDArray, K extends NDArray>(
95110
binary: GPGPUBinary, inputs: T[], output: K,
96-
customSetup?: (gpgpu: GPGPUContext) => void): void {
111+
customSetup?: (gpgpu: GPGPUContext, webGLProgram: WebGLProgram) =>
112+
void): void {
97113
validateBinaryAndProgram(binary.inShapeInfos, inputs);
98114
validateBinaryAndProgram([binary.outShapeInfo], [output]);
99115

@@ -104,10 +120,12 @@ export function runProgram<T extends NDArray, K extends NDArray>(
104120
gpgpu.setProgram(binary.webGLProgram);
105121
inputs.forEach((input, i) => {
106122
const tex = input.getTexture();
107-
gpgpu.setInputMatrixTexture(tex, binary.program.variableNames[i], i);
123+
const variableName = binary.program.variableNames[i];
124+
const variableUniformLocation = binary.uniformLocations[variableName];
125+
gpgpu.setInputMatrixTexture(tex, variableUniformLocation, i);
108126
});
109127
if (customSetup != null) {
110-
customSetup(gpgpu);
128+
customSetup(gpgpu, binary.webGLProgram);
111129
}
112130
gpgpu.executeProgram();
113131
}

0 commit comments

Comments
 (0)