| { |
| "domain": "ai.onnx", |
| "name": "SimplifiedLayerNormalization", |
| "conformance": "legacy-default-domain", |
| "sinceVersion": 1, |
| "inputs": { "x": { "onnx": "X", "dtype": "T" }, "scale": { "dtype": "V" } }, |
| "outputs": { |
| "y": { "onnx": "Y", "dtype": "V", "rank": "ranks.x", "shape": "shapes.x" }, |
| "invStdVar": { |
| "onnx": "inv_std_var", |
| "dtype": "U", |
| "rank": "ranks.x", |
| "optional": true, |
| "shape": "prefix(shapes.x, axisNorm) + fill(1, ranks.x - axisNorm)" |
| } |
| }, |
| "attributes": { |
| "axis": { "default": -1 }, |
| "epsilon": { "default": 0.00001 }, |
| "stash_type": { "default": 1 }, |
| "keep_dims": { "default": 1 } |
| }, |
| "attributeConstraints": { "stash_type": { "values": [1] }, "keep_dims": { "values": [1] } }, |
| "typeConstraints": { "T": ["float32", "float16"], "V": ["float32", "float16"], "U": ["float32"] }, |
| "tunables": { |
| "WORKGROUP_SIZE": { "default": 256 }, |
| "SPLIT_MAX_ROWS": { "default": 256 }, |
| "SPLIT_MIN_HIDDEN": { "default": 16384 }, |
| "SPLIT_TARGET_ELEMENTS": { "default": 4096 }, |
| "MAX_SPLITS": { "default": 64 } |
| }, |
| "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\"))", |
| "normMaxWorkgroup": "min(tunables.WORKGROUP_SIZE, deviceWorkgroupCap)", |
| "hasSubgroupId": "device.features.has(\"subgroups\") and device.wgslLanguageFeatures.has(\"subgroup_id\")", |
| "axisNorm": "attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.x", |
| "normRows": "outer(shapes.x, axisNorm)", |
| "normHidden": "dim(shapes.x, axisNorm) * inner(shapes.x, axisNorm)", |
| "normRowStride": "max(1, min(normRows, min(device.limits.maxComputeWorkgroupsPerDimension, 65535)))", |
| "rowWg": "min(normMaxWorkgroup, pow2ceil(max(1, normHidden)))", |
| "baseOk": "ranks.x >= 1 and sameShape(shapes.y, shapes.x) and ranks.scale >= 0 and ranks.scale <= ranks.x and broadcastable(shapes.scale, shapes.x) and attrs.axis + ranks.x >= 0 and attrs.axis < ranks.x and normHidden > 0 and attrs.stash_type == onnxDtypeCode(\"float32\") and f16Ok(dtypes.T) and f16Ok(dtypes.V)", |
| "lastAxisOk": "baseOk and (attrs.axis == -1 or attrs.axis == ranks.x - 1)", |
| "suffixAxisOk": "baseOk and ranks.x >= 2 and not (attrs.axis == -1 or attrs.axis == ranks.x - 1)", |
| "noStats": "not present.invStdVar", |
| "statsOk": "present.invStdVar and ranks.invStdVar == ranks.x and sameShape(prefix(shapes.invStdVar, axisNorm), prefix(shapes.x, axisNorm)) and numel(suffix(shapes.invStdVar, axisNorm)) == 1", |
| "sameDtype": "dtypes.T == dtypes.V", |
| "splitCount": "min(tunables.MAX_SPLITS, pow2ceil(ceilDiv(normHidden, tunables.SPLIT_TARGET_ELEMENTS)))", |
| "splitScratchBytes": "normRows * splitCount * 4", |
| "splitFits": "normRows <= tunables.SPLIT_MAX_ROWS and splitCount <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and splitScratchBytes <= device.limits.maxStorageBufferBindingSize and splitScratchBytes <= device.limits.maxBufferSize" |
| }, |
| "bindings": { |
| "x": { "buffer": "read-only-storage", "elementType": "$xElement" }, |
| "scale": { "buffer": "read-only-storage", "elementType": "$ioElement" }, |
| "y": { "buffer": "storage", "elementType": "$ioElement" }, |
| "params": { |
| "buffer": "uniform", |
| "struct": [ |
| { "name": "rows", "type": "u32", "value": "normRows" }, |
| { "name": "rowStride", "type": "u32", "value": "normRowStride" } |
| ] |
| }, |
| "inv_std_out": { "arg": "invStdVar", "buffer": "storage", "elementType": "f32" }, |
| "partials_2": { "name": "partials", "buffer": "read-only-storage", "elementType": "f32" } |
| }, |
| "variants": [ |
| { |
| "id": "last_axis", |
| "priority": 1, |
| "when": ["lastAxisOk", "noStats"], |
| "derive": { |
| "scalar": "dtypes.V", |
| "xElement": "dtypes.T", |
| "ioElement": "dtypes.V", |
| "hiddenSize": "normHidden", |
| "workgroupSize": "rowWg", |
| "epsilon": "attrs.epsilon" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "SimplifiedLayerNormalization.Row", |
| "shader": "rms-normalization.wgsl.jinja", |
| "derive": { |
| "xShape": "shapes.x", |
| "scaleShape": "shapes.scale", |
| "xRank": "ranks.x", |
| "scaleRank": "ranks.scale", |
| "writeStats": false, |
| "rmsScaleAfterCast": false |
| }, |
| "bindings": ["x", "scale", "y", "params"], |
| "dispatch": { "x": "min(normRows, 65535)", "y": "ceilDiv(normRows, 65535)", "z": 1 } |
| } |
| ] |
| }, |
| { |
| "id": "last_axis_stats", |
| "priority": 2, |
| "when": ["lastAxisOk", "statsOk"], |
| "derive": { |
| "scalar": "dtypes.V", |
| "xElement": "dtypes.T", |
| "ioElement": "dtypes.V", |
| "hiddenSize": "normHidden", |
| "workgroupSize": "rowWg", |
| "epsilon": "attrs.epsilon" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "SimplifiedLayerNormalization.Row", |
| "shader": "rms-normalization.wgsl.jinja", |
| "derive": { |
| "xShape": "shapes.x", |
| "scaleShape": "shapes.scale", |
| "xRank": "ranks.x", |
| "scaleRank": "ranks.scale", |
| "writeStats": true, |
| "rmsScaleAfterCast": false |
| }, |
| "bindings": ["x", "scale", "y", "inv_std_out", "params"], |
| "dispatch": { "x": "min(normRows, 65535)", "y": "ceilDiv(normRows, 65535)", "z": 1 } |
| } |
| ] |
| }, |
| { |
| "id": "suffix_axis", |
| "priority": 10, |
| "when": ["suffixAxisOk", "noStats"], |
| "derive": { |
| "scalar": "dtypes.V", |
| "xElement": "dtypes.T", |
| "ioElement": "dtypes.V", |
| "hiddenSize": "normHidden", |
| "workgroupSize": "rowWg", |
| "epsilon": "attrs.epsilon" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "SimplifiedLayerNormalization.Row", |
| "shader": "rms-normalization.wgsl.jinja", |
| "derive": { |
| "xShape": "shapes.x", |
| "scaleShape": "shapes.scale", |
| "xRank": "ranks.x", |
| "scaleRank": "ranks.scale", |
| "writeStats": false, |
| "rmsScaleAfterCast": false |
| }, |
| "bindings": ["x", "scale", "y", "params"], |
| "dispatch": { "x": "min(normRows, 65535)", "y": "ceilDiv(normRows, 65535)", "z": 1 } |
| } |
| ] |
| }, |
| { |
| "id": "suffix_axis_stats", |
| "priority": 11, |
| "when": ["suffixAxisOk", "statsOk"], |
| "derive": { |
| "scalar": "dtypes.V", |
| "xElement": "dtypes.T", |
| "ioElement": "dtypes.V", |
| "hiddenSize": "normHidden", |
| "workgroupSize": "rowWg", |
| "epsilon": "attrs.epsilon" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "SimplifiedLayerNormalization.Row", |
| "shader": "rms-normalization.wgsl.jinja", |
| "derive": { |
| "xShape": "shapes.x", |
| "scaleShape": "shapes.scale", |
| "xRank": "ranks.x", |
| "scaleRank": "ranks.scale", |
| "writeStats": true, |
| "rmsScaleAfterCast": false |
| }, |
| "bindings": ["x", "scale", "y", "inv_std_out", "params"], |
| "dispatch": { "x": "min(normRows, 65535)", "y": "ceilDiv(normRows, 65535)", "z": 1 } |
| } |
| ] |
| }, |
| { |
| "id": "suffix_axis_splitk", |
| "priority": 15, |
| "when": ["baseOk", "ranks.x >= 2", "noStats", "splitFits"], |
| "demoteWhen": ["reportedNonWave32Adapter", "normHidden < tunables.SPLIT_MIN_HIDDEN"], |
| "derive": { |
| "scalar": "dtypes.V", |
| "xElement": "dtypes.T", |
| "ioElement": "dtypes.V", |
| "hiddenSize": "normHidden", |
| "workgroupSize": "normMaxWorkgroup", |
| "split": "splitCount", |
| "epsilon": "attrs.epsilon", |
| "normalizeRows": "normRows" |
| }, |
| "intermediates": [{ "id": "partials", "dtype": "float32", "shape": "[normRows * splitCount]" }], |
| "passes": [ |
| { |
| "id": "partials", |
| "name": "SimplifiedLayerNormalization.SplitKPartials", |
| "shader": "rms-normalization-splitk-partials.wgsl.jinja", |
| "bindings": ["x", { "name": "partials", "buffer": "storage", "elementType": "f32" }, "params"], |
| "dispatch": { |
| "x": "min(normRows, DISPATCH_FOLD_WIDTH)", |
| "y": "ceilDiv(normRows, DISPATCH_FOLD_WIDTH)", |
| "z": "splitCount" |
| } |
| }, |
| { |
| "id": "normalize", |
| "name": "SimplifiedLayerNormalization.SplitKNormalize", |
| "shader": "rms-normalization-splitk-normalize.wgsl.jinja", |
| "derive": { |
| "xShape": "shapes.x", |
| "scaleShape": "shapes.scale", |
| "xRank": "ranks.x", |
| "scaleRank": "ranks.scale", |
| "writeStats": false, |
| "rmsScaleAfterCast": false, |
| "normalizeBlocks": "max(split, min(min(device.limits.maxComputeWorkgroupsPerDimension, 65535), ceilDiv(workgroupSize, max(1, normalizeRows)), ceilDiv(hiddenSize, workgroupSize * 4)))", |
| "normalizeChunk": "ceilDiv(hiddenSize, normalizeBlocks)" |
| }, |
| "bindings": ["x", "scale", "partials_2", "y", "params"], |
| "dispatch": { |
| "x": "min(normRows, DISPATCH_FOLD_WIDTH)", |
| "y": "ceilDiv(normRows, DISPATCH_FOLD_WIDTH)", |
| "z": "normalizeBlocks" |
| } |
| } |
| ] |
| }, |
| { |
| "id": "suffix_axis_splitk_stats", |
| "priority": 16, |
| "when": ["baseOk", "ranks.x >= 2", "statsOk", "splitFits"], |
| "demoteWhen": ["reportedNonWave32Adapter", "normHidden < tunables.SPLIT_MIN_HIDDEN"], |
| "derive": { |
| "scalar": "dtypes.V", |
| "xElement": "dtypes.T", |
| "ioElement": "dtypes.V", |
| "hiddenSize": "normHidden", |
| "workgroupSize": "normMaxWorkgroup", |
| "split": "splitCount", |
| "epsilon": "attrs.epsilon", |
| "normalizeRows": "normRows" |
| }, |
| "intermediates": [{ "id": "partials", "dtype": "float32", "shape": "[normRows * splitCount]" }], |
| "passes": [ |
| { |
| "id": "partials", |
| "name": "SimplifiedLayerNormalization.SplitKPartials", |
| "shader": "rms-normalization-splitk-partials.wgsl.jinja", |
| "bindings": ["x", { "name": "partials", "buffer": "storage", "elementType": "f32" }, "params"], |
| "dispatch": { |
| "x": "min(normRows, DISPATCH_FOLD_WIDTH)", |
| "y": "ceilDiv(normRows, DISPATCH_FOLD_WIDTH)", |
| "z": "splitCount" |
| } |
| }, |
| { |
| "id": "normalize", |
| "name": "SimplifiedLayerNormalization.SplitKNormalize", |
| "shader": "rms-normalization-splitk-normalize.wgsl.jinja", |
| "derive": { |
| "xShape": "shapes.x", |
| "scaleShape": "shapes.scale", |
| "xRank": "ranks.x", |
| "scaleRank": "ranks.scale", |
| "writeStats": true, |
| "rmsScaleAfterCast": false, |
| "normalizeBlocks": "max(split, min(min(device.limits.maxComputeWorkgroupsPerDimension, 65535), ceilDiv(workgroupSize, max(1, normalizeRows)), ceilDiv(hiddenSize, workgroupSize * 4)))", |
| "normalizeChunk": "ceilDiv(hiddenSize, normalizeBlocks)" |
| }, |
| "bindings": ["x", "scale", "partials_2", "y", "inv_std_out", "params"], |
| "dispatch": { |
| "x": "min(normRows, DISPATCH_FOLD_WIDTH)", |
| "y": "ceilDiv(normRows, DISPATCH_FOLD_WIDTH)", |
| "z": "normalizeBlocks" |
| } |
| } |
| ] |
| }, |
| { |
| "id": "last_axis_row_vec4", |
| "priority": 110, |
| "when": ["lastAxisOk", "sameDtype", "noStats", "ranks.scale >= 1", "numel(shapes.scale) == dim(shapes.x, -1)", "dim(shapes.scale, -1) == dim(shapes.x, -1)", "dim(shapes.x, -1) % 4 == 0"], |
| "derive": { |
| "scalar": "dtypes.T", |
| "xElement": "\"vec4<\" ~ dtypes.T ~ \">\"", |
| "ioElement": "\"vec4<\" ~ dtypes.T ~ \">\"" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "SimplifiedLayerNormalization.LastAxisRow", |
| "shader": "norm-row-stats.wgsl.jinja", |
| "derive": { |
| "modeSpec": "\"rms\"", |
| "vec4": true, |
| "writeStats": false, |
| "rmsScaleAfterCast": false, |
| "scalar": "dtypes.T", |
| "usesF16Spec": "dtypes.T == \"f16\"", |
| "hidden": "dim(shapes.x, -1)", |
| "wg": "min(normMaxWorkgroup, pow2ceil(max(1, dim(shapes.x, -1) / 4)))", |
| "epsilon": "attrs.epsilon", |
| "hiddenVec": "dim(shapes.x, -1) / 4", |
| "vecType": "\"vec4<\" ~ dtypes.T ~ \">\"", |
| "combineSubgroups": "hasSubgroupId" |
| }, |
| "bindings": ["x", "scale", "y", "params"], |
| "dispatch": { "x": "min(normRows, 65535)", "y": "ceilDiv(normRows, 65535)", "z": 1 }, |
| "subgroupCollectivesWidth": "portable" |
| } |
| ] |
| }, |
| { |
| "id": "last_axis_row", |
| "priority": 100, |
| "when": ["lastAxisOk", "sameDtype", "noStats", "ranks.scale >= 1", "numel(shapes.scale) == dim(shapes.x, -1)", "dim(shapes.scale, -1) == dim(shapes.x, -1)"], |
| "derive": { "scalar": "dtypes.T", "xElement": "dtypes.T", "ioElement": "dtypes.T" }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "SimplifiedLayerNormalization.LastAxisRow", |
| "shader": "norm-row-stats.wgsl.jinja", |
| "derive": { |
| "modeSpec": "\"rms\"", |
| "vec4": false, |
| "writeStats": false, |
| "rmsScaleAfterCast": false, |
| "scalar": "dtypes.T", |
| "usesF16Spec": "dtypes.T == \"f16\"", |
| "hidden": "dim(shapes.x, -1)", |
| "wg": "min(normMaxWorkgroup, pow2ceil(max(1, dim(shapes.x, -1))))", |
| "epsilon": "attrs.epsilon", |
| "hiddenVec": 1, |
| "vecType": "\"vec4<\" ~ dtypes.T ~ \">\"", |
| "combineSubgroups": "hasSubgroupId" |
| }, |
| "bindings": ["x", "scale", "y", "params"], |
| "dispatch": { "x": "min(normRows, 65535)", "y": "ceilDiv(normRows, 65535)", "z": 1 }, |
| "subgroupCollectivesWidth": "portable" |
| } |
| ] |
| }, |
| { |
| "id": "last_axis_row_vec4_stats", |
| "priority": 112, |
| "when": ["lastAxisOk", "sameDtype", "statsOk", "ranks.scale >= 1", "numel(shapes.scale) == dim(shapes.x, -1)", "dim(shapes.scale, -1) == dim(shapes.x, -1)", "dim(shapes.x, -1) % 4 == 0"], |
| "derive": { |
| "scalar": "dtypes.T", |
| "xElement": "\"vec4<\" ~ dtypes.T ~ \">\"", |
| "ioElement": "\"vec4<\" ~ dtypes.T ~ \">\"" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "SimplifiedLayerNormalization.LastAxisRow", |
| "shader": "norm-row-stats.wgsl.jinja", |
| "derive": { |
| "modeSpec": "\"rms\"", |
| "vec4": true, |
| "writeStats": true, |
| "rmsScaleAfterCast": false, |
| "scalar": "dtypes.T", |
| "usesF16Spec": "dtypes.T == \"f16\"", |
| "hidden": "dim(shapes.x, -1)", |
| "wg": "min(normMaxWorkgroup, pow2ceil(max(1, dim(shapes.x, -1) / 4)))", |
| "epsilon": "attrs.epsilon", |
| "hiddenVec": "dim(shapes.x, -1) / 4", |
| "vecType": "\"vec4<\" ~ dtypes.T ~ \">\"", |
| "combineSubgroups": "hasSubgroupId" |
| }, |
| "bindings": ["x", "scale", "y", "inv_std_out", "params"], |
| "dispatch": { "x": "min(normRows, 65535)", "y": "ceilDiv(normRows, 65535)", "z": 1 }, |
| "subgroupCollectivesWidth": "portable" |
| } |
| ] |
| }, |
| { |
| "id": "last_axis_row_stats", |
| "priority": 102, |
| "when": ["lastAxisOk", "sameDtype", "statsOk", "ranks.scale >= 1", "numel(shapes.scale) == dim(shapes.x, -1)", "dim(shapes.scale, -1) == dim(shapes.x, -1)"], |
| "derive": { "scalar": "dtypes.T", "xElement": "dtypes.T", "ioElement": "dtypes.T" }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "SimplifiedLayerNormalization.LastAxisRow", |
| "shader": "norm-row-stats.wgsl.jinja", |
| "derive": { |
| "modeSpec": "\"rms\"", |
| "vec4": false, |
| "writeStats": true, |
| "rmsScaleAfterCast": false, |
| "scalar": "dtypes.T", |
| "usesF16Spec": "dtypes.T == \"f16\"", |
| "hidden": "dim(shapes.x, -1)", |
| "wg": "min(normMaxWorkgroup, pow2ceil(max(1, dim(shapes.x, -1))))", |
| "epsilon": "attrs.epsilon", |
| "hiddenVec": 1, |
| "vecType": "\"vec4<\" ~ dtypes.T ~ \">\"", |
| "combineSubgroups": "hasSubgroupId" |
| }, |
| "bindings": ["x", "scale", "y", "inv_std_out", "params"], |
| "dispatch": { "x": "min(normRows, 65535)", "y": "ceilDiv(normRows, 65535)", "z": 1 }, |
| "subgroupCollectivesWidth": "portable" |
| } |
| ] |
| } |
| ] |
| } |
|
|
