feat: implement WebNNProgram class for graph management and tensor operations in WebNNView

Author: caperaven

app/web-nn/web-nn.js

@@ -1,16 +1,11 @@
 import { EventsManager } from "../../src/system/events-manager.js";
-import { 
-    WebNNModule, 
-    OperandDescriptorBuilder } from "../../src/modules/webnn.js";
+import { WebNNProgram } from "../../src/modules/webnn/webnn-program.js";
 
 export default class WebNNView extends HTMLElement {
     static tag = "webnn-view";
 
     #eventsManager = new EventsManager();
-    #graphBuilder
-    #graph;
-    #context;
-    #tensors;
+    #program = new WebNNProgram();
 
     constructor() {
         super();
@@ -22,65 +17,39 @@ export default class WebNNView extends HTMLElement {
 			url: import.meta.url,
 		});  
 
-        requestAnimationFrame(() => {
+        requestAnimationFrame(async () => {
             const button = this.shadowRoot.querySelector("button");
             this.#eventsManager.addEvent(button, "click", this.#submit.bind(this));
-        })
-        
-        this.#context = await WebNNModule.createContext();
-        this.#graphBuilder = await WebNNModule.createGraph( { context: this.#context } );
-        const operands = await this.#createGraph();
-        await this.#createTensors(operands.A, operands.B, operands.C);
+
+            // 1. initialize
+            await this.#program.init();
+            
+            // 2. build graph
+            const descriptor = {dataType: 'float32', shape: [1]};
+            const A = this.#program.addToGraph("input", "A", descriptor);
+            const B = this.#program.addToGraph("input", "B", descriptor);
+            const C = this.#program.addToGraph("add", A, B);
+            this.#program.build({C});
+
+            // 3. add input and output tensors
+            await this.#program.addInputTensor("A", A);
+            await this.#program.addInputTensor("B", B);
+            await this.#program.addOutputTensor("C", C);
+        })        
     }
 
     async disconnectedCallback() {
         this.#eventsManager = this.#eventsManager.dispose();
     }
 
-    async #createGraph() {
-        const descriptor = new OperandDescriptorBuilder().build();
-        const A = this.#graphBuilder.input("A", descriptor);
-        const B = this.#graphBuilder.input("B", descriptor);
-        const C = this.#graphBuilder.add(A, B);  
-
-        this.#graph = await this.#graphBuilder.build({C});
-        return { A, B, C };
-    }
-
-    async #createTensors(inputA, inputB, outputC) {
-        this.#tensors = { 
-            inputA: await this.#createTensor(this.#context, inputA, true, true), 
-            inputB: await this.#createTensor(this.#context, inputB, true, true), 
-            outputC: await this.#createTensor(this.#context, outputC, false, true)
-        };
-    }
-
-    async #createTensor(context, operand, writable, readable) {
-        return context.createTensor({
-            dataType: operand.dataType, shape: operand.shape, writable, readable
-        })
-    }
-
     async #submit(event) {
         const value1 = this.shadowRoot.querySelector("#value1").value;
         const value2 = this.shadowRoot.querySelector("#value2").value;
 
-        this.#context.writeTensor(this.#tensors.inputA, new Float32Array([value1]));
-        this.#context.writeTensor(this.#tensors.inputB, new Float32Array([value2]));
-
-        const inputs = {
-            'A': this.#tensors.inputA,
-            'B': this.#tensors.inputB
-        };
-
-        const outputs = {
-            'C': this.#tensors.outputC
-        };
+        await this.#program.set("A", [value1]);
+        await this.#program.set("B", [value2]);
 
-        this.#context.dispatch(this.#graph, inputs, outputs);
-        const output = await this.#context.readTensor(this.#tensors.outputC); 
-        const result = new Float32Array(output)[0];
-        this.shadowRoot.querySelector("#result").textContent = result;
+        this.shadowRoot.querySelector("#result").textContent = await this.#program.run();
     }
 }
 

src/modules/webnn/webnn-program.js

@@ -0,0 +1,66 @@
+import DeviceType from "./device-type.js";
+import PowerOptions from "./power-options.js";
+
+export class WebNNProgram {
+    #context;
+    #graphBuilder;
+    #graph;
+    #inputTensors = {};
+    #outputTensors = {};
+
+    async init(deviceType = DeviceType.GPU, powerPreference = PowerOptions.DEFAULT) {
+        const contextOptions = {
+            deviceType,
+            powerPreference
+        };
+
+        this.#context = await navigator.ml.createContext(contextOptions);
+        this.#graphBuilder = new MLGraphBuilder(this.#context);
+    }
+
+    /**
+     * Create tensor on context
+     * @param {*} context 
+     * @param {*} operand 
+     * @param {*} writable 
+     * @param {*} readable 
+     * @returns 
+     */
+    async #createTensor(operand, writable, readable) {
+        return await this.#context.createTensor({
+            dataType: operand.dataType, shape: operand.shape, writable, readable
+        })
+    }
+
+    async addInputTensor(name, operand) {
+        this.#inputTensors[name] = await this.#createTensor(operand, true, false);
+    }
+
+    async addOutputTensor(name, operand) {
+        this.#outputTensors[name] = await this.#createTensor(operand, false, true);
+    }
+
+    addToGraph(action, ...args) {
+        return this.#graphBuilder[action](...args);
+    }
+
+    async set(name, values) {
+        await this.#context.writeTensor(this.#inputTensors[name], new Float32Array(values));
+    }
+
+    async build(args) {
+        this.#graph = await this.#graphBuilder.build(args);
+    }
+
+    async run() {
+        this.#context.dispatch(this.#graph, this.#inputTensors, this.#outputTensors);
+
+        const outputKey = Object.keys(this.#outputTensors)[0];
+        const output = await this.#context.readTensor(this.#outputTensors[outputKey]);
+        return new Float32Array(output)[0];
+    }
+}
+
+// const program = new WebNNProgram();
+// program.buildGraph();
+// const result = program.run({a: 1, b: 2});