onnxruntime

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

4
import { GlslContext, GlslLib, GlslLibRoutineNode, TopologicalSortGlslRoutines } from './glsl-definitions';
5
import { replaceInlines } from './glsl-function-inliner';
6
import { glslRegistry } from './glsl-registered-libs';
7
import { getDefaultFragShaderMain, getFragShaderPreamble } from './glsl-source';
8
import { ProgramInfo, TextureLayout, VariableInfo } from './types';
9
import { WebGLContext } from './webgl-context';
10

11
/**
12
 * Preprocessor for the additions to the GLSL language
13
 * It deals with:
14
 *  @include directives
15
 *  @inline
16
 *  Loop unrolling (not implemented)
17
 *  Macro resolution (not implemented)
18
 */
19
export class GlslPreprocessor {
20
  readonly context: GlslContext;
21
  readonly libs: { [name: string]: GlslLib } = {};
22
  readonly glslLibRoutineDependencyGraph: { [routineName: string]: GlslLibRoutineNode } = {};
23

24
  constructor(
25
    glContext: WebGLContext,
26
    programInfo: ProgramInfo,
27
    inputTextureLayouts: TextureLayout[],
28
    outputTextureLayout: TextureLayout,
29
  ) {
30
    this.context = new GlslContext(glContext, programInfo, inputTextureLayouts, outputTextureLayout);
31

32
    // construct GlslLibs
33
    Object.keys(glslRegistry).forEach((name: string) => {
34
      const lib = new glslRegistry[name](this.context);
35
      this.libs[name] = lib;
36
    });
37

38
    // construct GlslRoutineDependencyGraph
39
    const map = this.glslLibRoutineDependencyGraph;
40
    for (const libName in this.libs) {
41
      const lib = this.libs[libName];
42
      const routinesInLib = lib.getFunctions();
43
      for (const routine in routinesInLib) {
44
        const key = libName + '.' + routine;
45
        let currentNode: GlslLibRoutineNode;
46
        if (map[key]) {
47
          currentNode = map[key];
48
          currentNode.routineBody = routinesInLib[routine].routineBody;
49
        } else {
50
          currentNode = new GlslLibRoutineNode(key, routinesInLib[routine].routineBody);
51
          map[key] = currentNode;
52
        }
53
        const dependencies = routinesInLib[routine].dependencies;
54
        if (dependencies) {
55
          for (let i = 0; i < dependencies.length; ++i) {
56
            if (!map[dependencies[i]]) {
57
              const node = new GlslLibRoutineNode(dependencies[i]);
58
              map[dependencies[i]] = node;
59
              currentNode.addDependency(node);
60
            } else {
61
              currentNode.addDependency(map[dependencies[i]]);
62
            }
63
          }
64
        }
65
      }
66
    }
67
  }
68

69
  preprocess(): string {
70
    const programInfo = this.context.programInfo;
71
    let source = programInfo.shaderSource;
72

73
    // append main() function
74
    if (!this.context.programInfo.hasMain) {
75
      source = `${source}
76
      ${getDefaultFragShaderMain(this.context.glContext.version, this.context.outputTextureLayout.shape.length)}`;
77
    }
78
    // replace inlines
79
    source = replaceInlines(source);
80

81
    // concat final source string
82
    return `${getFragShaderPreamble(this.context.glContext.version)}
83
    ${this.getUniforms(programInfo.inputNames, programInfo.variables)}
84
    ${this.getImports(source)}
85
    ${source}`;
86
  }
87

88
  protected getImports(script: string): string {
89
    const routinesIncluded = this.selectGlslLibRoutinesToBeIncluded(script);
90

91
    if (routinesIncluded.length === 0) {
92
      return '';
93
    }
94

95
    let routines = '';
96
    for (let i = 0; i < routinesIncluded.length; ++i) {
97
      if (routinesIncluded[i].routineBody) {
98
        routines += routinesIncluded[i].routineBody + '\n';
99
      } else {
100
        throw new Error(`Missing body for the Glsl Library routine: ${routinesIncluded[i].name}`);
101
      }
102
    }
103

104
    return routines;
105
  }
106
  private selectGlslLibRoutinesToBeIncluded(script: string): GlslLibRoutineNode[] {
107
    const nodes: GlslLibRoutineNode[] = [];
108

109
    Object.keys(this.glslLibRoutineDependencyGraph).forEach((classAndRoutine) => {
110
      const routine = classAndRoutine.split('.')[1];
111
      if (script.indexOf(routine) !== -1) {
112
        nodes.push(this.glslLibRoutineDependencyGraph[classAndRoutine]);
113
      }
114
    });
115

116
    return TopologicalSortGlslRoutines.returnOrderedNodes(nodes);
117
  }
118

119
  protected getUniforms(samplers?: string[], variables?: VariableInfo[]): string {
120
    const uniformLines: string[] = [];
121
    if (samplers) {
122
      for (const sampler of samplers) {
123
        uniformLines.push(`uniform sampler2D ${sampler};`);
124
      }
125
    }
126
    if (variables) {
127
      for (const variable of variables) {
128
        uniformLines.push(
129
          `uniform ${variable.type} ${variable.name}${variable.arrayLength ? `[${variable.arrayLength}]` : ''};`,
130
        );
131
      }
132
    }
133
    return uniformLines.join('\n');
134
  }
135
}
136