texture to tensor.

tfjs-backend-webgl/src/backend_webgl.ts

@@ -20,6 +20,7 @@ import './flags_webgl';
 
 import * as tf from '@tensorflow/tfjs-core';
 import {backend_util, BackendValues, buffer, DataId, DataStorage, DataToGPUWebGLOption, DataType, engine, env, GPUData, kernel_impls, KernelBackend, MemoryInfo, nextFrame, NumericDataType, Rank, RecursiveArray, scalar, ShapeMap, Tensor, Tensor2D, TensorBuffer, TensorInfo, tidy, TimingInfo, TypedArray, util} from '@tensorflow/tfjs-core';
+
 import {getWebGLContext} from './canvas_util';
 import {DecodeMatrixProgram} from './decode_matrix_gpu';
 import {DecodeMatrixPackedProgram} from './decode_matrix_packed_gpu';
@@ -176,6 +177,49 @@ export class MathBackendWebGL extends KernelBackend {
     return this.texData.numDataIds() - this.pendingDeletes;
   }
 
+  // Writes a new entry to the data store with a WebGL texture, and registers it
+  // to the texture manager.
+  writeTexture(
+      texture: WebGLTexture, shape: number[], dtype: DataType,
+      texShape: [number, number]): DataId {
+    const shapeAs3D = webgl_util.getShapeAs3D(shape);
+    const inData = {
+      shape,
+      dtype,
+      texture,
+      texShape,
+      usage: TextureUsage.RENDER,
+      isPacked: false,
+      refCount: 1
+    };
+    const inputTensorData = {shape, texData: inData, isUniform: false};
+    const output = this.makeTensorInfo(shape, dtype);
+    const outputData = this.texData.get(output.dataId);
+    const outputTensorData: TensorData = {
+      shape: output.shape,
+      texData: outputData,
+      isUniform: false
+    };
+
+    this.uploadToGPU(output);
+
+    const program = new EncodeMatrixProgram(shapeAs3D, false /* isByteArray */);
+
+    const key =
+        gpgpu_math.makeShaderKey(program, [inputTensorData], outputTensorData);
+
+
+    const binary = this.getAndSaveBinary(key, () => {
+      return gpgpu_math.compileProgram(
+          this.gpgpu, program, inputTensorData, outputTensorData);
+    });
+
+    gpgpu_math.runProgram(
+        this.gpgpu, binary, inputTensorData, outputTensorData,
+        null /* customUniformValues */);
+  }
+
+
   write(values: BackendValues, shape: number[], dtype: DataType): DataId {
     if (env().getBool('WEBGL_CHECK_NUMERICAL_PROBLEMS') ||
         env().getBool('DEBUG')) {

tfjs-backend-webgl/src/backend_webgl_test.ts

@@ -28,6 +28,7 @@ import {computeBytes} from './texture_manager';
 import {PhysicalTextureType} from './tex_util';
 import {WEBGL_ENVS} from './backend_webgl_test_registry';
 import {GPGPUContext} from './gpgpu_context';
+import {createTensorFromTexture} from './base';
 
 function decodeStrings(bytes: Uint8Array[]): string[] {
   return bytes.map(b => decodeString(b));
@@ -1108,3 +1109,291 @@ describeWithFlags('Parallel compilation', WEBGL_ENVS, () => {
     tf.env().set('WEBGL_CPU_FORWARD', savedWebGLCPUForward);
   });
 });
+
+describeWithFlags('create tensor from texture', WEBGL2_ENVS, () => {
+  it('basic usage', async () => {
+    // In this test we create a WebGL texture using the GL context from the
+    // WebGL backend. Then we create a tensor from that texture, and ensure that
+    // we can perform a TF operation on that tensor and get the expected result.
+
+    const gpgpu = new GPGPUContext();
+    const width = 3;
+    const height = 4;
+
+    const gl = gpgpu.gl;
+    const texture = gl.createTexture();
+    const tex2d = gl.TEXTURE_2D;
+    // tslint:disable-next-line:no-any
+    const glany = gl as any;
+    const internalFormat = glany.R32F;
+    const textureFormat = glany.RED;
+    const textureType = glany.FLOAT;
+    const dataForUpload =
+        new Float32Array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]);
+
+    gl.bindTexture(tex2d, texture);
+    gl.texParameteri(tex2d, gl.TEXTURE_WRAP_S, gl.CLAMP_TO_EDGE);
+    gl.texParameteri(tex2d, gl.TEXTURE_WRAP_T, gl.CLAMP_TO_EDGE);
+    gl.texParameteri(tex2d, gl.TEXTURE_MIN_FILTER, gl.NEAREST);
+    gl.texParameteri(tex2d, gl.TEXTURE_MAG_FILTER, gl.NEAREST);
+    gl.texImage2D(
+        tex2d, 0, internalFormat, width, height, 0, textureFormat, textureType,
+        dataForUpload);
+
+    const logicalShape = [height, width];
+    const physicalShape: [number, number] = [height, width];
+    const a = createTensorFromTexture({
+      texture,
+      shape: logicalShape,
+      dtype: 'float32',
+      texShapeRC: physicalShape,
+      internalFormat,
+      textureFormat,
+      textureType
+    });
+    const b = tf.mul(a, 2);
+
+    expect(b.shape).toEqual(logicalShape);
+    expectArraysClose(
+        await b.data(), [0, 2, 4, 6, 8, 10, 12, 14, 16, 18, 20, 22]);
+
+    gpgpu.dispose();
+  });
+
+  it('logical and physical shapes do not match', async () => {
+    // In this test we create a WebGL texture using the GL context from the
+    // WebGL backend. Then we create a tensor from that texture, and ensure that
+    // we can perform a TF operation on that tensor and get the expected result.
+
+    const gpgpu = new GPGPUContext();
+    const width = 3;
+    const height = 4;
+
+    const gl = gpgpu.gl;
+    const texture = gl.createTexture();
+    const tex2d = gl.TEXTURE_2D;
+    // tslint:disable-next-line:no-any
+    const glany = gl as any;
+    const internalFormat = glany.R32F;
+    const textureFormat = glany.RED;
+    const textureType = glany.FLOAT;
+    const dataForUpload =
+        new Float32Array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]);
+
+    gl.bindTexture(tex2d, texture);
+    gl.texParameteri(tex2d, gl.TEXTURE_WRAP_S, gl.CLAMP_TO_EDGE);
+    gl.texParameteri(tex2d, gl.TEXTURE_WRAP_T, gl.CLAMP_TO_EDGE);
+    gl.texParameteri(tex2d, gl.TEXTURE_MIN_FILTER, gl.NEAREST);
+    gl.texParameteri(tex2d, gl.TEXTURE_MAG_FILTER, gl.NEAREST);
+    gl.texImage2D(
+        tex2d, 0, internalFormat, width, height, 0, textureFormat, textureType,
+        dataForUpload);
+
+    const logicalShape = [2, 6];
+    const physicalShape: [number, number] = [height, width];
+    const a = createTensorFromTexture({
+      texture,
+      shape: logicalShape,
+      dtype: 'float32',
+      texShapeRC: physicalShape,
+      internalFormat,
+      textureFormat,
+      textureType
+    });
+    const b = tf.mul(a, 2);
+
+    expect(b.shape).toEqual(logicalShape);
+    expectArraysClose(
+        await b.data(), [0, 2, 4, 6, 8, 10, 12, 14, 16, 18, 20, 22]);
+
+    gpgpu.dispose();
+  });
+
+  it('physical texture has empty entries', async () => {
+    // In this test we create a WebGL texture using the GL context from the
+    // WebGL backend. Then we create a tensor from that texture, and ensure that
+    // we can perform a TF operation on that tensor and get the expected result.
+
+    const gpgpu = new GPGPUContext();
+    const width = 3;
+    const height = 5;
+
+    const gl = gpgpu.gl;
+    const texture = gl.createTexture();
+    const tex2d = gl.TEXTURE_2D;
+    // tslint:disable-next-line:no-any
+    const glany = gl as any;
+    const internalFormat = glany.R32F;
+    const textureFormat = glany.RED;
+    const textureType = glany.FLOAT;
+
+    const dataForUpload = new Float32Array(width * height);
+    const data = new Float32Array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]);
+    dataForUpload.set(data);
+
+    gl.bindTexture(tex2d, texture);
+    gl.texParameteri(tex2d, gl.TEXTURE_WRAP_S, gl.CLAMP_TO_EDGE);
+    gl.texParameteri(tex2d, gl.TEXTURE_WRAP_T, gl.CLAMP_TO_EDGE);
+    gl.texParameteri(tex2d, gl.TEXTURE_MIN_FILTER, gl.NEAREST);
+    gl.texParameteri(tex2d, gl.TEXTURE_MAG_FILTER, gl.NEAREST);
+    gl.texImage2D(
+        tex2d, 0, internalFormat, width, height, 0, textureFormat, textureType,
+        dataForUpload);
+
+    const logicalShape = [1, 13];
+    const physicalShape: [number, number] = [height, width];
+    const a = createTensorFromTexture({
+      texture,
+      shape: logicalShape,
+      dtype: 'float32',
+      texShapeRC: physicalShape,
+      internalFormat,
+      textureFormat,
+      textureType
+    });
+    const b = tf.mul(a, 2);
+
+    expect(b.shape).toEqual(logicalShape);
+    expectArraysClose(
+        await b.data(), [0, 2, 4, 6, 8, 10, 12, 14, 16, 18, 20, 22, 24]);
+
+    gpgpu.dispose();
+  });
+
+  it('force f16', async () => {
+    // Unlike in the basic usage test, rather than creating a texture from
+    // scratch, we must extract the output texture from an operation because we
+    // cannot upload Float16 data directly to the GPU.
+
+    // We clean up explicitly so that we have full control over the environment
+    // flags during texture initialization / disposal.
+    tf.engine().startScope();
+
+    const webglRenderF32EnabledFlagSaved =
+        tf.env().getBool('WEBGL_RENDER_FLOAT32_ENABLED');
+    const webglPackedFlagSaved = tf.env().getBool('WEBGL_PACK');
+    tf.env().set('WEBGL_RENDER_FLOAT32_ENABLED', false);
+
+    // We must set `WEBGL_PACK` to false because createTensorFromTexture only
+    // accepts unpacked textures so this ensures the output texture (`bTexture`
+    // below) is unpacked.
+    tf.env().set('WEBGL_PACK', false);
+
+    const gpgpu = new GPGPUContext();
+    const gl = gpgpu.gl;
+    // tslint:disable-next-line:no-any
+    const glany = gl as any;
+
+    const width = 3;
+    const height = 4;
+
+    const logicalShape: [number, number] = [height, width];
+    const physicalShape: [number, number] = [height, width];
+    const a = tf.tensor2d([0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11], logicalShape);
+    const b = tf.relu(a);
+    const bTexture = (tf.backend() as MathBackendWebGL).getTexture(b.dataId);
+    const c = createTensorFromTexture({
+      texture: bTexture,
+      shape: logicalShape,
+      dtype: 'float32',
+      texShapeRC: physicalShape,
+      internalFormat: glany.R16F,
+      textureFormat: glany.RED,
+      textureType: glany.HALF_FLOAT
+    });
+    const d = tf.mul(c, 2);
+
+    expect(d.shape).toEqual(logicalShape);
+    expectArraysClose(
+        await d.data(), [0, 2, 4, 6, 8, 10, 12, 14, 16, 18, 20, 22]);
+
+    gpgpu.dispose();
+
+    tf.engine().endScope();
+    tf.env().set(
+        'WEBGL_RENDER_FLOAT32_ENABLED', webglRenderF32EnabledFlagSaved);
+    tf.env().set('WEBGL_PACK', webglPackedFlagSaved);
+  });
+});
+
+describeWithFlags('create tensor from texture', WEBGL1_ENVS, () => {
+  it('basic usage', async () => {
+    const gpgpu = new GPGPUContext();
+    const width = 3;
+    const height = 4;
+
+    const logicalShape: [number, number] = [height, width];
+    const physicalShape: [number, number] = [height, width];
+
+    const gl = gpgpu.gl;
+    const texture = gl.createTexture();
+    const tex2d = gl.TEXTURE_2D;
+    const internalFormat = gl.RGBA;
+    const textureFormat = gl.RGBA;
+    const textureType = gl.FLOAT;
+    // WebGL 1 does not accept gl.RED as an internalFormat, so we have to
+    // upload values for the unused channels as well.
+    const dataForUpload = new Float32Array([
+      0, 0, 0, 0, 1, 0, 0, 0, 2, 0, 0, 0, 3, 0, 0, 0, 4,  0, 0, 0, 5,  0, 0, 0,
+      6, 0, 0, 0, 7, 0, 0, 0, 8, 0, 0, 0, 9, 0, 0, 0, 10, 0, 0, 0, 11, 0, 0, 0,
+    ]);
+
+    gl.bindTexture(tex2d, texture);
+    gl.texParameteri(tex2d, gl.TEXTURE_WRAP_S, gl.CLAMP_TO_EDGE);
+    gl.texParameteri(tex2d, gl.TEXTURE_WRAP_T, gl.CLAMP_TO_EDGE);
+    gl.texParameteri(tex2d, gl.TEXTURE_MIN_FILTER, gl.NEAREST);
+    gl.texParameteri(tex2d, gl.TEXTURE_MAG_FILTER, gl.NEAREST);
+    gl.texImage2D(
+        tex2d, 0, internalFormat, width, height, 0, textureFormat, textureType,
+        dataForUpload);
+
+    const a = createTensorFromTexture({
+      texture,
+      shape: logicalShape,
+      texShapeRC: physicalShape,
+      dtype: 'float32',
+      internalFormat,
+      textureFormat,
+      textureType
+    });
+    const b = tf.mul(a, 2);
+
+    expect(b.shape).toEqual(logicalShape);
+    expectArraysClose(
+        await b.data(), [0, 2, 4, 6, 8, 10, 12, 14, 16, 18, 20, 22]);
+    gpgpu.dispose();
+  });
+
+  it('chained', async () => {
+    const gpgpu = new GPGPUContext();
+    const gl = gpgpu.gl;
+
+    const webglPackedFlagSaved = tf.env().getBool('WEBGL_PACK');
+    tf.env().set('WEBGL_PACK', false);
+
+    const width = 3;
+    const height = 4;
+
+    const logicalShape: [number, number] = [height, width];
+    const physicalShape: [number, number] = [height, width];
+    const a = tf.tensor2d([0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11], logicalShape);
+    const b = tf.relu(a);
+    const bTexture = (tf.backend() as MathBackendWebGL).getTexture(b.dataId);
+    const c = createTensorFromTexture({
+      texture: bTexture,
+      shape: logicalShape,
+      texShapeRC: physicalShape,
+      dtype: 'float32',
+      internalFormat: gl.RGBA,
+      textureFormat: gl.RGBA,
+      textureType: gl.FLOAT
+    });
+    const d = tf.mul(c, 2);
+
+    expect(d.shape).toEqual(logicalShape);
+    expectArraysClose(
+        await d.data(), [0, 2, 4, 6, 8, 10, 12, 14, 16, 18, 20, 22]);
+
+    tf.env().set('WEBGL_PACK', webglPackedFlagSaved);
+  });
+});

tfjs-backend-webgl/src/base.ts

@@ -28,6 +28,10 @@ if (device_util.isBrowser()) {
 // Export webgl utilities
 export * from './webgl';
 
-// Export forceHalfFlost under webgl namespace for the union bundle.
-import {forceHalfFloat} from './webgl';
-export const webgl = {forceHalfFloat};
+// Export forceHalfFloat and createTensorFromTexture under webgl namespace for
+// the union bundle.
+import {forceHalfFloat, createTensorFromTexture} from './webgl';
+export const webgl = {
+  forceHalfFloat,
+  createTensorFromTexture
+};