onnxruntime

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

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

6
import { GlslContext, GlslLib, GlslLibRoutine } from './glsl-definitions';
7
import { getGlsl } from './glsl-source';
8
import { squeezeShape } from './texture-layout-strategy';
9
import { TextureLayout } from './types';
10
import {
11
  generateShaderFuncNameFromInputSamplerName,
12
  generateShaderFuncNameFromInputSamplerNameAtOutCoords,
13
  getCoordsDataType,
14
  getGlChannels,
15
  getSqueezedParams,
16
  squeezeInputShape,
17
} from './utils';
18

19
/**
20
 * GLSL Library responsible for data types and routines for manipulating
21
 * coordinates and mapping to/from tensor indices
22
 */
23
export class CoordsGlslLib extends GlslLib {
24
  returnType: string;
25

26
  constructor(context: GlslContext) {
27
    super(context);
28
  }
29
  getFunctions(): { [name: string]: GlslLibRoutine } {
30
    return {
31
      ...this.offsetToCoords(),
32
      ...this.coordsToOffset(),
33
      ...this.toVec(),
34
      ...this.valueFrom(),
35
      // TODO return these only when packing is enabled.
36
      ...this.getCommonUtilFuncs(),
37
      ...this.getInputsSamplingSnippets(),
38
      ...this.getOutputSamplingSnippet(),
39
    };
40
  }
41
  getCustomTypes() {
42
    return {};
43
  }
44
  /**
45
   * Produces a function that can map from
46
   * 2D normalzied coordinates (s,t) to a flat offset
47
   */
48
  protected offsetToCoords(): { [name: string]: GlslLibRoutine } {
49
    const funcName = 'offsetToCoords';
50
    return {
51
      offsetToCoords: new GlslLibRoutine(`
52
      vec2 ${funcName}(int offset, int width, int height) {
53
        int t = offset / width;
54
        int s = offset - t*width;
55
        vec2 coords = (vec2(s,t) + vec2(0.5,0.5)) / vec2(width, height);
56
        return coords;
57
      }
58
      `),
59
    };
60
  }
61

62
  /**
63
   * Produces a function that can map from
64
   * 2D normalzied coordinates (s,t) to a flat offset
65
   */
66
  protected coordsToOffset(): { [name: string]: GlslLibRoutine } {
67
    const funcName = 'coordsToOffset';
68
    return {
69
      coordsToOffset: new GlslLibRoutine(`
70
      int ${funcName}(vec2 coords, int width, int height) {
71
        float s = coords.s * float(width);
72
        float t = coords.t * float(height);
73
        int offset = int(t) * width + int(s);
74
        return offset;
75
      }
76
      `),
77
    };
78
  }
79

80
  /**
81
   * Generates code for output sampler.
82
   */
83

84
  protected getOutputSamplingSnippet(): { [name: string]: GlslLibRoutine } {
85
    const outputLayout = this.context.outputTextureLayout;
86
    if (outputLayout.isPacked) {
87
      return this.getPackedOutputSamplingSnippet(outputLayout);
88
    } else {
89
      return this.getUnpackedOutputSamplingSnippet(outputLayout);
90
    }
91
  }
92

93
  /**
94
   * Generates code for packed output sampler.
95
   */
96
  protected getPackedOutputSamplingSnippet(outputLayout: TextureLayout): { [name: string]: GlslLibRoutine } {
97
    const outShape = outputLayout.unpackedShape;
98
    const outTexShape = [outputLayout.width, outputLayout.height];
99
    const result: { [name: string]: GlslLibRoutine } = {};
100
    const funcName = 'getOutputCoords';
101
    switch (outShape.length) {
102
      case 0:
103
        result[funcName] = this.getOutputScalarCoords();
104
        break;
105
      case 1:
106
        result[funcName] = this.getOutputPacked1DCoords(outShape as [number], outTexShape as [number, number]);
107
        break;
108
      case 2:
109
        result[funcName] = this.getOutputPacked2DCoords(outShape as [number, number], outTexShape as [number, number]);
110
        break;
111
      case 3:
112
        result[funcName] = this.getOutputPacked3DCoords(
113
          outShape as [number, number, number],
114
          outTexShape as [number, number],
115
        );
116
        break;
117
      default:
118
        result[funcName] = this.getOutputPackedNDCoords(outShape, outTexShape as [number, number]);
119
    }
120
    const glsl = getGlsl(this.context.glContext.version);
121
    // TODO we need this to properly return a packed vec4 from kernels.
122
    // Replace all '{glsl.output} = result' with 'setOutput(result)' in all kernels.
123
    const floatTextureSetRGBASource = `
124
      void setOutput(vec4 val) {
125
        ${glsl.output} = val;
126
      }
127
    `;
128
    const floatTextureSetRGBAFuncName = 'floatTextureSetRGBA';
129
    result[floatTextureSetRGBAFuncName] = new GlslLibRoutine(floatTextureSetRGBASource);
130
    return result;
131
  }
132

133
  /**
134
   * Generates code for unpacked output sampler.
135
   */
136
  protected getUnpackedOutputSamplingSnippet(outputLayout: TextureLayout): { [name: string]: GlslLibRoutine } {
137
    const outShape = outputLayout.unpackedShape;
138
    const outTexShape = [outputLayout.width, outputLayout.height];
139
    const result: { [name: string]: GlslLibRoutine } = {};
140
    const funcName = 'getOutputCoords';
141
    switch (outShape.length) {
142
      case 0:
143
        result[funcName] = this.getOutputScalarCoords();
144
        break;
145
      case 1:
146
        result[funcName] = this.getOutputUnpacked1DCoords(outShape as [number], outTexShape as [number, number]);
147
        break;
148
      case 2:
149
        result[funcName] = this.getOutputUnpacked2DCoords(
150
          outShape as [number, number],
151
          outTexShape as [number, number],
152
        );
153
        break;
154
      case 3:
155
        result[funcName] = this.getOutputUnpacked3DCoords(
156
          outShape as [number, number, number],
157
          outTexShape as [number, number],
158
        );
159
        break;
160
      case 4:
161
        result[funcName] = this.getOutputUnpacked4DCoords(
162
          outShape as [number, number, number, number],
163
          outTexShape as [number, number],
164
        );
165
        break;
166
      case 5:
167
        result[funcName] = this.getOutputUnpacked5DCoords(
168
          outShape as [number, number, number, number, number],
169
          outTexShape as [number, number],
170
        );
171
        break;
172
      case 6:
173
        result[funcName] = this.getOutputUnpacked6DCoords(
174
          outShape as [number, number, number, number, number, number],
175
          outTexShape as [number, number],
176
        );
177
        break;
178
      default:
179
        throw new Error(`Unsupported output dimensionality: ${outShape.length}`);
180
    }
181
    const glsl = getGlsl(this.context.glContext.version);
182
    // TODO we need this to properly return a packed vec4 from kernels.
183
    // Replace all '{glsl.output} = result' with 'setOutput(result)' in all kernels.
184
    const floatTextureSetRSource = `
185
        void setOutput(float val) {
186
          ${glsl.output} = vec4(val, 0, 0, 0);
187
        }
188
    `;
189
    const floatTextureSetRFuncName = 'floatTextureSetR';
190
    result[floatTextureSetRFuncName] = new GlslLibRoutine(floatTextureSetRSource);
191
    return result;
192
  }
193

194
  /**
195
   * Scalar output coordinates.
196
   */
197
  protected getOutputScalarCoords(): GlslLibRoutine {
198
    return new GlslLibRoutine(`
199
      int getOutputCoords() {
200
        return 0;
201
      }
202
    `);
203
  }
204

205
  /**
206
   * 1D packed output coordinates.
207
   */
208
  protected getOutputPacked1DCoords(_shape: [number], texShape: [number, number]): GlslLibRoutine {
209
    const packedTexShape = texShape;
210
    let source = '';
211
    if (packedTexShape[0] === 1) {
212
      source = `
213
          int getOutputCoords() {
214
            return 2 * int(TexCoords.y * ${packedTexShape[1]}.0);
215
          }
216
        `;
217
      return new GlslLibRoutine(source);
218
    }
219

220
    if (packedTexShape[1] === 1) {
221
      source = `
222
          int getOutputCoords() {
223
            return 2 * int(TexCoords.x * ${packedTexShape[0]}.0);
224
          }
225
        `;
226
      return new GlslLibRoutine(source);
227
    }
228

229
    source = `
230
        int getOutputCoords() {
231
          ivec2 resTexRC = ivec2(TexCoords.xy *
232
                                 vec2(${packedTexShape[0]}, ${packedTexShape[1]}));
233
          return 2 * (resTexRC.y * ${packedTexShape[0]} + resTexRC.x);
234
        }
235
      `;
236
    return new GlslLibRoutine(source);
237
  }
238

239
  /**
240
   * 2D packed output coordinates.
241
   */
242
  protected getOutputPacked2DCoords(shape: [number, number], texShape: [number, number]): GlslLibRoutine {
243
    let source = '';
244
    if (ArrayUtil.arraysEqual(shape, texShape)) {
245
      source = `
246
        ivec2 getOutputCoords() {
247
          return 2 * ivec2(TexCoords.xy * vec2(${texShape[0]}, ${texShape[1]}));
248
        }
249
      `;
250
      return new GlslLibRoutine(source);
251
    }
252

253
    const packedTexShape = texShape;
254
    // texels needed to accommodate a logical row
255
    const texelsInLogicalRow = Math.ceil(shape[1] / 2);
256

257
    /**
258
     * getOutputCoords
259
     *
260
     * resTexRC: The rows and columns of the texels. If you move over one
261
     * texel to the right in the packed texture, you are moving over one column
262
     * (not two).
263
     *
264
     * index: The texel index
265
     */
266
    source = `
267
        ivec2 getOutputCoords() {
268
          ivec2 resTexRC = ivec2(TexCoords.xy *
269
                                vec2(${packedTexShape[0]}, ${packedTexShape[1]}));
270

271
          int index = resTexRC.y * ${packedTexShape[0]} + resTexRC.x;
272

273
          // reverse r and c order for packed texture
274
          int r = imod(index, ${texelsInLogicalRow}) * 2;
275
          int c = 2 * (index / ${texelsInLogicalRow});
276

277
          return ivec2(r, c);
278
        }
279
      `;
280
    return new GlslLibRoutine(source);
281
  }
282

283
  /**
284
   * 3D packed output coordinates.
285
   */
286
  protected getOutputPacked3DCoords(shape: [number, number, number], texShape: [number, number]): GlslLibRoutine {
287
    const packedTexShape = [texShape[0], texShape[1]];
288
    const texelsInLogicalRow = Math.ceil(shape[2] / 2);
289
    const texelsInBatch = texelsInLogicalRow * Math.ceil(shape[1] / 2);
290
    const source = `
291
        ivec3 getOutputCoords() {
292
          ivec2 resTexRC = ivec2(TexCoords.xy *
293
                                vec2(${packedTexShape[0]}, ${packedTexShape[1]}));
294
          int index = resTexRC.y * ${packedTexShape[0]} + resTexRC.x;
295

296
          int b = index / ${texelsInBatch};
297
          index -= b * ${texelsInBatch};
298

299
          // reverse r and c order for packed texture
300
          int r = imod(index, ${texelsInLogicalRow}) * 2;
301
          int c = 2 * (index / ${texelsInLogicalRow});
302

303
          return ivec3(b, r, c);
304
        }
305
      `;
306
    return new GlslLibRoutine(source);
307
  }
308

309
  /**
310
   * ND packed output coordinates.
311
   */
312
  protected getOutputPackedNDCoords(shape: readonly number[], texShape: [number, number]): GlslLibRoutine {
313
    const packedTexShape = [texShape[0], texShape[1]];
314

315
    const texelsInLogicalRow = Math.ceil(shape[shape.length - 1] / 2);
316
    const texelsInBatch = texelsInLogicalRow * Math.ceil(shape[shape.length - 2] / 2);
317
    let texelsInBatchN = texelsInBatch;
318
    let batches = '';
319
    let coords = 'b, r, c';
320

321
    for (let b = 2; b < shape.length - 1; b++) {
322
      texelsInBatchN *= shape[shape.length - b - 1];
323
      batches =
324
        `
325
      int b${b} = index / ${texelsInBatchN};
326
      index -= b${b} * ${texelsInBatchN};
327
    ` + batches;
328
      coords = `b${b}, ` + coords;
329
    }
330
    const source = `
331
      ivec${shape.length} getOutputCoords() {
332
        ivec2 resTexRC = ivec2(TexCoords.xy *
333
                              vec2(${packedTexShape[0]}, ${packedTexShape[1]}));
334
        int index = resTexRC.y * ${packedTexShape[0]} + resTexRC.x;
335

336
        ${batches}
337

338
        int b = index / ${texelsInBatch};
339
        index -= b * ${texelsInBatch};
340

341
        // reverse r and c order for packed texture
342
        int r = imod(index, ${texelsInLogicalRow}) * 2;
343
        int c = 2 * (index / ${texelsInLogicalRow});
344

345
        return ivec${shape.length}(${coords});
346
      }
347
    `;
348
    return new GlslLibRoutine(source);
349
  }
350

351
  /**
352
   * Unpacked 1D output coordinates.
353
   */
354
  protected getOutputUnpacked1DCoords(_shape: [number], texShape: [number, number]): GlslLibRoutine {
355
    const source = `
356
        int getOutputCoords() {
357
          ivec2 resTexRC = ivec2(TexCoords.xy *
358
                                vec2(${texShape[0]}, ${texShape[1]}));
359
          return resTexRC.y * ${texShape[0]} + resTexRC.x;
360
        }
361
      `;
362
    return new GlslLibRoutine(source);
363
  }
364

365
  /**
366
   * Unpacked 2D output coordinates.
367
   */
368
  protected getOutputUnpacked2DCoords(shape: [number, number], texShape: [number, number]): GlslLibRoutine {
369
    const source = `
370
        ivec2 getOutputCoords() {
371
          ivec2 resTexRC = ivec2(TexCoords.xy *
372
                                vec2(${texShape[0]}, ${texShape[1]}));
373
          int index = resTexRC.y * ${texShape[0]} + resTexRC.x;
374
          int r = index / ${shape[1]};
375
          int c = index - r * ${shape[1]};
376
          return ivec2(r, c);
377
        }
378
      `;
379
    return new GlslLibRoutine(source);
380
  }
381

382
  /**
383
   * Unpacked 3D output coordinates.
384
   */
385
  protected getOutputUnpacked3DCoords(shape: [number, number, number], texShape: [number, number]): GlslLibRoutine {
386
    let source = '';
387
    const rank = shape.length;
388

389
    let strides = null;
390
    if (rank < 2) {
391
      strides = [];
392
    }
393

394
    strides = new Array(rank - 1);
395
    strides[rank - 2] = shape[rank - 1];
396
    for (let i = rank - 3; i >= 0; --i) {
397
      strides[i] = strides[i + 1] * shape[i + 1];
398
    }
399
    const coordsToCompute = ['r', 'c', 'd'];
400
    const coordsFromIndexSnippet = strides
401
      .map((stride, i) => {
402
        const line1 = `int ${coordsToCompute[i]} = index / ${stride}`;
403
        const line2 =
404
          i === strides.length - 1
405
            ? `int ${coordsToCompute[i + 1]} = index - ${coordsToCompute[i]} * ${stride}`
406
            : `index -= ${coordsToCompute[i]} * ${stride}`;
407
        return `${line1}; ${line2};`;
408
      })
409
      .join('');
410

411
    source = `
412
        ivec3 getOutputCoords() {
413
          ivec2 resTexRC = ivec2(TexCoords.xy *
414
                                vec2(${texShape[0]}, ${texShape[1]}));
415
          int index = resTexRC.y * ${texShape[0]} + resTexRC.x;
416
          ${coordsFromIndexSnippet}
417
          return ivec3(r, c, d);
418
        }
419
      `;
420
    return new GlslLibRoutine(source);
421
  }
422

423
  /**
424
   * Unpacked 4D output coordinates.
425
   */
426
  protected getOutputUnpacked4DCoords(
427
    shape: [number, number, number, number],
428
    texShape: [number, number],
429
  ): GlslLibRoutine {
430
    let source = '';
431
    const rank = shape.length;
432

433
    let strides = null;
434
    if (rank < 2) {
435
      strides = [];
436
    }
437

438
    strides = new Array(rank - 1);
439
    strides[rank - 2] = shape[rank - 1];
440
    for (let i = rank - 3; i >= 0; --i) {
441
      strides[i] = strides[i + 1] * shape[i + 1];
442
    }
443
    const coordsToCompute = ['r', 'c', 'd', 'd2'];
444
    const coordsFromIndexSnippet = strides
445
      .map((stride, i) => {
446
        const line1 = `int ${coordsToCompute[i]} = index / ${stride}`;
447
        const line2 =
448
          i === strides.length - 1
449
            ? `int ${coordsToCompute[i + 1]} = index - ${coordsToCompute[i]} * ${stride}`
450
            : `index -= ${coordsToCompute[i]} * ${stride}`;
451
        return `${line1}; ${line2};`;
452
      })
453
      .join('');
454

455
    source = `
456
      ivec4 getOutputCoords() {
457
          ivec2 resTexRC = ivec2(TexCoords.xy *
458
                                vec2(${texShape[0]}, ${texShape[1]}));
459
          int index = resTexRC.y * ${texShape[0]} + resTexRC.x;
460
          ${coordsFromIndexSnippet}
461
          return ivec4(r, c, d, d2);
462
        }
463
      `;
464
    return new GlslLibRoutine(source);
465
  }
466

467
  /**
468
   * Unpacked 5D output coordinates.
469
   */
470
  protected getOutputUnpacked5DCoords(
471
    shape: [number, number, number, number, number],
472
    texShape: [number, number],
473
  ): GlslLibRoutine {
474
    let source = '';
475
    const rank = shape.length;
476

477
    let strides = null;
478
    if (rank < 2) {
479
      strides = [];
480
    }
481

482
    strides = new Array(rank - 1);
483
    strides[rank - 2] = shape[rank - 1];
484
    for (let i = rank - 3; i >= 0; --i) {
485
      strides[i] = strides[i + 1] * shape[i + 1];
486
    }
487
    const coordsToCompute = ['r', 'c', 'd', 'd2', 'd3'];
488
    const coordsFromIndexSnippet = strides
489
      .map((stride, i) => {
490
        const line1 = `int ${coordsToCompute[i]} = index / ${stride}`;
491
        const line2 =
492
          i === strides.length - 1
493
            ? `int ${coordsToCompute[i + 1]} = index - ${coordsToCompute[i]} * ${stride}`
494
            : `index -= ${coordsToCompute[i]} * ${stride}`;
495
        return `${line1}; ${line2};`;
496
      })
497
      .join('');
498

499
    source = `
500
      ivec5 getOutputCoords() {
501
          ivec2 resTexRC = ivec2(TexCoords.xy *
502
                                vec2(${texShape[0]}, ${texShape[1]}));
503
          int index = resTexRC.y * ${texShape[0]} + resTexRC.x;
504
          ${coordsFromIndexSnippet}
505
          return ivec5(r, c, d, d2, d3);
506
        }
507
      `;
508
    return new GlslLibRoutine(source);
509
  }
510

511
  /**
512
   * Unpacked 6D output coordinates.
513
   */
514
  protected getOutputUnpacked6DCoords(
515
    shape: [number, number, number, number, number, number],
516
    texShape: [number, number],
517
  ): GlslLibRoutine {
518
    let source = '';
519
    const rank = shape.length;
520

521
    let strides = null;
522
    if (rank < 2) {
523
      strides = [];
524
    }
525

526
    strides = new Array(rank - 1);
527
    strides[rank - 2] = shape[rank - 1];
528
    for (let i = rank - 3; i >= 0; --i) {
529
      strides[i] = strides[i + 1] * shape[i + 1];
530
    }
531
    const coordsToCompute = ['r', 'c', 'd', 'd2', 'd3', 'd4'];
532
    const coordsFromIndexSnippet = strides
533
      .map((stride, i) => {
534
        const line1 = `int ${coordsToCompute[i]} = index / ${stride}`;
535
        const line2 =
536
          i === strides.length - 1
537
            ? `int ${coordsToCompute[i + 1]} = index - ${coordsToCompute[i]} * ${stride}`
538
            : `index -= ${coordsToCompute[i]} * ${stride}`;
539
        return `${line1}; ${line2};`;
540
      })
541
      .join('');
542

543
    source = `
544
     ivec6 getOutputCoords() {
545
         ivec2 resTexRC = ivec2(TexCoords.xy *
546
                               vec2(${texShape[0]}, ${texShape[1]}));
547
         int index = resTexRC.y * ${texShape[0]} + resTexRC.x;
548
         ${coordsFromIndexSnippet}
549
         return ivec6(r, c, d, d2, d3, d4);
550
       }
551
     `;
552
    return new GlslLibRoutine(source);
553
  }
554

555
  /**
556
   * Generates code for common UV coords computation utility functions.
557
   */
558
  protected getCommonUtilFuncs(): { [name: string]: GlslLibRoutine } {
559
    const result: { [name: string]: GlslLibRoutine } = {};
560
    let funcName = 'uvFromFlat';
561
    result[funcName] = new GlslLibRoutine(`
562
    vec2 uvFromFlat(int texNumR, int texNumC, int index) {
563
      int texC = index / texNumR;
564
      int texR = index - texC * texNumR;
565
      // TODO: swap texR, texC order in following function so row is corresponding to u and column is corresponding to
566
      //       v.
567
      return (vec2(texR, texC) + halfCR) / vec2(texNumR, texNumC);
568
    }
569
    `);
570
    funcName = 'packedUVfrom1D';
571
    result[funcName] = new GlslLibRoutine(`
572
      vec2 packedUVfrom1D(int texNumR, int texNumC, int index) {
573
        int texelIndex = index / 2;
574
        int texR = texelIndex / texNumC;
575
        int texC = texelIndex - texR * texNumC;
576
        return (vec2(texC, texR) + halfCR) / vec2(texNumC, texNumR);
577
      }
578
      `);
579
    funcName = 'packedUVfrom2D';
580
    result[funcName] = new GlslLibRoutine(`
581
      vec2 packedUVfrom2D(int texNumR, int texNumC, int texelsInLogicalRow, int row, int col) {
582
        int texelIndex = (row / 2) * texelsInLogicalRow + (col / 2);
583
        int texR = texelIndex / texNumC;
584
        int texC = texelIndex - texR * texNumC;
585
        return (vec2(texC, texR) + halfCR) / vec2(texNumC, texNumR);
586
      }
587
      `);
588
    funcName = 'packedUVfrom3D';
589
    result[funcName] = new GlslLibRoutine(`
590
      vec2 packedUVfrom3D(int texNumR, int texNumC,
591
          int texelsInBatch, int texelsInLogicalRow, int b,
592
          int row, int col) {
593
        int index = b * texelsInBatch + (row / 2) * texelsInLogicalRow + (col / 2);
594
        int texR = index / texNumC;
595
        int texC = index - texR * texNumC;
596
        return (vec2(texC, texR) + halfCR) / vec2(texNumC, texNumR);
597
      }
598
      `);
599
    funcName = 'sampleTexture';
600
    const glsl = getGlsl(this.context.glContext.version);
601
    result[funcName] = new GlslLibRoutine(`
602
        float sampleTexture(sampler2D textureSampler, vec2 uv) {
603
            return ${glsl.texture2D}(textureSampler, uv).r;
604
        }`);
605
    return result;
606
  }
607

608
  /**
609
   * Constructing snippets for inputs
610
   */
611
  protected getInputsSamplingSnippets(): { [name: string]: GlslLibRoutine } {
612
    const result: { [name: string]: GlslLibRoutine } = {};
613
    const outputLayout = this.context.outputTextureLayout;
614
    this.context.programInfo.inputNames.forEach((samplerName, i) => {
615
      const inputLayout = this.context.inputTextureLayouts[i];
616
      const funcName = generateShaderFuncNameFromInputSamplerName(samplerName);
617
      if (inputLayout.isPacked) {
618
        result[funcName] = this.getPackedSamplerFromInput(funcName, samplerName, inputLayout);
619
      } else {
620
        result[funcName] = this.getUnpackedSamplerFromInput(funcName, samplerName, inputLayout);
621
      }
622

623
      const outCoordFuncName = generateShaderFuncNameFromInputSamplerNameAtOutCoords(samplerName);
624
      if (inputLayout.unpackedShape.length <= outputLayout.unpackedShape.length) {
625
        if (inputLayout.isPacked) {
626
          result[outCoordFuncName] = this.getPackedSamplerAtOutputCoords(
627
            outCoordFuncName,
628
            inputLayout,
629
            outputLayout,
630
            samplerName,
631
          );
632
        } else {
633
          result[outCoordFuncName] = this.getUnpackedSamplerAtOutputCoords(
634
            outCoordFuncName,
635
            inputLayout,
636
            outputLayout,
637
            samplerName,
638
          );
639
        }
640
      }
641
    });
642

643
    return result;
644
  }
645

646
  /**
647
   * Constructing snippets for output coordinates of samplers
648
   */
649
  protected getPackedSamplerAtOutputCoords(
650
    funcName: string,
651
    inputLayout: TextureLayout,
652
    outputLayout: TextureLayout,
653
    name: string,
654
  ): GlslLibRoutine {
655
    const inShape = inputLayout.unpackedShape;
656
    const outShape = outputLayout.unpackedShape;
657
    const texName = name;
658
    const texFuncSnippet = generateShaderFuncNameFromInputSamplerName(texName);
659

660
    const inRank = inShape.length;
661
    const outRank = outShape.length;
662

663
    const broadcastDims = BroadcastUtil.getBroadcastDims(inShape, outShape);
664

665
    const type = getCoordsDataType(outRank);
666
    const rankDiff = outRank - inRank;
667
    let coordsSnippet: string;
668
    const fields = getGlChannels();
669

670
    if (inRank === 0) {
671
      coordsSnippet = '';
672
    } else if (outRank < 2 && broadcastDims.length >= 1) {
673
      coordsSnippet = 'coords = 0;';
674
    } else {
675
      coordsSnippet = broadcastDims.map((d) => `coords.${fields[d + rankDiff]} = 0;`).join('\n');
676
    }
677
    let unpackedCoordsSnippet = '';
678
    if (outRank < 2 && inRank > 0) {
679
      unpackedCoordsSnippet = 'coords';
680
    } else {
681
      unpackedCoordsSnippet = inShape.map((_s, i) => `coords.${fields[i + rankDiff]}`).join(', ');
682
    }
683

684
    let output = 'return outputValue;';
685
    const inSize = ShapeUtil.size(inShape);
686
    const isInputScalar = inSize === 1;
687
    const outSize = ShapeUtil.size(outShape);
688
    const isOutputScalar = outSize === 1;
689

690
    if (inRank === 1 && !isInputScalar && !isOutputScalar) {
691
      output = `
692
        return vec4(outputValue.xy, outputValue.xy);
693
      `;
694
    } else if (isInputScalar && !isOutputScalar) {
695
      if (outRank === 1) {
696
        output = `
697
          return vec4(outputValue.x, outputValue.x, 0., 0.);
698
        `;
699
      } else {
700
        output = `
701
          return vec4(outputValue.x);
702
        `;
703
      }
704
    } else if (broadcastDims.length) {
705
      const rows = inRank - 2;
706
      const cols = inRank - 1;
707

708
      if (broadcastDims.indexOf(rows) > -1 && broadcastDims.indexOf(cols) > -1) {
709
        output = 'return vec4(outputValue.x);';
710
      } else if (broadcastDims.indexOf(rows) > -1) {
711
        output = 'return vec4(outputValue.x, outputValue.y, ' + 'outputValue.x, outputValue.y);';
712
      } else if (broadcastDims.indexOf(cols) > -1) {
713
        output = 'return vec4(outputValue.xx, outputValue.zz);';
714
      }
715
    }
716

717
    const swapLastDimsSnippet = `
718
        int lastDim = coords.${fields[outRank - 1]};
719
        coords.${fields[outRank - 1]} = coords.${fields[outRank - 2]};
720
        coords.${fields[outRank - 2]} = lastDim;
721
      `;
722
    const source = `
723
      vec4 ${funcName}() {
724
        ${type} coords = getOutputCoords();
725
        ${swapLastDimsSnippet}
726
        ${coordsSnippet}
727
        vec4 outputValue = ${texFuncSnippet}(${unpackedCoordsSnippet});
728
        ${output}
729
      }
730
    `;
731
    return new GlslLibRoutine(source, ['coordinates.getOutputCoords']);
732
  }
733

734
  /**
735
   * Constructing snippets for unpacked output coordinates of samplers
736
   */
737
  protected getUnpackedSamplerAtOutputCoords(
738
    funcName: string,
739
    inputLayout: TextureLayout,
740
    outputLayout: TextureLayout,
741
    name: string,
742
  ): GlslLibRoutine {
743
    const outTexShape = [outputLayout.width, outputLayout.height];
744
    const inTexShape = [inputLayout.width, inputLayout.height];
745
    const inRank = inputLayout.unpackedShape.length;
746
    const outRank = outputLayout.unpackedShape.length;
747
    const inShape = inputLayout.unpackedShape;
748
    const outShape = outputLayout.unpackedShape;
749
    const texFuncSnippet = generateShaderFuncNameFromInputSamplerName(name);
750

751
    if (inRank === outRank && ArrayUtil.arraysEqual(inTexShape, outTexShape)) {
752
      const source = `
753
          float ${funcName}() {
754
            return sampleTexture(${name}, TexCoords);
755
          }
756
        `;
757
      return new GlslLibRoutine(source, ['coordinates.sampleTexture']);
758
    }
759

760
    const type = getCoordsDataType(outRank);
761
    const broadcastDims = BroadcastUtil.getBroadcastDims(inShape, outShape);
762
    const rankDiff = outRank - inRank;
763
    let coordsSnippet: string;
764
    const fields = getGlChannels();
765

766
    if (inRank === 0) {
767
      coordsSnippet = '';
768
    } else if (outRank < 2 && broadcastDims.length >= 1) {
769
      coordsSnippet = 'coords = 0;';
770
    } else {
771
      coordsSnippet = broadcastDims.map((d) => `coords.${fields[d + rankDiff]} = 0;`).join('\n');
772
    }
773
    let unpackedCoordsSnippet = '';
774
    if (outRank < 2 && inRank > 0) {
775
      unpackedCoordsSnippet = 'coords';
776
    } else {
777
      unpackedCoordsSnippet = inputLayout.unpackedShape.map((_s, i) => `coords.${fields[i + rankDiff]}`).join(', ');
778
    }
779
    const source = `
780
        float ${funcName}() {
781
          ${type} coords = getOutputCoords();
782
          ${coordsSnippet}
783
          return ${texFuncSnippet}(${unpackedCoordsSnippet});
784
        }
785
      `;
786
    return new GlslLibRoutine(source, ['coordinates.getOutputCoords']);
787
  }
788

789
  /**
790
   * Constructing snippets for packed operations.
791
   */
792
  protected getPackedSamplerFromInput(funcName: string, name: string, inputLayout: TextureLayout): GlslLibRoutine {
793
    switch (inputLayout.unpackedShape.length) {
794
      case 0:
795
        return this.getPackedSamplerScalar(funcName, name);
796
      case 1:
797
        return this.getPackedSampler1D(funcName, name, inputLayout);
798
      case 2:
799
        return this.getPackedSampler2D(funcName, name, inputLayout);
800
      case 3:
801
        return this.getPackedSampler3D(funcName, name, inputLayout);
802
      default:
803
        return this.getPackedSamplerND(funcName, name, inputLayout);
804
    }
805
  }
806

807
  /**
808
   * Constructing snippets for unpacked operations.
809
   */
810
  protected getUnpackedSamplerFromInput(funcName: string, name: string, inputLayout: TextureLayout): GlslLibRoutine {
811
    const shape = inputLayout.unpackedShape;
812
    switch (shape.length) {
813
      case 0:
814
        return this.getUnpackedSamplerScalar(funcName, name, inputLayout);
815
      case 1:
816
        return this.getUnpackedSampler1D(funcName, name, inputLayout);
817
      case 2:
818
        return this.getUnpackedSampler2D(funcName, name, inputLayout);
819
      case 3:
820
        return this.getUnpackedSampler3D(funcName, name, inputLayout);
821
      case 4:
822
        return this.getUnpackedSampler4D(funcName, name, inputLayout);
823
      case 5:
824
        return this.getUnpackedSampler5D(funcName, name, inputLayout);
825
      case 6:
826
        return this.getUnpackedSampler6D(funcName, name, inputLayout);
827
      default:
828
        // TODO support more dimensionalities
829
        throw new Error(`Unsupported dimension ${shape.length}-D`);
830
    }
831
  }
832

833
  /**
834
   * Packed scalar snippet.
835
   */
836
  protected getPackedSamplerScalar(funcName: string, name: string): GlslLibRoutine {
837
    const glsl = getGlsl(this.context.glContext.version);
838
    const source = `
839
          vec4 ${funcName}() {
840
            return ${glsl.texture2D}(${name}, halfCR);
841
          }
842
        `;
843
    return new GlslLibRoutine(source);
844
  }
845

846
  /**
847
   * Packed 1D snippet.
848
   */
849
  protected getPackedSampler1D(funcName: string, name: string, inputLayout: TextureLayout): GlslLibRoutine {
850
    const texShape = [inputLayout.width, inputLayout.height];
851
    const packedTexShape = [texShape[1], texShape[0]];
852
    const glsl = getGlsl(this.context.glContext.version);
853

854
    const packedSampler = `vec4 ${funcName}(int index) {
855
      vec2 uv = packedUVfrom1D(
856
      ${packedTexShape[0]}, ${packedTexShape[1]}, index);
857
      return ${glsl.texture2D}(${name}, uv);
858
    }`;
859
    const source = packedSampler;
860
    return new GlslLibRoutine(source, ['coordinates.packedUVfrom1D']);
861
  }
862

863
  /**
864
   * Packed 2D snippet.
865
   */
866
  protected getPackedSampler2D(funcName: string, name: string, inputLayout: TextureLayout): GlslLibRoutine {
867
    const shape = inputLayout.unpackedShape;
868
    const texShape = [inputLayout.width, inputLayout.height];
869
    const glsl = getGlsl(this.context.glContext.version);
870
    const texNumR = texShape[0];
871
    const texNumC = texShape[1];
872

873
    if (texShape != null && ArrayUtil.arraysEqual(shape, texShape)) {
874
      const packedSampler = `vec4 ${funcName}(int row, int col) {
875
        vec2 uv = (vec2(col, row) + halfCR) / vec2(${texNumC}.0, ${texNumR}.0);
876
        return ${glsl.texture2D}(${name}, uv);
877
      }`;
878

879
      return new GlslLibRoutine(packedSampler);
880
    }
881
    const packedTexShape = texShape;
882
    const valuesPerRow = Math.ceil(shape[1] / 2);
883
    const packedSampler = `vec4 ${funcName}(int row, int col) {
884
      vec2 uv = packedUVfrom2D(${packedTexShape[1]}, ${packedTexShape[0]}, ${valuesPerRow}, row, col);
885
      return ${glsl.texture2D}(${name}, uv);
886
    }`;
887
    const source = packedSampler;
888
    return new GlslLibRoutine(source, ['coordinates.packedUVfrom2D']);
889
  }
890

891
  /**
892
   * Packed 3D snippet.
893
   */
894
  protected getPackedSampler3D(funcName: string, name: string, inputLayout: TextureLayout): GlslLibRoutine {
895
    const shape = inputLayout.unpackedShape;
896
    const texShape = [inputLayout.width, inputLayout.height];
897
    const packedTexShape = [texShape[0], texShape[1]];
898
    const glsl = getGlsl(this.context.glContext.version);
899

900
    if (shape[0] === 1) {
901
      const squeezedShape = shape.slice(1);
902
      const keptDims = [1, 2];
903
      const newInputShape = squeezeInputShape(shape, squeezedShape);
904
      const params = ['b', 'row', 'col'];
905
      // Deep copy of input texture layout.
906
      const newInputLayout: TextureLayout = JSON.parse(JSON.stringify(inputLayout));
907
      newInputLayout.unpackedShape = newInputShape;
908
      const samplerRoutine = this.getPackedSamplerFromInput(funcName, name, newInputLayout);
909
      const packedSampler = `${samplerRoutine.routineBody}
910
      vec4 ${funcName}(int b, int row, int col) {
911
        return ${funcName}(${getSqueezedParams(params, keptDims)});
912
      } `;
913
      const source = packedSampler;
914
      return new GlslLibRoutine(source, samplerRoutine.dependencies);
915
    }
916
    const texNumR = packedTexShape[0];
917
    const texNumC = packedTexShape[1];
918

919
    const valuesPerRow = Math.ceil(shape[2] / 2);
920
    const texelsInBatch = valuesPerRow * Math.ceil(shape[1] / 2);
921

922
    const packedSampler = `vec4 ${funcName}(int b, int row, int col) {
923
      vec2 uv = packedUVfrom3D(
924
        ${texNumC}, ${texNumR}, ${texelsInBatch}, ${valuesPerRow}, b, row, col);
925
      return ${glsl.texture2D}(${name}, uv);}`;
926
    const source = packedSampler;
927
    return new GlslLibRoutine(source, ['coordinates.packedUVfrom3D']);
928
  }
929
  /*
930
   * Packed ND snippet.
931
   */
932
  protected getPackedSamplerND(funcName: string, name: string, inputLayout: TextureLayout): GlslLibRoutine {
933
    const shape = inputLayout.unpackedShape;
934
    const rank = shape.length;
935
    const texShape = [inputLayout.width, inputLayout.height];
936
    const glsl = getGlsl(this.context.glContext.version);
937

938
    const packedTexShape = [texShape[0], texShape[1]];
939
    const texNumR = packedTexShape[1];
940
    const texNumC = packedTexShape[0];
941
    const valuesPerRow = Math.ceil(shape[rank - 1] / 2);
942
    let texelsInBatch = valuesPerRow * Math.ceil(shape[rank - 2] / 2);
943
    let params = 'int b, int row, int col';
944
    let index = `b * ${texelsInBatch} + (row / 2) * ${valuesPerRow} + (col / 2)`;
945
    for (let b = 2; b < rank - 1; b++) {
946
      params = `int b${b}, ` + params;
947
      texelsInBatch *= shape[rank - b - 1];
948
      index = `b${b} * ${texelsInBatch} + ` + index;
949
    }
950
    const packedSampler = `vec4 ${funcName}(${params}) {
951
      int index = ${index};
952
      int texR = index / ${texNumC};
953
      int texC = index - texR * ${texNumC};
954
      vec2 uv = (vec2(texC, texR) + halfCR) / vec2(${texNumC}, ${texNumR});
955
      return ${glsl.texture2D}(${name}, uv);
956
    }`;
957
    const source = packedSampler;
958
    return new GlslLibRoutine(source);
959
  }
960

961
  /**
962
   * Unpacked scalar snippet.
963
   */
964
  protected getUnpackedSamplerScalar(funcName: string, name: string, inputLayout: TextureLayout): GlslLibRoutine {
965
    const [texNumR, texNumC] = [inputLayout.width, inputLayout.height];
966
    if (texNumR === 1 && texNumC === 1) {
967
      const source = `
968
          float ${funcName}() {
969
            return sampleTexture(${name}, halfCR);
970
          }
971
        `;
972
      return new GlslLibRoutine(source, ['coordinates.sampleTexture']);
973
    }
974

975
    const source = `
976
        float ${funcName}() {
977
          int offset_${name} = coordsToOffset(TexCoords, ${texNumR}, ${texNumC});
978
          vec2 uv = uvFromFlat(${texNumR}, ${texNumC}, offset_${name});
979
          return sampleTexture(${name}, uv);
980
        }
981
      `;
982
    return new GlslLibRoutine(source, [
983
      'coordinates.uvFromFlat',
984
      'coordinates.sampleTexture',
985
      'coordinates.coordsToOffset',
986
    ]);
987
  }
988

989
  /**
990
   * Unpacked 1D snippet.
991
   */
992
  protected getUnpackedSampler1D(funcName: string, name: string, inputLayout: TextureLayout): GlslLibRoutine {
993
    const tNumR = inputLayout.width;
994
    const tNumC = inputLayout.height;
995

996
    if (tNumC === 1 && tNumR === 1) {
997
      const source = `
998
        float ${funcName}(int index) {
999
          return sampleTexture(${name}, halfCR);
1000
        }
1001
      `;
1002
      return new GlslLibRoutine(source, ['coordinates.sampleTexture']);
1003
    }
1004

1005
    if (tNumC === 1) {
1006
      const source = `
1007
          float ${funcName}(int index) {
1008
            vec2 uv = vec2((float(index) + 0.5) / ${tNumR}.0, 0.5);
1009
            return sampleTexture(${name}, uv);
1010
          }
1011
        `;
1012
      return new GlslLibRoutine(source, ['coordinates.sampleTexture']);
1013
    }
1014
    if (tNumR === 1) {
1015
      const source = `
1016
          float ${funcName}(int index) {
1017
            vec2 uv = vec2(0.5, (float(index) + 0.5) / ${tNumC}.0);
1018
            return sampleTexture(${name}, uv);
1019
          }
1020
        `;
1021
      return new GlslLibRoutine(source, ['coordinates.sampleTexture']);
1022
    }
1023
    const source = `
1024
        float ${funcName}(int index) {
1025
          vec2 uv = uvFromFlat(${tNumR}, ${tNumC}, index);
1026
          return sampleTexture(${name}, uv);
1027
        }
1028
      `;
1029
    return new GlslLibRoutine(source, ['coordinates.uvFromFlat', 'coordinates.sampleTexture']);
1030
  }
1031

1032
  /**
1033
   * Unpacked 2D snippet.
1034
   */
1035

1036
  protected getUnpackedSampler2D(funcName: string, name: string, inputLayout: TextureLayout): GlslLibRoutine {
1037
    const shape = inputLayout.unpackedShape;
1038

1039
    // TODO: modify row/col order for other dimensions.
1040
    const texShape = [inputLayout.height, inputLayout.width];
1041

1042
    if (texShape != null && ArrayUtil.arraysEqual(shape, texShape)) {
1043
      const texNumR = texShape[1];
1044
      const texNumC = texShape[0];
1045
      const source = `
1046
          float ${funcName}(int row, int col) {
1047
            vec2 uv = (vec2(row, col) + halfCR) / vec2(${texNumR}.0, ${texNumC}.0);
1048
            return sampleTexture(${name}, uv);
1049
          }
1050
        `;
1051
      return new GlslLibRoutine(source, ['coordinates.sampleTexture']);
1052
    }
1053

1054
    const { newShape, keptDims } = squeezeShape(shape as number[]);
1055
    const squeezedShape = newShape;
1056
    if (squeezedShape.length < shape.length) {
1057
      const newInputShape = squeezeInputShape(shape, squeezedShape);
1058
      // Deep copy of input texture layout.
1059
      const newInputLayout: TextureLayout = JSON.parse(JSON.stringify(inputLayout));
1060
      newInputLayout.unpackedShape = newInputShape;
1061

1062
      const params = ['col', 'row'];
1063
      const source = `
1064
          ${this.getUnpackedSamplerFromInput(funcName, name, newInputLayout).routineBody}
1065
          float ${funcName}(int row, int col) {
1066
            return ${funcName}(${getSqueezedParams(params, keptDims)});
1067
          }
1068
        `;
1069
      return new GlslLibRoutine(source, ['coordinates.sampleTexture']);
1070
    }
1071

1072
    const texNumR = texShape[1];
1073
    const texNumC = texShape[0];
1074
    if (texNumC === 1) {
1075
      const source = `
1076
          float ${funcName}(int row, int col) {
1077
            int offset_${name} = coordsToOffset(TexCoords, ${texNumR}, ${texNumC});
1078
            float index = dot(vec3(row, col, offset_${name}), vec3(${shape[1]}, 1, 1));
1079
            vec2 uv = vec2(0.5, (index + 0.5) / ${texNumR}.0);
1080
            return sampleTexture(${name}, uv);
1081
          }
1082
        `;
1083
      return new GlslLibRoutine(source, ['coordinates.sampleTexture', 'coordinates.coordsToOffset']);
1084
    }
1085

1086
    if (texNumR === 1) {
1087
      const source = `
1088
          float ${funcName}(int row, int col) {
1089
            int offset_${name} = coordsToOffset(TexCoords, ${texNumR}, ${texNumC});
1090
            float index = dot(vec3(row, col, offset_${name}), vec3(${shape[1]}, 1, 1));
1091
            vec2 uv = vec2((index + 0.5) / ${texNumC}.0, 0.5);
1092
            return sampleTexture(${name}, uv);
1093
          }
1094
        `;
1095
      return new GlslLibRoutine(source, ['coordinates.sampleTexture', 'coordinates.coordsToOffset']);
1096
    }
1097

1098
    const source = `
1099
        float ${funcName}(int row, int col) {
1100
          int index = col * ${shape[1]} + row;
1101
          vec2 uv = uvFromFlat(${texNumR}, ${texNumC}, index);
1102
          return sampleTexture(${name}, uv);
1103
        }
1104
      `;
1105
    return new GlslLibRoutine(source, [
1106
      'coordinates.uvFromFlat',
1107
      'coordinates.sampleTexture',
1108
      'coordinates.coordsToOffset',
1109
    ]);
1110
  }
1111

1112
  /**
1113
   * Unpacked 3D snippet.
1114
   */
1115

1116
  protected getUnpackedSampler3D(funcName: string, name: string, inputLayout: TextureLayout): GlslLibRoutine {
1117
    const shape = inputLayout.unpackedShape;
1118
    const stride0 = shape[1] * shape[2];
1119
    const stride1 = shape[2];
1120

1121
    const { newShape, keptDims } = squeezeShape(shape as number[]);
1122
    const squeezedShape = newShape;
1123
    if (squeezedShape.length < shape.length) {
1124
      const newInputShape = squeezeInputShape(shape, squeezedShape);
1125
      const params = ['batch', 'col', 'row'];
1126
      // Deep copy of input texture layout.
1127
      const newInputLayout: TextureLayout = JSON.parse(JSON.stringify(inputLayout));
1128
      newInputLayout.unpackedShape = newInputShape;
1129
      const routine = this.getUnpackedSamplerFromInput(funcName, name, newInputLayout);
1130
      // TODO: revisit the logic here to make it simpler
1131
      const revDims = keptDims.reverse();
1132
      const source = `
1133
          ${routine.routineBody}
1134
          float ${funcName}(int batch, int row, int col) {
1135
            return ${funcName}(${getSqueezedParams(params, revDims)});
1136
          }
1137
        `;
1138
      return new GlslLibRoutine(source, routine.dependencies);
1139
    }
1140

1141
    const texNumR = inputLayout.width;
1142
    const texNumC = inputLayout.height;
1143
    const source = `
1144
          float ${funcName}(int depth, int row, int col) {
1145
            // Explicitly use integer operations as dot() only works on floats.
1146
            int index = depth * ${stride0} + col * ${stride1} + row;
1147
            vec2 uv = uvFromFlat(${texNumR}, ${texNumC}, index);
1148
            return sampleTexture(${name}, uv);
1149
          }
1150
      `;
1151
    return new GlslLibRoutine(source, [
1152
      'coordinates.uvFromFlat',
1153
      'coordinates.sampleTexture',
1154
      'coordinates.coordsToOffset',
1155
    ]);
1156
  }
1157

1158
  /**
1159
   * Unpacked 4D snippet.
1160
   */
1161

1162
  protected getUnpackedSampler4D(funcName: string, name: string, inputLayout: TextureLayout): GlslLibRoutine {
1163
    const shape = inputLayout.unpackedShape;
1164
    const stride2 = shape[3];
1165
    const stride1 = shape[2] * stride2;
1166
    const stride0 = shape[1] * stride1;
1167

1168
    //
1169
    // TODO: re-enable this shortcut once the index calculation bug is fixed.
1170
    //
1171
    // const {newShape, keptDims} = squeezeShape(shape as number[]);
1172
    // if (newShape.length < shape.length) {
1173
    //   const newInputShape = squeezeInputShape(shape, newShape);
1174
    //   const params = ['row', 'col', 'depth', 'depth2'];
1175
    //   // Deep copy of input texture layout.
1176
    //   const newInputLayout: TextureLayout = JSON.parse(JSON.stringify(inputLayout));
1177
    //   newInputLayout.unpackedShape = newInputShape;
1178
    //   const source = `
1179
    //       ${this.getUnpackedSamplerFromInput(funcName, name, newInputLayout).routineBody}
1180
    //       float ${funcName}(int row, int col, int depth, int depth2) {
1181
    //         return ${funcName}(${getSqueezedParams(params, keptDims)});
1182
    //       }
1183
    //     `;
1184
    //   return new GlslLibRoutine(
1185
    //       source, ['coordinates.uvFromFlat', 'coordinates.sampleTexture', 'coordinates.coordsToOffset']);
1186
    // }
1187

1188
    const texNumR = inputLayout.width;
1189
    const texNumC = inputLayout.height;
1190
    const source = `
1191
        float ${funcName}(int row, int col, int depth, int depth2) {
1192
          int index = row * ${stride0} + col * ${stride1} +
1193
              depth2 * ${stride2} + depth;
1194
          vec2 uv = uvFromFlat(${texNumR}, ${texNumC}, index);
1195
          return sampleTexture(${name}, uv);
1196
        }
1197
      `;
1198
    return new GlslLibRoutine(source, ['coordinates.uvFromFlat', 'coordinates.sampleTexture']);
1199
  }
1200

1201
  /**
1202
   * Unpacked 5D snippet.
1203
   */
1204
  protected getUnpackedSampler5D(funcName: string, name: string, inputLayout: TextureLayout): GlslLibRoutine {
1205
    const shape = inputLayout.unpackedShape;
1206
    const stride3 = shape[4];
1207
    const stride2 = shape[3] * stride3;
1208
    const stride1 = shape[2] * stride2;
1209
    const stride0 = shape[1] * stride1;
1210

1211
    const { newShape, keptDims } = squeezeShape(shape as number[]);
1212
    if (newShape.length < shape.length) {
1213
      const newInputShape = squeezeInputShape(shape, newShape);
1214
      const params = ['row', 'col', 'depth', 'depth2', 'depth3'];
1215
      // Deep copy of input texture layout.
1216
      const newInputLayout: TextureLayout = JSON.parse(JSON.stringify(inputLayout));
1217
      newInputLayout.unpackedShape = newInputShape;
1218

1219
      const source = `
1220
          ${this.getUnpackedSamplerFromInput(funcName, name, newInputLayout).routineBody}
1221
          float ${funcName}(int row, int col, int depth, int depth2, int depth3) {
1222
            return ${funcName}(${getSqueezedParams(params, keptDims)});
1223
          }
1224
        `;
1225
      return new GlslLibRoutine(source, ['coordinates.sampleTexture', 'coordinates.uvFromFlat']);
1226
    }
1227

1228
    const texNumR = inputLayout.width;
1229
    const texNumC = inputLayout.height;
1230
    const source = `
1231
        float ${funcName}(int row, int col, int depth, int depth2, int depth3) {
1232
          int index = row * ${stride0} + col * ${stride1} + depth * ${stride2} +
1233
          depth3 * ${stride3} + depth2;
1234
          vec2 uv = uvFromFlat(${texNumR}, ${texNumC}, index);
1235
          return sampleTexture(${name}, uv);
1236
        }
1237
      `;
1238
    return new GlslLibRoutine(source, ['coordinates.sampleTexture', 'coordinates.uvFromFlat']);
1239
  }
1240

1241
  /**
1242
   * Unpacked 6D snippet.
1243
   */
1244
  protected getUnpackedSampler6D(funcName: string, name: string, inputLayout: TextureLayout): GlslLibRoutine {
1245
    const shape = inputLayout.unpackedShape;
1246
    const stride4 = shape[5];
1247
    const stride3 = shape[4] * stride4;
1248
    const stride2 = shape[3] * stride3;
1249
    const stride1 = shape[2] * stride2;
1250
    const stride0 = shape[1] * stride1;
1251

1252
    const { newShape, keptDims } = squeezeShape(shape as number[]);
1253
    if (newShape.length < shape.length) {
1254
      const newInputShape = squeezeInputShape(shape, newShape);
1255
      const params = ['row', 'col', 'depth', 'depth2', 'depth3', 'depth4'];
1256
      // Deep copy of input texture layout.
1257
      const newInputLayout: TextureLayout = JSON.parse(JSON.stringify(inputLayout));
1258
      newInputLayout.unpackedShape = newInputShape;
1259

1260
      const source = `
1261
            ${this.getUnpackedSamplerFromInput(funcName, name, newInputLayout).routineBody}
1262
            float ${funcName}(int row, int col, int depth,
1263
              int depth2, int depth3, int depth4) {
1264
              return ${funcName}(${getSqueezedParams(params, keptDims)});
1265
            }
1266
          `;
1267
      return new GlslLibRoutine(source, ['coordinates.sampleTexture', 'coordinates.uvFromFlat']);
1268
    }
1269

1270
    const texNumR = inputLayout.width;
1271
    const texNumC = inputLayout.height;
1272
    const source = `
1273
          float ${funcName}(int row, int col, int depth,
1274
            int depth2, int depth3, int depth4) {
1275
            int index = row * ${stride0} + col * ${stride1} + depth * ${stride2} +
1276
            depth2 * ${stride3} + depth3 * ${stride4} + depth4;
1277
            vec2 uv = uvFromFlat(${texNumR}, ${texNumC}, index);
1278
            return sampleTexture(${name}, uv);
1279
          }
1280
        `;
1281
    return new GlslLibRoutine(source, [
1282
      'coordinates.uvFromFlat',
1283
      'coordinates.sampleTexture',
1284
      'coordinates.coordsToOffset',
1285
    ]);
1286
  }
1287

1288
  /**
1289
   * This is the main function to map from the given texture coordinates (s,t)
1290
   * to logical indices for the output
1291
   * There will only be one single variation of this
1292
   * Also see coordsToOffset and offsetToIndices for input-specific versions
1293
   */
1294
  protected toVec(): { [name: string]: GlslLibRoutine } {
1295
    const output = this.context.outputTextureLayout;
1296
    const rank = output.shape.length;
1297
    const strides = output.strides;
1298
    const xScale = output.width;
1299
    const yScale = output.height;
1300

1301
    const stridesBlock = [];
1302
    for (let i = 0; i < rank - 1; ++i) {
1303
      stridesBlock.push(`
1304
        c[${i}] = offset / ${strides[i]};`);
1305
      stridesBlock.push(`
1306
        offset -= c[${i}] * ${strides[i]};`);
1307
    }
1308
    stridesBlock.push(`
1309
        c[${rank - 1}] = offset;`);
1310
    const body = `
1311
      void toVec(vec2 texCoords, out int c[${rank}]) {
1312
        int offset = coordsToOffset(texCoords, ${xScale}, ${yScale});
1313
        ${stridesBlock.join('')}
1314
      }
1315
      void toVec(int offset, out int c[${rank}]) {
1316
        ${stridesBlock.join('')}
1317
      }
1318
    `;
1319
    return { toVec: new GlslLibRoutine(body, ['coordinates.coordsToOffset']) };
1320
  }
1321
  /**
1322
   * These are value getter functions generated for each input
1323
   * Each function is hardwired to the name and dimensions of the input
1324
   * An '_T' variation is also produced which accesses values as if the
1325
   * input was transposed
1326
   */
1327
  protected valueFrom(): { [name: string]: GlslLibRoutine } {
1328
    const result: { [name: string]: GlslLibRoutine } = {};
1329
    this.context.programInfo.inputNames.forEach((name, i) => {
1330
      const layout = this.context.inputTextureLayouts[i];
1331
      const shape = layout.unpackedShape.length > 0 ? layout.unpackedShape : layout.shape;
1332
      const rank = shape.length;
1333
      let funcName = `_${name}`;
1334
      result[funcName] = new GlslLibRoutine(this.getValueFromSingle(name, rank, layout.width, layout.height, false), [
1335
        `shapeUtils.indicesToOffset${funcName}`,
1336
        'coordinates.offsetToCoords',
1337
        'fragcolor.getColorAsFloat',
1338
      ]);
1339
      funcName = funcName + '_T';
1340
      result[funcName] = new GlslLibRoutine(this.getValueFromSingle(name, rank, layout.width, layout.height, true), [
1341
        `shapeUtils.indicesToOffset${funcName}`,
1342
        'coordinates.offsetToCoords',
1343
        'fragcolor.getColorAsFloat',
1344
      ]);
1345
    });
1346
    return result;
1347
  }
1348
  /**
1349
   * Produces one value getter function for the name and rank given
1350
   * If a transpose is set proper offsetToCoords mapping will be used
1351
   * @param name name of the function
1352
   * @param rank rank of the input
1353
   * @param transpose whether or not should generate a transpose variation
1354
   */
1355
  protected getValueFromSingle(
1356
    varName: string,
1357
    rank: number,
1358
    width: number,
1359
    height: number,
1360
    transpose: boolean,
1361
  ): string {
1362
    let name = `_${varName}`;
1363
    if (transpose) {
1364
      name = name + '_T';
1365
    }
1366
    const glsl = getGlsl(this.context.glContext.version);
1367
    return `
1368
        float ${name}(int m[${rank}]) {
1369
          int offset = indicesToOffset${name}(m);
1370
          vec2 coords = offsetToCoords(offset, ${width}, ${height});
1371
          float value = getColorAsFloat(${glsl.texture2D}(${varName}, coords));
1372
          return value;
1373
        }
1374
        `;
1375
  }
1376

1377
  /**
1378
   * Produces a packed value getter function for the name and rank given
1379
   * If a transpose is set proper offsetToCoords mapping will be used
1380
   * @param name name of the function
1381
   * @param rank rank of the input
1382
   * @param transpose whether or not should generate a transpose variation
1383
   */
1384
  protected getPackedValueFrom(
1385
    varName: string,
1386
    rank: number,
1387
    width: number,
1388
    height: number,
1389
    transpose: boolean,
1390
  ): string {
1391
    let name = `_${varName}_Pack`;
1392
    if (transpose) {
1393
      name = name + '_T';
1394
    }
1395
    const glsl = getGlsl(this.context.glContext.version);
1396
    return `
1397
        vec4 ${name}(int m[${rank}]) {
1398
          int offset = indicesToOffset_${varName}(m);
1399
          vec2 coords = offsetToCoords(offset, ${width}, ${height});
1400
          return ${glsl.texture2D}(${varName}, coords);
1401
        }
1402
        `;
1403
  }
1404
}
1405