onnxruntime

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

4
import { GlslContext, GlslLib, GlslLibRoutine } from './glsl-definitions';
5

6
/**
7
 * GLSL Library responsible for vec routines
8
 * Vec is an varible length int array. The length is fixed at the time of
9
 * generating the library functions from the dimensions of the output.
10
 */
11
export class VecGlslLib extends GlslLib {
12
  constructor(context: GlslContext) {
13
    super(context);
14
  }
15
  getCustomTypes(): { [name: string]: string } {
16
    return {};
17
  }
18
  getFunctions(): { [name: string]: GlslLibRoutine } {
19
    return { ...this.binaryVecFunctions(), ...this.copyVec(), ...this.setVecItem(), ...this.getVecItem() };
20
  }
21
  protected binaryVecFunctions(): { [name: string]: GlslLibRoutine } {
22
    const outputLayout = this.context.outputTextureLayout;
23
    const rank = outputLayout.shape.length;
24
    const nameOp: { [name: string]: string } = { add: '+=', sub: '-=', mul: '*=', div: '/=' };
25
    const result: { [name: string]: GlslLibRoutine } = {};
26
    for (const name in nameOp) {
27
      const fname = `${name}Vec`;
28
      let assignmentBlock = '';
29
      for (let i = 0; i < rank; ++i) {
30
        assignmentBlock += `
31
          dest[${i}] ${nameOp[name]} src[${i}];
32
          `;
33
      }
34
      const body = `
35
        void ${fname}(int src[${rank}], out int dest[${rank}]) {
36
          ${assignmentBlock}
37
        }
38
        `;
39
      result[fname] = new GlslLibRoutine(body);
40
    }
41

42
    return result;
43
  }
44
  protected copyVec(): { [name: string]: GlslLibRoutine } {
45
    const outputLayout = this.context.outputTextureLayout;
46
    const rank = outputLayout.shape.length;
47
    let assignmentBlock = '';
48
    for (let i = 0; i < rank; ++i) {
49
      assignmentBlock += `
50
        dest[${i}] = src[${i}];
51
        `;
52
    }
53
    const body = `
54
      void copyVec(int src[${rank}], out int dest[${rank}]) {
55
        ${assignmentBlock}
56
      }
57
      `;
58
    return { copyVec: new GlslLibRoutine(body) };
59
  }
60

61
  protected setVecItem(): { [name: string]: GlslLibRoutine } {
62
    const outputLayout = this.context.outputTextureLayout;
63
    const rank = outputLayout.shape.length;
64
    let block = `
65
        if(index < 0)
66
            index =${rank} + index;
67
        if (index == 0)
68
            m[0] = value;
69
        `;
70
    for (let i = 1; i < rank - 1; ++i) {
71
      block += `
72
        else if (index == ${i})
73
            m[${i}] = value;
74
            `;
75
    }
76
    block += `
77
        else
78
            m[${rank - 1}] = value;
79
        `;
80
    const body = `
81
      void setVecItem(out int m[${rank}], int index, int value) {
82
        ${block}
83
      }
84
        `;
85
    return { setVecItem: new GlslLibRoutine(body) };
86
  }
87
  protected getVecItem(): { [name: string]: GlslLibRoutine } {
88
    const outputLayout = this.context.outputTextureLayout;
89
    const rank = outputLayout.shape.length;
90
    let block = `
91
        if(index < 0)
92
            index = ${rank} + index;
93
        if (index == 0)
94
            return m[0];
95
      `;
96
    for (let i = 1; i < rank - 1; ++i) {
97
      block += `
98
        else if (index == ${i})
99
            return m[${i}];
100
      `;
101
    }
102
    block += `
103
        else
104
            return m[${rank - 1}];
105
        `;
106
    const body = `
107
      int getVecItem(int m[${rank}], int index) {
108
        ${block}
109
      }
110
    `;
111
    return { getVecItem: new GlslLibRoutine(body) };
112
  }
113
}
114