[webgl] Modularize MaxPool, MaxPoolBackprop, AvgPool, AvgPoolBackprop webgl kernels

[jinjingforever]

This change is 

[tafsiri]

Lookin good. Just a couple of fixes to match the input validation strategy in cpu.

Reviewed 7 of 7 files at r1.
Reviewable status: complete! 1 of 1 approvals obtained (waiting on @annxingyuan, @jinjingforever, and @tafsiri)

---

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

  const dilations = 1;

  util.assert(

This looks like it might just belong in the op since the attrs it asserts are untransformed.

---

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

  const dilations = 1;

  util.assert(

same as above

[jinjingforever]

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

---

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

Previously, tafsiri (Yannick Assogba) wrote…

This looks like it might just belong in the op since the attrs it asserts are untransformed.

Thank you Yannick!

I asked @lina128 offline about this particular assert in my previous PR and we decided to leave it here because "dilations" (which is not coming from attrs) is technically mutated. I can imagine that the assert might be useful if for example somebody changed dilations to some other number here. Any thoughts/preferences on this?

Thanks!

[tafsiri]

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

---

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

Previously, jinjingforever (Jing Jin) wrote…

Thank you Yannick!

I asked @lina128 offline about this particular assert in my previous PR and we decided to leave it here because "dilations" (which is not coming from attrs) is technically mutated. I can imagine that the assert might be useful if for example somebody changed dilations to some other number here. Any thoughts/preferences on this?

Thanks!

Ah okay i see what you mean. No strong pref on my end so feel free to leave it in for documentation/robustness to future change purposes (though an alternative would be to assert dilation is 1 and have a comment for the rest). If we ever do a systematic exploration (or see them show up in profiles) of the cost of assertions and find them expensive we can revisit it.

[annxingyuan]

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

---

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

      () => 'Error in avgPool: Either strides or dilations must be 1. ' +
          `Got strides ${strides} and dilations '${dilations}'`);

IIUC, the avgPool op includes an optimization in the forwardFunc of cloning the input in the case of a 1x1 convolution. Should we try to include that in this kernel, or add a TODO item to include it? Same for them axPool kernel below.

---

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

  const convInfo = backend_util.computePool2DInfo(
      x.shape as [number, number, number, number], filterSize, strides,

I noticed that the avgPool op reshapes inputs to be 4D, and then shapes them back at the end of the op - should this be handled by the kernels instead? Same for the maxPool kernel below.

Sorry if this was already discussed in the avgPool CPU PR!

---

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

      x.shape as [number, number, number, number], filterSize, strides,
      dilations, pad, dimRoundingMode);
  const program = new Pool2DProgram(convInfo, 'max', false);

name this variable

[jinjingforever]

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

---

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

Previously, tafsiri (Yannick Assogba) wrote…

Ah okay i see what you mean. No strong pref on my end so feel free to leave it in for documentation/robustness to future change purposes (though an alternative would be to assert dilation is 1 and have a comment for the rest). If we ever do a systematic exploration (or see them show up in profiles) of the cost of assertions and find them expensive we can revisit it.

SG. Thank you!

---

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

Previously, annxingyuan (Ann Yuan) wrote…

IIUC, the avgPool op includes an optimization in the forwardFunc of cloning the input in the case of a 1x1 convolution. Should we try to include that in this kernel, or add a TODO item to include it? Same for them axPool kernel below.

Good catch! I included that in CPU backend but forgot here..

I added an Identity kernel (very similar to CPU) and used it here.

---

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

Previously, annxingyuan (Ann Yuan) wrote…

I noticed that the avgPool op reshapes inputs to be 4D, and then shapes them back at the end of the op - should this be handled by the kernels instead? Same for the maxPool kernel below.

Sorry if this was already discussed in the avgPool CPU PR!

Thanks! I had some discussion with @lina128 (please correct me if I am wrong). Looks like the official TF document says that it expects a 4D tensor as the input. To match that, our kernel should in theory expect a 4D tensor as well, which means it is probably better to put the 3D->4D transformation outside kernel.

---

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

Previously, annxingyuan (Ann Yuan) wrote…

name this variable

Done