Use MLTensor and dispatch() and deprecate compute()

code/.eslintrc.js

@@ -3,6 +3,7 @@ module.exports = {
     'CodeMirror': 'readonly',
     'executeCodeSnippet': 'readonly',
     'sizeOfShape': 'readonly',
+    'MLTensorUsage': 'readonly',
   },
   ignorePatterns: ['libs/'],
 };

code/samples/matmul.js

@@ -1,25 +1,30 @@
 // Step 0: Create a context and graph builder for 'gpu', 'cpu' or 'npu'.
 const context = await navigator.ml.createContext({deviceType: 'gpu'});
 const builder = new MLGraphBuilder(context);
+const descA = {dataType: 'float32', dimensions: [3, 4], shape: [3, 4]};
+const descB = {dataType: 'float32', dimensions: [4, 3], shape: [4, 3]};
 // Step 1: Create a computational graph calculating `c = a * b`.
-const a = builder.input('a', {
-  dataType: 'float32',
-  dimensions: [3, 4],
-  shape: [3, 4],
-});
-const b = builder.input('b', {
-  dataType: 'float32',
-  dimensions: [4, 3],
-  shape: [4, 3],
-});
+const a = builder.input('a', descA);
+const b = builder.input('b', descB);
 const c = builder.matmul(a, b);
 // Step 2: Compile it into an executable graph.
 const graph = await builder.build({c});
 // Step 3: Bind input and output buffers to the graph and execute.
 const bufferA = new Float32Array(3*4).fill(1.0);
 const bufferB = new Float32Array(4*3).fill(0.8);
-const bufferC = new Float32Array(3*3);
-const results = await context.compute(
-    graph, {'a': bufferA, 'b': bufferB}, {'c': bufferC});
+descA.usage = MLTensorUsage.WRITE;
+descB.usage = MLTensorUsage.WRITE;
+const tensorA = await context.createTensor(descA);
+const tensorB = await context.createTensor(descB);
+context.writeTensor(tensorA, bufferA);
+context.writeTensor(tensorB, bufferB);
+const tensorC = await context.createTensor({
+  dataType: 'float32',
+  dimensions: [3, 3],
+  shape: [3, 3],
+  usage: MLTensorUsage.READ,
+});
+context.dispatch(graph, {a: tensorA, b: tensorB}, {c: tensorC});
+const results = await context.readTensor(tensorC);
 // Step 4: Retrieve the results.
-console.log(`values: ${results.outputs.c}`);
+console.log(`values: ${new Float32Array(results)}`);

code/samples/mul_add.js

@@ -1,21 +1,30 @@
-const operandType = {dataType: 'float32', dimensions: [2, 2], shape: [2, 2]};
+const desc = {dataType: 'float32', dimensions: [2, 2], shape: [2, 2]};
 const context = await navigator.ml.createContext();
 const builder = new MLGraphBuilder(context);
 // 1. Create a computational graph 'C = 0.2 * A + B'.
 const constant = builder.constant(
     {dataType: 'float32'}, new Float32Array([0.2]));
-const A = builder.input('A', operandType);
-const B = builder.input('B', operandType);
+const A = builder.input('A', desc);
+const B = builder.input('B', desc);
 const C = builder.add(builder.mul(A, constant), B);
 // 2. Build the graph into an executable.
 const graph = await builder.build({'C': C});
 // 3. Bind inputs to the graph and execute for the result.
 const bufferA = new Float32Array(4).fill(1.0);
 const bufferB = new Float32Array(4).fill(0.8);
-const bufferC = new Float32Array(4);
-const inputs = {'A': bufferA, 'B': bufferB};
-const outputs = {'C': bufferC};
-const results = await context.compute(graph, inputs, outputs);
+desc.usage = MLTensorUsage.WRITE;
+const tensorA = await context.createTensor(desc);
+const tensorB = await context.createTensor(desc);
+context.writeTensor(tensorA, bufferA);
+context.writeTensor(tensorB, bufferB);
+const tensorC = await context.createTensor({
+  ...desc,
+  usage: MLTensorUsage.READ,
+});
+const inputs = {'A': tensorA, 'B': tensorB};
+const outputs = {'C': tensorC};
+context.dispatch(graph, inputs, outputs);
 // The computed result of [[1, 1], [1, 1]] is in the buffer associated with
 // the output operand.
-console.log('Output value: ' + results.outputs.C);
+const results = await context.readTensor(tensorC);
+console.log('Output value: ' + new Float32Array(results));

code/samples/simple_graph.js

@@ -50,15 +50,26 @@ const graph = await builder.build({'output': output});
 // Setup the input buffers with value 1.
 const inputBuffer1 = new Float32Array(TENSOR_SIZE).fill(1);
 const inputBuffer2 = new Float32Array(TENSOR_SIZE).fill(1);
-const outputBuffer = new Float32Array(TENSOR_SIZE);
+
+desc.usage = MLTensorUsage.WRITE;
+const inputTensor1 = await context.createTensor(desc);
+const inputTensor2 = await context.createTensor(desc);
+context.writeTensor(inputTensor1, inputBuffer1);
+context.writeTensor(inputTensor2, inputBuffer2);
+
+const outputTensor = await context.createTensor({
+  ...desc,
+  usage: MLTensorUsage.READ,
+});
 
 // Execute the compiled graph with the specified inputs.
 const inputs = {
-  'input1': inputBuffer1,
-  'input2': inputBuffer2,
+  'input1': inputTensor1,
+  'input2': inputTensor2,
 };
-const outputs = {'output': outputBuffer};
-const results = await context.compute(graph, inputs, outputs);
+const outputs = {'output': outputTensor};
+context.dispatch(graph, inputs, outputs);
 
-console.log('Output value: ' + results.outputs.output);
+const results = await context.readTensor(outputTensor);
+console.log('Output value: ' + new Float32Array(results));
 // Output value: 2.25,2.25,2.25,2.25,2.25,2.25,2.25,2.25

face_recognition/.eslintrc.js

@@ -1,6 +1,7 @@
 module.exports = {
   globals: {
     'MLGraphBuilder': 'readonly',
+    'MLTensorUsage': 'readonly',
     'tf': 'readonly',
   },
 };

face_recognition/facenet_nchw.js

@@ -12,6 +12,8 @@ export class FaceNetNchw {
     this.context_ = null;
     this.builder_ = null;
     this.graph_ = null;
+    this.inputTensor_ = null;
+    this.outputTensor_ = null;
     this.weightsUrl_ = weightsOrigin() +
       '/test-data/models/facenet_nchw/weights';
     this.inputOptions = {
@@ -25,6 +27,7 @@ export class FaceNetNchw {
       distanceMetric: 'euclidean',
       threshold: 1.26,
     };
+    this.outputShape_ = [1, 512];
   }
 
   async buildConv_(
@@ -138,10 +141,19 @@ export class FaceNetNchw {
   async load(contextOptions) {
     this.context_ = await navigator.ml.createContext(contextOptions);
     this.builder_ = new MLGraphBuilder(this.context_);
-    const input = this.builder_.input('input', {
+    const inputDesc = {
       dataType: 'float32',
       dimensions: this.inputOptions.inputShape,
       shape: this.inputOptions.inputShape,
+    };
+    const input = this.builder_.input('input', inputDesc);
+    inputDesc.usage = MLTensorUsage.WRITE;
+    this.inputTensor_ = await this.context_.createTensor(inputDesc);
+    this.outputTensor_ = await this.context_.createTensor({
+      dataType: 'float32',
+      dimensions: this.outputShape_,
+      shape: this.outputShape_,
+      usage: MLTensorUsage.READ,
     });
 
     const poolOptions = {windowDimensions: [3, 3], strides};
@@ -272,9 +284,12 @@ export class FaceNetNchw {
     }
   }
 
-  async compute(inputBuffer, outputs) {
-    const inputs = {'input': inputBuffer};
-    const results = await this.context_.compute(this.graph_, inputs, outputs);
-    return results;
+  async compute(inputBuffer) {
+    this.context_.writeTensor(this.inputTensor_, inputBuffer);
+    const inputs = {'input': this.inputTensor_};
+    const outputs = {'output': this.outputTensor_};
+    this.context_.dispatch(this.graph_, inputs, outputs);
+    const results = await this.context_.readTensor(this.outputTensor_);
+    return new Float32Array(results);
   }
 }

face_recognition/facenet_nhwc.js

@@ -12,6 +12,8 @@ export class FaceNetNhwc {
     this.context_ = null;
     this.builder_ = null;
     this.graph_ = null;
+    this.inputTensor_ = null;
+    this.outputTensor_ = null;
     this.weightsUrl_ = weightsOrigin() +
       '/test-data/models/facenet_nhwc/weights';
     this.inputOptions = {
@@ -25,6 +27,7 @@ export class FaceNetNhwc {
       distanceMetric: 'euclidean',
       threshold: 1.26,
     };
+    this.outputShape_ = [1, 512];
   }
 
   async buildConv_(input, namePrefix, options = undefined, relu = true) {
@@ -139,10 +142,19 @@ export class FaceNetNhwc {
   async load(contextOptions) {
     this.context_ = await navigator.ml.createContext(contextOptions);
     this.builder_ = new MLGraphBuilder(this.context_);
-    const input = this.builder_.input('input', {
+    const inputDesc = {
       dataType: 'float32',
       dimensions: this.inputOptions.inputShape,
       shape: this.inputOptions.inputShape,
+    };
+    const input = this.builder_.input('input', inputDesc);
+    inputDesc.usage = MLTensorUsage.WRITE;
+    this.inputTensor_ = await this.context_.createTensor(inputDesc);
+    this.outputTensor_ = await this.context_.createTensor({
+      dataType: 'float32',
+      dimensions: this.outputShape_,
+      shape: this.outputShape_,
+      usage: MLTensorUsage.READ,
     });
 
     const poolOptions = {windowDimensions: [3, 3], strides, layout: 'nhwc'};
@@ -241,9 +253,12 @@ export class FaceNetNhwc {
     }
   }
 
-  async compute(inputBuffer, outputs) {
-    const inputs = {'input': inputBuffer};
-    const results = await this.context_.compute(this.graph_, inputs, outputs);
-    return results;
+  async compute(inputBuffer) {
+    this.context_.writeTensor(this.inputTensor_, inputBuffer);
+    const inputs = {'input': this.inputTensor_};
+    const outputs = {'output': this.outputTensor_};
+    this.context_.dispatch(this.graph_, inputs, outputs);
+    const results = await this.context_.readTensor(this.outputTensor_);
+    return new Float32Array(results);
   }
 }

face_recognition/main.js

@@ -32,8 +32,6 @@ let stream = null;
 let loadTime = 0;
 let buildTime = 0;
 let computeTime = 0;
-let fdOutputs;
-let frOutputs;
 let deviceType = '';
 let lastdeviceType = '';
 let backend = '';
@@ -180,15 +178,14 @@ async function getEmbeddings(inputElem) {
   const fdInputBuffer = utils.getInputTensor(inputElem, fdInputOptions);
   let totalComputeTime = 0;
   let start = performance.now();
-  const results = await fdInstance.compute(fdInputBuffer, fdOutputs);
+  const results = await fdInstance.compute(fdInputBuffer);
   totalComputeTime = performance.now() - start;
-  fdOutputs = results.outputs;
   const strokedRects = [];
   const embeddings = [];
   const height = inputElem.naturalHeight || inputElem.height;
   const width = inputElem.naturalWidth || inputElem.width;
   const fdOutputArrary = [];
-  for (const output of Object.entries(fdOutputs)) {
+  for (const output of Object.entries(results)) {
     fdOutputArrary.push(output[1]);
   }
   const fdSsdOutputs = SsdDecoder.processSsdOutputTensor(
@@ -222,10 +219,9 @@ async function getEmbeddings(inputElem) {
     frInputOptions.drawOptions = drawOptions;
     const frInputBuffer = utils.getInputTensor(inputElem, frInputOptions);
     start = performance.now();
-    const results = await frInstance.compute(frInputBuffer, frOutputs);
+    const results = await frInstance.compute(frInputBuffer);
     totalComputeTime += performance.now() - start;
-    frOutputs = results.outputs;
-    const [...normEmbedding] = Float32Array.from(frOutputs.output);
+    const [...normEmbedding] = Float32Array.from(results);
     embeddings.push(normEmbedding);
   }
   return {computeTime: totalComputeTime, strokedRects, embeddings};
@@ -330,12 +326,6 @@ async function main() {
       frInstance = constructNetObject(frInstanceType);
       fdInputOptions = fdInstance.inputOptions;
       frInputOptions = frInstance.inputOptions;
-      fdOutputs = {};
-      for (const outputInfo of Object.entries(fdInstance.outputsInfo)) {
-        fdOutputs[outputInfo[0]] =
-            new Float32Array(utils.sizeOfShape(outputInfo[1]));
-      }
-      frOutputs = {'output': new Float32Array(utils.sizeOfShape([1, 512]))};
       isFirstTimeLoad = false;
       console.log(`- Model name: ${fdModelName}, Model layout: ${layout} -`);
       // UI shows model loading progress
@@ -374,12 +364,10 @@ async function main() {
       let medianComputeTime;
       console.log('- Computing... ');
       // Do warm up
-      const fdResults = await fdInstance.compute(new Float32Array(
-          utils.sizeOfShape(fdInputOptions.inputShape)), fdOutputs);
-      const frResults = await frInstance.compute(new Float32Array(
-          utils.sizeOfShape(frInputOptions.inputShape)), frOutputs);
-      fdOutputs = fdResults.outputs;
-      frOutputs = frResults.outputs;
+      await fdInstance.compute(new Float32Array(
+          utils.sizeOfShape(fdInputOptions.inputShape)));
+      await frInstance.compute(new Float32Array(
+          utils.sizeOfShape(frInputOptions.inputShape)));
       for (let i = 0; i < numRuns; i++) {
         if (numRuns > 1) {
           // clear all predicted embeddings for benckmarking

facial_landmark_detection/.eslintrc.js

@@ -1,6 +1,7 @@
 module.exports = {
   globals: {
     'MLGraphBuilder': 'readonly',
+    'MLTensorUsage': 'readonly',
     'tf': 'readonly',
   },
 };

facial_landmark_detection/face_landmark_nchw.js

@@ -8,12 +8,15 @@ export class FaceLandmarkNchw {
     this.context_ = null;
     this.builder_ = null;
     this.graph_ = null;
+    this.inputTensor_ = null;
+    this.outputTensor_ = null;
     this.weightsUrl_ = weightsOrigin() +
       '/test-data/models/face_landmark_nchw/weights';
     this.inputOptions = {
       inputLayout: 'nchw',
       inputShape: [1, 3, 128, 128],
     };
+    this.outputShape_ = [1, 136];
   }
 
   async buildMaxPool2d(input, options) {
@@ -69,10 +72,19 @@ export class FaceLandmarkNchw {
   async load(contextOptions) {
     this.context_ = await navigator.ml.createContext(contextOptions);
     this.builder_ = new MLGraphBuilder(this.context_);
-    const input = this.builder_.input('input', {
+    const inputDesc = {
       dataType: 'float32',
       dimensions: this.inputOptions.inputShape,
       shape: this.inputOptions.inputShape,
+    };
+    const input = this.builder_.input('input', inputDesc);
+    inputDesc.usage = MLTensorUsage.WRITE;
+    this.inputTensor_ = await this.context_.createTensor(inputDesc);
+    this.outputTensor_ = await this.context_.createTensor({
+      dataType: 'float32',
+      dimensions: this.outputShape_,
+      shape: this.outputShape_,
+      usage: MLTensorUsage.READ,
     });
 
     const poolOptions =
@@ -112,9 +124,12 @@ export class FaceLandmarkNchw {
     }
   }
 
-  async compute(inputBuffer, outputs) {
-    const inputs = {'input': inputBuffer};
-    const results = await this.context_.compute(this.graph_, inputs, outputs);
-    return results;
+  async compute(inputBuffer) {
+    this.context_.writeTensor(this.inputTensor_, inputBuffer);
+    const inputs = {'input': this.inputTensor_};
+    const outputs = {'output': this.outputTensor_};
+    this.context_.dispatch(this.graph_, inputs, outputs);
+    const results = await this.context_.readTensor(this.outputTensor_);
+    return new Float32Array(results);
   }
 }