File size: 5,496 Bytes
b5d19c4 d27ea5a b5d19c4 d27ea5a b5d19c4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 | {
"domain": "com.microsoft",
"name": "GatedAdd",
"sinceVersion": 1,
"description": "Adds `Y`, scaled by a per-row `gate`, to `X`: `output = X + round_to_T(Y * gate)`. `X` and `Y` have shape `(..., C)`; `gate` has the same rank with a trailing dimension of 1, so one value covers each row of `C` channels. Rounding the product to `T` before the addition preserves the semantics of a separate `Mul` followed by `Add`. Bfloat16 is not implemented.",
"inputs": [
{
"role": "X",
"dtype": "T",
"description": "Unscaled input with shape `(..., C)`. Any rank of at least 1 is accepted; only the trailing channel axis is distinguished."
},
{ "role": "Y", "dtype": "T", "description": "Input scaled by the gate, with the same shape as `X`." },
{
"role": "gate",
"dtype": "T",
"description": "Per-row gate with shape `(..., 1)`: the same rank and leading dimensions as `X`, with a trailing dimension of 1 that broadcasts over the `C` channels."
}
],
"outputs": [
{
"role": "output",
"dtype": "T",
"rank": "ranks.X",
"shape": "shapes.X",
"description": "Gated sum `X + round_to_T(Y * gate)`, with the same shape as `X`."
}
],
"typeConstraints": { "T": ["float32", "float16"] },
"tunables": { "WORKGROUP_SIZE": 256 },
"args": {
"X": { "kind": "tensor", "semantic": "X", "role": "input" },
"Y": { "kind": "tensor", "semantic": "Y", "role": "input" },
"gate": { "kind": "tensor", "semantic": "gate", "role": "input" },
"output": { "kind": "tensor", "semantic": "output", "role": "output" }
},
"derive": {
"channels": "dim(shapes.X, ranks.X - 1)",
"gateContract": "ranks.X >= 1 and channels > 0 and ranks.Y == ranks.X and ranks.gate == ranks.X and sameShape(shapes.Y, shapes.X) and sameShape(shapes.output, shapes.X) and dim(shapes.gate, ranks.gate - 1) == 1 and sameShape(prefix(shapes.gate, ranks.gate - 1), prefix(shapes.X, ranks.X - 1)) and f16Ok(dtypes.T)",
"vec4Rows": "channels % 4 == 0 and numel(shapes.X) % 4 == 0"
},
"constants": { "scalar": "dtypes.T", "usesF16": "dtypes.T == \"f16\"", "hidden": "channels if channels > 0 else 1" },
"variants": [
{
"id": "vec4",
"priority": 30,
"when": ["gateContract", "vec4Rows", "numel(shapes.X) > 0"],
"constants": { "vec4": true, "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"" },
"passes": [
{
"id": "main",
"name": "GatedAdd.vec4",
"shader": "gated-add.wgsl.jinja",
"bindings": [
{
"name": "x",
"arg": "X",
"semantic": "X",
"buffer": { "type": "read-only-storage" },
"elementType": "$vectorScalar"
},
{
"name": "y",
"arg": "Y",
"semantic": "Y",
"buffer": { "type": "read-only-storage" },
"elementType": "$vectorScalar"
},
{
"name": "gate",
"arg": "gate",
"semantic": "gate",
"buffer": { "type": "read-only-storage" },
"elementType": "$scalar"
},
{
"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.X) / 4" }]
}
}
],
"dispatch": { "threads": "numel(shapes.X) / 4", "workgroupSize": "tunables.WORKGROUP_SIZE" }
}
]
},
{
"id": "scalar",
"priority": 0,
"when": ["gateContract"],
"constants": { "vec4": false },
"passes": [
{
"id": "main",
"name": "GatedAdd.scalar",
"shader": "gated-add.wgsl.jinja",
"source": { "inputs": { "itemsPerInvocation": 4 } },
"bindings": [
{
"name": "x",
"arg": "X",
"semantic": "X",
"buffer": { "type": "read-only-storage" },
"elementType": "$scalar"
},
{
"name": "y",
"arg": "Y",
"semantic": "Y",
"buffer": { "type": "read-only-storage" },
"elementType": "$scalar"
},
{
"name": "gate",
"arg": "gate",
"semantic": "gate",
"buffer": { "type": "read-only-storage" },
"elementType": "$scalar"
},
{
"name": "output",
"arg": "output",
"semantic": "output",
"buffer": { "type": "storage" },
"elementType": "$scalar"
},
{
"name": "params",
"semantic": "kernel.params",
"buffer": { "type": "uniform" },
"struct": { "name": "Params", "fields": [{ "name": "count", "type": "u32", "value": "numel(shapes.X)" }] }
}
],
"dispatch": { "threads": "ceilDiv(numel(shapes.X), 4)", "workgroupSize": "tunables.WORKGROUP_SIZE" }
}
]
}
]
}
|