| 145 | } |
| 146 | |
| 147 | getUserCode(): string { |
| 148 | const userCode = ` |
| 149 | ${main('index')} { |
| 150 | if (index < uniforms.size) { |
| 151 | let outC = getCoordsFromIndex(index); |
| 152 | let batch = outC[0]; |
| 153 | let elemIdx = outC[1]; |
| 154 | // The output size is half of the previous size. |
| 155 | // If the previous sequence is | | | | _ _ _ _ | | | | _ _ _ _ |
| 156 | // (k=4), we only need to output the indices at positions |, the |
| 157 | // indices at positions _ can be thrown away, see Figure5(b) After |
| 158 | // Phase 2 (Merge phase) in the Bitonic Top K paper referenced |
| 159 | // above. |
| 160 | // For example, the paper shows we only need to output the orange |
| 161 | // bars. The output sequence should look like this | | | | | | | |. |
| 162 | // Because the sequence is halved, to map the output index back to |
| 163 | // the previous sequence to find the corresponding value, we need |
| 164 | // to double the index. When we double the index, we basically |
| 165 | // interpolate a position, so 2i looks like |
| 166 | // | _ | _ | _ | _ | _ | _ | _. We move the | to the first k |
| 167 | // position of each 2k positions by - elemIdx % k. E.g. for output |
| 168 | // at index 4,5,6,7, we want to get the corresponding element at |
| 169 | // original index 8,9,10,11, for output at index 8,9,10,11, |
| 170 | // we want to get the corresponding element at original index |
| 171 | // 16,17,18,19, so on and so forth. |
| 172 | |
| 173 | var i = 0; |
| 174 | if (elemIdx < uniforms.k) { |
| 175 | i = elemIdx; |
| 176 | } else { |
| 177 | i = elemIdx * 2 - elemIdx % uniforms.k; |
| 178 | } |
| 179 | var i0 = 0; |
| 180 | if (uniforms.firstPass == 1) { |
| 181 | i0 = i; |
| 182 | } else { |
| 183 | i0 = i32(getIndices(batch, i)); |
| 184 | } |
| 185 | var i1 = 0; |
| 186 | if (uniforms.firstPass == 1) { |
| 187 | i1 = i + uniforms.k; |
| 188 | } else { |
| 189 | i1 = i32(getIndices(batch, i + uniforms.k)); |
| 190 | } |
| 191 | |
| 192 | let x0 = getX(batch, i0); |
| 193 | var x1 = f32(0.0); |
| 194 | if (i1 < uniforms.inputSize) { |
| 195 | x1 = getX(batch, i1); |
| 196 | } else { |
| 197 | x1 = x0; |
| 198 | } |
| 199 | |
| 200 | if (x0 >= x1) { |
| 201 | setOutputAtIndex(index, f32(i0)); |
| 202 | } else { |
| 203 | setOutputAtIndex(index, f32(i1)); |
| 204 | } |