10000 upload strides and packedTexShape as uniforms · tensorflow/tfjs-core@94688af · GitHub
[go: up one dir, main page]
Skip to content
This repository was archived by the owner on Aug 15, 2019. It is now read-only.

Commit 94688af

Browse files
tafsiritafsiri
committed
upload strides and packedTexShape as uniforms
1 parent 13e3dec commit 94688af

File tree

2 files changed

+60
-11
lines changed

2 files changed

+60
-11
lines changed

‎src/backends/webgl/gpgpu_math.ts

Lines changed: 43 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -125,6 +125,9 @@ export function compileProgram<T extends Tensor, K extends Tensor>(
125125
if (inShapeInfo.logicalShape.length > 0) {
126126
const varShapeName = `shape${inputInfo.name}`;
127127
const varTexShapeName = `texShape${inputInfo.name}`;
128+
const varStridesName = `strides${inputInfo.name}`;
129+
const varPackedTexShapeName = `packedTexShape${inputInfo.name}`;
130+
128131
const shouldThrow = false;
129132
const shapeLocation =
130133
gpgpu.getUniformLocation(webGLProgram, varShapeName, shouldThrow);
@@ -137,6 +140,18 @@ export function compileProgram<T extends Tensor, K extends Tensor>(
137140
if (texShapeLocation != null) {
138141
uniformLocations[varTexShapeName] = texShapeLocation;
139142
}
143+
144+
const varStridesLocation =
145+
gpgpu.getUniformLocation(webGLProgram, varStridesName, shouldThrow);
146+
if (varStridesLocation != null) {
147+
uniformLocations[varStridesName] = varStridesLocation;
148+
}
149+
150+
const varPackedTexShapeLocation = gpgpu.getUniformLocation(
151+
webGLProgram, varPackedTexShapeName, shouldThrow);
152+
if (varPackedTexShapeLocation != null) {
153+
uniformLocations[varPackedTexShapeName] = varPackedTexShapeLocation;
154+
}
140155
}
141156
}
142157

@@ -255,14 +270,12 @@ export function runProgram<T extends Tensor, K extends Tensor>(
255270
}
256271
return;
257272
}
258-
// Upload the shape/texShape information as a uniform as well.
259-
const varShapeName = `shape${varName}`;
260-
const varShapeLoc = binary.uniformLocations[varShapeName];
273+
// Upload shape information as uniform.
274+
const varShapeLoc = binary.uniformLocations[`shape${varName}`];
261275
if (varShapeLoc != null) {
262276
let shape: number[]|Int32Array;
263277
if (binary.program.usesPackedTextures) {
264278
shape = util.packedShapeTransform(input.shape);
265-
// TODO yassogba@ should anything special happen for isPackShader?
266279
} else {
267280
// Call squeezeShape to match the shape used in the shader program
268281
const {newShape} = util.squeezeShape(input.shape);
@@ -275,8 +288,22 @@ export function runProgram<T extends Tensor, K extends Tensor>(
275288
gpgpu.gl.uniform1iv(varShapeLoc, shape);
276289
}
277290

278-
const varTexShapeName = `texShape${varName}`;
279-
const varTexShapeLoc = binary.uniformLocations[varTexShapeName];
291+
// Upload precomputed strides
292+
const varStridesLoc = binary.uniformLocations[`strides${varName}`];
293+
if (varStridesLoc != null) {
294+
let strides: number[]|Int32Array;
295+
const {newShape} = util.squeezeShape(input.shape);
296+
strides = util.computeStrides(newShape);
297+
298+
if (!(strides instanceof Int32Array)) {
299+
strides = new Int32Array(strides);
300+
}
301+
gpgpu.gl.uniform1iv(varStridesLoc, strides);
302+
}
303+
304+
305+
// Upload texShape/packedTexShape information as uniform.
306+
const varTexShapeLoc = binary.uniformLocations[`texShape${varName}`];
280307
// TODO(yassogba, nsthoat) rename/document these two shapes:
281308
// input.texData.shape and input.texData.texShape
282309
// to make it more apparent why they are both needed.
@@ -285,6 +312,15 @@ export function runProgram<T extends Tensor, K extends Tensor>(
285312
texShape = new Int32Array(texShape);
286313
}
287314
gpgpu.gl.uniform1iv(varTexShapeLoc, texShape);
315+
if (binary.program.usesPackedTextures) {
316+
const varPackedTexShapeLoc =
317+
binary.uniformLocations[`packedTexShape${varName}`];
318+
if (varPackedTexShapeLoc != null) {
319+
const packedTexShape =
320+
[Math.ceil(texShape[0] / 2), Math.ceil(texShape[1] / 2)];
321+
gpgpu.gl.uniform1iv(varPackedTexShapeLoc, packedTexShape);
322+
}
323+
}
288324

289325
// If the input was sliced, upload the flat offset index.
290326
if (input.texData.slice != null && varOffsetLoc != null) {
@@ -295,6 +331,7 @@ export function runProgram<T extends Tensor, K extends Tensor>(
295331
});
296332

297333
// Upload output shape uniforms
334+
// TODO yassogba upload outputStrides and outputPackedTexShape
298335
if (output.shape.length > 0) {
299336
const outputShapeName = `outputShape`;
300337
const outputShapeLoc = binary.uniformLocations[outputShapeName];

‎src/backends/webgl/shader_compiler.ts

Lines changed: 17 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -159,6 +159,8 @@ function getInputSamplingSnippet(
159159
const inputName = inInfo.name;
160160
const shapeUniform = `shape${inputName}`;
161161
const texShapeUniform = `texShape${inputName}`;
162+
const stridesUniform = `strides${inputName}`;
163+
const packedTexShapeUniform = `packedTexShape${inputName}`;
162164

163165
let sampler;
164166
if (usesPackedTextures) {
@@ -173,10 +175,16 @@ function getInputSamplingSnippet(
173175
// for a good cache hit).
174176
if (inInfo.shapeInfo.logicalShape.length > 0) {
175177
if (sampler.indexOf(shapeUniform) !== -1) {
176-
res += `uniform int shape${inputName}[${MAX_TENSOR_RANK}]; \n`;
178+
res += `uniform int ${shapeUniform}[${MAX_TENSOR_RANK}]; \n`;
179+
}
180+
if (sampler.indexOf(stridesUniform) !== -1) {
181+
res += `uniform int ${stridesUniform}[${MAX_TENSOR_RANK - 1}]; \n`;
177182
}
178183
if (sampler.indexOf(texShapeUniform) !== -1) {
179-
res += `uniform int texShape${inputName}[${TEXSHAPE_RANK}]; \n`;
184+
res += `uniform int ${texShapeUniform}[${TEXSHAPE_RANK}]; \n`;
185+
}
186+
if (sampler.indexOf(packedTexShapeUniform) !== -1) {
187+
res += `uniform int ${packedTexShapeUniform}[${TEXSHAPE_RANK}]; \n`;
180188
}
181189
}
182190
res += sampler;
@@ -807,15 +815,19 @@ function getPackedSampler2D(inputInfo: InputInfo): string {
807815
}
808816

809817
const shapeUniform = `shape${texName}`;
818+
const packedTexShapeUniform = `packedTexShape${texName}`;
810819
return `
811820
vec4 ${funcName}(int row, int col) {
812-
int packedTexShapeR = int(ceil(float(${texShapeUniform}[0]) * 0.5));
813-
int packedTexShapeC = int(ceil(float(${texShapeUniform}[1]) * 0.5));
821+
//int packedTexShapeR = int(ceil(float(${texShapeUniform}[0]) * 0.5));
822+
//int packedTexShapeC = int(ceil(float(${texShapeUniform}[1]) * 0.5));
823+
824+
int packedTexShapeR = ${packedTexShapeUniform}[0];
825+
int packedTexShapeC = ${packedTexShapeUniform}[1];
814826
815827
int valuesPerRow = int(ceil(float(${shapeUniform}[1]) * 0.5));
816828
817829
vec2 uv = packedUVfrom2D(valuesPerRow,
818-
packedTexShapeR, packedTexShapeC, row, col);
830+
${packedTexShapeUniform}[0], packedTexShapeC, row, col);
819831
return ${glsl.texture2D}(${texName}, uv);
820832
}
821833
`;

0 commit comments

Comments
 (0)
0