| { |
| "domain": "ai.onnx", |
| "name": "InstanceNormalization", |
| "sinceVersion": 6, |
| "description": "Applies instance normalization to the input: `y = scale * (x - mean) / sqrt(variance + epsilon) + B`, where `mean` and `variance` are computed per instance per channel over the spatial dimensions. Equivalent to batch normalization with a batch size of one per channel.", |
| "inputs": [ |
| { "role": "input", "dtype": "T", "description": "Input tensor of shape `(N x C x D1 x ... x Dn)`; at least 3-D." }, |
| { |
| "role": "scale", |
| "dtype": "T", |
| "rank": 1, |
| "description": "1-D scale tensor of size C, one scale factor per channel." |
| }, |
| { "role": "B", "dtype": "T", "rank": 1, "description": "1-D bias tensor of size C, one bias value per channel." } |
| ], |
| "outputs": [ |
| { |
| "role": "output", |
| "dtype": "T", |
| "rank": "ranks.input", |
| "description": "Normalized output tensor; same shape as the input.", |
| "shape": "shapes.input" |
| } |
| ], |
| "attributes": { "epsilon": 0.00001 }, |
| "attributeDescriptions": { |
| "epsilon": "Small constant added to the variance before taking the square root to avoid division by zero." |
| }, |
| "typeConstraints": { "T": ["float32", "float16"] }, |
| "args": { |
| "input": { "kind": "tensor", "semantic": "input", "role": "input" }, |
| "scale": { "kind": "tensor", "semantic": "scale", "role": "input" }, |
| "b": { "kind": "tensor", "semantic": "B", "role": "input" }, |
| "output": { "kind": "tensor", "semantic": "output", "role": "output" } |
| }, |
| "tunables": { |
| "WORKGROUP_SIZE": 256, |
| "MAX_STATS_SPLITS": 256, |
| "STATS_VALUES_PER_SPLIT": 2048, |
| "SPLIT_STATS_MIN_SPATIAL": 65536, |
| "SPLIT_STATS_MAX_PLANES": 256, |
| "COMBINE_WORKGROUP_SIZE": 64, |
| "BATCHED_MIN_PLANES_PER_WORKGROUP": 8 |
| }, |
| "derive": { |
| "deviceWorkgroupCap": "min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)", |
| "wave32Adapter": "has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize == 32 and device.adapterInfo.subgroupMaxSize == 32", |
| "reportedNonWave32Adapter": "not wave32Adapter and (has(device.adapterInfo, \"subgroupMinSize\") or has(device.adapterInfo, \"subgroupMaxSize\"))", |
| "instanceContractOk": "f16Ok(dtypes.T) and ranks.input >= 3 and ranks.output == ranks.input and sameShape(shapes.output, shapes.input) and ranks.scale == 1 and ranks.B == 1 and dim(shapes.scale, 0) == dim(shapes.input, 1) and dim(shapes.B, 0) == dim(shapes.input, 1)", |
| "instancePlanes": "dim(shapes.input, 0) * dim(shapes.input, 1)", |
| "instanceSpatial": "inner(shapes.input, 1)", |
| "normDeviceWorkgroupCap": "min(tunables.WORKGROUP_SIZE, deviceWorkgroupCap)", |
| "normWorkgroupCap": "max(1, pow2ceil(normDeviceWorkgroupCap + 1) / 2)", |
| "normSubgroupMin": "device.adapterInfo.subgroupMinSize if has(device.adapterInfo, \"subgroupMinSize\") else 1", |
| "normSubgroupMax": "device.adapterInfo.subgroupMaxSize if has(device.adapterInfo, \"subgroupMaxSize\") else 32", |
| "hasSubgroupId": "device.features.has(\"subgroups\") and device.wgslLanguageFeatures.has(\"subgroup_id\")", |
| "instanceScalarWorkgroup": "min(normWorkgroupCap, pow2ceil(instanceSpatial))", |
| "instanceVec4Workgroup": "min(normWorkgroupCap, pow2ceil(instanceSpatial / 4))", |
| "instanceVec4SubgroupEfficient": "not hasSubgroupId or instanceVec4Workgroup >= normSubgroupMin", |
| "instanceBatchedVec4Lanes": "instanceVec4Workgroup", |
| "instanceBatchedVec4PlanesPerWorkgroup": "max(1, floor(normWorkgroupCap / instanceBatchedVec4Lanes))", |
| "instanceBatchedVec4Workgroups": "ceilDiv(instancePlanes, instanceBatchedVec4PlanesPerWorkgroup)", |
| "instanceBatchedVec4StorageBytes": "normWorkgroupCap * 2 * 4 + instanceBatchedVec4PlanesPerWorkgroup * 4", |
| "instanceRowWorkgroupBytes": "normWorkgroupCap * 2 * 4", |
| "instanceStatsBytes": "instancePlanes * 2 * 4", |
| "instanceStatsFits": "instanceStatsBytes <= device.limits.maxStorageBufferBindingSize and instanceStatsBytes <= device.limits.maxBufferSize", |
| "instanceRowCovered": "instanceContractOk and instanceRowWorkgroupBytes <= device.limits.maxComputeWorkgroupStorageSize", |
| "instanceSplitCount": "min(tunables.MAX_STATS_SPLITS, device.limits.maxComputeWorkgroupsPerDimension, pow2ceil(ceilDiv(instanceSpatial, tunables.STATS_VALUES_PER_SPLIT)))", |
| "instancePartialBytes": "instancePlanes * instanceSplitCount * 2 * 4", |
| "splitStatsCovered": "instanceRowCovered and instanceStatsFits and instancePlanes <= tunables.SPLIT_STATS_MAX_PLANES and instancePlanes <= device.limits.maxComputeWorkgroupsPerDimension and instanceSpatial >= tunables.SPLIT_STATS_MIN_SPATIAL and instancePartialBytes <= device.limits.maxStorageBufferBindingSize and instancePartialBytes <= device.limits.maxBufferSize", |
| "splitStatsPreferred": "splitStatsCovered and instancePlanes < normSubgroupMax" |
| }, |
| "bindingSets": { |
| "planeIo": [ |
| { |
| "name": "x", |
| "arg": "input", |
| "semantic": "input", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$ioElement" |
| }, |
| { |
| "name": "scale", |
| "arg": "scale", |
| "semantic": "scale", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$T" |
| }, |
| { "name": "bias", "arg": "b", "semantic": "B", "buffer": { "type": "read-only-storage" }, "elementType": "$T" }, |
| { |
| "name": "y", |
| "arg": "output", |
| "semantic": "output", |
| "buffer": { "type": "storage" }, |
| "elementType": "$ioElement" |
| } |
| ], |
| "plane": [ |
| { |
| "name": "x", |
| "arg": "input", |
| "semantic": "input", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$ioElement" |
| }, |
| { |
| "name": "scale", |
| "arg": "scale", |
| "semantic": "scale", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$T" |
| }, |
| { "name": "bias", "arg": "b", "semantic": "B", "buffer": { "type": "read-only-storage" }, "elementType": "$T" }, |
| { |
| "name": "y", |
| "arg": "output", |
| "semantic": "output", |
| "buffer": { "type": "storage" }, |
| "elementType": "$ioElement" |
| }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [ |
| { "name": "rows", "type": "u32", "value": "dim(shapes.input, 0) * dim(shapes.input, 1)" }, |
| { |
| "name": "rowStride", |
| "type": "u32", |
| "value": "max(1, min(dim(shapes.input, 0) * dim(shapes.input, 1), device.limits.maxComputeWorkgroupsPerDimension))" |
| } |
| ] |
| } |
| } |
| ], |
| "planeBatched": [ |
| { |
| "name": "x", |
| "arg": "input", |
| "semantic": "input", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$ioElement" |
| }, |
| { |
| "name": "scale", |
| "arg": "scale", |
| "semantic": "scale", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$T" |
| }, |
| { "name": "bias", "arg": "b", "semantic": "B", "buffer": { "type": "read-only-storage" }, "elementType": "$T" }, |
| { |
| "name": "y", |
| "arg": "output", |
| "semantic": "output", |
| "buffer": { "type": "storage" }, |
| "elementType": "$ioElement" |
| }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [{ "name": "rows", "type": "u32", "value": "dim(shapes.input, 0) * dim(shapes.input, 1)" }] |
| } |
| } |
| ], |
| "applyScalar": [ |
| { |
| "name": "input", |
| "arg": "input", |
| "semantic": "input", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$T" |
| }, |
| { "name": "stats", "semantic": "stats", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, |
| { |
| "name": "scale", |
| "arg": "scale", |
| "semantic": "scale", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$T" |
| }, |
| { "name": "bias", "arg": "b", "semantic": "B", "buffer": { "type": "read-only-storage" }, "elementType": "$T" }, |
| { "name": "output", "arg": "output", "semantic": "output", "buffer": { "type": "storage" }, "elementType": "$T" }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [ |
| { "name": "count", "type": "u32", "value": "numel(shapes.output)" }, |
| { "name": "channels", "type": "u32", "value": "dim(shapes.input, 1)" }, |
| { "name": "spatial", "type": "u32", "value": "instanceSpatial" } |
| ] |
| } |
| } |
| ], |
| "applyVec4": [ |
| { |
| "name": "input", |
| "arg": "input", |
| "semantic": "input", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$vectorScalar" |
| }, |
| { "name": "stats", "semantic": "stats", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, |
| { |
| "name": "scale", |
| "arg": "scale", |
| "semantic": "scale", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$T" |
| }, |
| { "name": "bias", "arg": "b", "semantic": "B", "buffer": { "type": "read-only-storage" }, "elementType": "$T" }, |
| { |
| "name": "output", |
| "arg": "output", |
| "semantic": "output", |
| "buffer": { "type": "storage" }, |
| "elementType": "$vectorScalar" |
| }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [ |
| { "name": "count", "type": "u32", "value": "numel(shapes.output) / 4" }, |
| { "name": "channels", "type": "u32", "value": "dim(shapes.input, 1)" }, |
| { "name": "spatial", "type": "u32", "value": "instanceSpatial" } |
| ] |
| } |
| } |
| ], |
| "splitPartials": [ |
| { |
| "name": "input", |
| "arg": "input", |
| "semantic": "input", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$splitInputElement" |
| }, |
| { "name": "partials", "semantic": "partials", "buffer": { "type": "storage" }, "elementType": "f32" }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [ |
| { "name": "planes", "type": "u32", "value": "instancePlanes" }, |
| { "name": "spatial", "type": "u32", "value": "instanceSpatial" } |
| ] |
| } |
| } |
| ], |
| "splitCombine": [ |
| { "name": "partials", "semantic": "partials", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, |
| { "name": "stats", "semantic": "stats", "buffer": { "type": "storage" }, "elementType": "f32" }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [ |
| { "name": "planes", "type": "u32", "value": "instancePlanes" }, |
| { "name": "spatial", "type": "u32", "value": "instanceSpatial" }, |
| { "name": "epsilon", "type": "f32", "value": "attrs.epsilon" } |
| ] |
| } |
| } |
| ] |
| }, |
| "variants": [ |
| { |
| "id": "plane_batched_vec4", |
| "priority": 115, |
| "when": ["instanceRowCovered", "instanceSpatial % 4 == 0", "instanceSpatial >= 4", "instancePlanes >= normWorkgroupCap", "instanceBatchedVec4PlanesPerWorkgroup >= tunables.BATCHED_MIN_PLANES_PER_WORKGROUP", "instanceBatchedVec4StorageBytes <= device.limits.maxComputeWorkgroupStorageSize"], |
| "demoteWhen": ["reportedNonWave32Adapter and instancePlanes <= device.limits.maxComputeWorkgroupsPerDimension"], |
| "constants": { |
| "usesF16": "dtypes.T == \"f16\"", |
| "ioElement": "\"vec4<\" ~ dtypes.T ~ \">\"", |
| "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"", |
| "hidden": "instanceSpatial", |
| "hiddenVec": "instanceSpatial / 4", |
| "channels": "dim(shapes.input, 1) if dim(shapes.input, 1) > 0 else 1", |
| "epsilon": "attrs.epsilon", |
| "workgroupSize": "normWorkgroupCap", |
| "lanesPerPlane": "instanceBatchedVec4Lanes", |
| "planesPerWorkgroup": "instanceBatchedVec4PlanesPerWorkgroup" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "InstanceNormalization.PlaneBatchedVec4", |
| "shader": "instance-normalization-batched-planes-vec4.wgsl.jinja", |
| "bindings": "planeBatched", |
| "dispatch": { "workgroups": "instanceBatchedVec4Workgroups" } |
| } |
| ] |
| }, |
| { |
| "id": "plane_subgroup_vec4", |
| "priority": 110, |
| "requires": { "features": [] }, |
| "when": ["instanceRowCovered", "inner(shapes.input, 1) % 4 == 0", "instanceVec4SubgroupEfficient"], |
| "constants": { "ioElement": "\"vec4<\" ~ dtypes.T ~ \">\"", "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"" }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "InstanceNormalization.plane_subgroup_vec4", |
| "source": { |
| "shader": "norm-row-stats.wgsl.jinja", |
| "inputs": { |
| "mode": "\"instance\"", |
| "vec4": true, |
| "scalar": "dtypes.T", |
| "usesF16": "dtypes.T == \"f16\"", |
| "hidden": "instanceSpatial", |
| "wg": "instanceVec4Workgroup", |
| "epsilon": "attrs.epsilon", |
| "channels": "dim(shapes.input, 1) if dim(shapes.input, 1) > 0 else 1", |
| "hiddenVec": "instanceSpatial / 4", |
| "vecType": "\"vec4<\" ~ dtypes.T ~ \">\"", |
| "combineSubgroups": "hasSubgroupId" |
| } |
| }, |
| "subgroupCollectivesWidth": "portable", |
| "bindings": "plane", |
| "dispatch": { "workgroups": "instancePlanes" } |
| } |
| ] |
| }, |
| { |
| "id": "plane_subgroup_vec4_scalar_io", |
| "priority": 111, |
| "requires": { "features": ["subgroups"] }, |
| "when": ["dtypes.T == \"f32\"", "wave32Adapter", "device.wgslLanguageFeatures.has(\"subgroup_id\")", "instanceRowCovered", "instanceSpatial % 4 == 0", "instanceVec4SubgroupEfficient"], |
| "constants": { "ioElement": "dtypes.T" }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "InstanceNormalization.plane_subgroup_vec4_scalar_io", |
| "source": { |
| "shader": "norm-row-stats.wgsl.jinja", |
| "inputs": { |
| "mode": "\"instance\"", |
| "vec4": true, |
| "scalarIo": true, |
| "scalar": "dtypes.T", |
| "usesF16": false, |
| "hidden": "instanceSpatial", |
| "wg": "instanceVec4Workgroup", |
| "epsilon": "attrs.epsilon", |
| "channels": "dim(shapes.input, 1) if dim(shapes.input, 1) > 0 else 1", |
| "hiddenVec": "instanceSpatial / 4", |
| "vecType": "\"vec4<f32>\"", |
| "combineSubgroups": true |
| } |
| }, |
| "subgroupCollectivesWidth": "portable", |
| "bindings": "plane", |
| "dispatch": { "workgroups": "instancePlanes" } |
| } |
| ] |
| }, |
| { |
| "id": "plane_subgroup", |
| "priority": 100, |
| "requires": { "features": [] }, |
| "when": ["instanceRowCovered"], |
| "constants": { "ioElement": "dtypes.T" }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "InstanceNormalization.plane_subgroup", |
| "source": { |
| "shader": "norm-row-stats.wgsl.jinja", |
| "inputs": { |
| "mode": "\"instance\"", |
| "vec4": false, |
| "scalar": "dtypes.T", |
| "usesF16": "dtypes.T == \"f16\"", |
| "hidden": "instanceSpatial", |
| "wg": "instanceScalarWorkgroup", |
| "epsilon": "attrs.epsilon", |
| "channels": "dim(shapes.input, 1) if dim(shapes.input, 1) > 0 else 1", |
| "combineSubgroups": "hasSubgroupId" |
| } |
| }, |
| "subgroupCollectivesWidth": "portable", |
| "bindings": "plane", |
| "dispatch": { "workgroups": "instancePlanes" } |
| } |
| ] |
| }, |
| { |
| "id": "plane_splitk_vec4", |
| "priority": 121, |
| "when": ["splitStatsPreferred", "instanceSpatial % 4 == 0"], |
| "constants": { |
| "vectorized": true, |
| "usesF16": "dtypes.T == \"f16\"", |
| "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"", |
| "useSubgroups": "hasSubgroupId", |
| "splitInputElement": "\"vec4<\" ~ dtypes.T ~ \">\"", |
| "workgroupSize": "normWorkgroupCap", |
| "split": "instanceSplitCount", |
| "combineWorkgroupSize": "min(tunables.COMBINE_WORKGROUP_SIZE, normWorkgroupCap)", |
| "applyWorkgroupSize": "normWorkgroupCap" |
| }, |
| "intermediates": [ |
| { "id": "partials", "dtype": "float32", "shape": "[instancePlanes * instanceSplitCount, 2]" }, |
| { "id": "stats", "dtype": "float32", "shape": "[instancePlanes, 2]" } |
| ], |
| "passes": [ |
| { |
| "id": "partials", |
| "name": "InstanceNormalization.SplitKPartialsVec4", |
| "shader": "instance-normalization-splitk-partials.wgsl.jinja", |
| "subgroupCollectivesWidth": "portable", |
| "bindings": "splitPartials", |
| "dispatch": { "workgroups": "instancePlanes", "z": "instanceSplitCount" } |
| }, |
| { |
| "id": "combine", |
| "name": "InstanceNormalization.SplitKCombine", |
| "shader": "instance-normalization-splitk-combine.wgsl.jinja", |
| "bindings": "splitCombine", |
| "dispatch": { "threads": "instancePlanes", "workgroupSize": "constants.combineWorkgroupSize" } |
| }, |
| { |
| "id": "apply", |
| "name": "InstanceNormalization.ApplyVec4", |
| "shader": "instance-normalization-apply.wgsl.jinja", |
| "bindings": "applyVec4", |
| "dispatch": { "threads": "numel(shapes.output) / 4", "workgroupSize": "constants.applyWorkgroupSize" } |
| } |
| ] |
| }, |
| { |
| "id": "plane_splitk", |
| "priority": 120, |
| "when": ["splitStatsPreferred"], |
| "constants": { |
| "scalar": "dtypes.T", |
| "usesF16": "dtypes.T == \"f16\"", |
| "splitInputElement": "dtypes.T", |
| "workgroupSize": "normWorkgroupCap", |
| "split": "instanceSplitCount", |
| "combineWorkgroupSize": "min(tunables.COMBINE_WORKGROUP_SIZE, normWorkgroupCap)", |
| "applyWorkgroupSize": "normWorkgroupCap" |
| }, |
| "intermediates": [ |
| { "id": "partials", "dtype": "float32", "shape": "[instancePlanes * instanceSplitCount, 2]" }, |
| { "id": "stats", "dtype": "float32", "shape": "[instancePlanes, 2]" } |
| ], |
| "passes": [ |
| { |
| "id": "partials", |
| "name": "InstanceNormalization.SplitKPartials", |
| "shader": "instance-normalization-splitk-partials.wgsl.jinja", |
| "bindings": "splitPartials", |
| "dispatch": { "workgroups": "instancePlanes", "z": "instanceSplitCount" } |
| }, |
| { |
| "id": "combine", |
| "name": "InstanceNormalization.SplitKCombine", |
| "shader": "instance-normalization-splitk-combine.wgsl.jinja", |
| "bindings": "splitCombine", |
| "dispatch": { "threads": "instancePlanes", "workgroupSize": "constants.combineWorkgroupSize" } |
| }, |
| { |
| "id": "apply", |
| "name": "InstanceNormalization.Apply", |
| "shader": "instance-normalization-apply.wgsl.jinja", |
| "bindings": "applyScalar", |
| "dispatch": { "threads": "numel(shapes.output)", "workgroupSize": "constants.applyWorkgroupSize" } |
| } |
| ] |
| } |
| ] |
| } |
|
|
