| { |
| "domain": "ai.onnx", |
| "name": "GroupNormalization", |
| "sinceVersion": 21, |
| "inputs": { |
| "x": { "onnx": "X", "dtype": "T" }, |
| "scale": { "dtype": "T", "rank": 1 }, |
| "bias": { "dtype": "T", "rank": 1 } |
| }, |
| "outputs": { "y": { "onnx": "Y", "dtype": "T", "rank": "ranks.x", "shape": "shapes.x" } }, |
| "attributes": { "epsilon": { "default": 0.00001 }, "stash_type": { "default": 1 }, "num_groups": {} }, |
| "attributeConstraints": { "num_groups": { "required": true }, "stash_type": { "values": [1, 10] } }, |
| "typeConstraints": { "T": ["float32", "float16"] }, |
| "tunables": { |
| "WORKGROUP_SIZE": { "default": 256 }, |
| "MAX_STATS_SPLITS": { "default": 64 }, |
| "STATS_VALUES_PER_SPLIT": { "default": 4096 }, |
| "SPLIT_STATS_MIN_HIDDEN": { "default": 65536 }, |
| "SPLIT_STATS_MAX_ROWS": { "default": 256 } |
| }, |
| "derive": { |
| "deviceWorkgroupCap": "min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)", |
| "groupAttributesOk": "attrs.num_groups >= 1", |
| "groupShapeOk": "groupAttributesOk and f16Ok(dtypes.T) and ranks.x >= 3 and ranks.scale == 1 and ranks.bias == 1 and ranks.y == ranks.x and sameShape(shapes.y, shapes.x) and dim(shapes.scale, 0) == dim(shapes.x, 1) and dim(shapes.bias, 0) == dim(shapes.x, 1) and dim(shapes.x, 1) % attrs.num_groups == 0", |
| "groupContractOk": "groupShapeOk and attrs.stash_type == onnxDtypeCode(\"float32\")", |
| "groupStashF16Ok": "groupShapeOk and attrs.stash_type == onnxDtypeCode(\"float16\")", |
| "groupRows": "dim(shapes.x, 0) * attrs.num_groups if groupAttributesOk else 0", |
| "groupSpatial": "inner(shapes.x, 1)", |
| "groupChannelsPerGroup": "dim(shapes.x, 1) / attrs.num_groups if groupAttributesOk else 0", |
| "groupHidden": "groupChannelsPerGroup * groupSpatial", |
| "normDeviceWorkgroupCap": "min(tunables.WORKGROUP_SIZE, deviceWorkgroupCap)", |
| "normWorkgroupCap": "max(1, pow2ceil(normDeviceWorkgroupCap + 1) / 2)", |
| "groupScalarWorkgroup": "min(normWorkgroupCap, pow2ceil(groupHidden))", |
| "groupVec4Workgroup": "min(normWorkgroupCap, pow2ceil(groupHidden / 4))", |
| "hasSubgroupId": "device.features.has(\"subgroups\") and device.wgslLanguageFeatures.has(\"subgroup_id\")", |
| "groupRowWorkgroupBytes": "normWorkgroupCap * 2 * 4", |
| "groupRowCovered": "groupContractOk and groupRowWorkgroupBytes <= device.limits.maxComputeWorkgroupStorageSize", |
| "groupSplitCount": "min(tunables.MAX_STATS_SPLITS, min(device.limits.maxComputeWorkgroupsPerDimension, 65535), pow2ceil(ceilDiv(groupHidden, tunables.STATS_VALUES_PER_SPLIT)))", |
| "groupPartialBytes": "groupRows * groupSplitCount * 2 * 4", |
| "groupSplitCovered": "groupRowCovered and groupRows <= tunables.SPLIT_STATS_MAX_ROWS and groupRows <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and groupHidden >= tunables.SPLIT_STATS_MIN_HIDDEN and groupPartialBytes <= device.limits.maxStorageBufferBindingSize and groupPartialBytes <= device.limits.maxBufferSize" |
| }, |
| "bindings": { |
| "x": { "elementType": "$ioElement" }, |
| "scale": { "elementType": "$scalar" }, |
| "bias": { "elementType": "$scalar" }, |
| "y": { "elementType": "$ioElement" }, |
| "params": { |
| "struct": [ |
| { "name": "rows", "type": "u32", "value": "groupRows" }, |
| { |
| "name": "rowStride", |
| "type": "u32", |
| "value": "max(1, min(groupRows, min(device.limits.maxComputeWorkgroupsPerDimension, 65535)))" |
| } |
| ] |
| }, |
| "x_apply": { "name": "x", "elementType": "$scalar" }, |
| "params_rows": { "name": "params", "struct": [{ "name": "rows", "type": "u32", "value": "groupRows" }] } |
| }, |
| "variants": [ |
| { |
| "id": "group_stash_f16_serial", |
| "priority": 1000, |
| "when": ["groupStashF16Ok"], |
| "derive": { |
| "scalar": "dtypes.T", |
| "ioElement": "dtypes.T", |
| "hiddenSize": "groupHidden", |
| "spatial": "groupSpatial", |
| "channelsPerGroup": "groupChannelsPerGroup", |
| "numGroups": "attrs.num_groups", |
| "epsilon": "attrs.epsilon" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "GroupNormalization.StashF16Serial", |
| "shader": "group-normalization-stash-f16-serial.wgsl.jinja", |
| "bindings": ["x", "scale", "bias", "y", "params"], |
| "dispatch": { "x": "min(groupRows, 65535)", "y": "ceilDiv(groupRows, 65535)", "z": 1 } |
| } |
| ] |
| }, |
| { |
| "id": "group_splitk", |
| "priority": 120, |
| "when": ["groupSplitCovered"], |
| "derive": { |
| "scalar": "dtypes.T", |
| "hiddenSize": "groupHidden", |
| "spatial": "groupSpatial", |
| "channelsPerGroup": "groupChannelsPerGroup", |
| "numGroups": "attrs.num_groups", |
| "workgroupSize": "normWorkgroupCap", |
| "split": "groupSplitCount", |
| "epsilon": "attrs.epsilon" |
| }, |
| "intermediates": [{ "id": "partials", "dtype": "float32", "shape": "[groupRows * groupSplitCount, 2]" }], |
| "passes": [ |
| { |
| "id": "partials", |
| "name": "GroupNormalization.SplitKPartials", |
| "shader": "group-normalization-splitk-partials.wgsl.jinja", |
| "bindings": ["x_apply", { "name": "partials", "elementType": "vec2<f32>" }, "params_rows"], |
| "dispatch": { |
| "x": "min(groupRows, DISPATCH_FOLD_WIDTH)", |
| "y": "ceilDiv(groupRows, DISPATCH_FOLD_WIDTH)", |
| "z": "groupSplitCount" |
| } |
| }, |
| { |
| "id": "apply", |
| "name": "GroupNormalization.SplitKApply", |
| "shader": "group-normalization-splitk-apply.wgsl.jinja", |
| "bindings": [ |
| "x_apply", |
| "scale", |
| "bias", |
| { "name": "partials", "buffer": "read-only-storage", "elementType": "vec2<f32>" }, |
| { "arg": "y", "elementType": "$scalar" }, |
| "params_rows" |
| ], |
| "dispatch": { |
| "x": "min(groupRows, DISPATCH_FOLD_WIDTH)", |
| "y": "ceilDiv(groupRows, DISPATCH_FOLD_WIDTH)", |
| "z": "groupSplitCount" |
| } |
| } |
| ] |
| }, |
| { |
| "id": "group_subgroup_vec4", |
| "priority": 110, |
| "when": ["groupRowCovered", "groupSpatial % 4 == 0"], |
| "derive": { "scalar": "dtypes.T", "ioElement": "\"vec4<\" ~ dtypes.T ~ \">\"" }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "GroupNormalization.GroupSubgroupVec4", |
| "shader": "norm-row-stats.wgsl.jinja", |
| "derive": { |
| "vec4": true, |
| "scalar": "dtypes.T", |
| "hidden": "groupHidden", |
| "wg": "groupVec4Workgroup", |
| "epsilon": "attrs.epsilon", |
| "numGroupsSpec": "attrs.num_groups", |
| "cpg": "groupChannelsPerGroup", |
| "hiddenVec": "groupHidden / 4", |
| "vecType": "\"vec4<\" ~ dtypes.T ~ \">\"", |
| "spatialVec": "groupSpatial / 4", |
| "combineSubgroups": "hasSubgroupId" |
| }, |
| "bindings": ["x", "scale", "bias", "y", "params"], |
| "dispatch": { "x": "min(groupRows, 65535)", "y": "ceilDiv(groupRows, 65535)", "z": 1 } |
| } |
| ] |
| }, |
| { |
| "id": "group_subgroup", |
| "priority": 100, |
| "when": ["groupRowCovered"], |
| "derive": { "scalar": "dtypes.T", "ioElement": "dtypes.T" }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "GroupNormalization.GroupSubgroup", |
| "shader": "norm-row-stats.wgsl.jinja", |
| "derive": { |
| "vec4": false, |
| "scalar": "dtypes.T", |
| "hidden": "groupHidden", |
| "wg": "groupScalarWorkgroup", |
| "epsilon": "attrs.epsilon", |
| "numGroupsSpec": "attrs.num_groups", |
| "cpg": "groupChannelsPerGroup", |
| "spatial": "groupSpatial", |
| "combineSubgroups": "hasSubgroupId" |
| }, |
| "bindings": ["x", "scale", "bias", "y", "params"], |
| "dispatch": { "x": "min(groupRows, 65535)", "y": "ceilDiv(groupRows, 65535)", "z": 1 } |
| } |
| ] |
| } |
| ] |
| } |
|
|
