[wasm] Add min / max kernels.

[annxingyuan]

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

---

This change is 

[nsthorat]

Reviewed 8 of 8 files at r1.
Reviewable status: complete! 1 of 1 approvals obtained (waiting on @annxingyuan, @dsmilkov, and @nsthorat)

---

tfjs-backend-wasm/src/setup_test.ts, line 63 at r1 (raw file):

  'broadcast 6D + 2D',  // Broadcasting along inner dims not supported yet.

  // Min / max

can you separate these to min and max (separate comments)? im gonna send a follow up PR that actually does something smarter here, each include will have a list of excludes because we don't actually know if we're excluding the wrong things

---

tfjs-backend-wasm/src/cc/kernels/Max.cc, line 45 at r1 (raw file):

    float max = x_buf[offset];
    for (int j = 0; j < reduce_size; ++j) {
      float value = x_buf[offset + j];

if you're feeling brave you could do this with pointer arithmetic (you can operate on a float* pointer directly to create a new float*). This avoids double addition of offset, j and the beginning pointer of x_buf (so basically at the start of this inner loop you keep another float* that points to x_buf + offset and update that pointer).

Here and also for min

[nsthorat]

Wait for Daniel on this too :)

[annxingyuan]

Reviewable status: complete! 1 of 1 approvals obtained (waiting on @dsmilkov)

---

tfjs-backend-wasm/src/setup_test.ts, line 63 at r1 (raw file):

Previously, nsthorat (Nikhil Thorat) wrote…

can you separate these to min and max (separate comments)? im gonna send a follow up PR that actually does something smarter here, each include will have a list of excludes because we don't actually know if we're excluding the wrong things

Done

---

tfjs-backend-wasm/src/cc/kernels/Max.cc, line 45 at r1 (raw file):

Previously, nsthorat (Nikhil Thorat) wrote…

if you're feeling brave you could do this with pointer arithmetic (you can operate on a float* pointer directly to create a new float*). This avoids double addition of offset, j and the beginning pointer of x_buf (so basically at the start of this inner loop you keep another float* that points to x_buf + offset and update that pointer).

Here and also for min

Done - i think... is this what you mean?

[dsmilkov]

Nice work! Left a few comments. Biggest comment is about the exclude list in tests.

Reviewed 4 of 8 files at r1, 3 of 3 files at r2.
Reviewable status: complete! 1 of 1 approvals obtained (waiting on @annxingyuan and @dsmilkov)

---

tfjs-backend-wasm/src/setup_test.ts, line 64 at r2 (raw file):

  // max
  'logSoftmax ',  // Not yet implemented.

we try to be specific with the excluded test names, otherwise when we later add 'softmax' to the include list, we might think that we implemented it correctly because we forgot to remove it from the exclude list.

Let's change these to specific keywords that ignore only the failing tests.

---

tfjs-backend-wasm/src/cc/kernels/Max.cc, line 19 at r2 (raw file):

#endif

#include <math.h>

remove unused math.h and unused vector

---

tfjs-backend-wasm/src/cc/kernels/Max.cc, line 45 at r2 (raw file):

    float max = x_buf[offset];

    float* x_iter_end = x_info.buf.f32 + offset + reduce_size;

instead of summing x_info.buf.f32 + offset twice, have a single offset outside the two for loops that starts at x_info.buf.f32 and sum reduce_sum into it on every outer iteration. Smth like:

int offset = x_info.buf.f32;
for (...) {
  ...
  x_iter_end = offset + reduce_size;
  ...
  offset += reduce_size;
}

---

tfjs-backend-wasm/src/cc/kernels/Min.cc, line 19 at r2 (raw file):

#endif

#include <math.h>

remove unused math and vector

---

tfjs-core/src/ops/reduction_ops.ts, line 386 at r2 (raw file):

    save([xOrig, y]);
    return y;
  }, {x: $x}, grad, 'Max', {axes}, inputsToSave);

I'm surprised that a unit test didn't fail here since you didn't provide outputsToSave=[true] as the last param in runKernelFunc(). And the gradient above expects two things to exist in the saved list. Can you verify why a unit test didn't fail? Is saved[1] unused?

[annxingyuan]

Reviewable status: complete! 1 of 1 approvals obtained (waiting on @dsmilkov)

---

tfjs-backend-wasm/src/setup_test.ts, line 64 at r2 (raw file):

Previously, dsmilkov (Daniel Smilkov) wrote…

we try to be specific with the excluded test names, otherwise when we later add 'softmax' to the include list, we might think that we implemented it correctly because we forgot to remove it from the exclude list.

Let's change these to specific keywords that ignore only the failing tests.

Done

---

tfjs-backend-wasm/src/cc/kernels/Max.cc, line 19 at r2 (raw file):

Previously, dsmilkov (Daniel Smilkov) wrote…

remove unused math.h and unused vector

Done

---

tfjs-backend-wasm/src/cc/kernels/Max.cc, line 45 at r2 (raw file):

Previously, dsmilkov (Daniel Smilkov) wrote…

instead of summing x_info.buf.f32 + offset twice, have a single offset outside the two for loops that starts at x_info.buf.f32 and sum reduce_sum into it on every outer iteration. Smth like:

int offset = x_info.buf.f32;
for (...) {
  ...
  x_iter_end = offset + reduce_size;
  ...
  offset += reduce_size;
}

Done

---

tfjs-backend-wasm/src/cc/kernels/Min.cc, line 19 at r2 (raw file):

Previously, dsmilkov (Daniel Smilkov) wrote…

remove unused math and vector

Done

---

tfjs-core/src/ops/reduction_ops.ts, line 386 at r2 (raw file):

Previously, dsmilkov (Daniel Smilkov) wrote…

I'm surprised that a unit test didn't fail here since you didn't provide outputsToSave=[true] as the last param in runKernelFunc(). And the gradient above expects two things to exist in the saved list. Can you verify why a unit test didn't fail? Is saved[1] unused?

I added outputsToSave. This wasn't caught because we're excluding gradients in the EXCLUDE_LIST. But it doesn't work anyway because computing gradients in this case requires an equal kernel.

[dsmilkov]

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

---

tfjs-backend-wasm/src/setup_test.ts, line 64 at r2 (raw file):

Previously, annxingyuan (Ann Yuan) wrote…

Done

Thanks! Can you be more specific with the 'pool ' string below too?

---

tfjs-backend-wasm/src/cc/kernels/Max.cc, line 54 at r3 (raw file):

    }

    x_offset = x_offset + reduce_size;

use +=. In C++ this matters

---

tfjs-backend-wasm/src/cc/kernels/Min.cc, line 54 at r3 (raw file):

    }

    x_offset = x_offset + reduce_size;

+=

---

tfjs-core/src/ops/reduction_ops.ts, line 386 at r2 (raw file):

Previously, annxingyuan (Ann Yuan) wrote…

I added outputsToSave. This wasn't caught because we're excluding gradients in the EXCLUDE_LIST. But it doesn't work anyway because computing gradients in this case requires an equal kernel.

Got it. Thanks!

[dsmilkov]

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

---

tfjs-backend-wasm/src/cc/kernels/Min.cc, line 47 at r3 (raw file):

    float* x_iter_end = x_offset + reduce_size;

    for (float* x = x_offset; x != x_iter_end; ++x) {

also let's use < instead of != to make it a standard for loop which compilers can recognize and potentially optimize. here and in Max.cc

[annxingyuan]

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

---

tfjs-backend-wasm/src/setup_test.ts, line 64 at r2 (raw file):

Previously, dsmilkov (Daniel Smilkov) wrote…

Thanks! Can you be more specific with the 'pool ' string below too?

Done

---

tfjs-backend-wasm/src/cc/kernels/Max.cc, line 54 at r3 (raw file):

Previously, dsmilkov (Daniel Smilkov) wrote…

use +=. In C++ this matters

Done

---

tfjs-backend-wasm/src/cc/kernels/Min.cc, line 47 at r3 (raw file):

Previously, dsmilkov (Daniel Smilkov) wrote…

also let's use < instead of != to make it a standard for loop which compilers can recognize and potentially optimize. here and in Max.cc

Done

---

tfjs-backend-wasm/src/cc/kernels/Min.cc, line 54 at r3 (raw file):

Previously, dsmilkov (Daniel Smilkov) wrote…

+=

Done