[webgl] Modularize reshape, add refCounter.

[annxingyuan]

Changes

• Adds a ref counting mechanism to the WebGL backend modeled off the CPU backend. Paves the way for removing the existing custom slice ref counter in WebGL.
• Modularizes WebGL's reshape kernel.
• Adds helper for re-registering kernels for a custom backend.

To see the logs from the Cloud Build CI, please join either our discussion or announcement mailing list.

---

This change is 

[jinjingforever]

Reviewable status: 0 of 1 approvals obtained (waiting on @annxingyuan, @jinjingforever, @lina128, and @tafsiri)

---

tfjs-backend-webgl/src/backend_webgl.ts, line 287 at r1 (raw file):

TensorData

nit: update comment (here and decRef below)

---

tfjs-backend-webgl/src/kernels/Reshape.ts, line 37 at r1 (raw file):

$xSize

For my education: I see '$' is used many places in our code base. Is it some kind of convention in our team? Thanks!

[lina128]

Reviewable status: 0 of 1 approvals obtained (waiting on @annxingyuan, @jinjingforever, @lina128, and @tafsiri)

---

tfjs-backend-webgl/src/backend_webgl.ts, line 381 at r1 (raw file):

    const {values, shape, slice, dtype, complexTensors, isPacked} = texData;

    if (slice != null) {

Can we add some annotation on this special treatment?

---

tfjs-backend-webgl/src/backend_webgl.ts, line 2446 at r1 (raw file):

  }

  private packedReshape(input: TensorInfo, afterShape: number[]): TensorInfo {

Seems that we need to keep this private function to prevent circular dependency.

---

tfjs-backend-webgl/src/kernels/Max_impl.ts, line 37 at r1 (raw file):

  backend.disposeIntermediateTensorInfo(reshapedInput);

  return reshape({inputs: {x: reduced}, attrs: {shape: outShape}, backend});

Should reduced be disposed too?

const result = reshape(...);

backend.disposeIntermediateTensorInfo(reduced);

return result;

---

tfjs-backend-webgl/src/kernels/Reshape.ts, line 19 at r1 (raw file):

import {KernelFunc, Reshape, ReshapeAttrs, ReshapeInputs, TensorInfo} from '@tensorflow/tfjs-core';
import {KernelConfig, util} from '@tensorflow/tfjs-core';

Merge above two lines.

[tafsiri]

Excited to see this!

Reviewed 12 of 12 files at r1.
Reviewable status: 0 of 1 approvals obtained (waiting on @annxingyuan and @tafsiri)

---

tfjs-backend-webgl/src/backend_webgl.ts, line 1870 at r1 (raw file):

      const targetShape = isChannelsLast ? xShape[0] * xShape[1] * xShape[2] :
                                           xShape[0] * xShape[2] * xShape[3];
      const xReshaped = x.reshape([1, targetShape, convInfo.inChannels]);

use op instead of chaining api

---

tfjs-backend-webgl/src/backend_webgl.ts, line 1872 at r1 (raw file):

      const xReshaped = x.reshape([1, targetShape, convInfo.inChannels]);
      const filterReshaped =
          filter.reshape([1, convInfo.inChannels, convInfo.outChannels]);

use op instead of chaining api

---

tfjs-backend-webgl/src/backend_webgl.ts, line 1883 at r1 (raw file):

        preluActivationWeights
      });
      return result.reshape(convInfo.outShape);

use op instead of chaining api

---

tfjs-backend-webgl/src/backend_webgl.ts, line 1918 at r1 (raw file):

            xReshaped.shape} isn't free`);
    const filterReshaped =
        filter.reshape([1, convInfo.inChannels, convInfo.outChannels]);

remove chaining api

---

tfjs-backend-webgl/src/kernels/Reshape_test.ts, line 25 at r1 (raw file):

import {describeWithFlags, ALL_ENVS} from '@tensorflow/tfjs-core/dist/jasmine_util';

describeWithFlags('Reshape.', ALL_ENVS, () => {

could you add a test(s) that captures that tests usign both slice and reshape and check that memory is as expected.

---

tfjs-core/src/kernel_registry.ts, line 203 at r1 (raw file):

 * @param newBackendName New backend.
 */
export function reregisterKernelsFromExistingToNewBackend(

naming suggestion: copyRegisteredKernels (reregister... stood out to me as it makes me think the kernels had already been registered on the target backend already and were being reinitialized or something).

[annxingyuan]

Reviewable status: 0 of 1 approvals obtained (waiting on @jinjingforever, @lina128, and @tafsiri)

---

tfjs-backend-webgl/src/backend_webgl.ts, line 287 at r1 (raw file):

Previously, jinjingforever (Jing Jin) wrote…

TensorData

nit: update comment (here and decRef below)

Done

---

tfjs-backend-webgl/src/backend_webgl.ts, line 381 at r1 (raw file):

Previously, lina128 (Na Li) wrote…

Can we add some annotation on this special treatment?

Done

---

tfjs-backend-webgl/src/backend_webgl.ts, line 1870 at r1 (raw file):

Previously, tafsiri (Yannick Assogba) wrote…

use op instead of chaining api

Done.

---

tfjs-backend-webgl/src/backend_webgl.ts, line 1872 at r1 (raw file):

Previously, tafsiri (Yannick Assogba) wrote…

use op instead of chaining api

Done.

---

tfjs-backend-webgl/src/backend_webgl.ts, line 1883 at r1 (raw file):

Previously, tafsiri (Yannick Assogba) wrote…

use op instead of chaining api

Done.

---

tfjs-backend-webgl/src/backend_webgl.ts, line 1918 at r1 (raw file):

Previously, tafsiri (Yannick Assogba) wrote…

remove chaining api

Done.

---

tfjs-backend-webgl/src/backend_webgl.ts, line 2446 at r1 (raw file):

Previously, lina128 (Na Li) wrote…

Seems that we need to keep this private function to prevent circular dependency.

Done

---

tfjs-backend-webgl/src/kernels/Max_impl.ts, line 37 at r1 (raw file):

Previously, lina128 (Na Li) wrote…

Should reduced be disposed too?

const result = reshape(...);

backend.disposeIntermediateTensorInfo(reduced);

return result;

Done

---

tfjs-backend-webgl/src/kernels/Reshape.ts, line 19 at r1 (raw file):

Previously, lina128 (Na Li) wrote…

Merge above two lines.

Done.

---

tfjs-backend-webgl/src/kernels/Reshape.ts, line 37 at r1 (raw file):

Previously, jinjingforever (Jing Jin) wrote…

$xSize

For my education: I see '$' is used many places in our code base. Is it some kind of convention in our team? Thanks!

Hi Jing, yes it's often used in our op code when trying to avoid reassigning the arguments. I think most of the time it's used to name variables that hold the tensorized version of 'tensor-like' inputs (e.g. const $a = convertToTensor(a);.

---

tfjs-backend-webgl/src/kernels/Reshape_test.ts, line 25 at r1 (raw file):

Previously, tafsiri (Yannick Assogba) wrote…

could you add a test(s) that captures that tests usign both slice and reshape and check that memory is as expected.

Done

---

tfjs-core/src/kernel_registry.ts, line 203 at r1 (raw file):

Previously, tafsiri (Yannick Assogba) wrote…

naming suggestion: copyRegisteredKernels (reregister... stood out to me as it makes me think the kernels had already been registered on the target backend already and were being reinitialized or something).

Done.

[lina128]

This is great! Thank you Ann!

Reviewable status: complete! 1 of 1 approvals obtained (waiting on @jinjingforever, @lina128, and @tafsiri)

[tafsiri]

Thanks Ann! Great stuff. (left one naming suggestion)

Reviewed 8 of 8 files at r2.
Reviewable status: complete! 2 of 1 approvals obtained (waiting on @annxingyuan, @jinjingforever, @lina128, and @tafsiri)

---

tfjs-core/src/kernel_registry.ts, line 203 at r1 (raw file):

Previously, annxingyuan (Ann Yuan) wrote…

Done.

I think you can drop the "FromExistingToNewBackend" suffix.

[annxingyuan]

Reviewable status: complete! 2 of 1 approvals obtained (waiting on @jinjingforever, @lina128, and @tafsiri)

---

tfjs-core/src/kernel_registry.ts, line 203 at r1 (raw file):

Previously, tafsiri (Yannick Assogba) wrote…

I think you can drop the "FromExistingToNewBackend" suffix.

Done