onnxruntime

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

4
import { SessionHandler } from '../../backend';
5
import { Graph } from '../../graph';
6
import { Logger } from '../../instrument';
7
import { Operator } from '../../operators';
8
import { OpSet, resolveOperator } from '../../opset';
9
import { Session } from '../../session';
10
import { Tensor } from '../../tensor';
11
import { WebGLBackend } from '../backend-webgl';
12

13
import { WebGLInferenceHandler } from './inference-handler';
14
import { WEBGL_OP_RESOLVE_RULES } from './op-resolve-rules';
15
import { ProgramManager } from './program-manager';
16
import { PreferLogicalStrategy, TextureLayoutStrategy } from './texture-layout-strategy';
17
import { TextureManager } from './texture-manager';
18
import { TextureData } from './types';
19

20
export class WebGLSessionHandler implements SessionHandler {
21
  programManager: ProgramManager;
22
  textureManager: TextureManager;
23
  layoutStrategy: TextureLayoutStrategy;
24
  packedTextureDataCache: Map<Tensor.Id, TextureData>;
25
  unpackedTextureDataCache: Map<Tensor.Id, TextureData>;
26
  pack2unpackMap: Map<Tensor.Id, Tensor.Id>;
27
  unpack2packMap: Map<Tensor.Id, Tensor.Id>;
28
  initializers: Set<Tensor.Id>;
29
  pack?: boolean;
30

31
  constructor(
32
    public readonly backend: WebGLBackend,
33
    public readonly context: Session.Context,
34
  ) {
35
    this.layoutStrategy = new PreferLogicalStrategy(backend.glContext.maxTextureSize);
36
    this.programManager = new ProgramManager(this.context.profiler, backend.glContext, this.layoutStrategy);
37
    this.textureManager = new TextureManager(backend.glContext, this.layoutStrategy, this.context.profiler, {
38
      reuseTextures: backend.textureCacheMode === 'full',
39
    });
40
    this.packedTextureDataCache = new Map();
41
    this.unpackedTextureDataCache = new Map();
42
    this.pack = backend.pack;
43
    this.pack2unpackMap = new Map();
44
    this.unpack2packMap = new Map();
45
  }
46

47
  createInferenceHandler() {
48
    return new WebGLInferenceHandler(this);
49
  }
50
  onGraphInitialized(graph: Graph): void {
51
    const initializers = graph
52
      .getValues()
53
      .filter((v) => v.from === -1 && v.tensor)
54
      .map((v) => v.tensor!.dataId);
55
    this.initializers = new Set(initializers);
56
  }
57
  isInitializer(tensorId: Tensor.Id): boolean {
58
    return this.initializers ? this.initializers.has(tensorId) : false;
59
  }
60
  addInitializer(tensorId: Tensor.Id): void {
61
    this.initializers.add(tensorId);
62
  }
63
  getTextureData(tensorId: Tensor.Id, isPacked: boolean): TextureData | undefined {
64
    if (isPacked) {
65
      return this.packedTextureDataCache.get(tensorId);
66
    } else {
67
      return this.unpackedTextureDataCache.get(tensorId);
68
    }
69
  }
70
  setTextureData(tensorId: Tensor.Id, textureData: TextureData, isPacked = false): void {
71
    Logger.verbose('WebGLSessionHandler', 'Storing Texture data in cache');
72
    if (isPacked) {
73
      this.packedTextureDataCache.set(tensorId, textureData);
74
    } else {
75
      this.unpackedTextureDataCache.set(tensorId, textureData);
76
    }
77
  }
78
  dispose(): void {
79
    this.programManager.dispose();
80
    this.textureManager.clearActiveTextures();
81
    this.packedTextureDataCache.forEach((td) => this.textureManager.releaseTexture(td, true));
82
    this.packedTextureDataCache = new Map();
83
    this.unpackedTextureDataCache.forEach((td) => this.textureManager.releaseTexture(td, true));
84
    this.unpackedTextureDataCache = new Map();
85
  }
86
  resolve(node: Graph.Node, opsets: readonly OpSet[], graph: Graph): Operator {
87
    const op = resolveOperator(node, opsets, WEBGL_OP_RESOLVE_RULES);
88
    return { impl: op.opImpl, context: op.opInit ? op.opInit(node, graph) : node };
89
  }
90
}
91