| 49 | } |
| 50 | |
| 51 | getUserCode(): string { |
| 52 | const userCode = ` |
| 53 | ${main('index')} { |
| 54 | if (index < uniforms.size) { |
| 55 | let outC = getCoordsFromIndex(index); |
| 56 | let batch = outC[0]; |
| 57 | let elemIdx = outC[1]; |
| 58 | // We compare elements pair-wise within a group of size 2 * inc. |
| 59 | // The comparing rule for each group alternates between ascending |
| 60 | // and descending. Within each group, we compare each pair at |
| 61 | // positions i and i+inc. To decide whether an element at position i |
| 62 | // is x0 or x1, we mod it by 2 * inc, if the result is smaller than |
| 63 | // inc, it is in the first half of the group, we denote it as x0, |
| 64 | // otherwise we denote it as x1. |
| 65 | // For example, as shown in the Bitonic top K paper referenced |
| 66 | // above, Figure5(a) shows that element[1] is in the second half of |
| 67 | // the group when group size is 2, but it is in the first half of |
| 68 | // the group when group size is 4. |
| 69 | let isFirstInPair = elemIdx % (2 * uniforms.inc) < uniforms.inc; |
| 70 | var i = 0; |
| 71 | if (isFirstInPair) { |
| 72 | i = elemIdx; |
| 73 | } else { |
| 74 | i = elemIdx - uniforms.inc; |
| 75 | } |
| 76 | |
| 77 | var i0 = 0; |
| 78 | if (uniforms.firstPass == 1) { |
| 79 | i0 = i; |
| 80 | } else { |
| 81 | i0 = i32(getIndices(batch, i)); |
| 82 | } |
| 83 | |
| 84 | var i1 = 0; |
| 85 | if (uniforms.firstPass == 1) { |
| 86 | i1 = i + uniforms.inc; |
| 87 | } else { |
| 88 | i1 = i32(getIndices(batch, i + uniforms.inc)); |
| 89 | } |
| 90 | |
| 91 | var x0 = f32(0.0); |
| 92 | var x1 = f32(0.0); |
| 93 | if (i0 < uniforms.inputSize) { |
| 94 | x0 = getX(batch, i0); |
| 95 | } else { |
| 96 | x0 = uniforms.negativeInf; |
| 97 | } |
| 98 | if (i1 < uniforms.inputSize) { |
| 99 | x1 = getX(batch, i1); |
| 100 | } else { |
| 101 | x1 = uniforms.negativeInf; |
| 102 | } |
| 103 | |
| 104 | let reverse = elemIdx % (2 * uniforms.dir) >= uniforms.dir; |
| 105 | let isGreater = x0 > x1 || (x0 == x1 && i1 > i0); |
| 106 | if (reverse == isGreater) { |
| 107 | // Elements in opposite order of direction |
| 108 | let iTemp = i0; |