gx

chenyc

2025-02-12 ea42ff3ebee1eeb3fb29423aa848a249441db81c

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
import { getCoordsDataType } from './shader_compiler';
export class GatherNDProgram {
    constructor(sliceDim, strides, shape, paramsShape) {
        this.sliceDim = sliceDim;
        this.strides = strides;
        this.paramsShape = paramsShape;
        this.variableNames = ['x', 'indices'];
        this.outputShape = shape;
        const dtype = getCoordsDataType(shape.length);
        let mainLoop = `
    int index;`;
        for (let j = 0; j < this.sliceDim; j++) {
            mainLoop += `
          index = round(getIndices(coords[0], ${j}));
          out_of_bounds = out_of_bounds || index < 0;
          out_of_bounds = out_of_bounds || index >= ${this.paramsShape[j]};
          flattenIndex += index * ${this.strides[j]};`;
        }
        this.userCode = `
         void main() {
          ${dtype} coords = getOutputCoords();
          int flattenIndex = 0;
          bool out_of_bounds = false;
 
          ${mainLoop}
 
          setOutput(out_of_bounds ? 0.0 : getX(flattenIndex, coords[1]));
        }
      `;
    }
}
//# sourceMappingURL=data:application/json;base64,eyJ2ZXJzaW9uIjozLCJmaWxlIjoiZ2F0aGVyX25kX2dwdS5qcyIsInNvdXJjZVJvb3QiOiIiLCJzb3VyY2VzIjpbIi4uLy4uLy4uLy4uLy4uL3RmanMtYmFja2VuZC13ZWJnbC9zcmMvZ2F0aGVyX25kX2dwdS50cyJdLCJuYW1lcyI6W10sIm1hcHBpbmdzIjoiQUFpQkEsT0FBTyxFQUFDLGlCQUFpQixFQUFDLE1BQU0sbUJBQW1CLENBQUM7QUFFcEQsTUFBTSxPQUFPLGVBQWU7SUFJMUIsWUFDWSxRQUFnQixFQUFVLE9BQWlCLEVBQUUsS0FBZSxFQUM1RCxXQUFxQjtRQURyQixhQUFRLEdBQVIsUUFBUSxDQUFRO1FBQVUsWUFBTyxHQUFQLE9BQU8sQ0FBVTtRQUMzQyxnQkFBVyxHQUFYLFdBQVcsQ0FBVTtRQUxqQyxrQkFBYSxHQUFHLENBQUMsR0FBRyxFQUFFLFNBQVMsQ0FBQyxDQUFDO1FBTS9CLElBQUksQ0FBQyxXQUFXLEdBQUcsS0FBSyxDQUFDO1FBQ3pCLE1BQU0sS0FBSyxHQUFHLGlCQUFpQixDQUFDLEtBQUssQ0FBQyxNQUFNLENBQUMsQ0FBQztRQUU5QyxJQUFJLFFBQVEsR0FBRztlQUNKLENBQUM7UUFDWixLQUFLLElBQUksQ0FBQyxHQUFHLENBQUMsRUFBRSxDQUFDLEdBQUcsSUFBSSxDQUFDLFFBQVEsRUFBRSxDQUFDLEVBQUUsRUFBRTtZQUN0QyxRQUFRLElBQUk7Z0RBQzhCLENBQUM7O3NEQUVLLElBQUksQ0FBQyxXQUFXLENBQUMsQ0FBQyxDQUFDO29DQUNyQyxJQUFJLENBQUMsT0FBTyxDQUFDLENBQUMsQ0FBQyxHQUFHLENBQUM7U0FDbEQ7UUFFRCxJQUFJLENBQUMsUUFBUSxHQUFHOztZQUVSLEtBQUs7Ozs7WUFJTCxRQUFROzs7O09BSWIsQ0FBQztJQUNOLENBQUM7Q0FDRiIsInNvdXJjZXNDb250ZW50IjpbIi8qKlxuICogQGxpY2Vuc2VcbiAqIENvcHlyaWdodCAyMDE4IEdvb2dsZSBMTEMuIEFsbCBSaWdodHMgUmVzZXJ2ZWQuXG4gKiBMaWNlbnNlZCB1bmRlciB0aGUgQXBhY2hlIExpY2Vuc2UsIFZlcnNpb24gMi4wICh0aGUgXCJMaWNlbnNlXCIpO1xuICogeW91IG1heSBub3QgdXNlIHRoaXMgZmlsZSBleGNlcHQgaW4gY29tcGxpYW5jZSB3aXRoIHRoZSBMaWNlbnNlLlxuICogWW91IG1heSBvYnRhaW4gYSBjb3B5IG9mIHRoZSBMaWNlbnNlIGF0XG4gKlxuICogaHR0cDovL3d3dy5hcGFjaGUub3JnL2xpY2Vuc2VzL0xJQ0VOU0UtMi4wXG4gKlxuICogVW5sZXNzIHJlcXVpcmVkIGJ5IGFwcGxpY2FibGUgbGF3IG9yIGFncmVlZCB0byBpbiB3cml0aW5nLCBzb2Z0d2FyZVxuICogZGlzdHJpYnV0ZWQgdW5kZXIgdGhlIExpY2Vuc2UgaXMgZGlzdHJpYnV0ZWQgb24gYW4gXCJBUyBJU1wiIEJBU0lTLFxuICogV0lUSE9VVCBXQVJSQU5USUVTIE9SIENPTkRJVElPTlMgT0YgQU5ZIEtJTkQsIGVpdGhlciBleHByZXNzIG9yIGltcGxpZWQuXG4gKiBTZWUgdGhlIExpY2Vuc2UgZm9yIHRoZSBzcGVjaWZpYyBsYW5ndWFnZSBnb3Zlcm5pbmcgcGVybWlzc2lvbnMgYW5kXG4gKiBsaW1pdGF0aW9ucyB1bmRlciB0aGUgTGljZW5zZS5cbiAqID09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09XG4gKi9cbmltcG9ydCB7R1BHUFVQcm9ncmFtfSBmcm9tICcuL2dwZ3B1X21hdGgnO1xuaW1wb3J0IHtnZXRDb29yZHNEYXRhVHlwZX0gZnJvbSAnLi9zaGFkZXJfY29tcGlsZXInO1xuXG5leHBvcnQgY2xhc3MgR2F0aGVyTkRQcm9ncmFtIGltcGxlbWVudHMgR1BHUFVQcm9ncmFtIHtcbiAgdmFyaWFibGVOYW1lcyA9IFsneCcsICdpbmRpY2VzJ107XG4gIG91dHB1dFNoYXBlOiBudW1iZXJbXTtcbiAgdXNlckNvZGU6IHN0cmluZztcbiAgY29uc3RydWN0b3IoXG4gICAgICBwcml2YXRlIHNsaWNlRGltOiBudW1iZXIsIHByaXZhdGUgc3RyaWRlczogbnVtYmVyW10sIHNoYXBlOiBudW1iZXJbXSxcbiAgICAgIHByaXZhdGUgcGFyYW1zU2hhcGU6IG51bWJlcltdKSB7XG4gICAgdGhpcy5vdXRwdXRTaGFwZSA9IHNoYXBlO1xuICAgIGNvbnN0IGR0eXBlID0gZ2V0Q29vcmRzRGF0YVR5cGUoc2hhcGUubGVuZ3RoKTtcblxuICAgIGxldCBtYWluTG9vcCA9IGBcbiAgICBpbnQgaW5kZXg7YDtcbiAgICBmb3IgKGxldCBqID0gMDsgaiA8IHRoaXMuc2xpY2VEaW07IGorKykge1xuICAgICAgbWFpbkxvb3AgKz0gYFxuICAgICAgICAgIGluZGV4ID0gcm91bmQoZ2V0SW5kaWNlcyhjb29yZHNbMF0sICR7an0pKTtcbiAgICAgICAgICBvdXRfb2ZfYm91bmRzID0gb3V0X29mX2JvdW5kcyB8fCBpbmRleCA8IDA7XG4gICAgICAgICAgb3V0X29mX2JvdW5kcyA9IG91dF9vZl9ib3VuZHMgfHwgaW5kZXggPj0gJHt0aGlzLnBhcmFtc1NoYXBlW2pdfTtcbiAgICAgICAgICBmbGF0dGVuSW5kZXggKz0gaW5kZXggKiAke3RoaXMuc3RyaWRlc1tqXX07YDtcbiAgICB9XG5cbiAgICB0aGlzLnVzZXJDb2RlID0gYFxuICAgICAgICAgdm9pZCBtYWluKCkge1xuICAgICAgICAgICR7ZHR5cGV9IGNvb3JkcyA9IGdldE91dHB1dENvb3JkcygpO1xuICAgICAgICAgIGludCBmbGF0dGVuSW5kZXggPSAwO1xuICAgICAgICAgIGJvb2wgb3V0X29mX2JvdW5kcyA9IGZhbHNlO1xuXG4gICAgICAgICAgJHttYWluTG9vcH1cblxuICAgICAgICAgIHNldE91dHB1dChvdXRfb2ZfYm91bmRzID8gMC4wIDogZ2V0WChmbGF0dGVuSW5kZXgsIGNvb3Jkc1sxXSkpO1xuICAgICAgICB9XG4gICAgICBgO1xuICB9XG59XG4iXX0=