Skip to content

Commit f6db363

Browse files
authored
Merge branch 'master' into parallel_compilation
2 parents 643c81e + 745ecca commit f6db363

File tree

150 files changed

+908
-1079
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

150 files changed

+908
-1079
lines changed

tfjs-backend-webgl/src/backend_webgl.ts

Lines changed: 16 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,7 @@ import {simpleAbsImplCPU} from './kernel_utils/shared';
3535
import {PackProgram} from './pack_gpu';
3636
import {ReshapePackedProgram} from './reshape_packed_gpu';
3737
import * as tex_util from './tex_util';
38-
import {TextureData, TextureUsage} from './tex_util';
38+
import {Texture, TextureData, TextureUsage} from './tex_util';
3939
import {TextureManager} from './texture_manager';
4040
import * as unary_op from './unaryop_gpu';
4141
import {UnaryOpProgram} from './unaryop_gpu';
@@ -328,7 +328,7 @@ export class MathBackendWebGL extends KernelBackend {
328328
const tmpData = this.texData.get(tmpDownloadTarget.dataId);
329329

330330
buffer = this.gpgpu.createBufferFromTexture(
331-
tmpData.texture, ...tex_util.getDenseTexShape(shape));
331+
tmpData.texture.texture, ...tex_util.getDenseTexShape(shape));
332332
}
333333

334334
this.pendingRead.set(dataId, []);
@@ -419,10 +419,11 @@ export class MathBackendWebGL extends KernelBackend {
419419
if (env().getBool('WEBGL_DOWNLOAD_FLOAT_ENABLED')) {
420420
const tmpTarget = this.decode(dataId);
421421
const tmpData = this.texData.get(tmpTarget.dataId);
422-
const vals = this.gpgpu
423-
.downloadMatrixFromPackedTexture(
424-
tmpData.texture, ...tex_util.getDenseTexShape(shape))
425-
.subarray(0, size);
422+
const vals =
423+
this.gpgpu
424+
.downloadMatrixFromPackedTexture(
425+
tmpData.texture.texture, ...tex_util.getDenseTexShape(shape))
426+
.subarray(0, size);
426427

427428
this.disposeIntermediateTensorInfo(tmpTarget);
428429

@@ -439,11 +440,11 @@ export class MathBackendWebGL extends KernelBackend {
439440
const output = this.runWebGLProgram(
440441
program, [{shape: outputShape, dtype, dataId}], 'float32');
441442
const tmpData = this.texData.get(output.dataId);
442-
const vals =
443-
this.gpgpu
444-
.downloadByteEncodedFloatMatrixFromOutputTexture(
445-
tmpData.texture, tmpData.texShape[0], tmpData.texShape[1])
446-
.subarray(0, size);
443+
const vals = this.gpgpu
444+
.downloadByteEncodedFloatMatrixFromOutputTexture(
445+
tmpData.texture.texture, tmpData.texShape[0],
446+
tmpData.texShape[1])
447+
.subarray(0, size);
447448
this.disposeIntermediateTensorInfo(output);
448449

449450
return vals;
@@ -618,7 +619,7 @@ export class MathBackendWebGL extends KernelBackend {
618619

619620
getTexture(dataId: DataId): WebGLTexture {
620621
this.uploadToGPU(dataId);
621-
return this.texData.get(dataId).texture;
622+
return this.texData.get(dataId).texture.texture;
622623
}
623624

624625
/**
@@ -1001,6 +1002,8 @@ export class MathBackendWebGL extends KernelBackend {
10011002

10021003
let texShape = texData.texShape;
10031004
if (texShape == null) {
1005+
// This texShape may not be the final texture shape. For packed or dense
1006+
// textures, the texShape will be changed when textures are created.
10041007
texShape = webgl_util.getTextureShapeFromLogicalShape(shape, isPacked);
10051008
texData.texShape = texShape;
10061009
}
@@ -1094,7 +1097,7 @@ export class MathBackendWebGL extends KernelBackend {
10941097

10951098
private acquireTexture(
10961099
texShape: [number, number], texType: TextureUsage, dtype: DataType,
1097-
isPacked: boolean): WebGLTexture {
1100+
isPacked: boolean): Texture {
10981101
this.numBytesInGPU += this.computeBytes(texShape, dtype);
10991102
if (!this.warnedAboutMemory &&
11001103
this.numBytesInGPU > this.numMBBeforeWarning * 1024 * 1024) {

tfjs-backend-webgl/src/gpgpu_context.ts

Lines changed: 6 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@ import {env, PixelData, TypedArray, util} from '@tensorflow/tfjs-core';
2020
import {getWebGLContext, setWebGLContext} from './canvas_util';
2121
import * as gpgpu_util from './gpgpu_util';
2222
import * as tex_util from './tex_util';
23-
import {TextureConfig} from './tex_util';
23+
import {Texture, TextureConfig} from './tex_util';
2424
import {WebGL1DisjointQueryTimerExtension, WebGL2DisjointQueryTimerExtension, WebGLParallelCompilationExtension} from './webgl_types';
2525
import * as webgl_util from './webgl_util';
2626

@@ -138,22 +138,20 @@ export class GPGPUContext {
138138
this.disposed = true;
139139
}
140140

141-
public createFloat32MatrixTexture(rows: number, columns: number):
142-
WebGLTexture {
141+
public createFloat32MatrixTexture(rows: number, columns: number): Texture {
143142
this.throwIfDisposed();
144143
return gpgpu_util.createFloat32MatrixTexture(
145144
this.gl, rows, columns, this.textureConfig);
146145
}
147146

148-
public createFloat16MatrixTexture(rows: number, columns: number):
149-
WebGLTexture {
147+
public createFloat16MatrixTexture(rows: number, columns: number): Texture {
150148
this.throwIfDisposed();
151149
return gpgpu_util.createFloat16MatrixTexture(
152150
this.gl, rows, columns, this.textureConfig);
153151
}
154152

155153
public createUnsignedBytesMatrixTexture(rows: number, columns: number):
156-
WebGLTexture {
154+
Texture {
157155
this.throwIfDisposed();
158156
return gpgpu_util.createUnsignedBytesMatrixTexture(
159157
this.gl, rows, columns, this.textureConfig);
@@ -175,14 +173,13 @@ export class GPGPUContext {
175173
}
176174

177175
public createFloat16PackedMatrixTexture(rows: number, columns: number):
178-
WebGLTexture {
176+
Texture {
179177
this.throwIfDisposed();
180178
return gpgpu_util.createFloat16PackedMatrixTexture(
181179
this.gl, rows, columns, this.textureConfig);
182180
}
183181

184-
public createPackedMatrixTexture(rows: number, columns: number):
185-
WebGLTexture {
182+
public createPackedMatrixTexture(rows: number, columns: number): Texture {
186183
this.throwIfDisposed();
187184
return gpgpu_util.createPackedMatrixTexture(
188185
this.gl, rows, columns, this.textureConfig);

tfjs-backend-webgl/src/gpgpu_context_test.ts

Lines changed: 14 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -18,12 +18,15 @@
1818
import * as tf from '@tensorflow/tfjs-core';
1919
// tslint:disable-next-line: no-imports-from-dist
2020
import {describeWithFlags} from '@tensorflow/tfjs-core/dist/jasmine_util';
21+
// tslint:disable-next-line: no-imports-from-dist
22+
import {expectArraysEqual} from '@tensorflow/tfjs-core/dist/test_util';
2123

2224
import {WEBGL_ENVS} from './backend_webgl_test_registry';
2325
import * as canvas_util from './canvas_util';
2426
import {getGlslDifferences} from './glsl_version';
2527
import {GPGPUContext, linearSearchLastTrue} from './gpgpu_context';
2628
import * as tex_util from './tex_util';
29+
import {Texture} from './tex_util';
2730
import {createFragmentShader} from './webgl_util';
2831

2932
const DOWNLOAD_FLOAT_ENVS = {
@@ -54,15 +57,15 @@ describeWithFlags(
5457
});
5558

5659
it('sets the output texture property to the output texture', () => {
57-
texture = gpgpu.createFloat32MatrixTexture(1, 1);
60+
texture = gpgpu.createFloat32MatrixTexture(1, 1).texture;
5861
gpgpu.setOutputMatrixTexture(texture, 1, 1);
5962
expect(gpgpu.outputTexture).toBe(texture);
6063
});
6164

6265
it('sets the gl viewport to the output texture dimensions', () => {
6366
const columns = 456;
6467
const rows = 123;
65-
texture = gpgpu.createFloat32MatrixTexture(rows, columns);
68+
texture = gpgpu.createFloat32MatrixTexture(rows, columns).texture;
6669
gpgpu.setOutputMatrixTexture(texture, rows, columns);
6770
const expected = new Int32Array([0, 0, columns, rows]);
6871
expect(gpgpu.gl.getParameter(gpgpu.gl.VIEWPORT)).toEqual(expected);
@@ -72,7 +75,7 @@ describeWithFlags(
7275
describeWithFlags(
7376
'GPGPUContext setOutputPackedMatrixTexture', DOWNLOAD_FLOAT_ENVS, () => {
7477
let gpgpu: GPGPUContext;
75-
let texture: WebGLTexture;
78+
let texture: Texture;
7679
let gl: WebGLRenderingContext;
7780

7881
beforeEach(() => {
@@ -86,26 +89,28 @@ describeWithFlags(
8689

8790
afterEach(() => {
8891
if (texture != null) {
89-
gpgpu.deleteMatrixTexture(texture);
92+
gpgpu.deleteMatrixTexture(texture.texture);
9093
}
9194
gpgpu.dispose();
9295
});
9396

9497
it('sets the output texture property to the output texture', () => {
9598
texture = gpgpu.createPackedMatrixTexture(1, 1);
96-
gpgpu.setOutputPackedMatrixTexture(texture, 1, 1);
97-
expect(gpgpu.outputTexture).toBe(texture);
99+
expectArraysEqual(texture.texShape, [1, 1]);
100+
gpgpu.setOutputPackedMatrixTexture(texture.texture, 1, 1);
101+
expect(gpgpu.outputTexture).toBe(texture.texture);
98102
});
99103

100104
it('sets the gl viewport to the output packed texture dimensions', () => {
101105
const columns = 456;
102106
const rows = 123;
103107
texture = gpgpu.createPackedMatrixTexture(rows, columns);
104-
gpgpu.setOutputPackedMatrixTexture(texture, rows, columns);
108+
gpgpu.setOutputPackedMatrixTexture(texture.texture, rows, columns);
105109
const [width, height] =
106110
tex_util.getPackedMatrixTextureShapeWidthHeight(rows, columns);
107111
const expected = new Int32Array([0, 0, width, height]);
108112
expect(gpgpu.gl.getParameter(gpgpu.gl.VIEWPORT)).toEqual(expected);
113+
expectArraysEqual(texture.texShape, [height, width]);
109114
});
110115
});
111116

@@ -131,7 +136,7 @@ describeWithFlags(
131136
// tslint:disable-next-line: max-line-length
132137
const fragmentShader = createFragmentShader(gpgpu.gl, src);
133138
program = gpgpu.createProgram(fragmentShader);
134-
output = gpgpu.createPackedMatrixTexture(4, 4);
139+
output = gpgpu.createPackedMatrixTexture(4, 4).texture;
135140
gpgpu.uploadDenseMatrixToTexture(output, 2, 2, new Float32Array(16));
136141
gpgpu.setOutputMatrixTexture(output, 4, 4);
137142
gpgpu.setProgram(program);
@@ -181,7 +186,7 @@ describeWithFlags('GPGPUContext', DOWNLOAD_FLOAT_ENVS, () => {
181186
`;
182187
const fragmentShader = createFragmentShader(gpgpu.gl, src);
183188
const program = gpgpu.createProgram(fragmentShader);
184-
const result = gpgpu.createFloat32MatrixTexture(1, 1);
189+
const result = gpgpu.createFloat32MatrixTexture(1, 1).texture;
185190
gpgpu.setOutputMatrixTexture(result, 1, 1);
186191
gpgpu.setProgram(program);
187192
gpgpu.deleteMatrixTexture(result);

tfjs-backend-webgl/src/gpgpu_math.ts

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -242,9 +242,11 @@ export function runProgram<T extends Tensor, K extends Tensor>(
242242
const outTex = output.texData.texture;
243243
const outTexShape = output.texData.texShape;
244244
if (output.texData.isPacked) {
245-
gpgpu.setOutputPackedMatrixTexture(outTex, outTexShape[0], outTexShape[1]);
245+
gpgpu.setOutputPackedMatrixTexture(
246+
outTex.texture, outTexShape[0], outTexShape[1]);
246247
} else {
247-
gpgpu.setOutputMatrixTexture(outTex, outTexShape[0], outTexShape[1]);
248+
gpgpu.setOutputMatrixTexture(
249+
outTex.texture, outTexShape[0], outTexShape[1]);
248250
}
249251
gpgpu.setProgram(binary.webGLProgram);
250252

@@ -315,7 +317,7 @@ export function runProgram<T extends Tensor, K extends Tensor>(
315317
gpgpu.gl.uniform1i(varOffsetLoc, input.texData.slice.flatOffset);
316318
}
317319

318-
gpgpu.setInputMatrixTexture(input.texData.texture, varLoc, i);
320+
gpgpu.setInputMatrixTexture(input.texData.texture.texture, varLoc, i);
319321
});
320322

321323
const outShapeLoc = binary.outShapeLocation;

tfjs-backend-webgl/src/gpgpu_util.ts

Lines changed: 9 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@ import {env, PixelData, TypedArray} from '@tensorflow/tfjs-core';
1919

2020
import {getGlslDifferences} from './glsl_version';
2121
import * as tex_util from './tex_util';
22-
import {TextureConfig} from './tex_util';
22+
import {Texture, TextureConfig} from './tex_util';
2323
import * as webgl_util from './webgl_util';
2424

2525
export function createVertexShader(gl: WebGLRenderingContext): WebGLShader {
@@ -53,7 +53,7 @@ export function createIndexBuffer(gl: WebGLRenderingContext): WebGLBuffer {
5353
function createAndConfigureTexture(
5454
gl: WebGLRenderingContext, width: number, height: number,
5555
internalFormat: number, textureFormat: number,
56-
textureType: number): WebGLTexture {
56+
textureType: number): Texture {
5757
webgl_util.validateTextureSize(width, height);
5858
const texture = webgl_util.createTexture(gl);
5959

@@ -80,7 +80,8 @@ function createAndConfigureTexture(
8080
.texStorage2D(tex2d, 1, internalFormat, width, height));
8181
}
8282
webgl_util.callAndCheck(gl, () => gl.bindTexture(gl.TEXTURE_2D, null));
83-
return texture;
83+
84+
return {texture, texShape: [height, width]};
8485
}
8586

8687
export function getInternalFormatForFloat32MatrixTexture(
@@ -90,7 +91,7 @@ export function getInternalFormatForFloat32MatrixTexture(
9091

9192
export function createFloat32MatrixTexture(
9293
gl: WebGLRenderingContext, rows: number, columns: number,
93-
textureConfig: TextureConfig): WebGLTexture {
94+
textureConfig: TextureConfig): Texture {
9495
const [width, height] =
9596
tex_util.getUnpackedMatrixTextureShapeWidthHeight(rows, columns);
9697
return createAndConfigureTexture(
@@ -106,7 +107,7 @@ export function getInternalFormatForFloat16MatrixTexture(
106107

107108
export function createFloat16MatrixTexture(
108109
gl: WebGLRenderingContext, rows: number, columns: number,
109-
textureConfig: TextureConfig): WebGLTexture {
110+
textureConfig: TextureConfig): Texture {
110111
const [width, height] =
111112
tex_util.getUnpackedMatrixTextureShapeWidthHeight(rows, columns);
112113
return createAndConfigureTexture(
@@ -122,7 +123,7 @@ export function getInternalFormatForUnsignedBytesMatrixTexture(
122123

123124
export function createUnsignedBytesMatrixTexture(
124125
gl: WebGLRenderingContext, rows: number, columns: number,
125-
textureConfig: TextureConfig): WebGLTexture {
126+
textureConfig: TextureConfig): Texture {
126127
const [width, height] =
127128
tex_util.getUnpackedMatrixTextureShapeWidthHeight(rows, columns);
128129
return createAndConfigureTexture(
@@ -138,7 +139,7 @@ export function getInternalFormatForPackedMatrixTexture(
138139

139140
export function createPackedMatrixTexture(
140141
gl: WebGLRenderingContext, rows: number, columns: number,
141-
textureConfig: TextureConfig): WebGLTexture {
142+
textureConfig: TextureConfig): Texture {
142143
const [width, height] =
143144
tex_util.getPackedMatrixTextureShapeWidthHeight(rows, columns);
144145
return createAndConfigureTexture(
@@ -153,7 +154,7 @@ export function getInternalFormatForFloat16PackedMatrixTexture(
153154

154155
export function createFloat16PackedMatrixTexture(
155156
gl: WebGLRenderingContext, rows: number, columns: number,
156-
textureConfig: TextureConfig): WebGLTexture {
157+
textureConfig: TextureConfig): Texture {
157158
const [width, height] =
158159
tex_util.getPackedMatrixTextureShapeWidthHeight(rows, columns);
159160
return createAndConfigureTexture(

tfjs-backend-webgl/src/gpgpu_util_test.ts

Lines changed: 8 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -63,7 +63,8 @@ describeWithFlags('gpgpu_util createFloat32MatrixTexture', WEBGL_ENVS, () => {
6363
const gpgpu = new GPGPUContext();
6464
const textureConfig = tex_util.getTextureConfig(gpgpu.gl);
6565
const tex =
66-
gpgpu_util.createFloat32MatrixTexture(gpgpu.gl, 32, 32, textureConfig);
66+
gpgpu_util.createFloat32MatrixTexture(gpgpu.gl, 32, 32, textureConfig)
67+
.texture;
6768
gpgpu.gl.bindTexture(gpgpu.gl.TEXTURE_2D, tex);
6869
expect(
6970
gpgpu.gl.getTexParameter(gpgpu.gl.TEXTURE_2D, gpgpu.gl.TEXTURE_WRAP_S))
@@ -80,7 +81,8 @@ describeWithFlags('gpgpu_util createFloat32MatrixTexture', WEBGL_ENVS, () => {
8081
const gpgpu = new GPGPUContext();
8182
const textureConfig = tex_util.getTextureConfig(gpgpu.gl);
8283
const tex =
83-
gpgpu_util.createFloat32MatrixTexture(gpgpu.gl, 32, 32, textureConfig);
84+
gpgpu_util.createFloat32MatrixTexture(gpgpu.gl, 32, 32, textureConfig)
85+
.texture;
8486
gpgpu.gl.bindTexture(gpgpu.gl.TEXTURE_2D, tex);
8587
expect(gpgpu.gl.getTexParameter(
8688
gpgpu.gl.TEXTURE_2D, gpgpu.gl.TEXTURE_MIN_FILTER))
@@ -100,15 +102,15 @@ describeWithFlags('gpgpu_util createPackedMatrixTexture', WEBGL_ENVS, () => {
100102
const textureConfig = tex_util.getTextureConfig(gpgpu.gl);
101103
const tex =
102104
gpgpu_util.createPackedMatrixTexture(gpgpu.gl, 32, 32, textureConfig);
103-
gpgpu.gl.bindTexture(gpgpu.gl.TEXTURE_2D, tex);
105+
gpgpu.gl.bindTexture(gpgpu.gl.TEXTURE_2D, tex.texture);
104106
expect(
105107
gpgpu.gl.getTexParameter(gpgpu.gl.TEXTURE_2D, gpgpu.gl.TEXTURE_WRAP_S))
106108
.toEqual(gpgpu.gl.CLAMP_TO_EDGE);
107109
expect(
108110
gpgpu.gl.getTexParameter(gpgpu.gl.TEXTURE_2D, gpgpu.gl.TEXTURE_WRAP_T))
109111
.toEqual(gpgpu.gl.CLAMP_TO_EDGE);
110112
gpgpu.gl.bindTexture(gpgpu.gl.TEXTURE_2D, null);
111-
gpgpu.deleteMatrixTexture(tex);
113+
gpgpu.deleteMatrixTexture(tex.texture);
112114
gpgpu.dispose();
113115
});
114116

@@ -117,15 +119,15 @@ describeWithFlags('gpgpu_util createPackedMatrixTexture', WEBGL_ENVS, () => {
117119
const textureConfig = tex_util.getTextureConfig(gpgpu.gl);
118120
const tex =
119121
gpgpu_util.createPackedMatrixTexture(gpgpu.gl, 32, 32, textureConfig);
120-
gpgpu.gl.bindTexture(gpgpu.gl.TEXTURE_2D, tex);
122+
gpgpu.gl.bindTexture(gpgpu.gl.TEXTURE_2D, tex.texture);
121123
expect(gpgpu.gl.getTexParameter(
122124
gpgpu.gl.TEXTURE_2D, gpgpu.gl.TEXTURE_MIN_FILTER))
123125
.toEqual(gpgpu.gl.NEAREST);
124126
expect(gpgpu.gl.getTexParameter(
125127
gpgpu.gl.TEXTURE_2D, gpgpu.gl.TEXTURE_MAG_FILTER))
126128
.toEqual(gpgpu.gl.NEAREST);
127129
gpgpu.gl.bindTexture(gpgpu.gl.TEXTURE_2D, null);
128-
gpgpu.deleteMatrixTexture(tex);
130+
gpgpu.deleteMatrixTexture(tex.texture);
129131
gpgpu.dispose();
130132
});
131133
});

tfjs-backend-webgl/src/tex_util.ts

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -69,14 +69,18 @@ export enum PhysicalTextureType {
6969
PACKED_2X2_FLOAT16
7070
}
7171

72+
export interface Texture {
73+
texture: WebGLTexture;
74+
texShape: [number, number];
75+
}
7276
export interface TextureData {
7377
// Required.
7478
shape: number[];
7579
dtype: DataType;
7680

7781
// Optional.
7882
values?: backend_util.BackendValues;
79-
texture?: WebGLTexture;
83+
texture?: Texture;
8084
// For complex numbers, the real and imaginary parts are stored as their own
8185
// individual tensorInfos, with a parent joining the two with the
8286
// complexTensors field. When this is defined, texture will be null.

0 commit comments

Comments
 (0)