| { |
| "domain": "com.microsoft", |
| "name": "GroupNorm", |
| "sinceVersion": 1, |
| "inputs": { |
| "X": { "dtype": "T", "rank": 4 }, |
| "gamma": { "dtype": "M", "rank": 1 }, |
| "beta": { "dtype": "M", "rank": 1 } |
| }, |
| "outputs": { "Y": { "dtype": "T", "rank": 4, "shape": "shapes.X" } }, |
| "attributes": { "activation": {}, "groups": {}, "channels_last": { "default": 1 }, "epsilon": { "default": 0.00001 } }, |
| "attributeConstraints": { "activation": { "required": true, "values": [0, 1] }, "groups": { "required": true } }, |
| "typeConstraints": { "T": ["float32", "float16"], "M": ["float32", "float16"] }, |
| "tunables": { |
| "WORKGROUP_SIZE": { "default": 256 }, |
| "MAX_SPLITS": { "default": 256 }, |
| "VALUES_PER_SPLIT": { "default": 4096 }, |
| "SPLIT_MIN_ELEMENTS": { "default": 262144 }, |
| "APPLY_SLAB_ELEMENTS": { "default": 8192 }, |
| "FUSED_MIN_ROWS": { "default": 64 }, |
| "FUSED_MAX_GROUP_BYTES": { "default": 131072 } |
| }, |
| "derive": { |
| "attrsPresent": "has(attrs, \"groups\") and has(attrs, \"activation\")", |
| "numGroups": "attrs.groups if attrsPresent else 0", |
| "activationKind": "attrs.activation if attrsPresent else 0", |
| "channelsLast": "attrs.channels_last != 0", |
| "rankOk": "ranks.X == 4 and ranks.Y == 4 and ranks.gamma == 1 and ranks.beta == 1", |
| "dimN": "dim(shapes.X, 0) if rankOk else 0", |
| "dimC": "(dim(shapes.X, 3) if channelsLast else dim(shapes.X, 1)) if rankOk else 0", |
| "dimHW": "(dim(shapes.X, 1) * dim(shapes.X, 2) if channelsLast else dim(shapes.X, 2) * dim(shapes.X, 3)) if rankOk else 0", |
| "groupsOk": "numGroups >= 1 and dimC >= 1 and dimC % numGroups == 0", |
| "channelsPerGroup": "dimC / numGroups if groupsOk else 0", |
| "planeCount": "dimN * dimC", |
| "groupSize": "channelsPerGroup * dimHW", |
| "normRows": "dimN * numGroups", |
| "elementCount": "numel(shapes.X) if rankOk else 0", |
| "shapeOk": "rankOk and groupsOk and dimN > 0 and dimHW > 0 and groupSize > 0 and sameShape(shapes.Y, shapes.X) and dim(shapes.gamma, 0) == dimC and dim(shapes.beta, 0) == dimC", |
| "deviceWorkgroupCap": "min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)", |
| "storageBufferLimit": "min(device.limits.maxStorageBufferBindingSize, device.limits.maxBufferSize)", |
| "foldedDispatchCapacity": "min(device.limits.maxComputeWorkgroupsPerDimension, 65535) * min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", |
| "statsVec4Ok": "(channelsPerGroup % 4 == 0 and dimC % 4 == 0) if channelsLast else (dimHW % 4 == 0)", |
| "applyVec4Ok": "(dimC % 4 == 0) if channelsLast else (dimHW % 4 == 0)", |
| "fusedVecWidth": "4 if statsVec4Ok else 1", |
| "statsVecWidth": "4 if statsVec4Ok else (2 if channelsLast and channelsPerGroup % 2 == 0 else 1)", |
| "applyVecWidth": "4 if applyVec4Ok else 1", |
| "workgroupCap": "max(1, floor(pow2ceil(min(tunables.WORKGROUP_SIZE, deviceWorkgroupCap) + 1) / 2))", |
| "statsChunks": "max(1, floor(dimC / statsVecWidth))", |
| "fusedRowChunks": "max(1, floor(dimC / fusedVecWidth))", |
| "statsSegment": "max(1, floor(channelsPerGroup / statsVecWidth))", |
| "statsGroupsPerBlock": "max(1, min(numGroups, floor(workgroupCap / statsSegment))) if channelsLast and statsChunks > deviceWorkgroupCap and statsSegment <= workgroupCap else max(1, numGroups)", |
| "statsGroupBlocks": "ceilDiv(max(1, numGroups), statsGroupsPerBlock)", |
| "statsTileChunks": "statsGroupsPerBlock * statsSegment if statsGroupBlocks > 1 else statsChunks", |
| "applyChunks": "max(1, floor(dimC / applyVecWidth))", |
| "statsRowsPerStep": "max(1, floor(workgroupCap / statsTileChunks))", |
| "applyRowsPerStep": "max(1, floor(workgroupCap / applyChunks))", |
| "planeVecs": "max(1, floor(dimHW / statsVecWidth))", |
| "slabCount": "ceilDiv(dimHW, ceilDiv(dimHW, max(1, min(tunables.MAX_SPLITS, dimHW, ceilDiv(dimHW * dimC, tunables.VALUES_PER_SPLIT)))))", |
| "planeSplits": "max(1, min(tunables.MAX_SPLITS, planeVecs, ceilDiv(dimHW, tunables.VALUES_PER_SPLIT)))", |
| "splitCount": "slabCount if channelsLast else planeSplits", |
| "partialsPerRow": "slabCount if channelsLast else max(1, channelsPerGroup) * planeSplits", |
| "partialElements": "max(1, normRows * partialsPerRow)", |
| "statsElements": "max(1, normRows)", |
| "statsWorkgroup": "(statsTileChunks * statsRowsPerStep) if channelsLast else min(workgroupCap, max(32, pow2ceil(ceilDiv(planeVecs, planeSplits))))", |
| "applyWorkgroup": "(applyChunks * applyRowsPerStep) if channelsLast else statsWorkgroup", |
| "applySlabCount": "ceilDiv(dimHW, ceilDiv(dimHW, max(1, min(dimHW, min(device.limits.maxComputeWorkgroupsPerDimension, 65535), ceilDiv(elementCount, tunables.APPLY_SLAB_ELEMENTS)))))", |
| "applySplit": "applySlabCount if channelsLast else planeSplits", |
| "applyUnits": "(dimN * applySlabCount) if channelsLast else (planeCount * planeSplits)", |
| "workUnits": "(dimN * slabCount * statsGroupBlocks) if channelsLast else (planeCount * planeSplits)", |
| "splitFits": "statsWorkgroup <= deviceWorkgroupCap and applyWorkgroup <= deviceWorkgroupCap and statsWorkgroup * 8 <= device.limits.maxComputeWorkgroupStorageSize and partialElements * 8 <= storageBufferLimit and workUnits > 0 and workUnits <= foldedDispatchCapacity and applyUnits <= foldedDispatchCapacity", |
| "fusedChunks": "max(1, floor(channelsPerGroup / fusedVecWidth)) if channelsLast else 1", |
| "fusedWorkgroup": "(fusedChunks * max(1, floor(workgroupCap / fusedChunks))) if channelsLast else min(workgroupCap, max(32, pow2ceil(max(1, floor(groupSize / fusedVecWidth)))))", |
| "fusedLanes": "pow2ceil(fusedWorkgroup)", |
| "fusedFits": "fusedWorkgroup <= deviceWorkgroupCap and fusedLanes * 8 <= device.limits.maxComputeWorkgroupStorageSize and normRows <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) * min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", |
| "combineWorkgroup": "min(workgroupCap, 64)", |
| "statsElement": "(\"vec\" ~ statsVecWidth ~ \"<\" ~ dtypes.T ~ \">\") if statsVecWidth > 1 else dtypes.T", |
| "fusedElement": "(\"vec4<\" ~ dtypes.T ~ \">\") if statsVec4Ok else dtypes.T", |
| "applyElement": "(\"vec4<\" ~ dtypes.T ~ \">\") if applyVec4Ok else dtypes.T", |
| "mScalar": "dtypes.M", |
| "tScalar": "dtypes.T", |
| "contractOk": "shapeOk and f16Ok(dtypes.T) and f16Ok(dtypes.M)" |
| }, |
| "when": ["contractOk"], |
| "bindings": { |
| "xIn": { "arg": "X", "name": "X", "elementType": "$statsElement" }, |
| "srcX": { "arg": "X", "name": "src", "elementType": "$statsElement" }, |
| "gamma": { "elementType": "$mScalar" }, |
| "beta": { "elementType": "$mScalar" }, |
| "yOut": { "arg": "Y", "name": "Y", "elementType": "$applyElement" }, |
| "partialsW": { "name": "partials", "buffer": "storage", "elementType": "vec2<f32>" }, |
| "partialsR": { "name": "partials", "buffer": "read-only-storage", "elementType": "vec2<f32>" }, |
| "statsW": { "name": "stats", "buffer": "storage", "elementType": "vec2<f32>" }, |
| "statsR": { "name": "stats", "buffer": "read-only-storage", "elementType": "vec2<f32>" }, |
| "srcApply": { "arg": "X", "name": "src", "elementType": "$applyElement" }, |
| "yFused": { "arg": "Y", "name": "Y", "elementType": "$statsElement" } |
| }, |
| "variants": [ |
| { |
| "id": "split", |
| "priority": 20, |
| "when": ["splitFits"], |
| "demoteWhen": ["elementCount < tunables.SPLIT_MIN_ELEMENTS * statsGroupBlocks", "(not channelsLast) and normRows >= tunables.FUSED_MIN_ROWS and groupSize * dtypeBytes(tensorDtypes.X) <= tunables.FUSED_MAX_GROUP_BYTES"], |
| "derive": { |
| "phaseLayout": "channelsLast", |
| "hasSkip": "false", |
| "hasBias": "false", |
| "broadcastSkip": "false", |
| "vecSpec": "statsVecWidth", |
| "dimCSpec": "dimC", |
| "dimHWSpec": "dimHW", |
| "cpgSpec": "channelsPerGroup", |
| "groupsSpec": "numGroups", |
| "groupSizeSpec": "groupSize", |
| "normRowsSpec": "normRows", |
| "epsilonSpec": "attrs.epsilon", |
| "activationSpec": "activationKind", |
| "chunksSpec": "statsChunks", |
| "rowsPerStepSpec": "statsRowsPerStep", |
| "splitSpec": "splitCount", |
| "rowsPerSlabSpec": "ceilDiv(dimHW, slabCount)", |
| "planeVecsSpec": "planeVecs", |
| "chunkVecsSpec": "ceilDiv(planeVecs, planeSplits)", |
| "partialsPerRowSpec": "partialsPerRow", |
| "wgSpec": "statsWorkgroup", |
| "combineWgSpec": "combineWorkgroup", |
| "vecF32": "(\"vec\" ~ statsVecWidth ~ \"<f32>\") if statsVecWidth > 1 else \"f32\"", |
| "storeElement": "statsElement", |
| "applySrc": "\"src\"", |
| "groupUniform": "true" |
| }, |
| "intermediates": [ |
| { "id": "partials", "dtype": "float32", "shape": "[partialElements, 2]" }, |
| { "id": "stats", "dtype": "float32", "shape": "[statsElements, 2]" } |
| ], |
| "passes": [ |
| { |
| "id": "stats", |
| "name": "GroupNorm.Partials", |
| "shader": "group-norm-slab.wgsl.jinja", |
| "derive": { |
| "phase": "\"stats\"", |
| "unitsSpec": "workUnits", |
| "statsBlockedSpec": "channelsLast and statsGroupBlocks > 1", |
| "statsGroupsPerBlockSpec": "statsGroupsPerBlock", |
| "statsGroupBlocksSpec": "statsGroupBlocks", |
| "statsTileChunksSpec": "statsTileChunks" |
| }, |
| "bindings": ["xIn", "partialsW"], |
| "dispatch": { "x": "min(workUnits, 65535)", "y": "ceilDiv(workUnits, 65535)", "z": 1 } |
| }, |
| { |
| "id": "combine", |
| "name": "GroupNorm.Combine", |
| "shader": "group-norm-combine.wgsl.jinja", |
| "bindings": ["srcX", "partialsR", "statsW"], |
| "dispatch": { |
| "x": "min(ceilDiv((normRows), (combineWorkgroup)), 65535)", |
| "y": "ceilDiv(ceilDiv((normRows), (combineWorkgroup)), 65535)", |
| "z": 1 |
| } |
| }, |
| { |
| "id": "apply", |
| "name": "GroupNorm.Apply", |
| "shader": "group-norm-slab.wgsl.jinja", |
| "derive": { |
| "phase": "\"apply\"", |
| "vecSpec": "applyVecWidth", |
| "chunksSpec": "applyChunks", |
| "rowsPerStepSpec": "applyRowsPerStep", |
| "wgSpec": "applyWorkgroup", |
| "splitSpec": "applySplit", |
| "rowsPerSlabSpec": "ceilDiv(dimHW, applySlabCount)", |
| "groupUniform": "(channelsPerGroup % applyVecWidth == 0) if channelsLast else true", |
| "storeElement": "applyElement", |
| "vecF32": "\"vec4<f32>\" if applyVec4Ok else \"f32\"", |
| "unitsSpec": "applyUnits" |
| }, |
| "bindings": ["srcApply", "gamma", "beta", "statsR", "yOut"], |
| "dispatch": { "x": "min(applyUnits, 65535)", "y": "ceilDiv(applyUnits, 65535)", "z": 1 } |
| } |
| ] |
| }, |
| { |
| "id": "fused", |
| "priority": 10, |
| "when": ["fusedFits"], |
| "derive": { |
| "statsElement": "fusedElement", |
| "phaseLayout": "channelsLast", |
| "hasSkip": "false", |
| "hasBias": "false", |
| "broadcastSkip": "false", |
| "useVec4": "statsVec4Ok", |
| "vecSpec": "fusedVecWidth", |
| "dimCSpec": "dimC", |
| "dimHWSpec": "dimHW", |
| "cpgSpec": "channelsPerGroup", |
| "groupsSpec": "numGroups", |
| "groupSizeSpec": "groupSize", |
| "normRowsSpec": "normRows", |
| "epsilonSpec": "attrs.epsilon", |
| "activationSpec": "activationKind", |
| "chunksSpec": "fusedRowChunks", |
| "fusedChunksSpec": "fusedChunks", |
| "fusedRowsSpec": "max(1, floor(fusedWorkgroup / fusedChunks))", |
| "spanVecsSpec": "max(1, floor(groupSize / fusedVecWidth))", |
| "wgSpec": "fusedWorkgroup", |
| "ldsSpec": "fusedLanes", |
| "vecF32": "\"vec4<f32>\" if statsVec4Ok else \"f32\"", |
| "storeElement": "fusedElement", |
| "fusedSrc": "\"X\"", |
| "planeVecsSpec": "planeVecs" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "GroupNorm.Fused", |
| "shader": "group-norm-fused.wgsl.jinja", |
| "bindings": ["xIn", "gamma", "beta", "yFused"], |
| "dispatch": { "x": "min(normRows, 65535)", "y": "ceilDiv(normRows, 65535)", "z": 1 } |
| } |
| ] |
| } |
| ] |
| } |
|
|
