| 24 | userCode: string; |
| 25 | |
| 26 | constructor(reduceInfo: backend_util.ReduceInfo, divisor?: number) { |
| 27 | const {windowSize, batchSize, inSize, outSize} = reduceInfo; |
| 28 | this.outputShape = [batchSize, outSize]; |
| 29 | |
| 30 | const windowSizeNearestVec4 = Math.floor(windowSize / 4) * 4; |
| 31 | const windowSizeVec4Remainder = windowSize % 4; |
| 32 | |
| 33 | let updateSnippet = `sumValue += dot(values, ones);`; |
| 34 | if (divisor != null) { |
| 35 | const denominator = 1 / divisor; |
| 36 | updateSnippet = `sumValue += dot(values * ${ |
| 37 | util.isInt(denominator) ? denominator.toPrecision(2) : |
| 38 | denominator}, ones);`; |
| 39 | } |
| 40 | |
| 41 | let checkOutOfBounds = ''; |
| 42 | if (inSize % windowSize > 0) { |
| 43 | checkOutOfBounds = ` |
| 44 | if (inIdx < 0 || inIdx >= ${inSize}) { |
| 45 | return 0.0; |
| 46 | } |
| 47 | `; |
| 48 | } |
| 49 | |
| 50 | this.userCode = ` |
| 51 | const vec4 ones = vec4(1.0, 1.0, 1.0, 1.0); |
| 52 | |
| 53 | float getValue(int batch, int inIdx) { |
| 54 | ${checkOutOfBounds} |
| 55 | return getX(batch, inIdx); |
| 56 | } |
| 57 | |
| 58 | void main() { |
| 59 | ivec2 coords = getOutputCoords(); |
| 60 | int batch = coords[0]; |
| 61 | int outIdx = coords[1]; |
| 62 | int inOffset = outIdx * ${windowSize}; |
| 63 | |
| 64 | float sumValue = 0.0; |
| 65 | |
| 66 | for (int i = 0; i < ${windowSizeNearestVec4}; i += 4) { |
| 67 | int inIdx = inOffset + i; |
| 68 | vec4 values = vec4( |
| 69 | getValue(batch, inIdx), |
| 70 | getValue(batch, inIdx + 1), |
| 71 | getValue(batch, inIdx + 2), |
| 72 | getValue(batch, inIdx + 3) |
| 73 | ); |
| 74 | |
| 75 | ${updateSnippet} |
| 76 | } |
| 77 | |
| 78 | int inIdx = inOffset + ${windowSizeNearestVec4}; |
| 79 | if (${windowSizeVec4Remainder === 1}) { |
| 80 | vec4 values = vec4(getValue(batch, inIdx), 0.0, 0.0, 0.0); |
| 81 | |
| 82 | ${updateSnippet} |
| 83 | } else if (${windowSizeVec4Remainder === 2}) { |