(op: BinaryOpType, aShape: number[], bShape: number[])
| 39 | private type: string; |
| 40 | |
| 41 | constructor(op: BinaryOpType, aShape: number[], bShape: number[]) { |
| 42 | this.outputShape = backend_util.assertAndGetBroadcastShape(aShape, bShape); |
| 43 | this.dispatchLayout = flatDispatchLayout(this.outputShape); |
| 44 | this.op = op; |
| 45 | |
| 46 | this.useSharedMemoryWithA = |
| 47 | aShape.length <= 1 && bShape.length > 1 && aShape[0] < 128; |
| 48 | this.useSharedMemoryWithB = |
| 49 | bShape.length <= 1 && aShape.length > 1 && bShape[0] < 128; |
| 50 | |
| 51 | if (this.useSharedMemoryWithA || this.useSharedMemoryWithB) { |
| 52 | this.outputComponent = 1; |
| 53 | this.variableComponents = [1, 1]; |
| 54 | // lastDimensionSize is used as sharedBuf array size, so can not be |
| 55 | // used as uniform. |
| 56 | this.lastDimensionSize = |
| 57 | this.useSharedMemoryWithB ? bShape[0] : aShape[0]; |
| 58 | this.shaderKey = `binary_${op}_${this.lastDimensionSize}`; |
| 59 | this.type = 'shared'; |
| 60 | // This is an experimental value when using shared memory. |
| 61 | // Note that the maximum of workgroup X dimension is 256. |
| 62 | this.workgroupSize = [256, 1, 1]; |
| 63 | } else { |
| 64 | const aDivisibleBy4 = |
| 65 | aShape.length > 0 && aShape[aShape.length - 1] % 4 === 0; |
| 66 | const bDivisibleBy4 = |
| 67 | bShape.length > 0 && bShape[bShape.length - 1] % 4 === 0; |
| 68 | if (aDivisibleBy4 && bDivisibleBy4) { |
| 69 | this.outputComponent = 4; |
| 70 | this.variableComponents = [4, 4]; |
| 71 | } else if ( |
| 72 | (aDivisibleBy4 && |
| 73 | (util.isScalarShape(bShape) || bShape[bShape.length - 1] === 1)) || |
| 74 | (bDivisibleBy4 && |
| 75 | (util.isScalarShape(aShape) || aShape[aShape.length - 1] === 1))) { |
| 76 | this.outputComponent = 4; |
| 77 | this.variableComponents = aDivisibleBy4 ? [4, 1] : [1, 4]; |
| 78 | } else { |
| 79 | this.outputComponent = 1; |
| 80 | this.variableComponents = [1, 1]; |
| 81 | } |
| 82 | this.type = 'nonshared'; |
| 83 | this.shaderKey = `binary_${op}_${this.variableComponents}`; |
| 84 | // TODO(jiajia.qin@intel.com): Heuristically select a good work group |
| 85 | // size. |
| 86 | this.workgroupSize = [128, 1, 1]; |
| 87 | } |
| 88 | this.dispatch = computeDispatch( |
| 89 | this.dispatchLayout, this.outputShape, this.workgroupSize, |
| 90 | [this.outputComponent, 1, 1]); |
| 91 | } |
| 92 | |
| 93 | getUserCode(): string { |
| 94 | let userCode; |
nothing calls this directly
no test coverage detected