Xenova's picture
Xenova HF Staff
sync 91d990483a17
f1a8138 verified
Raw
History Blame
22 kB
{
"domain": "com.microsoft",
"name": "SkipSimplifiedLayerNormalization",
"sinceVersion": 1,
"inputs": {
"inputT": { "onnx": "input", "dtype": "T" },
"skipT": { "onnx": "skip", "dtype": "T" },
"gammaT": { "onnx": "gamma", "dtype": "T", "rank": 1 },
"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))))",
"rowDispatchFits": "rowCount <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) * min(device.limits.maxComputeWorkgroupsPerDimension, 65535)",
"normResourcesFit": "skipWg * 8 <= device.limits.maxComputeWorkgroupStorageSize and skipWgVec4 * 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",
"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",
"hasSubgroups": "device.features.has(\"subgroups\") and device.wgslLanguageFeatures.has(\"subgroup_id\")",
"hasF16": "device.features.has(\"shader-f16\")",
"no_bias_contract": "not present.biasT",
"f32_bias_contract": "false if not present.biasT else (ranks.biasT == 1 and tensorDtypes.biasT == \"float32\" and dim(shapes.biasT, 0) == hiddenSize)",
"f16_bias_contract": "false if not present.biasT else (ranks.biasT == 1 and tensorDtypes.biasT == \"float16\" and dim(shapes.biasT, 0) == hiddenSize)",
"f32_no_bias_residual_contract": "coreContract and residualOutputContract and f32ResidualDtypes and no_bias_contract",
"f32_bias_residual_contract": "coreContract and residualOutputContract and f32ResidualDtypes and f32_bias_contract",
"f16_no_bias_residual_contract": "hasF16 and coreContract and residualOutputContract and f16ResidualDtypes and no_bias_contract",
"f16_bias_residual_contract": "hasF16 and coreContract and residualOutputContract and f16ResidualDtypes and f16_bias_contract",
"f32_no_bias_output_contract": "coreContract and outputOnlyContract and f32MainDtypes and no_bias_contract",
"f32_bias_output_contract": "coreContract and outputOnlyContract and f32MainDtypes and f32_bias_contract",
"f16_no_bias_output_contract": "hasF16 and coreContract and outputOnlyContract and f16MainDtypes and no_bias_contract",
"f16_bias_output_contract": "hasF16 and coreContract and outputOnlyContract and f16MainDtypes and f16_bias_contract"
},
"when": ["normResourcesFit", "rowDispatchFits"],
"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" },
"output": { "arg": "outputT", "buffer": "storage", "elementType": "$vectorScalar" },
"input_skip_bias_sum": { "arg": "residualT", "buffer": "storage", "elementType": "$vectorScalar" },
"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" }
]
},
"bias": { "arg": "biasT", "buffer": "read-only-storage", "elementType": "$vectorScalar", "length": "$HIDDEN_LEN" },
"input_2": { "arg": "inputT", "name": "input", "buffer": "read-only-storage", "elementType": "$scalar" },
"skip_2": { "arg": "skipT", "name": "skip", "buffer": "read-only-storage", "elementType": "$scalar" },
"gamma_2": {
"arg": "gammaT",
"name": "gamma",
"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_2": {
"name": "params",
"buffer": "uniform",
"struct": [
{ "name": "rows", "type": "u32", "value": "rowCount" },
{ "name": "epsilon", "type": "f32", "value": "attrs.epsilon" }
]
},
"bias_2": {
"arg": "biasT",
"name": "bias",
"buffer": "read-only-storage",
"elementType": "$scalar",
"length": "$HIDDEN_LEN"
}
},
"variants": [
{
"id": "no_bias_vec4_f16",
"priority": 21,
"when": ["f16_no_bias_residual_contract", "vec4Aligned"],
"derive": {
"scalar": "\"f16\"",
"vectorScalar": "\"vec4<f16>\"",
"hasBias": "\"no_bias\" == \"bias\"",
"HIDDEN_LEN": "hiddenSize / 4"
},
"passes": [
{
"id": "main",
"name": "SkipSimplifiedLayerNormalization.Vec4",
"shader": "norm-skip-row-vec4.wgsl.jinja",
"derive": {
"simplified": true,
"hasBias": "\"no_bias\" == \"bias\"",
"hasBeta": false,
"writeResidualSum": true,
"usesF16Spec": true,
"hidden": "hiddenSize",
"hiddenVec": "hiddenSize / 4",
"wg": "skipWgVec4",
"vecType": "\"vec4<f16>\"",
"useSubgroups": "hasSubgroups"
},
"bindings": ["input", "skip", "gamma", "output", "input_skip_bias_sum", "params"],
"dispatch": { "x": "min(rowCount, 65535)", "y": "ceilDiv(rowCount, 65535)", "z": 1 },
"subgroupCollectivesWidth": "portable"
}
]
},
{
"id": "no_bias_vec4",
"priority": 20,
"when": ["f32_no_bias_residual_contract", "vec4Aligned"],
"derive": {
"scalar": "\"f32\"",
"vectorScalar": "\"vec4<f32>\"",
"hasBias": "\"no_bias\" == \"bias\"",
"HIDDEN_LEN": "hiddenSize / 4"
},
"passes": [
{
"id": "main",
"name": "SkipSimplifiedLayerNormalization.Vec4",
"shader": "norm-skip-row-vec4.wgsl.jinja",
"derive": {
"simplified": true,
"hasBias": "\"no_bias\" == \"bias\"",
"hasBeta": false,
"writeResidualSum": true,
"usesF16Spec": false,
"hidden": "hiddenSize",
"hiddenVec": "hiddenSize / 4",
"wg": "skipWgVec4",
"vecType": "\"vec4<f32>\"",
"useSubgroups": "hasSubgroups"
},
"bindings": ["input", "skip", "gamma", "output", "input_skip_bias_sum", "params"],
"dispatch": { "x": "min(rowCount, 65535)", "y": "ceilDiv(rowCount, 65535)", "z": 1 },
"subgroupCollectivesWidth": "portable"
}
]
},
{
"id": "no_bias_output_only_vec4",
"priority": 20,
"when": ["f32_no_bias_output_contract", "vec4Aligned"],
"derive": {
"scalar": "\"f32\"",
"vectorScalar": "\"vec4<f32>\"",
"hasBias": "\"no_bias\" == \"bias\"",
"HIDDEN_LEN": "hiddenSize / 4"
},
"passes": [
{
"id": "main",
"name": "SkipSimplifiedLayerNormalization.Vec4OutputOnly",
"shader": "norm-skip-row-vec4.wgsl.jinja",
"derive": {
"simplified": true,
"hasBias": "\"no_bias\" == \"bias\"",
"hasBeta": false,
"writeResidualSum": false,
"usesF16Spec": false,
"hidden": "hiddenSize",
"hiddenVec": "hiddenSize / 4",
"wg": "skipWgVec4",
"vecType": "\"vec4<f32>\"",
"useSubgroups": "hasSubgroups"
},
"bindings": ["input", "skip", "gamma", "output", "params"],
"dispatch": { "x": "min(rowCount, 65535)", "y": "ceilDiv(rowCount, 65535)", "z": 1 },
"subgroupCollectivesWidth": "portable"
}
]
},
{
"id": "no_bias_output_only_vec4_f16",
"priority": 21,
"when": ["f16_no_bias_output_contract", "vec4Aligned"],
"derive": {
"scalar": "\"f16\"",
"vectorScalar": "\"vec4<f16>\"",
"hasBias": "\"no_bias\" == \"bias\"",
"HIDDEN_LEN": "hiddenSize / 4"
},
"passes": [
{
"id": "main",
"name": "SkipSimplifiedLayerNormalization.Vec4OutputOnly",
"shader": "norm-skip-row-vec4.wgsl.jinja",
"derive": {
"simplified": true,
"hasBias": "\"no_bias\" == \"bias\"",
"hasBeta": false,
"writeResidualSum": false,
"usesF16Spec": true,
"hidden": "hiddenSize",
"hiddenVec": "hiddenSize / 4",
"wg": "skipWgVec4",
"vecType": "\"vec4<f16>\"",
"useSubgroups": "hasSubgroups"
},
"bindings": ["input", "skip", "gamma", "output", "params"],
"dispatch": { "x": "min(rowCount, 65535)", "y": "ceilDiv(rowCount, 65535)", "z": 1 },
"subgroupCollectivesWidth": "portable"
}
]
},
{
"id": "no_bias",
"priority": 0,
"when": ["f32_no_bias_residual_contract"],
"derive": {
"simplified": true,
"useSubgroups": false,
"hasBeta": false,
"writeResidualSum": true,
"hasBias": "\"no_bias\" == \"bias\"",
"scalar": "\"f32\"",
"workgroupSize": "skipWg",
"HIDDEN_LEN": "hiddenSize"
},
"passes": [
{
"id": "main",
"name": "SkipSimplifiedLayerNormalization",
"shader": "norm-skip-row.wgsl.jinja",
"bindings": ["input_2", "skip_2", "gamma_2", "output_2", "input_skip_bias_sum_2", "params_2"],
"dispatch": { "x": "min(rowCount, 65535)", "y": "ceilDiv(rowCount, 65535)", "z": 1 }
}
]
},
{
"id": "no_bias_f16",
"priority": 0,
"when": ["f16_no_bias_residual_contract"],
"requires": { "features": ["shader-f16"] },
"derive": {
"simplified": true,
"useSubgroups": false,
"hasBeta": false,
"writeResidualSum": true,
"hasBias": "\"no_bias\" == \"bias\"",
"scalar": "\"f16\"",
"usesF16": true,
"workgroupSize": "skipWg",
"HIDDEN_LEN": "hiddenSize"
},
"passes": [
{
"id": "main",
"name": "SkipSimplifiedLayerNormalization",
"shader": "norm-skip-row.wgsl.jinja",
"bindings": ["input_2", "skip_2", "gamma_2", "output_2", "input_skip_bias_sum_2", "params_2"],
"dispatch": { "x": "min(rowCount, 65535)", "y": "ceilDiv(rowCount, 65535)", "z": 1 }
}
]
},
{
"id": "no_bias_output_only_f16",
"priority": 0,
"when": ["f16_no_bias_output_contract"],
"requires": { "features": ["shader-f16"] },
"derive": {
"simplified": true,
"useSubgroups": false,
"hasBeta": false,
"writeResidualSum": false,
"hasBias": "\"no_bias\" == \"bias\"",
"scalar": "\"f16\"",
"usesF16": true,
"workgroupSize": "skipWg",
"HIDDEN_LEN": "hiddenSize"
},
"passes": [
{
"id": "main",
"name": "SkipSimplifiedLayerNormalization.OutputOnly",
"shader": "norm-skip-row.wgsl.jinja",
"bindings": ["input_2", "skip_2", "gamma_2", "output_2", "params_2"],
"dispatch": { "x": "min(rowCount, 65535)", "y": "ceilDiv(rowCount, 65535)", "z": 1 }
}
]
},
{
"id": "no_bias_output_only",
"priority": 0,
"when": ["f32_no_bias_output_contract"],
"derive": {
"simplified": true,
"useSubgroups": false,
"hasBeta": false,
"writeResidualSum": false,
"hasBias": "\"no_bias\" == \"bias\"",
"scalar": "\"f32\"",
"workgroupSize": "skipWg",
"HIDDEN_LEN": "hiddenSize"
},
"passes": [
{
"id": "main",
"name": "SkipSimplifiedLayerNormalization.OutputOnly",
"shader": "norm-skip-row.wgsl.jinja",
"bindings": ["input_2", "skip_2", "gamma_2", "output_2", "params_2"],
"dispatch": { "x": "min(rowCount, 65535)", "y": "ceilDiv(rowCount, 65535)", "z": 1 }
}
]
},
{
"id": "bias_vec4_f16",
"priority": 21,
"when": ["f16_bias_residual_contract", "vec4Aligned"],
"derive": {
"scalar": "\"f16\"",
"vectorScalar": "\"vec4<f16>\"",
"hasBias": "\"bias\" == \"bias\"",
"HIDDEN_LEN": "hiddenSize / 4"
},
"passes": [
{
"id": "main",
"name": "SkipSimplifiedLayerNormalization.Vec4",
"shader": "norm-skip-row-vec4.wgsl.jinja",
"derive": {
"simplified": true,
"hasBias": "\"bias\" == \"bias\"",
"hasBeta": false,
"writeResidualSum": true,
"usesF16Spec": true,
"hidden": "hiddenSize",
"hiddenVec": "hiddenSize / 4",
"wg": "skipWgVec4",
"vecType": "\"vec4<f16>\"",
"useSubgroups": "hasSubgroups"
},
"bindings": ["input", "skip", "gamma", "bias", "output", "input_skip_bias_sum", "params"],
"dispatch": { "x": "min(rowCount, 65535)", "y": "ceilDiv(rowCount, 65535)", "z": 1 },
"subgroupCollectivesWidth": "portable"
}
]
},
{
"id": "bias_vec4",
"priority": 20,
"when": ["f32_bias_residual_contract", "vec4Aligned"],
"derive": {
"scalar": "\"f32\"",
"vectorScalar": "\"vec4<f32>\"",
"hasBias": "\"bias\" == \"bias\"",
"HIDDEN_LEN": "hiddenSize / 4"
},
"passes": [
{
"id": "main",
"name": "SkipSimplifiedLayerNormalization.Vec4",
"shader": "norm-skip-row-vec4.wgsl.jinja",
"derive": {
"simplified": true,
"hasBias": "\"bias\" == \"bias\"",
"hasBeta": false,
"writeResidualSum": true,
"usesF16Spec": false,
"hidden": "hiddenSize",
"hiddenVec": "hiddenSize / 4",
"wg": "skipWgVec4",
"vecType": "\"vec4<f32>\"",
"useSubgroups": "hasSubgroups"
},
"bindings": ["input", "skip", "gamma", "bias", "output", "input_skip_bias_sum", "params"],
"dispatch": { "x": "min(rowCount, 65535)", "y": "ceilDiv(rowCount, 65535)", "z": 1 },
"subgroupCollectivesWidth": "portable"
}
]
},
{
"id": "bias_output_only_vec4",
"priority": 20,
"when": ["f32_bias_output_contract", "vec4Aligned"],
"derive": {
"scalar": "\"f32\"",
"vectorScalar": "\"vec4<f32>\"",
"hasBias": "\"bias\" == \"bias\"",
"HIDDEN_LEN": "hiddenSize / 4"
},
"passes": [
{
"id": "main",
"name": "SkipSimplifiedLayerNormalization.Vec4OutputOnly",
"shader": "norm-skip-row-vec4.wgsl.jinja",
"derive": {
"simplified": true,
"hasBias": "\"bias\" == \"bias\"",
"hasBeta": false,
"writeResidualSum": false,
"usesF16Spec": false,
"hidden": "hiddenSize",
"hiddenVec": "hiddenSize / 4",
"wg": "skipWgVec4",
"vecType": "\"vec4<f32>\"",
"useSubgroups": "hasSubgroups"
},
"bindings": ["input", "skip", "gamma", "bias", "output", "params"],
"dispatch": { "x": "min(rowCount, 65535)", "y": "ceilDiv(rowCount, 65535)", "z": 1 },
"subgroupCollectivesWidth": "portable"
}
]
},
{
"id": "bias_output_only_vec4_f16",
"priority": 21,
"when": ["f16_bias_output_contract", "vec4Aligned"],
"derive": {
"scalar": "\"f16\"",
"vectorScalar": "\"vec4<f16>\"",
"hasBias": "\"bias\" == \"bias\"",
"HIDDEN_LEN": "hiddenSize / 4"
},
"passes": [
{
"id": "main",
"name": "SkipSimplifiedLayerNormalization.Vec4OutputOnly",
"shader": "norm-skip-row-vec4.wgsl.jinja",
"derive": {
"simplified": true,
"hasBias": "\"bias\" == \"bias\"",
"hasBeta": false,
"writeResidualSum": false,
"usesF16Spec": true,
"hidden": "hiddenSize",
"hiddenVec": "hiddenSize / 4",
"wg": "skipWgVec4",
"vecType": "\"vec4<f16>\"",
"useSubgroups": "hasSubgroups"
},
"bindings": ["input", "skip", "gamma", "bias", "output", "params"],
"dispatch": { "x": "min(rowCount, 65535)", "y": "ceilDiv(rowCount, 65535)", "z": 1 },
"subgroupCollectivesWidth": "portable"
}
]
},
{
"id": "bias",
"priority": 0,
"when": ["f32_bias_residual_contract"],
"derive": {
"simplified": true,
"useSubgroups": false,
"hasBeta": false,
"writeResidualSum": true,
"hasBias": "\"bias\" == \"bias\"",
"scalar": "\"f32\"",
"workgroupSize": "skipWg",
"HIDDEN_LEN": "hiddenSize"
},
"passes": [
{
"id": "main",
"name": "SkipSimplifiedLayerNormalization",
"shader": "norm-skip-row.wgsl.jinja",
"bindings": ["input_2", "skip_2", "gamma_2", "bias_2", "output_2", "input_skip_bias_sum_2", "params_2"],
"dispatch": { "x": "min(rowCount, 65535)", "y": "ceilDiv(rowCount, 65535)", "z": 1 }
}
]
},
{
"id": "bias_f16",
"priority": 0,
"when": ["f16_bias_residual_contract"],
"requires": { "features": ["shader-f16"] },
"derive": {
"simplified": true,
"useSubgroups": false,
"hasBeta": false,
"writeResidualSum": true,
"hasBias": "\"bias\" == \"bias\"",
"scalar": "\"f16\"",
"usesF16": true,
"workgroupSize": "skipWg",
"HIDDEN_LEN": "hiddenSize"
},
"passes": [
{
"id": "main",
"name": "SkipSimplifiedLayerNormalization",
"shader": "norm-skip-row.wgsl.jinja",
"bindings": ["input_2", "skip_2", "gamma_2", "bias_2", "output_2", "input_skip_bias_sum_2", "params_2"],
"dispatch": { "x": "min(rowCount, 65535)", "y": "ceilDiv(rowCount, 65535)", "z": 1 }
}
]
},
{
"id": "bias_output_only_f16",
"priority": 0,
"when": ["f16_bias_output_contract"],
"requires": { "features": ["shader-f16"] },
"derive": {
"simplified": true,
"useSubgroups": false,
"hasBeta": false,
"writeResidualSum": false,
"hasBias": "\"bias\" == \"bias\"",
"scalar": "\"f16\"",
"usesF16": true,
"workgroupSize": "skipWg",
"HIDDEN_LEN": "hiddenSize"
},
"passes": [
{
"id": "main",
"name": "SkipSimplifiedLayerNormalization.OutputOnly",
"shader": "norm-skip-row.wgsl.jinja",
"bindings": ["input_2", "skip_2", "gamma_2", "bias_2", "output_2", "params_2"],
"dispatch": { "x": "min(rowCount, 65535)", "y": "ceilDiv(rowCount, 65535)", "z": 1 }
}
]
},
{
"id": "bias_output_only",
"priority": 0,
"when": ["f32_bias_output_contract"],
"derive": {
"simplified": true,
"useSubgroups": false,
"hasBeta": false,
"writeResidualSum": false,
"hasBias": "\"bias\" == \"bias\"",
"scalar": "\"f32\"",
"workgroupSize": "skipWg",
"HIDDEN_LEN": "hiddenSize"
},
"passes": [
{
"id": "main",
"name": "SkipSimplifiedLayerNormalization.OutputOnly",
"shader": "norm-skip-row.wgsl.jinja",
"bindings": ["input_2", "skip_2", "gamma_2", "bias_2", "output_2", "params_2"],
"dispatch": { "x": "min(rowCount, 65535)", "y": "ceilDiv(rowCount, 65535)", "z": 1 }
}
]
}
]
}