| 66 | } |
| 67 | |
| 68 | getUserCode(): string { |
| 69 | const workgroupSizeX = this.workgroupSize[0]; |
| 70 | const getInputShapeLastDim = () => { |
| 71 | if (this.inputShape.length === 1) { |
| 72 | return 'uniforms.xShape'; |
| 73 | } else { |
| 74 | return `uniforms.xShape.${getCoordsXYZ(this.inputShape.length - 1)}`; |
| 75 | } |
| 76 | }; |
| 77 | |
| 78 | const splitOutputCoords = () => { |
| 79 | let snippet = ''; |
| 80 | if (this.outputShape.length === 1) { |
| 81 | if (this.inputShape.length !== 1) { |
| 82 | snippet += 'outputCoords,'; |
| 83 | } |
| 84 | } else { |
| 85 | for (let i = 0; i < this.outputShape.length; i++) { |
| 86 | snippet += `outputCoords.${getCoordsXYZ(i)},`; |
| 87 | } |
| 88 | } |
| 89 | return snippet; |
| 90 | }; |
| 91 | |
| 92 | if (this.type === 'shared') { |
| 93 | const sharedMemorySnippet = ` |
| 94 | var<workgroup> xBestIndices : array<i32, ${workgroupSizeX}>; |
| 95 | var<workgroup> xBestValues : array<f32, ${workgroupSizeX}>; |
| 96 | `; |
| 97 | const userCode = ` |
| 98 | fn DIV_CEIL(a : u32, b : u32) -> u32 { |
| 99 | return ((a - 1u) / b + 1u); |
| 100 | } |
| 101 | |
| 102 | ${sharedMemorySnippet} |
| 103 | |
| 104 | ${main('index')} { |
| 105 | let outputIndex = index / ${workgroupSizeX}; |
| 106 | let reduceLength = ${getInputShapeLastDim()}; |
| 107 | |
| 108 | var bestIndex = i32(localId.x); |
| 109 | var bestValue = uniforms.infinityValue; |
| 110 | let outputCoords = getCoordsFromIndex(outputIndex); |
| 111 | for (var k = i32(localId.x); k < reduceLength && outputIndex < uniforms.size; |
| 112 | k = k + ${workgroupSizeX}) { |
| 113 | let candidate = getX(${splitOutputCoords()} k); |
| 114 | if (!isnan(candidate) && candidate ${this.op} bestValue) { |
| 115 | bestValue = candidate; |
| 116 | bestIndex = k; |
| 117 | } |
| 118 | } |
| 119 | xBestValues[localId.x] = bestValue; |
| 120 | xBestIndices[localId.x] = bestIndex; |
| 121 | workgroupBarrier(); |
| 122 | |
| 123 | var reduceSize = min(u32(reduceLength), ${workgroupSizeX}u); |
| 124 | for (var currentSize = reduceSize / 2u; reduceSize > 1u; |
| 125 | currentSize = reduceSize / 2u) { |