| { |
| "domain": "com.microsoft", |
| "name": "SkipLayerNormalization", |
| "sinceVersion": 1, |
| "inputs": { |
| "inputT": { "onnx": "input", "dtype": "T" }, |
| "skipT": { "onnx": "skip", "dtype": "T" }, |
| "gammaT": { "onnx": "gamma", "dtype": "T", "rank": 1 }, |
| "betaT": { "onnx": "beta", "dtype": "T", "rank": 1, "optional": true }, |
| "biasT": { "onnx": "bias", "dtype": "T", "rank": 1, "optional": true } |
| }, |
| "outputs": { |
| "outputT": { "onnx": "output", "dtype": "T", "rank": "ranks.inputT", "shape": "shapes.inputT" }, |
| "residualT": { |
| "onnx": "input_skip_bias_sum", |
| "dtype": "T", |
| "rank": "ranks.inputT", |
| "optional": true, |
| "shape": "shapes.inputT" |
| } |
| }, |
| "attributes": { "epsilon": { "default": 9.999999960041972e-13 } }, |
| "typeConstraints": { "T": ["float32", "float16"] }, |
| "tunables": { "MAX_WORKGROUP_SIZE": { "default": 256 } }, |
| "derive": { |
| "rowCount": "numel(shapes.inputT) / max(1, dim(shapes.inputT, -1))", |
| "hiddenSize": "dim(shapes.inputT, -1)", |
| "skipWg": "max(1, min(tunables.MAX_WORKGROUP_SIZE, device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX, pow2ceil(hiddenSize)))", |
| "skipWgVec4": "max(1, min(tunables.MAX_WORKGROUP_SIZE, device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX, pow2ceil(ceilDiv(hiddenSize, 4))))", |
| "portableWideExecution": "not has(device.adapterInfo, \"subgroupMinSize\") or device.adapterInfo.subgroupMinSize >= 32", |
| "broadcastRows": "dim(shapes.inputT, 0) * dim(shapes.inputT, 1)", |
| "broadcastHiddenSize": "dim(shapes.inputT, 2)", |
| "broadcastSkipWgVec4": "max(1, min(tunables.MAX_WORKGROUP_SIZE, device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX, pow2ceil(ceilDiv(broadcastHiddenSize, 4))))", |
| "rowDispatchFits": "rowCount <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) * min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", |
| "broadcastDispatchFits": "broadcastRows <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) * min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", |
| "normResourcesFit": "skipWg * 8 <= device.limits.maxComputeWorkgroupStorageSize and skipWgVec4 * 8 <= device.limits.maxComputeWorkgroupStorageSize", |
| "broadcastResourcesFit": "broadcastSkipWgVec4 * 8 <= device.limits.maxComputeWorkgroupStorageSize", |
| "epsilonOk": "attrs.epsilon >= 0", |
| "coreContract": "epsilonOk and (ranks.inputT == 2 or ranks.inputT == 3) and ranks.skipT == ranks.inputT and ranks.gammaT == 1 and ranks.outputT == ranks.inputT and sameShape(shapes.inputT, shapes.skipT) and sameShape(shapes.outputT, shapes.inputT) and dim(shapes.inputT, -1) > 0 and dim(shapes.gammaT, 0) == dim(shapes.inputT, -1)", |
| "residualOutputContract": "present.residualT and sameShape(shapes.residualT, shapes.inputT)", |
| "outputOnlyContract": "not present.residualT", |
| "betaContract": "false if not present.betaT else (ranks.betaT == 1 and dim(shapes.betaT, 0) == dim(shapes.inputT, -1))", |
| "noBetaContract": "not present.betaT", |
| "f32MainDtypes": "tensorDtypes.inputT == \"float32\" and tensorDtypes.skipT == \"float32\" and tensorDtypes.gammaT == \"float32\" and tensorDtypes.outputT == \"float32\"", |
| "f16MainDtypes": "tensorDtypes.inputT == \"float16\" and tensorDtypes.skipT == \"float16\" and tensorDtypes.gammaT == \"float16\" and tensorDtypes.outputT == \"float16\"", |
| "f32ResidualDtypes": "f32MainDtypes and tensorDtypes.residualT == \"float32\" if present.residualT else false", |
| "f16ResidualDtypes": "f16MainDtypes and tensorDtypes.residualT == \"float16\" if present.residualT else false", |
| "vec4Aligned": "dim(shapes.inputT, -1) % 4 == 0", |
| "broadcastSkipShapeOk": "(ranks.skipT == 2 and dim(shapes.skipT, 0) == dim(shapes.inputT, 1) and dim(shapes.skipT, 1) == dim(shapes.inputT, 2)) or (ranks.skipT == 3 and ((dim(shapes.skipT, 0) == 1 and dim(shapes.skipT, 1) == dim(shapes.inputT, 1) and dim(shapes.skipT, 2) == dim(shapes.inputT, 2)) or sameShape(shapes.skipT, shapes.inputT)))", |
| "broadcastOutputOnlyContract": "false if ranks.inputT != 3 or not present.betaT else (epsilonOk and not present.biasT and not present.residualT and dim(shapes.inputT, 2) % 4 == 0 and broadcastSkipShapeOk and ranks.gammaT == 1 and ranks.betaT == 1 and ranks.outputT == 3 and tensorDtypes.inputT == \"float32\" and tensorDtypes.skipT == \"float32\" and tensorDtypes.gammaT == \"float32\" and tensorDtypes.betaT == \"float32\" and tensorDtypes.outputT == \"float32\" and dim(shapes.inputT, 2) > 0 and dim(shapes.gammaT, 0) == dim(shapes.inputT, 2) and dim(shapes.betaT, 0) == dim(shapes.inputT, 2) and sameShape(shapes.outputT, shapes.inputT))", |
| "hasSubgroups": "device.features.has(\"subgroups\") and device.wgslLanguageFeatures.has(\"subgroup_id\")", |
| "hasF16": "device.features.has(\"shader-f16\")", |
| "f32_beta_no_bias_residual_contract": "coreContract and residualOutputContract and betaContract and f32ResidualDtypes and not present.biasT and tensorDtypes.betaT == \"float32\"", |
| "f32_beta_bias_residual_contract": "false if not present.biasT else (coreContract and residualOutputContract and betaContract and f32ResidualDtypes and ranks.biasT == 1 and tensorDtypes.betaT == \"float32\" and tensorDtypes.biasT == \"float32\" and dim(shapes.biasT, 0) == hiddenSize)", |
| "f16_beta_bias_residual_contract": "false if not present.biasT else (hasF16 and coreContract and residualOutputContract and betaContract and f16ResidualDtypes and ranks.biasT == 1 and tensorDtypes.betaT == \"float16\" and tensorDtypes.biasT == \"float16\" and dim(shapes.biasT, 0) == hiddenSize)", |
| "f32_no_beta_output_contract": "coreContract and outputOnlyContract and noBetaContract and f32MainDtypes and not present.biasT", |
| "f32_beta_no_bias_output_only_contract": "coreContract and outputOnlyContract and betaContract and f32MainDtypes and not present.biasT and tensorDtypes.betaT == \"float32\"", |
| "f32_beta_bias_output_only_contract": "false if not present.biasT else (coreContract and outputOnlyContract and betaContract and f32MainDtypes and ranks.biasT == 1 and tensorDtypes.betaT == \"float32\" and tensorDtypes.biasT == \"float32\" and dim(shapes.biasT, 0) == hiddenSize)", |
| "f16_no_beta_output_contract": "hasF16 and coreContract and outputOnlyContract and noBetaContract and f16MainDtypes and not present.biasT", |
| "f16_beta_no_bias_output_only_contract": "hasF16 and coreContract and outputOnlyContract and betaContract and f16MainDtypes and not present.biasT and tensorDtypes.betaT == \"float16\"", |
| "f16_beta_bias_output_only_contract": "false if not present.biasT else (hasF16 and coreContract and outputOnlyContract and betaContract and f16MainDtypes and ranks.biasT == 1 and tensorDtypes.betaT == \"float16\" and tensorDtypes.biasT == \"float16\" and dim(shapes.biasT, 0) == hiddenSize)" |
| }, |
| "bindings": { |
| "input": { "arg": "inputT", "buffer": "read-only-storage", "elementType": "$vectorScalar" }, |
| "skip": { "arg": "skipT", "buffer": "read-only-storage", "elementType": "$vectorScalar" }, |
| "gamma": { "arg": "gammaT", "buffer": "read-only-storage", "elementType": "$vectorScalar", "length": "$HIDDEN_LEN" }, |
| "beta": { "arg": "betaT", "buffer": "read-only-storage", "elementType": "$vectorScalar", "length": "$HIDDEN_LEN" }, |
| "output": { "arg": "outputT", "buffer": "storage", "elementType": "$vectorScalar" }, |
| "bias": { "arg": "biasT", "buffer": "read-only-storage", "elementType": "$vectorScalar", "length": "$HIDDEN_LEN" }, |
| "input_skip_bias_sum": { "arg": "residualT", "buffer": "storage", "elementType": "$vectorScalar" }, |
| "params_2": { |
| "name": "params", |
| "buffer": "uniform", |
| "struct": [ |
| { "name": "rows", "type": "u32", "value": "rowCount" }, |
| { |
| "name": "rowStride", |
| "type": "u32", |
| "value": "max(1, min(rowCount, min(device.limits.maxComputeWorkgroupsPerDimension, 65535)))" |
| }, |
| { "name": "epsilon", "type": "f32", "value": "attrs.epsilon" } |
| ] |
| }, |
| "input_2": { "arg": "inputT", "name": "input", "buffer": "read-only-storage", "elementType": "$scalar" }, |
| "skip_2": { "arg": "skipT", "name": "skip", "buffer": "read-only-storage", "elementType": "$scalar" }, |
| "bias_2": { |
| "arg": "biasT", |
| "name": "bias", |
| "buffer": "read-only-storage", |
| "elementType": "$scalar", |
| "length": "$HIDDEN_LEN" |
| }, |
| "gamma_2": { |
| "arg": "gammaT", |
| "name": "gamma", |
| "buffer": "read-only-storage", |
| "elementType": "$scalar", |
| "length": "$HIDDEN_LEN" |
| }, |
| "beta_2": { |
| "arg": "betaT", |
| "name": "beta", |
| "buffer": "read-only-storage", |
| "elementType": "$scalar", |
| "length": "$HIDDEN_LEN" |
| }, |
| "output_2": { "arg": "outputT", "name": "output", "buffer": "storage", "elementType": "$scalar" }, |
| "input_skip_bias_sum_2": { |
| "arg": "residualT", |
| "name": "input_skip_bias_sum", |
| "buffer": "storage", |
| "elementType": "$scalar" |
| }, |
| "params_3": { |
| "name": "params", |
| "buffer": "uniform", |
| "struct": [ |
| { "name": "rows", "type": "u32", "value": "rowCount" }, |
| { "name": "epsilon", "type": "f32", "value": "attrs.epsilon" } |
| ] |
| } |
| }, |
| "variants": [ |
| { |
| "id": "beta_output_only_vec4_broadcast", |
| "priority": 19, |
| "when": ["broadcastOutputOnlyContract", "broadcastResourcesFit", "broadcastDispatchFits"], |
| "derive": { "scalar": "\"f32\"", "vectorScalar": "\"vec4<f32>\"", "HIDDEN_LEN": "broadcastHiddenSize / 4" }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "SkipLayerNormalization.BroadcastSkip", |
| "shader": "norm-skip-row-vec4.wgsl.jinja", |
| "derive": { |
| "simplified": false, |
| "hasBias": false, |
| "hasBeta": true, |
| "writeResidualSum": false, |
| "usesF16Spec": false, |
| "broadcastSkip": true, |
| "hidden": "broadcastHiddenSize", |
| "hiddenVec": "broadcastHiddenSize / 4", |
| "wg": "broadcastSkipWgVec4", |
| "vecType": "\"vec4<f32>\"", |
| "useSubgroups": "hasSubgroups" |
| }, |
| "bindings": [ |
| "input", |
| "skip", |
| "gamma", |
| "beta", |
| "output", |
| { |
| "name": "params", |
| "struct": [ |
| { "name": "rows", "type": "u32", "value": "dim(shapes.inputT, 0) * dim(shapes.inputT, 1)" }, |
| { |
| "name": "rowStride", |
| "type": "u32", |
| "value": "max(1, min(broadcastRows, min(device.limits.maxComputeWorkgroupsPerDimension, 65535)))" |
| }, |
| { "name": "epsilon", "type": "f32", "value": "attrs.epsilon" }, |
| { "name": "skipRows", "type": "u32", "value": "numel(shapes.skipT) / broadcastHiddenSize" } |
| ] |
| } |
| ], |
| "dispatch": { "x": "min(broadcastRows, 65535)", "y": "ceilDiv(broadcastRows, 65535)", "z": 1 }, |
| "subgroupCollectivesWidth": "portable" |
| } |
| ] |
| }, |
| { |
| "id": "beta_bias_vec4", |
| "priority": 15, |
| "when": ["f32_beta_bias_residual_contract", "vec4Aligned", "hasSubgroups or \"bias\" == \"bias\" or portableWideExecution", "normResourcesFit", "rowDispatchFits"], |
| "derive": { |
| "scalar": "\"f32\"", |
| "vectorScalar": "\"vec4<f32>\"", |
| "hasBias": "\"bias\" == \"bias\"", |
| "workgroupSize": "skipWg", |
| "HIDDEN_LEN": "hiddenSize / 4" |
| }, |
| "passes": [ |
| { |
| "id": "normalize", |
| "name": "SkipLayerNormalization.Vec4.Normalize", |
| "shader": "norm-skip-row-vec4.wgsl.jinja", |
| "derive": { |
| "simplified": false, |
| "hasBias": "\"bias\" == \"bias\"", |
| "hasBeta": true, |
| "writeResidualSum": true, |
| "usesF16Spec": false, |
| "hidden": "hiddenSize", |
| "hiddenVec": "hiddenSize / 4", |
| "wg": "skipWgVec4", |
| "vecType": "\"vec4<f32>\"", |
| "useSubgroups": "hasSubgroups" |
| }, |
| "bindings": ["input", "skip", "bias", "gamma", "beta", "output", "input_skip_bias_sum", "params_2"], |
| "dispatch": { "x": "min(rowCount, 65535)", "y": "ceilDiv(rowCount, 65535)", "z": 1 }, |
| "subgroupCollectivesWidth": "portable" |
| } |
| ] |
| }, |
| { |
| "id": "beta_bias_row", |
| "priority": 5, |
| "when": ["f32_beta_bias_residual_contract", "normResourcesFit", "rowDispatchFits"], |
| "derive": { |
| "simplified": false, |
| "useSubgroups": "hasSubgroups", |
| "hasBeta": true, |
| "writeResidualSum": true, |
| "hasBias": "\"bias\" == \"bias\"", |
| "scalar": "\"f32\"", |
| "workgroupSize": "skipWg", |
| "HIDDEN_LEN": "hiddenSize" |
| }, |
| "passes": [ |
| { |
| "id": "normalize", |
| "name": "SkipLayerNormalization.Row.Normalize", |
| "shader": "norm-skip-row.wgsl.jinja", |
| "derive": { "writeResidualSum": true }, |
| "bindings": ["input_2", "skip_2", "bias_2", "gamma_2", "beta_2", "output_2", "input_skip_bias_sum_2", "params_3"], |
| "dispatch": { "x": "min(rowCount, 65535)", "y": "ceilDiv(rowCount, 65535)", "z": 1 }, |
| "subgroupCollectivesWidth": "portable" |
| } |
| ] |
| }, |
| { |
| "id": "beta_bias_vec4_f16", |
| "priority": 21, |
| "when": ["f16_beta_bias_residual_contract", "vec4Aligned", "normResourcesFit", "rowDispatchFits"], |
| "derive": { "scalar": "\"f16\"", "vectorScalar": "\"vec4<f16>\"", "HIDDEN_LEN": "hiddenSize / 4" }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "SkipLayerNormalization.Vec4", |
| "shader": "norm-skip-row-vec4.wgsl.jinja", |
| "derive": { |
| "simplified": false, |
| "hasBias": true, |
| "hasBeta": true, |
| "writeResidualSum": true, |
| "usesF16Spec": true, |
| "hidden": "hiddenSize", |
| "hiddenVec": "hiddenSize / 4", |
| "wg": "skipWgVec4", |
| "vecType": "\"vec4<f16>\"", |
| "useSubgroups": "hasSubgroups" |
| }, |
| "bindings": ["input", "skip", "bias", "gamma", "beta", "output", "input_skip_bias_sum", "params_2"], |
| "dispatch": { "x": "min(rowCount, 65535)", "y": "ceilDiv(rowCount, 65535)", "z": 1 }, |
| "subgroupCollectivesWidth": "portable" |
| } |
| ] |
| }, |
| { |
| "id": "no_beta_output_only_vec4", |
| "priority": 20, |
| "when": ["f32_no_beta_output_contract", "vec4Aligned", "normResourcesFit", "rowDispatchFits"], |
| "derive": { "scalar": "\"f32\"", "vectorScalar": "\"vec4<f32>\"", "HIDDEN_LEN": "hiddenSize / 4" }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "SkipLayerNormalization.NoBetaOutputOnly.Vec4", |
| "shader": "norm-skip-row-vec4.wgsl.jinja", |
| "derive": { |
| "simplified": false, |
| "hasBias": false, |
| "hasBeta": false, |
| "writeResidualSum": false, |
| "usesF16Spec": false, |
| "hidden": "hiddenSize", |
| "hiddenVec": "hiddenSize / 4", |
| "wg": "skipWgVec4", |
| "vecType": "\"vec4<f32>\"", |
| "useSubgroups": "hasSubgroups" |
| }, |
| "bindings": ["input", "skip", "gamma", "output", "params_2"], |
| "dispatch": { "x": "min(rowCount, 65535)", "y": "ceilDiv(rowCount, 65535)", "z": 1 }, |
| "subgroupCollectivesWidth": "portable" |
| } |
| ] |
| }, |
| { |
| "id": "no_beta_output_only_row", |
| "priority": 10, |
| "when": ["f32_no_beta_output_contract", "normResourcesFit", "rowDispatchFits"], |
| "derive": { |
| "simplified": false, |
| "hasBias": false, |
| "hasBeta": false, |
| "writeResidualSum": false, |
| "useSubgroups": "hasSubgroups", |
| "scalar": "\"f32\"", |
| "workgroupSize": "skipWg", |
| "HIDDEN_LEN": "hiddenSize" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "SkipLayerNormalization.NoBetaOutputOnly.Row", |
| "shader": "norm-skip-row.wgsl.jinja", |
| "bindings": ["input_2", "skip_2", "gamma_2", "output_2", "params_3"], |
| "dispatch": { "x": "min(rowCount, 65535)", "y": "ceilDiv(rowCount, 65535)", "z": 1 }, |
| "subgroupCollectivesWidth": "portable" |
| } |
| ] |
| }, |
| { |
| "id": "no_beta_output_only_vec4_f16", |
| "priority": 20, |
| "when": ["f16_no_beta_output_contract", "vec4Aligned", "normResourcesFit", "rowDispatchFits"], |
| "derive": { "scalar": "\"f16\"", "vectorScalar": "\"vec4<f16>\"", "HIDDEN_LEN": "hiddenSize / 4" }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "SkipLayerNormalization.NoBetaOutputOnly.Vec4.F16", |
| "shader": "norm-skip-row-vec4.wgsl.jinja", |
| "derive": { |
| "simplified": false, |
| "hasBias": false, |
| "hasBeta": false, |
| "writeResidualSum": false, |
| "usesF16Spec": true, |
| "hidden": "hiddenSize", |
| "hiddenVec": "hiddenSize / 4", |
| "wg": "skipWgVec4", |
| "vecType": "\"vec4<f16>\"", |
| "useSubgroups": "hasSubgroups" |
| }, |
| "bindings": ["input", "skip", "gamma", "output", "params_2"], |
| "dispatch": { "x": "min(rowCount, 65535)", "y": "ceilDiv(rowCount, 65535)", "z": 1 }, |
| "subgroupCollectivesWidth": "portable" |
| } |
| ] |
| }, |
| { |
| "id": "no_beta_output_only_row_f16", |
| "priority": 10, |
| "when": ["f16_no_beta_output_contract", "normResourcesFit", "rowDispatchFits"], |
| "derive": { |
| "simplified": false, |
| "hasBias": false, |
| "hasBeta": false, |
| "writeResidualSum": false, |
| "useSubgroups": "hasSubgroups", |
| "scalar": "\"f16\"", |
| "usesF16": true, |
| "workgroupSize": "skipWg", |
| "HIDDEN_LEN": "hiddenSize" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "SkipLayerNormalization.NoBetaOutputOnly.Row.F16", |
| "shader": "norm-skip-row.wgsl.jinja", |
| "bindings": ["input_2", "skip_2", "gamma_2", "output_2", "params_3"], |
| "dispatch": { "x": "min(rowCount, 65535)", "y": "ceilDiv(rowCount, 65535)", "z": 1 }, |
| "subgroupCollectivesWidth": "portable" |
| } |
| ] |
| }, |
| { |
| "id": "beta_no_bias_vec4", |
| "priority": 20, |
| "when": ["f32_beta_no_bias_residual_contract", "vec4Aligned", "hasSubgroups or portableWideExecution", "normResourcesFit", "rowDispatchFits"], |
| "derive": { |
| "scalar": "\"f32\"", |
| "vectorScalar": "\"vec4<f32>\"", |
| "hasBias": "\"no_bias\" == \"bias\"", |
| "workgroupSize": "skipWg", |
| "HIDDEN_LEN": "hiddenSize / 4" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "SkipLayerNormalization.Vec4", |
| "shader": "norm-skip-row-vec4.wgsl.jinja", |
| "derive": { |
| "simplified": false, |
| "hasBias": "\"no_bias\" == \"bias\"", |
| "hasBeta": true, |
| "writeResidualSum": true, |
| "usesF16Spec": false, |
| "hidden": "hiddenSize", |
| "hiddenVec": "hiddenSize / 4", |
| "wg": "skipWgVec4", |
| "vecType": "\"vec4<f32>\"", |
| "useSubgroups": "hasSubgroups" |
| }, |
| "bindings": ["input", "skip", "gamma", "beta", "output", "input_skip_bias_sum", "params_2"], |
| "dispatch": { "x": "min(rowCount, 65535)", "y": "ceilDiv(rowCount, 65535)", "z": 1 }, |
| "subgroupCollectivesWidth": "portable" |
| } |
| ] |
| }, |
| { |
| "id": "beta_no_bias_row", |
| "priority": 10, |
| "when": ["f32_beta_no_bias_residual_contract", "normResourcesFit", "rowDispatchFits"], |
| "derive": { |
| "simplified": false, |
| "useSubgroups": "hasSubgroups", |
| "hasBeta": true, |
| "writeResidualSum": true, |
| "hasBias": "\"no_bias\" == \"bias\"", |
| "scalar": "\"f32\"", |
| "workgroupSize": "skipWg", |
| "HIDDEN_LEN": "hiddenSize" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "SkipLayerNormalization.Row", |
| "shader": "norm-skip-row.wgsl.jinja", |
| "bindings": ["input_2", "skip_2", "gamma_2", "beta_2", "output_2", "input_skip_bias_sum_2", "params_3"], |
| "dispatch": { "x": "min(rowCount, 65535)", "y": "ceilDiv(rowCount, 65535)", "z": 1 }, |
| "subgroupCollectivesWidth": "portable" |
| } |
| ] |
| }, |
| { |
| "id": "beta_no_bias_output_only_vec4", |
| "priority": 20, |
| "when": ["f32_beta_no_bias_output_only_contract", "vec4Aligned", "hasSubgroups or \"no_bias\" == \"bias\" or portableWideExecution", "normResourcesFit", "rowDispatchFits"], |
| "derive": { |
| "scalar": "\"f32\"", |
| "vectorScalar": "\"vec4<f32>\"", |
| "hasBias": "\"no_bias\" == \"bias\"", |
| "workgroupSize": "skipWg", |
| "HIDDEN_LEN": "hiddenSize / 4" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "SkipLayerNormalization.Vec4", |
| "shader": "norm-skip-row-vec4.wgsl.jinja", |
| "derive": { |
| "simplified": false, |
| "hasBias": "\"no_bias\" == \"bias\"", |
| "hasBeta": true, |
| "writeResidualSum": false, |
| "usesF16Spec": false, |
| "hidden": "hiddenSize", |
| "hiddenVec": "hiddenSize / 4", |
| "wg": "skipWgVec4", |
| "vecType": "\"vec4<f32>\"", |
| "useSubgroups": "hasSubgroups" |
| }, |
| "bindings": ["input", "skip", "gamma", "beta", "output", "params_2"], |
| "dispatch": { "x": "min(rowCount, 65535)", "y": "ceilDiv(rowCount, 65535)", "z": 1 }, |
| "subgroupCollectivesWidth": "portable" |
| } |
| ] |
| }, |
| { |
| "id": "beta_no_bias_output_only_row", |
| "priority": 10, |
| "when": ["f32_beta_no_bias_output_only_contract", "normResourcesFit", "rowDispatchFits"], |
| "derive": { |
| "simplified": false, |
| "useSubgroups": "hasSubgroups", |
| "hasBeta": true, |
| "writeResidualSum": false, |
| "hasBias": "\"no_bias\" == \"bias\"", |
| "scalar": "\"f32\"", |
| "workgroupSize": "skipWg", |
| "HIDDEN_LEN": "hiddenSize" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "SkipLayerNormalization.Row", |
| "shader": "norm-skip-row.wgsl.jinja", |
| "bindings": ["input_2", "skip_2", "gamma_2", "beta_2", "output_2", "params_3"], |
| "dispatch": { "x": "min(rowCount, 65535)", "y": "ceilDiv(rowCount, 65535)", "z": 1 }, |
| "subgroupCollectivesWidth": "portable" |
| } |
| ] |
| }, |
| { |
| "id": "beta_bias_output_only_vec4", |
| "priority": 20, |
| "when": ["f32_beta_bias_output_only_contract", "vec4Aligned", "hasSubgroups or \"bias\" == \"bias\" or portableWideExecution", "normResourcesFit", "rowDispatchFits"], |
| "derive": { |
| "scalar": "\"f32\"", |
| "vectorScalar": "\"vec4<f32>\"", |
| "hasBias": "\"bias\" == \"bias\"", |
| "workgroupSize": "skipWg", |
| "HIDDEN_LEN": "hiddenSize / 4" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "SkipLayerNormalization.Vec4", |
| "shader": "norm-skip-row-vec4.wgsl.jinja", |
| "derive": { |
| "simplified": false, |
| "hasBias": "\"bias\" == \"bias\"", |
| "hasBeta": true, |
| "writeResidualSum": false, |
| "usesF16Spec": false, |
| "hidden": "hiddenSize", |
| "hiddenVec": "hiddenSize / 4", |
| "wg": "skipWgVec4", |
| "vecType": "\"vec4<f32>\"", |
| "useSubgroups": "hasSubgroups" |
| }, |
| "bindings": ["input", "skip", "gamma", "beta", "bias", "output", "params_2"], |
| "dispatch": { "x": "min(rowCount, 65535)", "y": "ceilDiv(rowCount, 65535)", "z": 1 }, |
| "subgroupCollectivesWidth": "portable" |
| } |
| ] |
| }, |
| { |
| "id": "beta_bias_output_only_row", |
| "priority": 10, |
| "when": ["f32_beta_bias_output_only_contract", "normResourcesFit", "rowDispatchFits"], |
| "derive": { |
| "simplified": false, |
| "useSubgroups": "hasSubgroups", |
| "hasBeta": true, |
| "writeResidualSum": false, |
| "hasBias": "\"bias\" == \"bias\"", |
| "scalar": "\"f32\"", |
| "workgroupSize": "skipWg", |
| "HIDDEN_LEN": "hiddenSize" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "SkipLayerNormalization.Row", |
| "shader": "norm-skip-row.wgsl.jinja", |
| "bindings": ["input_2", "skip_2", "bias_2", "gamma_2", "beta_2", "output_2", "params_3"], |
| "dispatch": { "x": "min(rowCount, 65535)", "y": "ceilDiv(rowCount, 65535)", "z": 1 }, |
| "subgroupCollectivesWidth": "portable" |
| } |
| ] |
| }, |
| { |
| "id": "beta_no_bias_output_only_vec4_f16", |
| "priority": 20, |
| "when": ["f16_beta_no_bias_output_only_contract", "vec4Aligned", "hasSubgroups or \"no_bias\" == \"bias\" or portableWideExecution", "normResourcesFit", "rowDispatchFits"], |
| "derive": { |
| "scalar": "\"f16\"", |
| "vectorScalar": "\"vec4<f16>\"", |
| "hasBias": "\"no_bias\" == \"bias\"", |
| "workgroupSize": "skipWg", |
| "HIDDEN_LEN": "hiddenSize / 4" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "SkipLayerNormalization.Vec4.F16", |
| "shader": "norm-skip-row-vec4.wgsl.jinja", |
| "derive": { |
| "simplified": false, |
| "hasBias": "\"no_bias\" == \"bias\"", |
| "hasBeta": true, |
| "writeResidualSum": false, |
| "usesF16Spec": true, |
| "hidden": "hiddenSize", |
| "hiddenVec": "hiddenSize / 4", |
| "wg": "skipWgVec4", |
| "vecType": "\"vec4<f16>\"", |
| "useSubgroups": "hasSubgroups" |
| }, |
| "bindings": ["input", "skip", "gamma", "beta", "output", "params_2"], |
| "dispatch": { "x": "min(rowCount, 65535)", "y": "ceilDiv(rowCount, 65535)", "z": 1 }, |
| "subgroupCollectivesWidth": "portable" |
| } |
| ] |
| }, |
| { |
| "id": "beta_no_bias_output_only_row_f16", |
| "priority": 10, |
| "when": ["f16_beta_no_bias_output_only_contract", "normResourcesFit", "rowDispatchFits"], |
| "derive": { |
| "simplified": false, |
| "useSubgroups": "hasSubgroups", |
| "hasBeta": true, |
| "writeResidualSum": false, |
| "hasBias": "\"no_bias\" == \"bias\"", |
| "scalar": "\"f16\"", |
| "usesF16": true, |
| "workgroupSize": "skipWg", |
| "HIDDEN_LEN": "hiddenSize" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "SkipLayerNormalization.Row.F16", |
| "shader": "norm-skip-row.wgsl.jinja", |
| "bindings": ["input_2", "skip_2", "gamma_2", "beta_2", "output_2", "params_3"], |
| "dispatch": { "x": "min(rowCount, 65535)", "y": "ceilDiv(rowCount, 65535)", "z": 1 }, |
| "subgroupCollectivesWidth": "portable" |
| } |
| ] |
| }, |
| { |
| "id": "beta_bias_output_only_vec4_f16", |
| "priority": 20, |
| "when": ["f16_beta_bias_output_only_contract", "vec4Aligned", "hasSubgroups or \"bias\" == \"bias\" or portableWideExecution", "normResourcesFit", "rowDispatchFits"], |
| "derive": { |
| "scalar": "\"f16\"", |
| "vectorScalar": "\"vec4<f16>\"", |
| "hasBias": "\"bias\" == \"bias\"", |
| "workgroupSize": "skipWg", |
| "HIDDEN_LEN": "hiddenSize / 4" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "SkipLayerNormalization.Vec4.F16", |
| "shader": "norm-skip-row-vec4.wgsl.jinja", |
| "derive": { |
| "simplified": false, |
| "hasBias": "\"bias\" == \"bias\"", |
| "hasBeta": true, |
| "writeResidualSum": false, |
| "usesF16Spec": true, |
| "hidden": "hiddenSize", |
| "hiddenVec": "hiddenSize / 4", |
| "wg": "skipWgVec4", |
| "vecType": "\"vec4<f16>\"", |
| "useSubgroups": "hasSubgroups" |
| }, |
| "bindings": ["input", "skip", "gamma", "beta", "bias", "output", "params_2"], |
| "dispatch": { "x": "min(rowCount, 65535)", "y": "ceilDiv(rowCount, 65535)", "z": 1 }, |
| "subgroupCollectivesWidth": "portable" |
| } |
| ] |
| }, |
| { |
| "id": "beta_bias_output_only_row_f16", |
| "priority": 10, |
| "when": ["f16_beta_bias_output_only_contract", "normResourcesFit", "rowDispatchFits"], |
| "derive": { |
| "simplified": false, |
| "useSubgroups": "hasSubgroups", |
| "hasBeta": true, |
| "writeResidualSum": false, |
| "hasBias": "\"bias\" == \"bias\"", |
| "scalar": "\"f16\"", |
| "usesF16": true, |
| "workgroupSize": "skipWg", |
| "HIDDEN_LEN": "hiddenSize" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "SkipLayerNormalization.Row.F16", |
| "shader": "norm-skip-row.wgsl.jinja", |
| "bindings": ["input_2", "skip_2", "bias_2", "gamma_2", "beta_2", "output_2", "params_3"], |
| "dispatch": { "x": "min(rowCount, 65535)", "y": "ceilDiv(rowCount, 65535)", "z": 1 }, |
| "subgroupCollectivesWidth": "portable" |
| } |
| ] |
| } |
| ] |
| } |
|
|
