NNotepad: Add ability to show additional tensors in the output (#311)

Adds `output()` helper which can be used to list an identifier (or
multiple identifiers) that will be shown in the results pane before
the result of the final expression or assignment.

A few notes:

- Ending with an `output()` is supported, as a convenience.

- The `output()` helper takes an identifier, not an expression. In the
  future this could change, either to taking an expression or showing
  the identifier name in the results pane.

- Mixing `split()` (which returns a list of tensors) and `output()` is
  probably a bit confusing, since split already gives you multiple
  tensors, you're gonna have to pay attention to how many are coming
  from each. Sorry!

Also fixes a minor issue with the call to asyncInit in js/tests.js -
it wasn't awaited, so if the first test case was a test with a call it
would fail.

Author: inexorabletash

nnotepad/README.md

@@ -50,6 +50,7 @@ In addition to WebNN [`MLGraphBuilder`](https://webmachinelearning.github.io/web
 
 * **load(_url_, _shape_, _dataType_)** - fetch a tensor resource. Must be served with appropriate [CORS](https://developer.mozilla.org/en-US/docs/Web/HTTP/CORS) headers. Example: `load('https://www.random.org/cgi-bin/randbyte?nbytes=256', [16, 16], 'uint8')`
 * **zeros(_shape_, _dataType_)** - constant zero-filled tensor of the given shape. Example: `zeros([2,2,2,2], 'int8')`
+* **output(_identifier_, ...)** - show the named variable(s) as an additional output, in addition to the last expression result. Example: `T = [1,2]  output(T)  mul(T,3)`
 
 
 # Details & Gotchas

nnotepad/js/nnotepad.js

@@ -432,22 +432,31 @@ export class NNotepad {
     // Generates WebNN code as the body of a function. `_` is passed as the
     // `MLGraphBuilder`. The output of the last expression is returned.
 
+    const kUtilArgName = '_util_';
+    const kOutputsArgName = '_outputs_';
+
     const src = lines
-        .map(
-            (line, index) =>
-              serializeLine(line, index === lines.length - 1))
+        .map((line, index) =>
+          serializeLine(line, index === lines.length - 1))
         .map((line) => line + ';\n')
         .join('');
     const AsyncFunction = async function() {}.constructor;
-    return [new AsyncFunction(['_', 'Util'], src), src];
+    return [new AsyncFunction(['_', kUtilArgName, kOutputsArgName], src), src];
 
     function serializeLine(line, last) {
       const expr = serializeExpr(line.expr);
+
+      // If the last thing is an `output()` call, don't wrap it.
+      const isOutputCall = line.type === 'expression' &&
+            line.expr.type === 'call' && line.expr.identifier === 'output';
+      const wrapAsOutput = last && !isOutputCall;
+
       switch (line.type) {
         case 'assignment':
-          return last ? `return ${expr}` : `const ${line.identifier} = ${expr}`;
+          return wrapAsOutput ? `${kOutputsArgName}.push(${expr})` :
+                                `const ${line.identifier} = ${expr}`;
         case 'expression':
-          return last ? `return ${expr}` : expr;
+          return wrapAsOutput ? `${kOutputsArgName}.push(${expr})` : expr;
       }
       throw new Error(`unexpected line type: ${line.type}`);
     }
@@ -588,7 +597,7 @@ export class NNotepad {
         return `_.constant({dataType: "${dataType.value}", dimensions: ${
           Util.stringify(dims)}, shape: ${
           Util.stringify(dims)}}, new ${
-          ctor.name}(await Util.loadBuffer(${
+          ctor.name}(await ${kUtilArgName}.${Util.loadBuffer.name}(${
           Util.stringify(url.value)})).buffer)`;
       }
 
@@ -609,6 +618,15 @@ export class NNotepad {
           ctor.name}(${len}).buffer)`;
       }
 
+      if (name === 'output') {
+        return args.map((arg) => {
+          if (arg.type !== 'identifier') {
+            throw new TypeError('output(): expected identifier');
+          }
+          return `${kOutputsArgName}.push(${arg.value})`;
+        }).join(';\n');
+      }
+
       return `_.${name}(${
         args.map(
             (arg, index) =>
@@ -630,21 +648,21 @@ export class NNotepad {
     const builder = new self.MLGraphBuilder(context);
 
     const outputOperands = [];
-    let output = await builderFunc(builder, Util);
-    if (output instanceof self.MLOperand) {
-      // TODO: remove try/catch once all back-ends support `identity()`.
-      try {
-        // In case `output` is a constant.
-        output = builder.identity(output);
-      } catch (ex) {
-        // Just live with it for now.
+    const outputs = [];
+    await builderFunc(builder, Util, outputs);
+    for (const output of outputs.flat()) {
+      if (output instanceof self.MLOperand) {
+        // TODO: remove try/catch once all back-ends support `identity()`.
+        try {
+          // In case `output` is a constant.
+          outputOperands.push(builder.identity(output));
+        } catch (ex) {
+          // Just live with it for now.
+          outputOperands.push(output);
+        }
+      } else {
+        throw new ParseError(`Non-MLOperand output: ${output}`);
       }
-      outputOperands.push(output);
-    } else if (Array.isArray(output)) {
-      outputOperands.push(...output);
-      // no-op
-    } else {
-      throw new ParseError(`Non-MLOperand output: ${output}`);
     }
 
     const namedOutputs = {};

nnotepad/js/tests.js

@@ -1,5 +1,5 @@
-import {Harness} from './testharness.js';
 import {NNotepad} from './nnotepad.js';
+import {Harness} from './testharness.js';
 
 // ============================================================
 // Helper for NNotepad-specific tests
@@ -60,7 +60,7 @@ async function testThrows(expr) {
 // ============================================================
 
 document.addEventListener('DOMContentLoaded', async (e) => {
-  NNotepad.asyncInit();
+  await NNotepad.asyncInit();
 
   Harness.section('Numbers');
   await test('125', {dataType: 'float32', shape: [], buffer: [125]});
@@ -158,6 +158,29 @@ document.addEventListener('DOMContentLoaded', async (e) => {
     {dataType: 'float32', shape: [2], buffer: [1, 2]},
     {dataType: 'float32', shape: [2], buffer: [3, 4]},
   ]);
+  await test(`A = [1,2]  output(A)  B = [3,4]`, [
+    {dataType: 'float32', shape: [2], buffer: [1, 2]},
+    {dataType: 'float32', shape: [2], buffer: [3, 4]},
+  ]);
+  await test(`A = [1,2]  output(A)  B = [3,4]  output(B)`, [
+    {dataType: 'float32', shape: [2], buffer: [1, 2]},
+    {dataType: 'float32', shape: [2], buffer: [3, 4]},
+  ]);
+  await test(`A = [1,2]  output(A)  split([1,2,3,4], 2)`, [
+    {dataType: 'float32', shape: [2], buffer: [1, 2]},
+    {dataType: 'float32', shape: [2], buffer: [1, 2]},
+    {dataType: 'float32', shape: [2], buffer: [3, 4]},
+  ]);
+  await test(`array = split([1,2,3,4], 2)  output(array)`, [
+    {dataType: 'float32', shape: [2], buffer: [1, 2]},
+    {dataType: 'float32', shape: [2], buffer: [3, 4]},
+  ]);
+  await test(
+      `A = [[1,7],[2,4]]  B = [[3,3],[5,2]]  output(A, B)  matmul(A,B)`, [
+        {dataType: 'float32', shape: [2, 2], buffer: [1, 7, 2, 4]},
+        {dataType: 'float32', shape: [2, 2], buffer: [3, 3, 5, 2]},
+        {dataType: 'float32', shape: [2, 2], buffer: [38, 17, 26, 14]},
+      ]);
 
   Harness.section('Non-operand arguments: array of operands');
   await test(

nnotepad/res/docs.html

@@ -110,6 +110,7 @@ <h1>Helpers</h1>
 <ul>
 <li><strong>load(<em>url</em>, <em>shape</em>, <em>dataType</em>)</strong> - fetch a tensor resource. Must be served with appropriate <a href="https://developer.mozilla.org/en-US/docs/Web/HTTP/CORS">CORS</a> headers. Example: <code>load('https://www.random.org/cgi-bin/randbyte?nbytes=256', [16, 16], 'uint8')</code></li>
 <li><strong>zeros(<em>shape</em>, <em>dataType</em>)</strong> - constant zero-filled tensor of the given shape. Example: <code>zeros([2,2,2,2], 'int8')</code></li>
+<li><strong>output(<em>identifier</em>, ...)</strong> - show the named variable(s) as an additional output, in addition to the last expression result. Example: <code>T = [1,2]  output(T)  mul(T,3)</code></li>
 </ul>
 <h1>Details &amp; Gotchas</h1>
 <ul>