Add async versions of dynamic ops: nonMaxSuppressionAsync, whereAsync

[dsmilkov]

Some ops in TF are dynamic in the sense that the output shape depends on the input's values (e.g. nonMaxSupression, where(cond)). For those dynamic ops, we add an async version to avoid blocking the main UI thread, since their regular implementation must call .dataSync().

Details:

• Add tf.nonMaxSuppressionAsync and tf.whereAsync since their regular counterparts are blocking the main UI thread.
• The regular version of those ops still works with a webgl backend, but will warn with a message: tf.opName() in WebGL locks the UI thread. Use tf.opNameAsync() instead.
• For both ops, since the implementation is shared across backends, move the common code to kernels/${opname}_impl.ts
• Split backend.where into backend.select and backend.where to match the kernel names in TF C.
• tf.where() will still require 2nd and 3rd params (x and y) and will call into backend.select() to match the Select kernel in C.
• tf.whereAsync() on the other hand, will take only 1 param (condition) and match the Where kernel in C.
• The plan is for tfjs-converter to notify the users to use frozenModel.predictAsync() whenever the model has dynamic ops. predictAsync() will internally call into the async versions of the dynamic ops.

FEATURE

---

This change is 

[nsthorat]

Reviewed 13 of 13 files at r1.
Reviewable status: 0 of 1 approvals obtained (waiting on @dsmilkov, @nsthorat, and @pyu10055)

---

src/kernels/non_max_suppression_impl.ts, line 62 at r1 (raw file):

}

function IOU(boxes: TypedArray, i: number, j: number) {

how about intersectionOverUnion

---

src/ops/image_ops.ts, line 203 at r1 (raw file):

function nonMaxSuppSanityCheck(
    boxes: Tensor2D, scores: Tensor1D, maxOutputSize: number,
    iouThreshold: number, scoreThreshold: number) {

return type

---

src/ops/topk.ts, line 59 at r1 (raw file):

/** This is the async version of `topk` */
async function topkAsync_<T extends Tensor>(

no need for this, lets just dataSync in topk and fix this on the GPU later.

[pyu10055]

Reviewed 13 of 13 files at r1.
Reviewable status: complete! 1 of 1 approvals obtained (waiting on @dsmilkov and @pyu10055)

---

src/kernels/backend.ts, line 246 at r1 (raw file):

  nonMaxSuppression(
      boxes: Tensor2D, scores: Tensor1D, maxOutputSize: number,
      iouThreshold: number, scoreThreshold: number): Tensor1D;

can you make the scoreThreshold optional? It is only defined in the NonMaxSuppresionV3, but not in NonMaxSuppressionV2. Thanks

[dsmilkov]

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

---

src/kernels/backend.ts, line 246 at r1 (raw file):

Previously, pyu10055 (Ping Yu) wrote…

can you make the scoreThreshold optional? It is only defined in the NonMaxSuppresionV3, but not in NonMaxSuppressionV2. Thanks

Done.

---

src/kernels/non_max_suppression_impl.ts, line 62 at r1 (raw file):

Previously, nsthorat (Nikhil Thorat) wrote…

how about intersectionOverUnion

Done. This is copied verbatim from existing code.

---

src/ops/image_ops.ts, line 203 at r1 (raw file):

Previously, nsthorat (Nikhil Thorat) wrote…

return type

Done.

---

src/ops/topk.ts, line 59 at r1 (raw file):

Previously, nsthorat (Nikhil Thorat) wrote…

no need for this, lets just dataSync in topk and fix this on the GPU later.

Done