src: allow operators to access graph in initialization (#227)

Author: fs-eire

lib/backend.ts

@@ -32,16 +32,16 @@ export interface SessionHandler {
 
   /**
    * Resolves the operator from the name and opset version; backend specific
-   * @param node
-   * @param opsets
+   * @param node the node to resolve
+   * @param opsets a list of opsets that exported from the model
+   * @param graph the completely initialized graph
    */
+  resolve(node: Graph.Node, opsets: ReadonlyArray<OpSet>, graph: Graph): Operator;
 
-  resolve(node: Graph.Node, opsets: ReadonlyArray<OpSet>): Operator;
   /**
    * This method let's the sessionHandler know that the graph initialization is complete
    * @param graph the completely initialized graph
    */
-
   onGraphInitialized?(graph: Graph): void;
 
   /**

lib/backends/cpu/session-handler.ts

@@ -19,9 +19,9 @@ export class CpuSessionHandler implements SessionHandler {
 
   dispose(): void {}
 
-  resolve(node: Graph.Node, opsets: ReadonlyArray<OpSet>): Operator {
+  resolve(node: Graph.Node, opsets: ReadonlyArray<OpSet>, graph: Graph): Operator {
     const op = resolveOperator(node, opsets, CPU_OP_RESOLVE_RULES);
-    op.initialize(node.attributes);
+    op.initialize(node.attributes, node, graph);
     return op;
   }
 }

lib/backends/wasm/session-handler.ts

@@ -23,9 +23,9 @@ export class WasmSessionHandler implements SessionHandler {
 
   dispose(): void {}
 
-  resolve(node: Graph.Node, opsets: ReadonlyArray<OpSet>): Operator {
+  resolve(node: Graph.Node, opsets: ReadonlyArray<OpSet>, graph: Graph): Operator {
     const op = resolveOperator(node, opsets, this.opResolveRules);
-    op.initialize(node.attributes);
+    op.initialize(node.attributes, node, graph);
     return op;
   }
 }

lib/backends/webgl/session-handler.ts

@@ -56,9 +56,9 @@ export class WebGLSessionHandler implements SessionHandler {
     this.textureDataCache.forEach(td => this.textureManager.releaseTexture(td, true));
     this.textureDataCache = new Map();
   }
-  resolve(node: Graph.Node, opsets: ReadonlyArray<OpSet>): Operator {
+  resolve(node: Graph.Node, opsets: ReadonlyArray<OpSet>, graph: Graph): Operator {
     const op = resolveOperator(node, opsets, WEBGL_OP_RESOLVE_RULES);
-    op.initialize(node.attributes);
+    op.initialize(node.attributes, node, graph);
     return op;
   }
 }

lib/operators.ts

@@ -3,10 +3,11 @@
 
 import {Attribute} from './attribute';
 import {InferenceHandler} from './backend';
+import {Graph} from './graph';
 import {Tensor} from './tensor';
 
 export interface Operator {
-  initialize(attributes: Attribute): void;
+  initialize(attributes: Attribute, node: Graph.Node, graph: Graph): void;
   checkInputs(inputs: Tensor[]): boolean;
   run(inferenceHandler: InferenceHandler, inputs: Tensor[]): Tensor[]|Promise<Tensor[]>;
 }

lib/session.ts

@@ -232,7 +232,7 @@ export class Session {
     this._ops = new Array(nodes.length);
 
     for (let i = 0; i < nodes.length; i++) {
-      this._ops[i] = this.sessionHandler.resolve(nodes[i], this._model.opsets);
+      this._ops[i] = this.sessionHandler.resolve(nodes[i], this._model.opsets, graph);
     }
   }
 

test/test-runner.ts

@@ -16,7 +16,7 @@ import {Logger, Profiler} from '../lib/instrument';
 import {Operator} from '../lib/operators';
 import {Tensor} from '../lib/tensor';
 
-import {base64toBuffer} from './test-shared';
+import {base64toBuffer, createMockGraph} from './test-shared';
 import {Test} from './test-types';
 
 // the threshold that used to compare 2 float numbers. See above for TensorResultValidator.floatEqual().
@@ -301,7 +301,8 @@ function initializeOperator(
     opsetImports: ReadonlyArray<Test.OperatorTestOpsetImport>): Operator {
   const attributes = new Attribute(undefined);
   attributeValues.forEach(value => attributes.set(value.name, value.type, value.data));
-  return sessionHandler.resolve({name: '', opType, inputs: [], outputs: [], attributes}, opsetImports);
+  const graph = createMockGraph(opType, attributes);
+  return sessionHandler.resolve(graph.getNodes()[0], opsetImports, graph);
 }
 
 /**

test/test-shared.ts

@@ -4,6 +4,9 @@
 import * as fs from 'fs';
 import {promisify} from 'util';
 
+import {Attribute} from '../lib/attribute';
+import {Graph} from '../lib/graph';
+
 export function base64toBuffer(data: string): Uint8Array {
   return Buffer.from(data, 'base64');
 }
@@ -28,3 +31,18 @@ async function readFile(file: string) {
     return Buffer.from(buffer);
   }
 }
+
+/**
+ * create a single-node graph for unit test purpose
+ */
+export function createMockGraph(opType: string, attributes: Attribute): Graph {
+  const node: Graph.Node = {name: '', opType, inputs: [], outputs: [], attributes};
+  return {
+    getInputIndices: () => [],
+    getInputNames: () => [],
+    getOutputIndices: () => [],
+    getOutputNames: () => [],
+    getNodes: () => [node],
+    getValues: () => []
+  };
+}

test/unittests/backends/webgl/test_conv_new.ts

@@ -8,6 +8,7 @@ import {CpuConv} from '../../../../lib/backends/cpu/ops/conv';
 import {Profiler} from '../../../../lib/instrument';
 import {Tensor} from '../../../../lib/tensor';
 import {TensorResultValidator} from '../../../test-runner';
+import {createMockGraph} from '../../../test-shared';
 
 const validator = new TensorResultValidator('webgl');
 let webglBackend: Backend|undefined;
@@ -28,8 +29,8 @@ function webglConv(
     attributes.set('pads', 'ints', pads);
   }
   attributes.set('strides', 'ints', strides);
-  const op = webglSessionhandler!.resolve(
-      {opType: 'Conv', attributes, inputs: [], outputs: [], name: `Conv`}, [{domain: '', version: 7}]);
+  const graph = createMockGraph('Conv', attributes);
+  const op = webglSessionhandler!.resolve(graph.getNodes()[0], [{domain: '', version: 7}], graph);
   if (!op.checkInputs([inputTensor, kernelTensor])) {
     throw new Error('Invalid inputs');
   }