onnxruntime

1
// Copyright (c) Microsoft Corporation. All rights reserved.
2
// Licensed under the MIT License.
3

4
import { WebGLBackend } from './backends/backend-webgl';
5
import { Graph } from './graph';
6
import { Operator } from './operators';
7
import { OpSet } from './opset';
8
import { Session } from './session';
9

10
export interface InferenceHandler {
11
  /**
12
   * dispose the inference handler. it will be called as the last step in Session.run()
13
   */
14
  dispose(): void;
15
}
16

17
export interface SessionHandler {
18
  /**
19
   * transform the graph at initialization time
20
   * @param graphTransformer the graph transformer to manipulate the model graph
21
   */
22
  transformGraph?(graphTransformer: Graph.Transformer): void;
23

24
  /**
25
   * create an instance of InferenceHandler to use in a Session.run() call
26
   */
27
  createInferenceHandler(): InferenceHandler;
28

29
  /**
30
   * dispose the session handler. it will be called when a session is being disposed explicitly
31
   */
32
  dispose(): void;
33

34
  /**
35
   * Resolves the operator from the name and opset version; backend specific
36
   * @param node the node to resolve
37
   * @param opsets a list of opsets that exported from the model
38
   * @param graph the completely initialized graph
39
   */
40
  resolve(node: Graph.Node, opsets: readonly OpSet[], graph: Graph): Operator;
41

42
  /**
43
   * This method let's the sessionHandler know that the graph initialization is complete
44
   * @param graph the completely initialized graph
45
   */
46
  onGraphInitialized?(graph: Graph): void;
47

48
  /**
49
   * a reference to the corresponding backend
50
   */
51
  readonly backend: Backend;
52

53
  /**
54
   * a reference to the session context
55
   */
56
  readonly context: Session.Context;
57
}
58

59
export interface Backend {
60
  /**
61
   * initialize the backend. will be called only once, when the first time the
62
   * backend it to be used
63
   */
64
  initialize(): boolean | Promise<boolean>;
65

66
  /**
67
   * create an instance of SessionHandler to use in a Session object's lifecycle
68
   */
69
  createSessionHandler(context: Session.Context): SessionHandler;
70

71
  /**
72
   * dispose the backend. currently this will not be called
73
   */
74
  dispose(): void;
75
}
76

77
// caches all initialized backend instances
78
const backendsCache: Map<string, Backend> = new Map();
79

80
export const backend: { [name: string]: Backend } = {
81
  webgl: new WebGLBackend(),
82
};
83

84
/**
85
 * Resolve a reference to the backend. If a hint is specified, the corresponding
86
 * backend will be used.
87
 */
88
export async function resolveBackend(hint?: string | readonly string[]): Promise<Backend> {
89
  if (!hint) {
90
    return resolveBackend(['webgl']);
91
  } else {
92
    const hints = typeof hint === 'string' ? [hint] : hint;
93

94
    for (const backendHint of hints) {
95
      const cache = backendsCache.get(backendHint);
96
      if (cache) {
97
        return cache;
98
      }
99

100
      const backend = await tryLoadBackend(backendHint);
101
      if (backend) {
102
        return backend;
103
      }
104
    }
105
  }
106

107
  throw new Error('no available backend to use');
108
}
109

110
async function tryLoadBackend(backendHint: string): Promise<Backend | undefined> {
111
  const backendObj = backend;
112

113
  if (typeof backendObj[backendHint] !== 'undefined' && isBackend(backendObj[backendHint])) {
114
    const backend = backendObj[backendHint];
115
    let init = backend.initialize();
116
    if (typeof init === 'object' && 'then' in init) {
117
      init = await init;
118
    }
119
    if (init) {
120
      backendsCache.set(backendHint, backend);
121
      return backend;
122
    }
123
  }
124

125
  return undefined;
126
}
127

128
function isBackend(obj: unknown) {
129
  // eslint-disable-next-line @typescript-eslint/no-explicit-any
130
  const o = obj as any;
131

132
  // check if an object is a Backend instance
133
  if (
134
    'initialize' in o &&
135
    typeof o.initialize === 'function' && // initialize()
136
    'createSessionHandler' in o &&
137
    typeof o.createSessionHandler === 'function' && // createSessionHandler()
138
    'dispose' in o &&
139
    typeof o.dispose === 'function' // dispose()
140
  ) {
141
    return true;
142
  }
143

144
  return false;
145
}
146

147
export type BackendType = Backend;
148
export type SessionHandlerType = ReturnType<BackendType['createSessionHandler']>;
149
export type InferenceHandlerType = ReturnType<SessionHandlerType['createInferenceHandler']>;
150