| 101 | } |
| 102 | |
| 103 | private getSetup(dims: string[]): string { |
| 104 | if (this.rank === 1) { |
| 105 | return ''; |
| 106 | } |
| 107 | |
| 108 | const innerDims = dims.slice(-2); |
| 109 | const col = this.enableShapeUniforms ? `outShape[${this.rank} - 1]` : |
| 110 | this.outputShape[this.rank - 1]; |
| 111 | const row = this.enableShapeUniforms ? `outShape[${this.rank} - 2]` : |
| 112 | this.outputShape[this.rank - 2]; |
| 113 | |
| 114 | return ` |
| 115 | int r = ${innerDims[0]}; |
| 116 | int c = ${innerDims[1]}; |
| 117 | int rp1 = r + 1; |
| 118 | int cp1 = c + 1; |
| 119 | |
| 120 | bool cEdge = cp1 >= ${col}; |
| 121 | bool rEdge = rp1 >= ${row}; |
| 122 | `; |
| 123 | } |
| 124 | |
| 125 | private getOutput(dims: string[]): string { |
| 126 | const sourceCoords = this.getSourceCoordsArr(dims); |