{ "domain": "ai.onnx", "name": "GroupNormalization", "sinceVersion": 21, "description": "Applies group normalization to the input: `y = scale * (x - mean) / sqrt(variance + epsilon) + bias`, where mean and variance are computed per instance per group of channels. The number of groups `num_groups` must divide the channel count `C` evenly; when `num_groups == C` this is equivalent to InstanceNormalization, and when `num_groups == 1` it is equivalent to LayerNormalization. The normalization stage supports TensorProto `stash_type` values `1` (float32) and `10` (float16).", "inputs": [ { "role": "X", "dtype": "T", "description": "Input data tensor of shape `(N x C x D1 x ... x Dn)` where `N` is batch size and `C` is the number of channels." }, { "role": "scale", "dtype": "T", "rank": 1, "description": "Scale tensor of shape `(C)`, one value per channel." }, { "role": "bias", "dtype": "T", "rank": 1, "description": "Bias tensor of shape `(C)`, one value per channel." } ], "outputs": [ { "role": "Y", "dtype": "T", "rank": "ranks.X", "description": "Normalized output tensor of the same shape as `X`.", "shape": "shapes.X" } ], "attributes": { "epsilon": 0.00001, "stash_type": 1 }, "attributeDescriptions": { "epsilon": "Small value added to the variance denominator to avoid division by zero.", "num_groups": "Required number of groups to divide the channels into; must be a divisor of `C`.", "stash_type": "TensorProto element type used for the normalization stage: `1` computes in float32, while `10` computes in float16. Normalized values are cast back to the input type before scale and bias are applied." }, "attributeConstraints": { "num_groups": { "required": true }, "stash_type": { "values": [1, 10] } }, "typeConstraints": { "T": ["float32", "float16"] }, "args": { "x": { "kind": "tensor", "semantic": "X", "role": "input" }, "scale": { "kind": "tensor", "semantic": "scale", "role": "input" }, "bias": { "kind": "tensor", "semantic": "bias", "role": "input" }, "y": { "kind": "tensor", "semantic": "Y", "role": "output" } }, "tunables": { "WORKGROUP_SIZE": 256, "MAX_STATS_SPLITS": 64, "STATS_VALUES_PER_SPLIT": 4096, "SPLIT_STATS_MIN_HIDDEN": 65536, "SPLIT_STATS_MAX_ROWS": 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, device.limits.maxComputeWorkgroupsPerDimension, pow2ceil(ceilDiv(groupHidden, tunables.STATS_VALUES_PER_SPLIT)))", "groupPartialBytes": "groupRows * groupSplitCount * 2 * 4", "groupSplitCovered": "groupRowCovered and groupRows <= tunables.SPLIT_STATS_MAX_ROWS and groupRows <= device.limits.maxComputeWorkgroupsPerDimension and groupHidden >= tunables.SPLIT_STATS_MIN_HIDDEN and groupPartialBytes <= device.limits.maxStorageBufferBindingSize and groupPartialBytes <= device.limits.maxBufferSize" }, "bindingSets": { "norm": [ { "name": "x", "arg": "x", "semantic": "X", "buffer": { "type": "read-only-storage" }, "elementType": "$ioElement" }, { "name": "scale", "arg": "scale", "semantic": "scale", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" }, { "name": "bias", "arg": "bias", "semantic": "bias", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" }, { "name": "y", "arg": "y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$ioElement" }, { "name": "params", "semantic": "kernel.params", "buffer": { "type": "uniform" }, "struct": { "name": "Params", "fields": [ { "name": "rows", "type": "u32", "value": "groupRows" }, { "name": "rowStride", "type": "u32", "value": "max(1, min(groupRows, device.limits.maxComputeWorkgroupsPerDimension))" } ] } } ], "splitPartials": [ { "name": "x", "arg": "x", "semantic": "X", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" }, { "name": "partials", "semantic": "partials", "buffer": { "type": "storage" }, "elementType": "vec2" }, { "name": "params", "semantic": "kernel.params", "buffer": { "type": "uniform" }, "struct": { "name": "Params", "fields": [{ "name": "rows", "type": "u32", "value": "groupRows" }] } } ], "splitApply": [ { "name": "x", "arg": "x", "semantic": "X", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" }, { "name": "scale", "arg": "scale", "semantic": "scale", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" }, { "name": "bias", "arg": "bias", "semantic": "bias", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" }, { "name": "partials", "semantic": "partials", "buffer": { "type": "read-only-storage" }, "elementType": "vec2" }, { "name": "y", "arg": "y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$scalar" }, { "name": "params", "semantic": "kernel.params", "buffer": { "type": "uniform" }, "struct": { "name": "Params", "fields": [{ "name": "rows", "type": "u32", "value": "groupRows" }] } } ] }, "variants": [ { "id": "group_stash_f16_serial", "priority": 1000, "when": "groupStashF16Ok", "constants": { "scalar": "dtypes.T", "ioElement": "dtypes.T", "usesF16": "dtypes.T == \"f16\"", "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": "norm", "dispatch": { "workgroups": "groupRows" } } ] }, { "id": "group_splitk", "priority": 120, "when": ["groupSplitCovered"], "constants": { "scalar": "dtypes.T", "usesF16": "dtypes.T == \"f16\"", "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": "splitPartials", "dispatch": { "workgroups": "groupRows", "z": "groupSplitCount" } }, { "id": "apply", "name": "GroupNormalization.SplitKApply", "shader": "group-normalization-splitk-apply.wgsl.jinja", "bindings": "splitApply", "dispatch": { "workgroups": "groupRows", "z": "groupSplitCount" } } ] }, { "id": "group_subgroup_vec4", "priority": 110, "when": ["groupRowCovered", "groupSpatial % 4 == 0"], "constants": { "scalar": "dtypes.T", "ioElement": "\"vec4<\" ~ dtypes.T ~ \">\"", "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"" }, "passes": [ { "id": "main", "name": "GroupNormalization.group_subgroup_vec4", "source": { "shader": "norm-row-stats.wgsl.jinja", "inputs": { "mode": "\"group\"", "vec4": true, "scalar": "dtypes.T", "usesF16": "dtypes.T == \"f16\"", "hidden": "groupHidden", "wg": "groupVec4Workgroup", "epsilon": "attrs.epsilon", "numGroups": "attrs.num_groups", "cpg": "groupChannelsPerGroup", "hiddenVec": "groupHidden / 4", "vecType": "\"vec4<\" ~ dtypes.T ~ \">\"", "spatialVec": "groupSpatial / 4", "combineSubgroups": "hasSubgroupId" } }, "subgroupCollectivesWidth": "portable", "bindings": "norm", "dispatch": { "workgroups": "groupRows" } } ] }, { "id": "group_subgroup", "priority": 100, "when": ["groupRowCovered"], "constants": { "scalar": "dtypes.T", "ioElement": "dtypes.T" }, "passes": [ { "id": "main", "name": "GroupNormalization.group_subgroup", "source": { "shader": "norm-row-stats.wgsl.jinja", "inputs": { "mode": "\"group\"", "vec4": false, "scalar": "dtypes.T", "usesF16": "dtypes.T == \"f16\"", "hidden": "groupHidden", "wg": "groupScalarWorkgroup", "epsilon": "attrs.epsilon", "numGroups": "attrs.num_groups", "cpg": "groupChannelsPerGroup", "spatial": "groupSpatial", "combineSubgroups": "hasSubgroupId" } }, "subgroupCollectivesWidth": "portable", "bindings": "norm", "dispatch": { "workgroups": "groupRows" } } ] } ] }