| { |
| "domain": "ai.onnx", |
| "name": "AveragePool", |
| "sinceVersion": 19, |
| "inputs": { "x": { "onnx": "X", "dtype": "T" } }, |
| "outputs": { "y": { "onnx": "Y", "dtype": "T", "rank": "ranks.x", "shape": "poolOutputShape" } }, |
| "attributes": { |
| "auto_pad": { "default": "NOTSET" }, |
| "count_include_pad": { "default": 0 }, |
| "ceil_mode": { "default": 0 }, |
| "kernel_shape": {}, |
| "strides": {}, |
| "pads": {}, |
| "dilations": {} |
| }, |
| "attributeConstraints": { |
| "auto_pad": { "values": ["NOTSET", "SAME_UPPER", "SAME_LOWER", "VALID"] }, |
| "ceil_mode": { "values": [0, 1] }, |
| "count_include_pad": { "values": [0, 1] }, |
| "kernel_shape": { "required": true } |
| }, |
| "typeConstraints": { "T": ["float32", "float16"] }, |
| "tunables": { |
| "WORKGROUP_SIZE": { "default": 256 }, |
| "SEPARABLE_MIN_WINDOW": { "default": 25 }, |
| "WINDOW_PARALLEL_MIN_TAPS": { "default": 128 }, |
| "WINDOW_PARALLEL_MAX_SERIAL_GROUPS": { "default": 64 } |
| }, |
| "derive": { |
| "spatialRank": "ranks.x - 2", |
| "narrowSubgroupRange": "has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize < device.adapterInfo.subgroupMaxSize and device.adapterInfo.subgroupMaxSize <= 16", |
| "variableSubgroup16To32": "device.features.has(\"subgroups\") and has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize == 16 and device.adapterInfo.subgroupMaxSize == 32", |
| "kernelShapeLengthOk": "not has(attrs, \"kernel_shape\") or (attrs.kernel_shape | length) == spatialRank", |
| "stridesLengthOk": "not has(attrs, \"strides\") or (attrs.strides | length) == spatialRank", |
| "dilationsLengthOk": "not has(attrs, \"dilations\") or (attrs.dilations | length) == spatialRank", |
| "padsLengthOk": "not has(attrs, \"pads\") or (attrs.pads | length) == 2 * spatialRank", |
| "kernelD": "attrs.kernel_shape[0] if kernelShapeLengthOk and has(attrs, \"kernel_shape\") and spatialRank == 3 else 1", |
| "kernelH": "attrs.kernel_shape[spatialRank - 2] if kernelShapeLengthOk and has(attrs, \"kernel_shape\") and spatialRank >= 2 else 1", |
| "kernelW": "attrs.kernel_shape[spatialRank - 1] if kernelShapeLengthOk and has(attrs, \"kernel_shape\") and spatialRank >= 1 else 1", |
| "strideD": "attrs.strides[0] if stridesLengthOk and has(attrs, \"strides\") and spatialRank == 3 else 1", |
| "strideH": "attrs.strides[spatialRank - 2] if stridesLengthOk and has(attrs, \"strides\") and spatialRank >= 2 else 1", |
| "strideW": "attrs.strides[spatialRank - 1] if stridesLengthOk and has(attrs, \"strides\") and spatialRank >= 1 else 1", |
| "dilationD": "attrs.dilations[0] if dilationsLengthOk and has(attrs, \"dilations\") and spatialRank == 3 else 1", |
| "dilationH": "attrs.dilations[spatialRank - 2] if dilationsLengthOk and has(attrs, \"dilations\") and spatialRank >= 2 else 1", |
| "dilationW": "attrs.dilations[spatialRank - 1] if dilationsLengthOk and has(attrs, \"dilations\") and spatialRank >= 1 else 1", |
| "padFront": "attrs.pads[0] if padsLengthOk and has(attrs, \"pads\") and spatialRank == 3 else 0", |
| "padTop": "attrs.pads[spatialRank - 2] if padsLengthOk and has(attrs, \"pads\") and spatialRank >= 2 else 0", |
| "padLeft": "attrs.pads[spatialRank - 1] if padsLengthOk and has(attrs, \"pads\") and spatialRank >= 1 else 0", |
| "padBack": "attrs.pads[spatialRank] if padsLengthOk and has(attrs, \"pads\") and spatialRank == 3 else 0", |
| "padBottom": "attrs.pads[2 * spatialRank - 2] if padsLengthOk and has(attrs, \"pads\") and spatialRank >= 2 else 0", |
| "padRight": "attrs.pads[2 * spatialRank - 1] if padsLengthOk and has(attrs, \"pads\") and spatialRank >= 1 else 0", |
| "explicitOutW": "min(ceil((dim(shapes.x, ranks.x - 1) + padLeft + padRight - ((kernelW - 1) * dilationW + 1)) / strideW) + 1, ceil((dim(shapes.x, ranks.x - 1) + padLeft) / strideW)) if attrs.ceil_mode != 0 else floor((dim(shapes.x, ranks.x - 1) + padLeft + padRight - ((kernelW - 1) * dilationW + 1)) / strideW) + 1", |
| "explicitOutH": "min(ceil((dim(shapes.x, ranks.x - 2) + padTop + padBottom - ((kernelH - 1) * dilationH + 1)) / strideH) + 1, ceil((dim(shapes.x, ranks.x - 2) + padTop) / strideH)) if attrs.ceil_mode != 0 else floor((dim(shapes.x, ranks.x - 2) + padTop + padBottom - ((kernelH - 1) * dilationH + 1)) / strideH) + 1", |
| "explicitOutD": "(min(ceil((dim(shapes.x, 2) + padFront + padBack - ((kernelD - 1) * dilationD + 1)) / strideD) + 1, ceil((dim(shapes.x, 2) + padFront) / strideD)) if attrs.ceil_mode != 0 else floor((dim(shapes.x, 2) + padFront + padBack - ((kernelD - 1) * dilationD + 1)) / strideD) + 1) if ranks.x == 5 else 0", |
| "inferredOutW": "ceil(dim(shapes.x, ranks.x - 1) / strideW) if attrs.auto_pad == \"SAME_UPPER\" or attrs.auto_pad == \"SAME_LOWER\" else explicitOutW", |
| "inferredOutH": "ceil(dim(shapes.x, ranks.x - 2) / strideH) if attrs.auto_pad == \"SAME_UPPER\" or attrs.auto_pad == \"SAME_LOWER\" else explicitOutH", |
| "inferredOutD": "ceil(dim(shapes.x, 2) / strideD) if attrs.auto_pad == \"SAME_UPPER\" or attrs.auto_pad == \"SAME_LOWER\" else explicitOutD", |
| "poolOutputShape": "[dim(shapes.x, 0), dim(shapes.x, 1), inferredOutW] if ranks.x == 3 else ([dim(shapes.x, 0), dim(shapes.x, 1), inferredOutH, inferredOutW] if ranks.x == 4 else [dim(shapes.x, 0), dim(shapes.x, 1), inferredOutD, inferredOutH, inferredOutW])", |
| "autoPadTotalD": "max(0, (dim(shapes.y, 2) - 1) * strideD + ((kernelD - 1) * dilationD + 1) - dim(shapes.x, 2)) if ranks.x == 5 else 0", |
| "autoPadTotalH": "max(0, (dim(shapes.y, ranks.x - 2) - 1) * strideH + ((kernelH - 1) * dilationH + 1) - dim(shapes.x, ranks.x - 2)) if ranks.x >= 4 else 0", |
| "autoPadTotalW": "max(0, (dim(shapes.y, ranks.x - 1) - 1) * strideW + ((kernelW - 1) * dilationW + 1) - dim(shapes.x, ranks.x - 1))", |
| "effectivePadD": "(floor(autoPadTotalD / 2) if attrs.auto_pad == \"SAME_UPPER\" else autoPadTotalD - floor(autoPadTotalD / 2)) if (attrs.auto_pad == \"SAME_UPPER\" or attrs.auto_pad == \"SAME_LOWER\") else padFront", |
| "effectivePadH": "(floor(autoPadTotalH / 2) if attrs.auto_pad == \"SAME_UPPER\" else autoPadTotalH - floor(autoPadTotalH / 2)) if (attrs.auto_pad == \"SAME_UPPER\" or attrs.auto_pad == \"SAME_LOWER\") else padTop", |
| "effectivePadW": "(floor(autoPadTotalW / 2) if attrs.auto_pad == \"SAME_UPPER\" else autoPadTotalW - floor(autoPadTotalW / 2)) if (attrs.auto_pad == \"SAME_UPPER\" or attrs.auto_pad == \"SAME_LOWER\") else padLeft", |
| "windowTaps": "kernelD * kernelH * kernelW", |
| "windowParallelEligible": "numel(shapes.y) > 0 and windowTaps >= tunables.WINDOW_PARALLEL_MIN_TAPS and ceilDiv(numel(shapes.y), tunables.WORKGROUP_SIZE) < tunables.WINDOW_PARALLEL_MAX_SERIAL_GROUPS and numel(shapes.y) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and numel(shapes.y) <= windowTaps", |
| "separableTmpElements": "dim(shapes.x, 0) * dim(shapes.x, 1) * dim(shapes.x, 2) * dim(shapes.y, 3)", |
| "separableFits": "separableTmpElements * 4 <= min(device.limits.maxStorageBufferBindingSize, device.limits.maxBufferSize) and ceilDiv(separableTmpElements, tunables.WORKGROUP_SIZE) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) * min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", |
| "poolBaseOk": "f16Ok(dtypes.T) and ranks.x >= 3 and ranks.x <= 5 and ranks.y == ranks.x and has(attrs, \"kernel_shape\") and kernelShapeLengthOk and stridesLengthOk and dilationsLengthOk and padsLengthOk and kernelD > 0 and kernelH > 0 and kernelW > 0 and strideD > 0 and strideH > 0 and strideW > 0 and dilationD > 0 and dilationH > 0 and dilationW > 0 and padFront >= 0 and padTop >= 0 and padLeft >= 0 and padBack >= 0 and padBottom >= 0 and padRight >= 0 and (attrs.auto_pad == \"NOTSET\" or not has(attrs, \"pads\")) and dim(shapes.y, 0) == dim(shapes.x, 0) and dim(shapes.y, 1) == dim(shapes.x, 1)", |
| "effectivePadBack": "autoPadTotalD - effectivePadD if (attrs.auto_pad == \"SAME_UPPER\" or attrs.auto_pad == \"SAME_LOWER\") else padBack", |
| "effectivePadBottom": "autoPadTotalH - effectivePadH if (attrs.auto_pad == \"SAME_UPPER\" or attrs.auto_pad == \"SAME_LOWER\") else padBottom", |
| "effectivePadRight": "autoPadTotalW - effectivePadW if (attrs.auto_pad == \"SAME_UPPER\" or attrs.auto_pad == \"SAME_LOWER\") else padRight", |
| "narrowF32W3UnrolledPreferred": "narrowSubgroupRange and dtypes.T == \"f32\"", |
| "globalWindowParallelCovered": "poolBaseOk and numel(shapes.y) == dim(shapes.x, 0) * dim(shapes.x, 1) and numel(shapes.y) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and attrs.auto_pad == \"NOTSET\" and strideD == 1 and strideH == 1 and strideW == 1 and dilationD == 1 and dilationH == 1 and dilationW == 1 and padFront == 0 and padTop == 0 and padLeft == 0 and padBack == 0 and padBottom == 0 and padRight == 0 and ((ranks.x == 3 and kernelW == dim(shapes.x, 2)) or (ranks.x == 4 and kernelH == dim(shapes.x, 2) and kernelW == dim(shapes.x, 3)) or (ranks.x == 5 and kernelD == dim(shapes.x, 2) and kernelH == dim(shapes.x, 3) and kernelW == dim(shapes.x, 4))) and inner(shapes.x, 1) >= 256", |
| "wholeWindowDenominator": "attrs.ceil_mode == 0 and (attrs.count_include_pad == 1 or (effectivePadH == 0 and effectivePadW == 0 and effectivePadBottom == 0 and effectivePadRight == 0))", |
| "scalar": "dtypes.T", |
| "inH": "dim(shapes.x, 2)", |
| "inPlane": "dim(shapes.x, 2) * dim(shapes.x, 3)", |
| "inW": "dim(shapes.x, 3)", |
| "outH": "dim(shapes.y, 2)", |
| "outW": "dim(shapes.y, 3)" |
| }, |
| "bindings": { |
| "y_main": { "name": "y", "elementType": "$outputElement" }, |
| "params_window_ncl1d": { |
| "name": "params", |
| "struct": [ |
| { "name": "channels", "type": "u32", "value": "dim(shapes.x, 1)" }, |
| { "name": "inW", "type": "u32", "value": "dim(shapes.x, 2)" }, |
| { "name": "outW", "type": "u32", "value": "dim(shapes.y, 2)" }, |
| { "name": "strideW", "type": "u32", "value": "strideW" }, |
| { "name": "dilationW", "type": "u32", "value": "dilationW" }, |
| { "name": "padW", "type": "i32", "value": "effectivePadW" } |
| ] |
| }, |
| "params_window_nchw2d": { |
| "name": "params", |
| "struct": [ |
| { "name": "channels", "type": "u32", "value": "dim(shapes.x, 1)" }, |
| { "name": "inH", "type": "u32", "value": "dim(shapes.x, 2)" }, |
| { "name": "inW", "type": "u32", "value": "dim(shapes.x, 3)" }, |
| { "name": "outH", "type": "u32", "value": "dim(shapes.y, 2)" }, |
| { "name": "outW", "type": "u32", "value": "dim(shapes.y, 3)" }, |
| { "name": "strideH", "type": "u32", "value": "strideH" }, |
| { "name": "strideW", "type": "u32", "value": "strideW" }, |
| { "name": "dilationH", "type": "u32", "value": "dilationH" }, |
| { "name": "dilationW", "type": "u32", "value": "dilationW" }, |
| { "name": "padH", "type": "i32", "value": "effectivePadH" }, |
| { "name": "padW", "type": "i32", "value": "effectivePadW" } |
| ] |
| }, |
| "params_window_ncdhw3d": { |
| "name": "params", |
| "struct": [ |
| { "name": "channels", "type": "u32", "value": "dim(shapes.x, 1)" }, |
| { "name": "inD", "type": "u32", "value": "dim(shapes.x, 2)" }, |
| { "name": "inH", "type": "u32", "value": "dim(shapes.x, 3)" }, |
| { "name": "inW", "type": "u32", "value": "dim(shapes.x, 4)" }, |
| { "name": "outD", "type": "u32", "value": "dim(shapes.y, 2)" }, |
| { "name": "outH", "type": "u32", "value": "dim(shapes.y, 3)" }, |
| { "name": "outW", "type": "u32", "value": "dim(shapes.y, 4)" }, |
| { "name": "strideD", "type": "u32", "value": "strideD" }, |
| { "name": "strideH", "type": "u32", "value": "strideH" }, |
| { "name": "strideW", "type": "u32", "value": "strideW" }, |
| { "name": "dilationD", "type": "u32", "value": "dilationD" }, |
| { "name": "dilationH", "type": "u32", "value": "dilationH" }, |
| { "name": "dilationW", "type": "u32", "value": "dilationW" }, |
| { "name": "padD", "type": "i32", "value": "effectivePadD" }, |
| { "name": "padH", "type": "i32", "value": "effectivePadH" }, |
| { "name": "padW", "type": "i32", "value": "effectivePadW" } |
| ] |
| } |
| }, |
| "variants": [ |
| { |
| "id": "global_window_parallel", |
| "priority": 50, |
| "when": ["globalWindowParallelCovered"], |
| "derive": { |
| "xElem": "(\"vec4<f16>\" if dtypes.T == \"f16\" else \"vec4<f32>\") if inner(shapes.x, 1) % 4 == 0 else dtypes.T" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "AveragePool.GlobalWindowParallel", |
| "shader": "pool-global-reduction.wgsl.jinja", |
| "derive": { |
| "op": "\"average\"", |
| "vec4": "inner(shapes.x, 1) % 4 == 0", |
| "workgroupSizeSpec": "tunables.WORKGROUP_SIZE", |
| "useSubgroups": "device.features.has(\"subgroups\")" |
| }, |
| "bindings": [ |
| { "arg": "x", "elementType": "$xElem" }, |
| { "arg": "y", "elementType": "$scalar" }, |
| { "name": "params", "struct": [{ "name": "spatial", "type": "u32", "value": "inner(shapes.x, 1)" }] } |
| ], |
| "dispatch": { "x": "numel(shapes.y)" } |
| } |
| ] |
| }, |
| { |
| "id": "window_parallel_ncl1d", |
| "priority": 28, |
| "when": ["poolBaseOk", "ranks.x == 3", "dim(shapes.y, 2) == inferredOutW", "attrs.count_include_pad == 0 or attrs.ceil_mode == 0", "windowParallelEligible"], |
| "supersededBy": ["global_window_parallel"], |
| "passes": [ |
| { |
| "id": "main", |
| "name": "AveragePool.WindowParallelNcl1d", |
| "shader": "pool-window-reduction.wgsl.jinja", |
| "derive": { |
| "op": "\"average\"", |
| "spatialRankSpec": 1, |
| "kernelDSpec": "kernelD", |
| "kernelHSpec": "kernelH", |
| "kernelWSpec": "kernelW", |
| "workgroupSizeSpec": "tunables.WORKGROUP_SIZE", |
| "useSubgroups": "device.features.has(\"subgroups\")", |
| "countIncludePad": "attrs.count_include_pad != 0" |
| }, |
| "bindings": ["x", "y", "params_window_ncl1d"], |
| "dispatch": { "x": "numel(shapes.y)" } |
| } |
| ] |
| }, |
| { |
| "id": "window_parallel_nchw2d", |
| "priority": 28, |
| "when": ["poolBaseOk", "ranks.x == 4", "dim(shapes.y, 2) == inferredOutH", "dim(shapes.y, 3) == inferredOutW", "attrs.count_include_pad == 0 or attrs.ceil_mode == 0", "windowParallelEligible"], |
| "supersededBy": ["global_window_parallel"], |
| "passes": [ |
| { |
| "id": "main", |
| "name": "AveragePool.WindowParallelNchw2d", |
| "shader": "pool-window-reduction.wgsl.jinja", |
| "derive": { |
| "op": "\"average\"", |
| "spatialRankSpec": 2, |
| "kernelDSpec": "kernelD", |
| "kernelHSpec": "kernelH", |
| "kernelWSpec": "kernelW", |
| "workgroupSizeSpec": "tunables.WORKGROUP_SIZE", |
| "useSubgroups": "device.features.has(\"subgroups\")", |
| "countIncludePad": "attrs.count_include_pad != 0" |
| }, |
| "bindings": ["x", "y", "params_window_nchw2d"], |
| "dispatch": { "x": "numel(shapes.y)" } |
| } |
| ] |
| }, |
| { |
| "id": "window_parallel_ncdhw3d", |
| "priority": 28, |
| "when": ["poolBaseOk", "ranks.x == 5", "dim(shapes.y, 2) == inferredOutD", "dim(shapes.y, 3) == inferredOutH", "dim(shapes.y, 4) == inferredOutW", "attrs.count_include_pad == 0 or attrs.ceil_mode == 0", "windowParallelEligible"], |
| "supersededBy": ["global_window_parallel"], |
| "passes": [ |
| { |
| "id": "main", |
| "name": "AveragePool.WindowParallelNcdhw3d", |
| "shader": "pool-window-reduction.wgsl.jinja", |
| "derive": { |
| "op": "\"average\"", |
| "spatialRankSpec": 3, |
| "kernelDSpec": "kernelD", |
| "kernelHSpec": "kernelH", |
| "kernelWSpec": "kernelW", |
| "workgroupSizeSpec": "tunables.WORKGROUP_SIZE", |
| "useSubgroups": "device.features.has(\"subgroups\")", |
| "countIncludePad": "attrs.count_include_pad != 0" |
| }, |
| "bindings": ["x", "y", "params_window_ncdhw3d"], |
| "dispatch": { "x": "numel(shapes.y)" } |
| } |
| ] |
| }, |
| { |
| "id": "nchw2d_k7s1_reuse_tail", |
| "priority": 34, |
| "when": ["poolBaseOk", "ranks.x == 4", "kernelH == 7", "kernelW == 7", "strideH == 1", "strideW == 1", "dilationH == 1", "dilationW == 1", "effectivePadH == 3", "effectivePadBottom == 3", "effectivePadW == 3", "effectivePadRight == 3", "attrs.count_include_pad == 0", "attrs.ceil_mode == 0", "dim(shapes.x, 3) % 4 != 0", "dim(shapes.y, 2) == dim(shapes.x, 2)", "dim(shapes.y, 3) == dim(shapes.x, 3)", "dim(shapes.y, 2) > 0", "dim(shapes.y, 3) > 0"], |
| "derive": { "outputElement": "dtypes.T" }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "AveragePool.Nchw2dK7s1ReuseTail", |
| "shader": "average-pool2d-nchw-horizontal-reuse.wgsl.jinja", |
| "derive": { |
| "kernelHSpec": 7, |
| "kernelWSpec": 7, |
| "strideHSpec": 1, |
| "strideWSpec": 1, |
| "dilationHSpec": 1, |
| "dilationWSpec": 1, |
| "padTopSpec": 3, |
| "padLeftSpec": 3, |
| "outW4": "ceilDiv(dim(shapes.y, 3), 4)", |
| "count4": "dim(shapes.y, 0) * dim(shapes.y, 1) * dim(shapes.y, 2) * ceilDiv(dim(shapes.y, 3), 4)", |
| "scalarOutput": true |
| }, |
| "bindings": ["x", "y_main"], |
| "dispatch": { |
| "x": "min(ceilDiv((dim(shapes.y, 0) * dim(shapes.y, 1) * dim(shapes.y, 2) * ceilDiv(dim(shapes.y, 3), 4)), (tunables.WORKGROUP_SIZE)), 65535)", |
| "y": "ceilDiv(ceilDiv((dim(shapes.y, 0) * dim(shapes.y, 1) * dim(shapes.y, 2) * ceilDiv(dim(shapes.y, 3), 4)), (tunables.WORKGROUP_SIZE)), 65535)", |
| "z": 1 |
| } |
| } |
| ] |
| }, |
| { |
| "id": "nchw2d_w3s1_reuse", |
| "priority": 30, |
| "when": ["poolBaseOk", "ranks.x == 4", "kernelW == 3", "strideW == 1", "dilationW == 1", "effectivePadW == 1", "effectivePadRight == 1", "kernelH >= 1", "kernelH <= 3", "strideH == 1", "dilationH >= 1", "attrs.count_include_pad == 0", "attrs.ceil_mode == 0", "dim(shapes.x, 3) % 4 == 0", "dim(shapes.y, 2) == dim(shapes.x, 2) + effectivePadH + effectivePadBottom - ((kernelH - 1) * dilationH + 1) + 1", "dim(shapes.y, 3) == dim(shapes.x, 3)", "dim(shapes.y, 2) > 0", "dim(shapes.y, 3) > 0"], |
| "demoteWhen": ["narrowF32W3UnrolledPreferred", "variableSubgroup16To32 and dtypes.T == \"f16\""], |
| "passes": [ |
| { |
| "id": "main", |
| "name": "AveragePool.Nchw2dW3s1Reuse", |
| "shader": "average-pool2d-nchw-w3s1-reuse.wgsl.jinja", |
| "derive": { |
| "kernelHSpec": "kernelH", |
| "dilationHSpec": "dilationH", |
| "padTopSpec": "effectivePadH", |
| "outWQuarter": "floor(dim(shapes.y, 3) / 4)" |
| }, |
| "bindings": [ |
| "x", |
| "y", |
| { "name": "params", "struct": [{ "name": "tileCount", "type": "u32", "value": "numel(shapes.y) / 4" }] } |
| ], |
| "dispatch": { |
| "x": "min(ceilDiv((numel(shapes.y) / 4), (tunables.WORKGROUP_SIZE)), 65535)", |
| "y": "ceilDiv(ceilDiv((numel(shapes.y) / 4), (tunables.WORKGROUP_SIZE)), 65535)", |
| "z": 1 |
| } |
| } |
| ] |
| }, |
| { |
| "id": "nchw2d_w3s1_reuse_tail", |
| "priority": 29, |
| "when": ["poolBaseOk", "ranks.x == 4", "kernelW == 3", "strideW == 1", "dilationW == 1", "effectivePadW == 1", "effectivePadRight == 1", "kernelH >= 1", "kernelH <= 3", "strideH == 1", "dilationH >= 1", "attrs.count_include_pad == 0", "attrs.ceil_mode == 0", "dim(shapes.x, 3) % 4 != 0", "dim(shapes.y, 2) == dim(shapes.x, 2) + effectivePadH + effectivePadBottom - ((kernelH - 1) * dilationH + 1) + 1", "dim(shapes.y, 3) == dim(shapes.x, 3)", "dim(shapes.y, 2) > 0", "dim(shapes.y, 3) > 0"], |
| "derive": { "outputElement": "dtypes.T" }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "AveragePool.Nchw2dW3s1ReuseTail", |
| "shader": "average-pool2d-nchw-horizontal-reuse.wgsl.jinja", |
| "derive": { |
| "kernelHSpec": "kernelH", |
| "kernelWSpec": 3, |
| "strideHSpec": 1, |
| "strideWSpec": 1, |
| "dilationHSpec": "dilationH", |
| "dilationWSpec": 1, |
| "padTopSpec": "effectivePadH", |
| "padLeftSpec": 1, |
| "outW4": "ceilDiv(dim(shapes.y, 3), 4)", |
| "count4": "dim(shapes.y, 0) * dim(shapes.y, 1) * dim(shapes.y, 2) * ceilDiv(dim(shapes.y, 3), 4)", |
| "scalarOutput": true |
| }, |
| "bindings": ["x", "y_main"], |
| "dispatch": { |
| "x": "min(ceilDiv((dim(shapes.y, 0) * dim(shapes.y, 1) * dim(shapes.y, 2) * ceilDiv(dim(shapes.y, 3), 4)), (tunables.WORKGROUP_SIZE)), 65535)", |
| "y": "ceilDiv(ceilDiv((dim(shapes.y, 0) * dim(shapes.y, 1) * dim(shapes.y, 2) * ceilDiv(dim(shapes.y, 3), 4)), (tunables.WORKGROUP_SIZE)), 65535)", |
| "z": 1 |
| } |
| } |
| ] |
| }, |
| { |
| "id": "nchw2d_separable", |
| "priority": 22, |
| "when": ["poolBaseOk", "ranks.x == 4", "ranks.y == 4", "dim(shapes.y, 2) == inferredOutH", "dim(shapes.y, 3) == inferredOutW", "dim(shapes.y, 2) > 0", "dim(shapes.y, 3) > 0", "kernelH > 1", "kernelW > 1", "kernelH * kernelW >= tunables.SEPARABLE_MIN_WINDOW", "strideH < kernelH", "wholeWindowDenominator", "separableFits"], |
| "derive": { "scalar": "dtypes.T" }, |
| "intermediates": [{ "id": "tmp", "dtype": "float32", "shape": "[separableTmpElements]" }], |
| "passes": [ |
| { |
| "id": "reduce_w", |
| "name": "AveragePool.SeparableW", |
| "shader": "pool2d-nchw-separable.wgsl.jinja", |
| "derive": { "axisSpec": "\"w\"", "op": "\"avg\"", "windowArea": "kernelH * kernelW" }, |
| "bindings": [ |
| "x", |
| { "name": "tmp", "elementType": "f32" }, |
| { |
| "name": "params", |
| "struct": [ |
| { "name": "inH", "type": "u32", "value": "dim(shapes.x, 2)" }, |
| { "name": "inW", "type": "u32", "value": "dim(shapes.x, 3)" }, |
| { "name": "outW", "type": "u32", "value": "dim(shapes.y, 3)" }, |
| { "name": "kernelW", "type": "u32", "value": "kernelW" }, |
| { "name": "strideW", "type": "u32", "value": "strideW" }, |
| { "name": "dilationW", "type": "u32", "value": "dilationW" }, |
| { "name": "padW", "type": "i32", "value": "effectivePadW" }, |
| { "name": "tmpCount", "type": "u32", "value": "separableTmpElements" } |
| ] |
| } |
| ], |
| "dispatch": { |
| "x": "min(ceilDiv((separableTmpElements), (tunables.WORKGROUP_SIZE)), 65535)", |
| "y": "ceilDiv(ceilDiv((separableTmpElements), (tunables.WORKGROUP_SIZE)), 65535)", |
| "z": 1 |
| } |
| }, |
| { |
| "id": "reduce_h", |
| "name": "AveragePool.SeparableH", |
| "shader": "pool2d-nchw-separable.wgsl.jinja", |
| "derive": { "axisSpec": "\"h\"", "op": "\"avg\"", "windowArea": "kernelH * kernelW" }, |
| "bindings": [ |
| { "name": "tmp", "buffer": "read-only-storage", "elementType": "f32" }, |
| "y", |
| { |
| "name": "params", |
| "struct": [ |
| { "name": "inH", "type": "u32", "value": "dim(shapes.x, 2)" }, |
| { "name": "outH", "type": "u32", "value": "dim(shapes.y, 2)" }, |
| { "name": "outW", "type": "u32", "value": "dim(shapes.y, 3)" }, |
| { "name": "kernelH", "type": "u32", "value": "kernelH" }, |
| { "name": "strideH", "type": "u32", "value": "strideH" }, |
| { "name": "dilationH", "type": "u32", "value": "dilationH" }, |
| { "name": "padH", "type": "i32", "value": "effectivePadH" }, |
| { "name": "count", "type": "u32", "value": "numel(shapes.y)" } |
| ] |
| } |
| ], |
| "dispatch": { |
| "x": "min(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), 65535)", |
| "y": "ceilDiv(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), 65535)", |
| "z": 1 |
| } |
| } |
| ] |
| }, |
| { |
| "id": "nchw2d_unrolled", |
| "priority": 25, |
| "when": ["poolBaseOk", "ranks.x == 4", "dim(shapes.y, 2) == inferredOutH", "dim(shapes.y, 3) == inferredOutW", "kernelH >= 1", "kernelH <= 3", "kernelW >= 1", "kernelW <= 3", "strideH >= 1", "strideW >= 1", "dilationH >= 1", "dilationW >= 1", "dim(shapes.y, 2) > 0", "dim(shapes.y, 3) > 0"], |
| "passes": [ |
| { |
| "id": "main", |
| "name": "AveragePool.Nchw2dUnrolled", |
| "shader": "pool-window-unroll.wgsl.jinja", |
| "derive": { |
| "op": "\"average\"", |
| "kernelHSpec": "kernelH", |
| "kernelWSpec": "kernelW", |
| "strideHSpec": "strideH", |
| "strideWSpec": "strideW", |
| "dilationHSpec": "dilationH", |
| "dilationWSpec": "dilationW", |
| "padTopSpec": "effectivePadH", |
| "padLeftSpec": "effectivePadW", |
| "inH": "dim(shapes.x, 2)", |
| "inW": "dim(shapes.x, 3)", |
| "outH": "dim(shapes.y, 2)", |
| "outW": "dim(shapes.y, 3)", |
| "padBottomSpec": "effectivePadBottom", |
| "padRightSpec": "effectivePadRight", |
| "countIncludePad": "attrs.count_include_pad != 0" |
| }, |
| "bindings": [ |
| "x", |
| "y", |
| { "name": "params", "struct": [{ "name": "count", "type": "u32", "value": "numel(shapes.y)" }] } |
| ], |
| "dispatch": { |
| "x": "min(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), 65535)", |
| "y": "ceilDiv(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), 65535)", |
| "z": 1 |
| } |
| } |
| ] |
| }, |
| { |
| "id": "ncl1d", |
| "priority": 10, |
| "when": ["poolBaseOk", "ranks.x == 3", "dim(shapes.y, 2) == inferredOutW"], |
| "passes": [ |
| { |
| "id": "main", |
| "name": "AveragePool.Ncl1d", |
| "shader": "pool-window-nd.wgsl.jinja", |
| "derive": { "op": "\"average\"", "spatialRankSpec": 1 }, |
| "bindings": [ |
| "x", |
| "y", |
| { |
| "name": "params", |
| "struct": [ |
| { "name": "channels", "type": "u32", "value": "dim(shapes.x, 1)" }, |
| { "name": "inW", "type": "u32", "value": "dim(shapes.x, 2)" }, |
| { "name": "outW", "type": "u32", "value": "dim(shapes.y, 2)" }, |
| { "name": "kernelW", "type": "u32", "value": "kernelW" }, |
| { "name": "strideW", "type": "u32", "value": "strideW" }, |
| { "name": "dilationW", "type": "u32", "value": "dilationW" }, |
| { "name": "padLeft", "type": "i32", "value": "effectivePadW" }, |
| { "name": "padRight", "type": "i32", "value": "effectivePadRight" }, |
| { "name": "countIncludePad", "type": "u32", "value": "attrs.count_include_pad" }, |
| { "name": "count", "type": "u32", "value": "numel(shapes.y)" } |
| ] |
| } |
| ], |
| "dispatch": { |
| "x": "min(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), 65535)", |
| "y": "ceilDiv(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), 65535)", |
| "z": 1 |
| } |
| } |
| ] |
| }, |
| { |
| "id": "nchw2d", |
| "when": ["poolBaseOk", "ranks.x == 4", "((dim(shapes.y, 2) == explicitOutH and dim(shapes.y, 3) == explicitOutW) or ((attrs.auto_pad == \"SAME_UPPER\" or attrs.auto_pad == \"SAME_LOWER\") and dim(shapes.y, 2) == ceil(dim(shapes.x, 2) / strideH) and dim(shapes.y, 3) == ceil(dim(shapes.x, 3) / strideW)))"], |
| "passes": [ |
| { |
| "id": "main", |
| "name": "AveragePool.Nchw2d", |
| "shader": "pool-window-nd.wgsl.jinja", |
| "derive": { "op": "\"average\"", "spatialRankSpec": 2 }, |
| "bindings": [ |
| "x", |
| "y", |
| { |
| "name": "params", |
| "struct": [ |
| { "name": "channels", "type": "u32", "value": "dim(shapes.x, 1)" }, |
| { "name": "inH", "type": "u32", "value": "dim(shapes.x, 2)" }, |
| { "name": "inW", "type": "u32", "value": "dim(shapes.x, 3)" }, |
| { "name": "outH", "type": "u32", "value": "dim(shapes.y, 2)" }, |
| { "name": "outW", "type": "u32", "value": "dim(shapes.y, 3)" }, |
| { "name": "kernelH", "type": "u32", "value": "kernelH" }, |
| { "name": "kernelW", "type": "u32", "value": "kernelW" }, |
| { "name": "strideH", "type": "u32", "value": "strideH" }, |
| { "name": "strideW", "type": "u32", "value": "strideW" }, |
| { "name": "dilationH", "type": "u32", "value": "dilationH" }, |
| { "name": "dilationW", "type": "u32", "value": "dilationW" }, |
| { "name": "padTop", "type": "i32", "value": "effectivePadH" }, |
| { "name": "padLeft", "type": "i32", "value": "effectivePadW" }, |
| { "name": "padBottom", "type": "i32", "value": "effectivePadBottom" }, |
| { "name": "padRight", "type": "i32", "value": "effectivePadRight" }, |
| { "name": "countIncludePad", "type": "u32", "value": "attrs.count_include_pad" }, |
| { "name": "count", "type": "u32", "value": "numel(shapes.y)" } |
| ] |
| } |
| ], |
| "dispatch": { |
| "x": "min(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), 65535)", |
| "y": "ceilDiv(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), 65535)", |
| "z": 1 |
| } |
| } |
| ] |
| }, |
| { |
| "id": "ncdhw3d", |
| "priority": 10, |
| "when": ["poolBaseOk", "ranks.x == 5", "dim(shapes.y, 2) == inferredOutD", "dim(shapes.y, 3) == inferredOutH", "dim(shapes.y, 4) == inferredOutW"], |
| "supersededBy": ["global_window_parallel"], |
| "passes": [ |
| { |
| "id": "main", |
| "name": "AveragePool.Ncdhw3d", |
| "shader": "pool-window-nd.wgsl.jinja", |
| "derive": { "op": "\"average\"", "spatialRankSpec": 3 }, |
| "bindings": [ |
| "x", |
| "y", |
| { |
| "name": "params", |
| "struct": [ |
| { "name": "channels", "type": "u32", "value": "dim(shapes.x, 1)" }, |
| { "name": "inD", "type": "u32", "value": "dim(shapes.x, 2)" }, |
| { "name": "inH", "type": "u32", "value": "dim(shapes.x, 3)" }, |
| { "name": "inW", "type": "u32", "value": "dim(shapes.x, 4)" }, |
| { "name": "outD", "type": "u32", "value": "dim(shapes.y, 2)" }, |
| { "name": "outH", "type": "u32", "value": "dim(shapes.y, 3)" }, |
| { "name": "outW", "type": "u32", "value": "dim(shapes.y, 4)" }, |
| { "name": "kernelD", "type": "u32", "value": "kernelD" }, |
| { "name": "kernelH", "type": "u32", "value": "kernelH" }, |
| { "name": "kernelW", "type": "u32", "value": "kernelW" }, |
| { "name": "strideD", "type": "u32", "value": "strideD" }, |
| { "name": "strideH", "type": "u32", "value": "strideH" }, |
| { "name": "strideW", "type": "u32", "value": "strideW" }, |
| { "name": "dilationD", "type": "u32", "value": "dilationD" }, |
| { "name": "dilationH", "type": "u32", "value": "dilationH" }, |
| { "name": "dilationW", "type": "u32", "value": "dilationW" }, |
| { "name": "padFront", "type": "i32", "value": "effectivePadD" }, |
| { "name": "padTop", "type": "i32", "value": "effectivePadH" }, |
| { "name": "padLeft", "type": "i32", "value": "effectivePadW" }, |
| { "name": "padBack", "type": "i32", "value": "effectivePadBack" }, |
| { "name": "padBottom", "type": "i32", "value": "effectivePadBottom" }, |
| { "name": "padRight", "type": "i32", "value": "effectivePadRight" }, |
| { "name": "countIncludePad", "type": "u32", "value": "attrs.count_include_pad" }, |
| { "name": "count", "type": "u32", "value": "numel(shapes.y)" } |
| ] |
| } |
| ], |
| "dispatch": { |
| "x": "min(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), 65535)", |
| "y": "ceilDiv(ceilDiv((numel(shapes.y)), (tunables.WORKGROUP_SIZE)), 65535)", |
| "z": 1 |
| } |
| } |
| ] |
| }, |
| { |
| "id": "nchw2d_k7s2_reuse", |
| "priority": 36, |
| "when": ["poolBaseOk", "ranks.x == 4", "kernelH == 7", "kernelW == 7", "strideH == 2", "strideW == 2", "dilationH == 1", "dilationW == 1", "effectivePadH == 3", "effectivePadBottom == 3", "effectivePadW == 3", "effectivePadRight == 3", "attrs.count_include_pad == 0", "attrs.ceil_mode == 0", "dim(shapes.y, 2) == inferredOutH", "dim(shapes.y, 3) == inferredOutW", "dim(shapes.y, 3) % 4 == 0", "dim(shapes.y, 2) > 0", "dim(shapes.y, 3) > 0"], |
| "derive": { "outputElement": "\"vec4<\" ~ dtypes.T ~ \">\"", "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"" }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "AveragePool.Nchw2dK7s2Reuse", |
| "shader": "average-pool2d-nchw-horizontal-reuse.wgsl.jinja", |
| "derive": { |
| "kernelHSpec": 7, |
| "kernelWSpec": 7, |
| "strideHSpec": 2, |
| "strideWSpec": 2, |
| "dilationHSpec": 1, |
| "dilationWSpec": 1, |
| "padTopSpec": 3, |
| "padLeftSpec": 3, |
| "outW4": "dim(shapes.y, 3) / 4", |
| "count4": "numel(shapes.y) / 4" |
| }, |
| "bindings": ["x", "y_main"], |
| "dispatch": { |
| "x": "min(ceilDiv((numel(shapes.y) / 4), (tunables.WORKGROUP_SIZE)), 65535)", |
| "y": "ceilDiv(ceilDiv((numel(shapes.y) / 4), (tunables.WORKGROUP_SIZE)), 65535)", |
| "z": 1 |
| } |
| } |
| ] |
| }, |
| { |
| "id": "nchw2d_k7s1_reuse", |
| "priority": 35, |
| "when": ["poolBaseOk", "ranks.x == 4", "kernelH == 7", "kernelW == 7", "strideH == 1", "strideW == 1", "dilationH == 1", "dilationW == 1", "effectivePadH == 3", "effectivePadBottom == 3", "effectivePadW == 3", "effectivePadRight == 3", "attrs.count_include_pad == 0", "attrs.ceil_mode == 0", "dim(shapes.x, 3) % 4 == 0", "dim(shapes.y, 2) == dim(shapes.x, 2)", "dim(shapes.y, 3) == dim(shapes.x, 3)", "dim(shapes.y, 2) > 0", "dim(shapes.y, 3) > 0"], |
| "derive": { "outputElement": "\"vec4<\" ~ dtypes.T ~ \">\"", "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"" }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "AveragePool.Nchw2dK7s1Reuse", |
| "shader": "average-pool2d-nchw-horizontal-reuse.wgsl.jinja", |
| "derive": { |
| "kernelHSpec": 7, |
| "kernelWSpec": 7, |
| "strideHSpec": 1, |
| "strideWSpec": 1, |
| "dilationHSpec": 1, |
| "dilationWSpec": 1, |
| "padTopSpec": 3, |
| "padLeftSpec": 3, |
| "outW4": "dim(shapes.y, 3) / 4", |
| "count4": "numel(shapes.y) / 4" |
| }, |
| "bindings": ["x", "y_main"], |
| "dispatch": { |
| "x": "min(ceilDiv((numel(shapes.y) / 4), (tunables.WORKGROUP_SIZE)), 65535)", |
| "y": "ceilDiv(ceilDiv((numel(shapes.y) / 4), (tunables.WORKGROUP_SIZE)), 65535)", |
| "z": 1 |
| } |
| } |
| ] |
| } |
| ] |
| } |
|
|
