| { |
| "domain": "ai.onnx", |
| "name": "GlobalMaxPool", |
| "sinceVersion": 1, |
| "inputs": { "x": { "onnx": "X", "dtype": "T" } }, |
| "outputs": { "y": { "onnx": "Y", "dtype": "T", "rank": "ranks.x" } }, |
| "typeConstraints": { "T": ["float32", "float16"] }, |
| "tunables": { |
| "WORKGROUP_SIZE": { "default": 256 }, |
| "PLANE_PARALLEL_MIN_SPATIAL": { "default": 64 }, |
| "PLANE_PARALLEL_AMORTIZED_SPATIAL": { "default": 768 }, |
| "PLANE_PARALLEL_STARVED_PLANES": { "default": 1024 } |
| }, |
| "derive": { |
| "baseContract": "f16Ok(dtypes.T) and ranks.x >= 3 and ranks.y == ranks.x and dim(shapes.y, 0) == dim(shapes.x, 0) and dim(shapes.y, 1) == dim(shapes.x, 1) and inner(shapes.y, 1) == 1", |
| "planeParallelEligible": "inner(shapes.x, 1) >= tunables.PLANE_PARALLEL_MIN_SPATIAL and (inner(shapes.x, 1) >= tunables.PLANE_PARALLEL_AMORTIZED_SPATIAL or numel(shapes.y) <= tunables.PLANE_PARALLEL_STARVED_PLANES) and numel(shapes.y) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", |
| "scalar": "dtypes.T" |
| }, |
| "when": ["baseContract"], |
| "variants": [ |
| { |
| "id": "wg_plane_parallel", |
| "priority": 30, |
| "when": ["planeParallelEligible"], |
| "derive": { "planeInput": "\"vec4<\" ~ dtypes.T ~ \">\" if inner(shapes.x, 1) % 4 == 0 else dtypes.T" }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "GlobalMaxPool.WgPlaneParallel", |
| "shader": "pool-global-reduction.wgsl.jinja", |
| "derive": { |
| "vec4": "inner(shapes.x, 1) % 4 == 0", |
| "workgroupSizeSpec": "tunables.WORKGROUP_SIZE", |
| "useSubgroups": "device.features.has(\"subgroups\")" |
| }, |
| "bindings": [ |
| { "arg": "x", "elementType": "$planeInput" }, |
| "y", |
| { "name": "params", "struct": [{ "name": "spatial", "type": "u32", "value": "inner(shapes.x, 1)" }] } |
| ], |
| "dispatch": { "x": "numel(shapes.y)" } |
| } |
| ] |
| }, |
| { |
| "id": "nchw", |
| "when": ["not planeParallelEligible"], |
| "passes": [ |
| { |
| "id": "main", |
| "name": "GlobalMaxPool.Nchw", |
| "shader": "pool-global-serial.wgsl.jinja", |
| "derive": { "outputsPerInvocation": 1 }, |
| "bindings": [ |
| "x", |
| "y", |
| { |
| "name": "params", |
| "struct": [ |
| { "name": "spatial", "type": "u32", "value": "inner(shapes.x, 1)" }, |
| { "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 |
| } |
| } |
| ] |
| } |
| ], |
| "bindings": {} |
| } |
|
|
