Modularize MaxPool, MaxPoolBackprop, AvgPool, AvgPoolBackprop cpu kernels.

[jinjingforever]

This change is 

[tafsiri]

Reviewed 6 of 6 files at r1.
Reviewable status: 0 of 1 approvals obtained (waiting on @jinjingforever, @lina128, and @tafsiri)

---

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

  if (convInfo.filterWidth === 1 && convInfo.filterHeight === 1 &&
      util.arraysEqual(convInfo.inShape, convInfo.outShape)) {
    res = {...x};

why this over res = x?

---

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

      util.arraysEqual(convInfo.inShape, convInfo.outShape)) {
    res = {...x};
    backend.incRef(x.dataId);

should this (both these lines) go through identity instead? I think keeping the incRef calls in fewer places might be better to make it easier to change memory management strategies.

---

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

  if (convInfo.filterWidth === 1 && convInfo.filterHeight === 1 &&
      util.arraysEqual(convInfo.inShape, convInfo.outShape)) {
    res = {...x};

same comment as above on using identity here

[jinjingforever]

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

---

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

Previously, tafsiri (Yannick Assogba) wrote…

why this over res = x?

Replaced with identity

---

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

Previously, tafsiri (Yannick Assogba) wrote…

should this (both these lines) go through identity instead? I think keeping the incRef calls in fewer places might be better to make it easier to change memory management strategies.

Ah, good point. I forgot about identity. Updated.

---

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

Previously, tafsiri (Yannick Assogba) wrote…

same comment as above on using identity here

Done

[lina128]

Hi Jing, pool and conv operations are some of the most complicated, thank you for starting this! I left some comment below.

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

---

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

  const xRank = x.shape.length;
  util.assert(

This assertion is at op level, I think we agreed to not repeat it at kernel level because x is not transformed. Is that right, @tafsiri ?

---

tfjs-backend-cpu/src/kernels/AvgPool.ts, line 45 at r2 (raw file):

          x.shape.length}}.`);
  if (dimRoundingMode != null) {
    util.assert(

This assertion is also at op level, and no transformation of this attribute, so no need to assert at kernel level. Otherwise these assertions always run twice if call starts from op level, which has cost. We had the discussion of whether it is safe to do so because call may start from other kernels, which may not have any checking. I think the consensus was, if there's no transformation in this kernel, we don't do the assertion, and kernels basically trust other kernels to pass in the right values. You may ask: Why we need to do the assertion at op level? Because op is public api, the data is user-provided, so we need to validate. Why we sometimes do assertion at kernel level if we trust kernel implementation? Most of them are for convenience, actually we should use unit tests to test all the possible cases instead of relying on assertion and bug report. We can consider converting them into unit tests once we have bandwidth. Jing, sorry I know this is sort of confusing about where to do transformation and where to assert. Above is basically what we discovered and agreed upon along the way. Hope it makes sense. :) Adding @tafsiri to double confirm above summary.

---

tfjs-backend-cpu/src/kernels/AvgPool.ts, line 62 at r2 (raw file):

    const xValues = backend.data.get(x.dataId).values as TypedArray;
    const strides = util.computeStrides(x.shape);
    const buffer = pool(xValues, x.shape, x.dtype, strides, convInfo, 'avg');

This uses op buffer and returns a tensorBuffer. Can we create a new util, say poolV2 and avoid using buffer and only return TypedArray? Alternatively, you can just revise pool and modularize all the affected kernels in this PR. Either way is fine by me, but I probably will prefer to not creating poolV2, because these versions generally align with tf kernel names' name.

---

tfjs-backend-cpu/src/kernels/AvgPoolBackprop.ts, line 35 at r2 (raw file):

  const inputRank = input.shape.length;
  const dyRank = dy.shape.length;
  util.assert(

No need to check rank here, this is type checked in Avgpool_grad.ts. Also if avgPool only work with 4D, it's not likely the backward propagation can get anything not 4D. So we can save an assertion here.

---

tfjs-backend-cpu/src/kernels/AvgPoolBackprop.ts, line 58 at r2 (raw file):

  const padTop = effectiveFilterHeight - 1 - convInfo.padInfo.top;
  const dx =
      buffer<Rank.R4>(x.shape as [number, number, number, number], 'float32');

This uses op buffer again. You can use create a typedArray for dx.

---

tfjs-backend-cpu/src/kernels/AvgPoolBackprop.ts, line 63 at r2 (raw file):

  const dyData = backend.data.get(dy.dataId).values as Float32Array;
  const dyBuf = buffer<Rank.R4>(

Same as above.

---

tfjs-backend-cpu/src/kernels/MaxPool.ts, line 35 at r2 (raw file):

  const xRank = x.shape.length;
  util.assert(

Same as above. No need for assertion.

---

tfjs-backend-cpu/src/kernels/MaxPool.ts, line 61 at r2 (raw file):

    const xValues = backend.data.get(x.dataId).values as TypedArray;
    const strides = util.computeStrides(x.shape);
    const buffer = pool(xValues, x.shape, x.dtype, strides, convInfo, 'max');

Same as above about changing pool implementation.

---

tfjs-backend-cpu/src/kernels/MaxPoolBackprop.ts, line 36 at r2 (raw file):

  const inputRank = input.shape.length;
  const dyRank = dy.shape.length;
  util.assert(

Same, no need to assert.

---

tfjs-backend-cpu/src/kernels/MaxPoolBackprop.ts, line 60 at r2 (raw file):

      1 /* dilations */, pad, dimRoundingMode);
  const xValues = backend.data.get(x.dataId).values as TypedArray;
  const maxPosBuf = buffer(

Same as above, remove buffer.

---

tfjs-backend-cpu/src/kernels/MaxPoolBackprop.ts, line 72 at r2 (raw file):

  const padTop = effectiveFilterHeight - 1 - convInfo.padInfo.top;
  const dx =
      buffer<Rank.R4>(x.shape as [number, number, number, number], 'float32');

Same as above, remove buffer.

---

tfjs-backend-cpu/src/kernels/MaxPoolBackprop.ts, line 75 at r2 (raw file):

  const dyData = backend.data.get(dy.dataId).values as Float32Array;
  const dyBuf = buffer<Rank.R4>(

Same as above, remove buffer.

[tafsiri]

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

---

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

Previously, lina128 (Na Li) wrote…

This assertion is at op level, I think we agreed to not repeat it at kernel level because x is not transformed. Is that right, @tafsiri ?

+1

---

tfjs-backend-cpu/src/kernels/AvgPool.ts, line 45 at r2 (raw file):

Previously, lina128 (Na Li) wrote…

This assertion is also at op level, and no transformation of this attribute, so no need to assert at kernel level. Otherwise these assertions always run twice if call starts from op level, which has cost. We had the discussion of whether it is safe to do so because call may start from other kernels, which may not have any checking. I think the consensus was, if there's no transformation in this kernel, we don't do the assertion, and kernels basically trust other kernels to pass in the right values. You may ask: Why we need to do the assertion at op level? Because op is public api, the data is user-provided, so we need to validate. Why we sometimes do assertion at kernel level if we trust kernel implementation? Most of them are for convenience, actually we should use unit tests to test all the possible cases instead of relying on assertion and bug report. We can consider converting them into unit tests once we have bandwidth. Jing, sorry I know this is sort of confusing about where to do transformation and where to assert. Above is basically what we discovered and agreed upon along the way. Hope it makes sense. :) Adding @tafsiri to double confirm above summary.

+1.

---

tfjs-backend-cpu/src/kernels/AvgPool.ts, line 62 at r2 (raw file):

Previously, lina128 (Na Li) wrote…

This uses op buffer and returns a tensorBuffer. Can we create a new util, say poolV2 and avoid using buffer and only return TypedArray? Alternatively, you can just revise pool and modularize all the affected kernels in this PR. Either way is fine by me, but I probably will prefer to not creating poolV2, because these versions generally align with tf kernel names' name.

A general comment about buffer (relevant here and below). I would say it is safe to use buffer because it is not really an op in the tensorflow definition of op (it involves no tensors). It is just a wrapper around the constructor for TensorBuffer.

TensorBuffer gives a nice interface for indexing into tensors, the logic for that would need to be reimplemented if you just have a typed array. So in terms of not using buffer(), i would ignore the 'kernels should not call ops' rule. However if there are other reasons to not use it we should discuss those, but I don't think buffer being an 'op' is a blocker in these cases.

[lina128]

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

---

tfjs-backend-cpu/src/kernels/AvgPool.ts, line 62 at r2 (raw file):

Previously, tafsiri (Yannick Assogba) wrote…

A general comment about buffer (relevant here and below). I would say it is safe to use buffer because it is not really an op in the tensorflow definition of op (it involves no tensors). It is just a wrapper around the constructor for TensorBuffer.

TensorBuffer gives a nice interface for indexing into tensors, the logic for that would need to be reimplemented if you just have a typed array. So in terms of not using buffer(), i would ignore the 'kernels should not call ops' rule. However if there are other reasons to not use it we should discuss those, but I don't think buffer being an 'op' is a blocker in these cases.

Make sense.

[jinjingforever]

Thanks Na and Yannick! I removed the unnecessary asserts and kept the "buffer" calls as Yannick suggested. PTAL. Thanks!

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

---

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

Previously, tafsiri (Yannick Assogba) wrote…

+1

Done

---

tfjs-backend-cpu/src/kernels/AvgPool.ts, line 45 at r2 (raw file):

Previously, tafsiri (Yannick Assogba) wrote…

+1.

Done

---

tfjs-backend-cpu/src/kernels/AvgPool.ts, line 62 at r2 (raw file):

Previously, lina128 (Na Li) wrote…

Make sense.

It makes sense. Thanks! Keep it as it is for now.

---

tfjs-backend-cpu/src/kernels/AvgPoolBackprop.ts, line 35 at r2 (raw file):

Previously, lina128 (Na Li) wrote…

No need to check rank here, this is type checked in Avgpool_grad.ts. Also if avgPool only work with 4D, it's not likely the backward propagation can get anything not 4D. So we can save an assertion here.

Done

---

tfjs-backend-cpu/src/kernels/MaxPool.ts, line 35 at r2 (raw file):

Previously, lina128 (Na Li) wrote…

Same as above. No need for assertion.

Done.

---

tfjs-backend-cpu/src/kernels/MaxPoolBackprop.ts, line 36 at r2 (raw file):

Previously, lina128 (Na Li) wrote…

Same, no need to assert.

Done

[lina128]

Thanks Jing!

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

[tafsiri]

Reviewed 1 of 4 files at r3.
Reviewable status: complete! 2 of 1 approvals obtained (waiting on @lina128 and @tafsiri)