| 41 | } |
| 42 | |
| 43 | getUserCode(): string { |
| 44 | // TODO(WGSL): below code can be merged with getUserCode. |
| 45 | let cCoords; |
| 46 | let abCoords; |
| 47 | if (this.rank > 4) { |
| 48 | throw Error(`Where for rank ${this.rank} is not yet supported`); |
| 49 | } |
| 50 | |
| 51 | if (this.rank === 1) { |
| 52 | abCoords = `resRC`; |
| 53 | cCoords = `resRC`; |
| 54 | } else { |
| 55 | const currentCoords = ['resRC.x', 'resRC.y', 'resRC.z', 'resRC.w']; |
| 56 | const cCoordVars = []; |
| 57 | const abCoordVars = []; |
| 58 | for (let i = 0; i < this.outputShape.length; i++) { |
| 59 | abCoordVars.push(`${currentCoords[i]}`); |
| 60 | if (i < this.cRank) { |
| 61 | cCoordVars.push(`${currentCoords[i]}`); |
| 62 | } |
| 63 | } |
| 64 | cCoords = cCoordVars.join(); |
| 65 | abCoords = abCoordVars.join(); |
| 66 | } |
| 67 | |
| 68 | const userCode = ` |
| 69 | ${main('index')} { |
| 70 | if (index < uniforms.size) { |
| 71 | let resRC = getCoordsFromIndex(index); |
| 72 | let cVal = getC(${cCoords}); |
| 73 | if (cVal >= 1.0) { |
| 74 | setOutputAtIndex(index, getA(${abCoords})); |
| 75 | } else { |
| 76 | setOutputAtIndex(index, getB(${abCoords})); |
| 77 | } |
| 78 | } |
| 79 | } |
| 80 | `; |
| 81 | return userCode; |
| 82 | } |
| 83 | } |