[wasm] Add pow, fusedMatMul.

[annxingyuan]

Changes

• Create Pow kernel (unblocks KNN model)
• Create FusedBatchMatMul kernel which supports 1D bias (unblocks speech-commands model)
• Create batch_mat_mul_impl to be shared among FusedBatchMatMul, BatchMatMul
• Move FusableActivation to backend so it can be shared among fused kernels

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

---

This change is 

[nsthorat]

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

---

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

#include <cstddef>

#include "src/cc/backend.h"

do you need this import?

---

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

#include <cstddef>

#include "src/cc/backend.h"

do you need this

---

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

  switch (dtype) {
    case DType::float32:
      binary_f32(a_id, b_id, out_id, power<float>);

can you not pass pow here directly?

---

tfjs-core/src/ops/binary_ops.ts, line 289 at r1 (raw file):

    save([$base, $exp, y]);
    return y;
  }, {a: $base, b: $exp}, grad, 'Pow', attrs, inputsToSave, outputsToSave) as T;

can we do base and exp instead of a and b?

[dsmilkov]

Reviewed 12 of 26 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/cc/BUILD, line 327 at r2 (raw file):

    hdrs = ["kernels/Conv2D.h"],
    deps = [
        ":backend",

add adding a dep on "backend" since nothing changed in those files, here and below

---

tfjs-backend-wasm/src/cc/batch_mat_mul_impl.h, line 1 at r2 (raw file):

/* Copyright 2019 Google Inc. All Rights Reserved.

2020 here and elsewhere

---

tfjs-backend-wasm/src/cc/batch_mat_mul_impl.h, line 23 at r2 (raw file):

namespace wasm {

void batchMatMul(const size_t a_id, const size_t* a_shape_ptr,

rename to "fused_batch_matmul" ? (more explicit what it does). also since this is an internal c++ function, use snake_casing (c++ style guide)

---

tfjs-backend-wasm/src/cc/batch_mat_mul_impl.cc, line 16 at r2 (raw file):

#ifdef __EMSCRIPTEN__
#include <emscripten.h>

remove import of emcsripten.h ?

---

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

#include <cstddef>

#include "src/cc/backend.h"

do you need backend.h?

---

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

  const res = ENGINE.runKernelFunc(
      (backend, save) => {

can you fit in 1 line

---

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

        return y;
      },
      inputs, grad, 'FusedBatchMatMul', {transposeA, transposeB, activation},

is this going to be the official kernel name? I see the converter uses '_FusedMatMul'

[annxingyuan]

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

---

tfjs-backend-wasm/src/cc/BUILD, line 327 at r2 (raw file):

Previously, dsmilkov (Daniel Smilkov) wrote…

add adding a dep on "backend" since nothing changed in those files, here and below

I think I do need this because Conv2D requires the enum of FusableActivations, which now lives in backend.h

---

tfjs-backend-wasm/src/cc/batch_mat_mul_impl.h, line 1 at r2 (raw file):

Previously, dsmilkov (Daniel Smilkov) wrote…

2020 here and elsewhere

Done

---

tfjs-backend-wasm/src/cc/batch_mat_mul_impl.h, line 23 at r2 (raw file):

Previously, dsmilkov (Daniel Smilkov) wrote…

rename to "fused_batch_matmul" ? (more explicit what it does). also since this is an internal c++ function, use snake_casing (c++ style guide)

Done

---

tfjs-backend-wasm/src/cc/batch_mat_mul_impl.cc, line 16 at r2 (raw file):

Previously, dsmilkov (Daniel Smilkov) wrote…

remove import of emcsripten.h ?

Done

---

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

Previously, dsmilkov (Daniel Smilkov) wrote…

do you need backend.h?

Yes that is where the FusableActivation enum is now defined.

---

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

Previously, nsthorat (Nikhil Thorat) wrote…

do you need this import?

Yes that is where the FusableActivation enum is now defined.

---

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

Previously, nsthorat (Nikhil Thorat) wrote…

do you need this

Yes that is where the FusableActivation enum is now defined.

---

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

Previously, nsthorat (Nikhil Thorat) wrote…

can you not pass pow here directly?

If I pass pow directly it complains about no overload of 'pow' matching the different dtypes.

---

tfjs-core/src/ops/binary_ops.ts, line 289 at r1 (raw file):

Previously, nsthorat (Nikhil Thorat) wrote…

can we do base and exp instead of a and b?

'a' and 'b' are the names assumed by the shared binary_kernel.ts

---

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

Previously, dsmilkov (Daniel Smilkov) wrote…

  const res = ENGINE.runKernelFunc(
      (backend, save) => {

can you fit in 1 line

clang doesn't seem to want me to do that?

---

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

Previously, dsmilkov (Daniel Smilkov) wrote…

is this going to be the official kernel name? I see the converter uses '_FusedMatMul'

Done

changed to _FusedMatMul

[dsmilkov]

LGTM! One optional comment about pow

Reviewed 9 of 26 files at r1, 12 of 12 files at r3.
Reviewable status: complete! 2 of 1 approvals obtained (waiting on @annxingyuan and @nsthorat)

---

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

Previously, annxingyuan (Ann Yuan) wrote…

If I pass pow directly it complains about no overload of 'pow' matching the different dtypes.

Is pow generic? That is, can you pass pow<float> directly here?

[annxingyuan]

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

---

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

Previously, dsmilkov (Daniel Smilkov) wrote…

Is pow generic? That is, can you pass pow<float> directly here?

Done