(
reduceInfo: backend_util.ReduceInfo,
reduceType: 'all'|'any'|'max'|'mean'|'min'|'prod'|'sum',
maxComputeWorkgroupSizeX: number)
| 32 | size = true; |
| 33 | |
| 34 | constructor( |
| 35 | reduceInfo: backend_util.ReduceInfo, |
| 36 | reduceType: 'all'|'any'|'max'|'mean'|'min'|'prod'|'sum', |
| 37 | maxComputeWorkgroupSizeX: number) { |
| 38 | this.inputShape = [reduceInfo.batchSize, reduceInfo.inSize]; |
| 39 | const [outputShape, ] = |
| 40 | backend_util.computeOutAndReduceShapes(this.inputShape, [1]); |
| 41 | this.outputShape = outputShape.length === 0 ? [1] : outputShape; |
| 42 | // If reduceSize |reduceInfo.inSize| is very large, the I/O accessing will |
| 43 | // become the bottleneck. Increasing workgroupSize can reduce the times of |
| 44 | // accessing global memory. The threshold value is just to make sure the |
| 45 | // reduceSize is large enough for a bigger workgroupSize. |
| 46 | if (reduceInfo.inSize >= 32768 && maxComputeWorkgroupSizeX >= 512) { |
| 47 | this.workgroupSize = [512, 1, 1]; |
| 48 | } else if (reduceInfo.inSize >= 4096) { |
| 49 | this.workgroupSize = [256, 1, 1]; |
| 50 | } else { |
| 51 | this.workgroupSize = [64, 1, 1]; |
| 52 | } |
| 53 | this.dispatchLayout = flatDispatchLayout(this.outputShape); |
| 54 | // A work group only outputs a data, so we transfer [1, 1, 1] to compute |
| 55 | // dispatch size. |
| 56 | this.dispatch = |
| 57 | computeDispatch(this.dispatchLayout, this.outputShape, [1, 1, 1]); |
| 58 | |
| 59 | this.reduceType = reduceType; |
| 60 | this.shaderKey = `reduce_${reduceType}`; |
| 61 | } |
| 62 | |
| 63 | getUserCode(): string { |
| 64 | let reduceOp = ``; |
nothing calls this directly
no test coverage detected