| 66 | } |
| 67 | |
| 68 | getUserCode(): string { |
| 69 | const userCode = ` |
| 70 | ${this.binaryOutput ? binaryWriteSnippet : writeSnippet} |
| 71 | ${main('index')} { |
| 72 | ${ |
| 73 | this.rank === 1 ? |
| 74 | `if (index < uniforms.xShape) { |
| 75 | let indexVal = i32(getX(index)); |
| 76 | if (indexVal < uniforms.binCountSize) { |
| 77 | let value = ${ |
| 78 | this.binaryOutput ? 1. : |
| 79 | (this.hasWeights ? 'getW(index)' : '1.')}; |
| 80 | bincount_write(indexVal, value); |
| 81 | } |
| 82 | }` : |
| 83 | `let coord = getCoordsFromIndex(index); |
| 84 | if (coordsInBounds2D(coord, uniforms.xShape)) { |
| 85 | let indexVal = i32(getX(coord[0], coord[1])); |
| 86 | if (indexVal < uniforms.binCountSize) { |
| 87 | let value = ${ |
| 88 | this.binaryOutput ? |
| 89 | 1. : |
| 90 | (this.hasWeights ? 'getW(coord[0], coord[1])' : '1.')}; |
| 91 | bincount_write(coord.x * uniforms.binCountSize + indexVal, value); |
| 92 | } |
| 93 | }`} |
| 94 | } |
| 95 | `; |
| 96 | return userCode; |
| 97 | } |
| 98 | } |