onnxruntime

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

4
import { Logger, Profiler } from '../../instrument';
5
import { Tensor } from '../../tensor';
6

7
import { Encoder, EncoderUsage } from './texture-data-encoder';
8
import { TextureLayoutStrategy } from './texture-layout-strategy';
9
import { TextureData, TextureLayout } from './types';
10
import { WebGLContext } from './webgl-context';
11

12
export interface TextureManagerConfig {
13
  reuseTextures?: boolean;
14
}
15

16
/**
17
 * TextureManager is the mainly responsible for caching Textures
18
 * Textures are cached in 2 levels:
19
 *   1. the texures which are associated with a dataId (from Tensor)
20
 *    Caching these is crucial to performance. These are In-use Textures
21
 *   2. textures which are not in use by any current ProgramInfo/Tensor
22
 *     These are called Free Textures
23
 * TextureManager is also used to help creating textures. For this it
24
 * uses WebGLContext and TextureLayoutStrategy
25
 */
26
export class TextureManager {
27
  private readonly inUseTextures: Map<string, WebGLTexture[]>;
28
  private readonly idleTextures: Map<string, WebGLTexture[]>;
29
  private readonly textureLookup: Map<WebGLTexture, string>;
30
  private readonly pendingRead: Map<Tensor.Id, Array<(arr: Tensor.NumberType) => void>> = new Map();
31

32
  constructor(
33
    public glContext: WebGLContext,
34
    public layoutStrategy: TextureLayoutStrategy,
35
    public profiler: Readonly<Profiler>,
36
    private config: TextureManagerConfig,
37
  ) {
38
    if (config.reuseTextures) {
39
      this.inUseTextures = new Map();
40
      this.idleTextures = new Map();
41
      this.textureLookup = new Map();
42
    }
43
  }
44
  createTextureFromLayout(
45
    dataType: Tensor.DataType,
46
    layout: TextureLayout,
47
    data?: Tensor.NumberType,
48
    usage?: EncoderUsage,
49
  ) {
50
    const textureDataType = this.toEncoderType(dataType);
51

52
    const encoder = this.glContext.getEncoder(textureDataType, layout.channels || 1, usage);
53
    if (layout.isPacked && usage === EncoderUsage.UploadOnly) {
54
      throw new Error('not implemented');
55
    }
56
    const width = layout.width;
57
    const height = layout.height;
58

59
    let key: string | undefined;
60
    let inUseTextures: WebGLTexture[] | undefined;
61
    if (this.config.reuseTextures) {
62
      key = `${width}x${height}_${encoder.format}_${encoder.internalFormat}_${encoder.textureType}`;
63
      inUseTextures = this.inUseTextures.get(key);
64
      if (!inUseTextures) {
65
        inUseTextures = [];
66
        this.inUseTextures.set(key, inUseTextures);
67
      }
68

69
      const idleTextures = this.idleTextures.get(key);
70
      if (idleTextures && idleTextures.length > 0) {
71
        const texture = idleTextures.pop()!;
72
        inUseTextures.push(texture);
73
        if (usage === EncoderUsage.UploadOnly) {
74
          this.glContext.updateTexture(texture, width, height, encoder, this.toTextureData(dataType, data)!);
75
        }
76
        return texture;
77
      }
78
    }
79

80
    Logger.verbose('TextureManager', `Creating new texture of size ${layout.width}x${layout.height}`);
81
    const texture = this.glContext.allocateTexture(width, height, encoder, this.toTextureData(dataType, data));
82

83
    if (this.config.reuseTextures) {
84
      inUseTextures!.push(texture);
85
      this.textureLookup.set(texture, key!);
86
    }
87
    return texture;
88
  }
89
  readTexture(td: TextureData, dataType: Tensor.DataType, channels?: number): Tensor.NumberType {
90
    if (!channels) {
91
      channels = 1;
92
    }
93
    return this.profiler.event('backend', 'TextureManager.readTexture', () => {
94
      const dataSize = td.shape.reduce((a, b) => a * b) * channels!;
95
      const data = this.glContext.readTexture(
96
        td.texture,
97
        td.width,
98
        td.height,
99
        dataSize,
100
        this.toEncoderType(dataType),
101
        channels!,
102
      );
103
      return this.toTensorData(dataType, data);
104
    });
105
  }
106
  async readTextureAsync(td: TextureData, dataType: Tensor.DataType, channels?: number): Promise<Tensor.NumberType> {
107
    const dataId = td.tensor.dataId;
108
    if (!channels) {
109
      channels = 1;
110
    }
111
    if (this.pendingRead.has(dataId)) {
112
      const subscribers = this.pendingRead.get(dataId);
113
      return new Promise<Tensor.NumberType>((resolve) => subscribers?.push(resolve));
114
    }
115
    return this.profiler.event('backend', 'TextureManager.readTextureAsync', async () => {
116
      this.pendingRead.set(dataId, []);
117
      const dataSize = td.shape.reduce((a, b) => a * b) * channels!;
118
      // add a fence waiting for the data to be ready
119
      await this.glContext.createAndWaitForFence();
120
      const data = this.glContext.readTexture(
121
        td.texture,
122
        td.width,
123
        td.height,
124
        dataSize,
125
        this.toEncoderType(dataType),
126
        channels!,
127
      );
128
      const tensorData = this.toTensorData(dataType, data);
129
      const subscribers = this.pendingRead.get(dataId);
130
      this.pendingRead.delete(dataId);
131
      subscribers?.forEach((resolve) => resolve(tensorData));
132
      return tensorData;
133
    });
134
  }
135
  readUint8TextureAsFloat(td: TextureData): Float32Array {
136
    return this.profiler.event('backend', 'TextureManager.readUint8TextureAsFloat', () => {
137
      const dataSize = td.shape.reduce((a, b) => a * b);
138
      const data = this.glContext.readTexture(td.texture, td.width, td.height, dataSize * 4, 'byte', 4);
139
      return new Float32Array(data.buffer, data.byteOffset, dataSize);
140
    });
141
  }
142
  releaseTexture(textureData: TextureData, deleteTexture?: boolean): void {
143
    let key: string | undefined;
144
    if (this.config.reuseTextures) {
145
      key = this.textureLookup.get(textureData.texture);
146
      if (key) {
147
        if (deleteTexture) {
148
          this.textureLookup.delete(key);
149
        }
150
        const inUseTextures = this.inUseTextures.get(key);
151
        if (inUseTextures) {
152
          const index = inUseTextures.indexOf(textureData.texture);
153
          if (index !== -1) {
154
            inUseTextures.splice(index, 1);
155
            let idleTextures = this.idleTextures.get(key);
156
            if (!idleTextures) {
157
              idleTextures = [];
158
              this.idleTextures.set(key, idleTextures);
159
            }
160
            idleTextures.push(textureData.texture);
161
          }
162
        }
163
      }
164
    }
165

166
    if (!key || deleteTexture) {
167
      Logger.verbose('TextureManager', `Deleting texture of size ${textureData.width}x${textureData.height}`);
168
      this.glContext.deleteTexture(textureData.texture);
169
    }
170
  }
171
  toTensorData(dataType: Tensor.DataType, data: Encoder.DataArrayType): Tensor.NumberType {
172
    switch (dataType) {
173
      case 'int16':
174
        return data instanceof Int16Array ? data : Int16Array.from(data);
175
      case 'int32':
176
        return data instanceof Int32Array ? data : Int32Array.from(data);
177
      case 'int8':
178
        return data instanceof Int8Array ? data : Int8Array.from(data);
179
      case 'uint16':
180
        return data instanceof Uint16Array ? data : Uint16Array.from(data);
181
      case 'uint32':
182
        return data instanceof Uint32Array ? data : Uint32Array.from(data);
183
      case 'uint8':
184
      case 'bool':
185
        return data instanceof Uint8Array ? data : Uint8Array.from(data);
186
      case 'float32':
187
        return data instanceof Float32Array ? data : Float32Array.from(data);
188
      case 'float64':
189
        return data instanceof Float64Array ? data : Float64Array.from(data);
190
      default:
191
        throw new Error(`TensorData type ${dataType} is not supported`);
192
    }
193
  }
194
  toTextureData(_dataType: Tensor.DataType, data: Tensor.NumberType | undefined): Encoder.DataArrayType | undefined {
195
    if (!data) {
196
      return undefined;
197
    }
198
    return data instanceof Float32Array ? data : new Float32Array(data);
199
    /*
200
    switch (dataType) {
201
      case 'int16':
202
      case 'int32':
203
      case 'uint16':
204
      case 'uint32':
205
        return (data.constructor === Uint32Array) ? data as Uint32Array : new Uint32Array(data);
206
      case 'int8':
207
      case 'uint8':
208
      case 'bool':
209
        return (data.constructor === Uint8Array) ? data as Uint8Array : new Uint8Array(data);
210
      case 'float32':
211
      case 'float64':
212
        return (data.constructor === Float32Array) ? data as Float32Array : new Float32Array(data);
213
      default:
214
        throw new Error(`TensorData type ${dataType} is not supported`);
215
    }
216
    */
217
  }
218
  toEncoderType(_dataType: Tensor.DataType): Encoder.DataType {
219
    return 'float';
220
    // switch (dataType) {
221
    //   case 'int16':
222
    //   case 'int32':
223
    //   case 'uint16':
224
    //   case 'uint32':
225
    //     return 'int';
226
    //   case 'uint8':
227
    //   case 'bool':
228
    //     return 'byte';
229
    //   case 'float32':
230
    //   case 'float64':
231
    //     return 'float';
232
    //   default:
233
    //     throw new Error(`TensorData type ${dataType} is not supported`);
234
    // }
235
  }
236
  clearActiveTextures(): void {
237
    this.glContext.clearActiveTextures();
238
  }
239
}
240