onnxruntime

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

4
import { ShapeUtil } from '../../util';
5

6
import { TextureLayoutStrategy, WidthHeightPrefs } from './texture-layout-strategy';
7
import { TextureLayout, TextureType } from './types';
8

9
export const createTextureLayoutFromTextureType = (
10
  textureLayoutStrategy: TextureLayoutStrategy,
11
  shape: readonly number[],
12
  textureType: TextureType,
13
): TextureLayout => {
14
  const channel = textureType === TextureType.unpacked || textureType === TextureType.unpackedReversed ? 1 : 4;
15
  const isPacked = textureType === TextureType.packed;
16
  const reverseWH = textureType === TextureType.unpackedReversed || textureType === TextureType.packed;
17
  const breakAxis = textureType === TextureType.packedLastDimension ? shape.length - 1 : undefined;
18
  const unpackedShape =
19
    textureType === TextureType.packedLastDimension
20
      ? shape.map((d, i) => (i === shape.length - 1 ? d * 4 : d))
21
      : undefined;
22
  return createTextureLayoutFromShape(textureLayoutStrategy, shape, channel, unpackedShape, {
23
    isPacked,
24
    reverseWH,
25
    breakAxis,
26
  });
27
};
28

29
export const calculateTextureWidthAndHeight = (
30
  textureLayoutStrategy: TextureLayoutStrategy,
31
  shape: readonly number[],
32
  textureType: TextureType,
33
): [number, number] => {
34
  const layout = createTextureLayoutFromTextureType(textureLayoutStrategy, shape, textureType);
35
  return [layout.width, layout.height];
36
};
37

38
/**
39
 * Create a TextureLayout object from shape.
40
 */
41
export const createTextureLayoutFromShape = (
42
  textureLayoutStrategy: TextureLayoutStrategy,
43
  shape: readonly number[],
44
  channels: 1 | 4 = 1,
45
  unpackedShape?: readonly number[],
46
  prefs?: WidthHeightPrefs,
47
): TextureLayout => {
48
  const isPacked = !!(prefs && prefs.isPacked);
49
  const [width, height] = textureLayoutStrategy.computeTextureWH(isPacked ? unpackedShape || shape : shape, prefs);
50
  const rank = shape.length;
51
  let inferredDims = shape.slice(0);
52
  if (rank === 0) {
53
    inferredDims = [1];
54
  }
55
  if (channels === 1) {
56
    // unpackedShape will take `shape` and not `inferredDims` so as to create a scalar Tensor if need be
57
    unpackedShape = shape;
58
  } else if (isPacked) {
59
    if (channels !== 4) {
60
      throw new Error('a packed texture must be 4-channel');
61
    }
62
    unpackedShape = shape;
63
    if (rank > 0) {
64
      inferredDims[rank - 1] = Math.ceil(inferredDims[rank - 1] / 2);
65
    }
66
    if (rank > 1) {
67
      inferredDims[rank - 2] = Math.ceil(inferredDims[rank - 2] / 2);
68
    }
69
  } else if (!unpackedShape) {
70
    throw new Error('Unpacked shape is needed when using channels > 1');
71
  }
72
  return {
73
    width,
74
    height,
75
    channels,
76
    isPacked,
77
    shape: inferredDims,
78
    strides: ShapeUtil.computeStrides(inferredDims),
79
    unpackedShape,
80
    reversedWH: prefs && prefs.reverseWH,
81
  };
82
};
83