| { |
| "domain": "com.microsoft", |
| "name": "CausalConvWithState", |
| "sinceVersion": 1, |
| "description": "Legacy Microsoft contrib form of stateful 1-D causal depthwise convolution. Each channel uses its own `(channels, 1, kernel)` weight over current and past positions, with optional activation and `past_state`/`present_state` tensors for incremental decoding. The contrib-only `state_window` attribute may retain several rollback states. This inference implementation preserves the existing contrib ABI with `ndim = 1`, float16 or float32 tensors, and float32 accumulation; spatial ranks 2 and 3 and bfloat16 are not implemented.", |
| "inputs": [ |
| { |
| "role": "input", |
| "dtype": "T", |
| "rank": 3, |
| "description": "Channels-first input tensor with shape `(batch_size, channels, sequence_length)` for the supported 1-D mode." |
| }, |
| { |
| "role": "weight", |
| "dtype": "T", |
| "rank": 3, |
| "description": "Depthwise convolution kernel with shape `(channels, 1, kernel_size)` for the supported 1-D mode." |
| }, |
| { |
| "role": "bias", |
| "dtype": "T", |
| "rank": 1, |
| "optional": true, |
| "description": "Optional per-channel bias with shape `(channels,)`." |
| }, |
| { |
| "role": "past_state", |
| "dtype": "T", |
| "rank": "3 if attrs.state_window == 0 else 4", |
| "optional": true, |
| "description": "Carry state from the previous step; shape `(batch_size, channels, k_1 - 1)`, or `(W, batch_size, channels, k_1 - 1)` when `state_window = W > 0`, in which case only slot `W - 1` is read. If absent, the left-side padding is zero." |
| } |
| ], |
| "outputs": [ |
| { |
| "role": "output", |
| "dtype": "T", |
| "rank": 3, |
| "shape": "shapes.input", |
| "description": "Convolution output with the same shape as `input`." |
| }, |
| { |
| "role": "present_state", |
| "dtype": "T", |
| "rank": "3 if attrs.state_window == 0 else 4", |
| "shape": "[dim(shapes.input, 0), dim(shapes.input, 1), dim(shapes.weight, 2) - 1] if attrs.state_window == 0 else [attrs.state_window, dim(shapes.input, 0), dim(shapes.input, 1), dim(shapes.weight, 2) - 1]", |
| "description": "Updated carry state; shape `(batch_size, channels, k_1 - 1)`, or `(W, batch_size, channels, k_1 - 1)` when `state_window = W > 0`. Slot `W - 1` holds the last `k - 1` values along the causal axis; slot `j` holds the same for the prefix ending at position `seq_len - W + j`." |
| } |
| ], |
| "attributes": { "activation": "none", "ndim": 1, "state_window": 0 }, |
| "attributeConstraints": { "activation": { "values": ["none", "silu", "swish"] }, "ndim": { "values": [1] } }, |
| "attributeDescriptions": { |
| "activation": "Activation applied after convolution and bias. Defaults to `none`; `swish` is an alias of SiLU.", |
| "ndim": "Number of spatial dimensions. This implementation supports the contrib 1D mode (`ndim = 1`).", |
| "state_window": "Contrib extension selecting the number of rollback state slots to retain, in the range 0 through 8. Defaults to 0." |
| }, |
| "typeConstraints": { "T": ["float32", "float16"] }, |
| "args": { |
| "inputT": { "kind": "tensor", "semantic": "input", "role": "input" }, |
| "weightT": { "kind": "tensor", "semantic": "weight", "role": "input" }, |
| "biasT": { "kind": "tensor", "semantic": "bias", "role": "input", "required": false }, |
| "pastStateT": { "kind": "tensor", "semantic": "past_state", "role": "input", "required": false }, |
| "outputT": { "kind": "tensor", "semantic": "output", "role": "output" }, |
| "presentStateT": { "kind": "tensor", "semantic": "present_state", "role": "output" } |
| }, |
| "tunables": { "workgroupSize": 256, "tiledWorkgroupSize": 128 }, |
| "derive": { |
| "stateWindow": "attrs.state_window", |
| "windowed": "stateWindow > 0", |
| "stateWindowOk": "stateWindow >= 0 and stateWindow <= 8", |
| "kernelSize": "dim(shapes.weightT, ranks.weightT - 1)", |
| "kernelSizePadded": "ceilDiv(kernelSize, 4) * 4", |
| "weightRankOk": "ranks.weightT == 3 and dim(shapes.weightT, 1) == 1", |
| "stateLength": "kernelSize - 1", |
| "stateSlotStride": "dim(shapes.inputT, 0) * dim(shapes.inputT, 1) * stateLength", |
| "windowedLengthOk": "not windowed or dim(shapes.inputT, 2) > 0", |
| "presentStateOk": "(ranks.presentStateT == 3 and dim(shapes.presentStateT, 0) == dim(shapes.inputT, 0) and dim(shapes.presentStateT, 1) == dim(shapes.inputT, 1) and dim(shapes.presentStateT, 2) == stateLength) if not windowed else (ranks.presentStateT == 4 and dim(shapes.presentStateT, 0) == stateWindow and dim(shapes.presentStateT, 1) == dim(shapes.inputT, 0) and dim(shapes.presentStateT, 2) == dim(shapes.inputT, 1) and dim(shapes.presentStateT, 3) == stateLength)", |
| "pastStateShapeOk": "present.pastStateT and ((ranks.pastStateT == 3 and dim(shapes.pastStateT, 0) == dim(shapes.inputT, 0) and dim(shapes.pastStateT, 1) == dim(shapes.inputT, 1) and dim(shapes.pastStateT, 2) == stateLength) if not windowed else (ranks.pastStateT == 4 and dim(shapes.pastStateT, 0) == stateWindow and dim(shapes.pastStateT, 1) == dim(shapes.inputT, 0) and dim(shapes.pastStateT, 2) == dim(shapes.inputT, 1) and dim(shapes.pastStateT, 3) == stateLength))", |
| "commonContract": "ranks.inputT == 3 and weightRankOk and ranks.outputT == 3 and (tensorDtypes.inputT == \"float32\" or tensorDtypes.inputT == \"float16\") and tensorDtypes.weightT == tensorDtypes.inputT and tensorDtypes.outputT == tensorDtypes.inputT and tensorDtypes.presentStateT == tensorDtypes.inputT and f16Ok(dtypes.T) and dim(shapes.inputT, 1) == dim(shapes.weightT, 0) and dim(shapes.outputT, 0) == dim(shapes.inputT, 0) and dim(shapes.outputT, 1) == dim(shapes.inputT, 1) and dim(shapes.outputT, 2) == dim(shapes.inputT, 2) and stateWindowOk and windowedLengthOk and presentStateOk", |
| "zeroStateContract": "commonContract and not present.pastStateT and not present.biasT", |
| "biasNoStateContract": "commonContract and not present.pastStateT and present.biasT and ranks.biasT == 1 and tensorDtypes.biasT == tensorDtypes.inputT and dim(shapes.biasT, 0) == dim(shapes.inputT, 1)", |
| "stateNoBiasContract": "commonContract and present.pastStateT and not present.biasT and tensorDtypes.pastStateT == tensorDtypes.inputT and pastStateShapeOk", |
| "stateBiasContract": "commonContract and present.pastStateT and present.biasT and ranks.biasT == 1 and tensorDtypes.pastStateT == tensorDtypes.inputT and tensorDtypes.biasT == tensorDtypes.inputT and pastStateShapeOk and dim(shapes.biasT, 0) == dim(shapes.inputT, 1)" |
| }, |
| "bindingSets": { |
| "zeroScalar": [ |
| { |
| "name": "input", |
| "arg": "inputT", |
| "semantic": "input", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$inputScalar" |
| }, |
| { |
| "name": "weight", |
| "arg": "weightT", |
| "semantic": "weight", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$inputScalar" |
| }, |
| { |
| "name": "output", |
| "arg": "outputT", |
| "semantic": "output", |
| "buffer": { "type": "storage" }, |
| "elementType": "$outputScalar" |
| }, |
| { |
| "name": "present_state", |
| "arg": "presentStateT", |
| "semantic": "present_state", |
| "buffer": { "type": "storage" }, |
| "elementType": "$outputScalar" |
| }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [ |
| { "name": "batchSize", "type": "u32", "value": "dim(shapes.inputT, 0)" }, |
| { "name": "channels", "type": "u32", "value": "dim(shapes.inputT, 1)" }, |
| { "name": "length", "type": "u32", "value": "dim(shapes.inputT, 2)" }, |
| { "name": "kernelSize", "type": "u32", "value": "kernelSize" }, |
| { "name": "stateWindow", "type": "u32", "value": "stateWindow" }, |
| { "name": "stateSlotStride", "type": "u32", "value": "stateSlotStride" } |
| ] |
| } |
| } |
| ], |
| "zeroVec4": [ |
| { |
| "name": "input", |
| "arg": "inputT", |
| "semantic": "input", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$inputVec4" |
| }, |
| { |
| "name": "weight", |
| "arg": "weightT", |
| "semantic": "weight", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$weightElem" |
| }, |
| { |
| "name": "output", |
| "arg": "outputT", |
| "semantic": "output", |
| "buffer": { "type": "storage" }, |
| "elementType": "$outputVec4" |
| }, |
| { |
| "name": "present_state", |
| "arg": "presentStateT", |
| "semantic": "present_state", |
| "buffer": { "type": "storage" }, |
| "elementType": "$outputScalar" |
| }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [ |
| { "name": "batchSize", "type": "u32", "value": "dim(shapes.inputT, 0)" }, |
| { "name": "channels", "type": "u32", "value": "dim(shapes.inputT, 1)" }, |
| { "name": "length", "type": "u32", "value": "dim(shapes.inputT, 2)" }, |
| { "name": "stateWindow", "type": "u32", "value": "stateWindow" }, |
| { "name": "stateSlotStride", "type": "u32", "value": "stateSlotStride" } |
| ] |
| } |
| } |
| ], |
| "biasNoState": [ |
| { |
| "name": "input", |
| "arg": "inputT", |
| "semantic": "input", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$inputScalar" |
| }, |
| { |
| "name": "weight", |
| "arg": "weightT", |
| "semantic": "weight", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$inputScalar" |
| }, |
| { |
| "name": "bias", |
| "arg": "biasT", |
| "semantic": "bias", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$inputScalar" |
| }, |
| { |
| "name": "output", |
| "arg": "outputT", |
| "semantic": "output", |
| "buffer": { "type": "storage" }, |
| "elementType": "$outputScalar" |
| }, |
| { |
| "name": "present_state", |
| "arg": "presentStateT", |
| "semantic": "present_state", |
| "buffer": { "type": "storage" }, |
| "elementType": "$outputScalar" |
| }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [ |
| { "name": "batchSize", "type": "u32", "value": "dim(shapes.inputT, 0)" }, |
| { "name": "channels", "type": "u32", "value": "dim(shapes.inputT, 1)" }, |
| { "name": "length", "type": "u32", "value": "dim(shapes.inputT, 2)" }, |
| { "name": "kernelSize", "type": "u32", "value": "kernelSize" }, |
| { "name": "stateWindow", "type": "u32", "value": "stateWindow" }, |
| { "name": "stateSlotStride", "type": "u32", "value": "stateSlotStride" } |
| ] |
| } |
| } |
| ], |
| "stateNoBias": [ |
| { |
| "name": "input", |
| "arg": "inputT", |
| "semantic": "input", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$inputScalar" |
| }, |
| { |
| "name": "weight", |
| "arg": "weightT", |
| "semantic": "weight", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$inputScalar" |
| }, |
| { |
| "name": "past_state", |
| "arg": "pastStateT", |
| "semantic": "past_state", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$inputScalar" |
| }, |
| { |
| "name": "output", |
| "arg": "outputT", |
| "semantic": "output", |
| "buffer": { "type": "storage" }, |
| "elementType": "$outputScalar" |
| }, |
| { |
| "name": "present_state", |
| "arg": "presentStateT", |
| "semantic": "present_state", |
| "buffer": { "type": "storage" }, |
| "elementType": "$outputScalar" |
| }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [ |
| { "name": "batchSize", "type": "u32", "value": "dim(shapes.inputT, 0)" }, |
| { "name": "channels", "type": "u32", "value": "dim(shapes.inputT, 1)" }, |
| { "name": "length", "type": "u32", "value": "dim(shapes.inputT, 2)" }, |
| { "name": "kernelSize", "type": "u32", "value": "kernelSize" }, |
| { "name": "stateWindow", "type": "u32", "value": "stateWindow" }, |
| { "name": "stateSlotStride", "type": "u32", "value": "stateSlotStride" } |
| ] |
| } |
| } |
| ], |
| "stateBias": [ |
| { |
| "name": "input", |
| "arg": "inputT", |
| "semantic": "input", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$inputScalar" |
| }, |
| { |
| "name": "weight", |
| "arg": "weightT", |
| "semantic": "weight", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$inputScalar" |
| }, |
| { |
| "name": "bias", |
| "arg": "biasT", |
| "semantic": "bias", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$inputScalar" |
| }, |
| { |
| "name": "past_state", |
| "arg": "pastStateT", |
| "semantic": "past_state", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$inputScalar" |
| }, |
| { |
| "name": "output", |
| "arg": "outputT", |
| "semantic": "output", |
| "buffer": { "type": "storage" }, |
| "elementType": "$outputScalar" |
| }, |
| { |
| "name": "present_state", |
| "arg": "presentStateT", |
| "semantic": "present_state", |
| "buffer": { "type": "storage" }, |
| "elementType": "$outputScalar" |
| }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [ |
| { "name": "batchSize", "type": "u32", "value": "dim(shapes.inputT, 0)" }, |
| { "name": "channels", "type": "u32", "value": "dim(shapes.inputT, 1)" }, |
| { "name": "length", "type": "u32", "value": "dim(shapes.inputT, 2)" }, |
| { "name": "kernelSize", "type": "u32", "value": "kernelSize" }, |
| { "name": "stateWindow", "type": "u32", "value": "stateWindow" }, |
| { "name": "stateSlotStride", "type": "u32", "value": "stateSlotStride" } |
| ] |
| } |
| } |
| ], |
| "zeroScalarIo": [ |
| { |
| "name": "input", |
| "arg": "inputT", |
| "semantic": "input", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$inputScalar" |
| }, |
| { |
| "name": "weight", |
| "arg": "weightT", |
| "semantic": "weight", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$inputScalar" |
| }, |
| { |
| "name": "output", |
| "arg": "outputT", |
| "semantic": "output", |
| "buffer": { "type": "storage" }, |
| "elementType": "$outputScalar" |
| }, |
| { |
| "name": "present_state", |
| "arg": "presentStateT", |
| "semantic": "present_state", |
| "buffer": { "type": "storage" }, |
| "elementType": "$outputScalar" |
| } |
| ], |
| "biasNoStateIo": [ |
| { |
| "name": "input", |
| "arg": "inputT", |
| "semantic": "input", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$inputScalar" |
| }, |
| { |
| "name": "weight", |
| "arg": "weightT", |
| "semantic": "weight", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$inputScalar" |
| }, |
| { |
| "name": "bias", |
| "arg": "biasT", |
| "semantic": "bias", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$inputScalar" |
| }, |
| { |
| "name": "output", |
| "arg": "outputT", |
| "semantic": "output", |
| "buffer": { "type": "storage" }, |
| "elementType": "$outputScalar" |
| }, |
| { |
| "name": "present_state", |
| "arg": "presentStateT", |
| "semantic": "present_state", |
| "buffer": { "type": "storage" }, |
| "elementType": "$outputScalar" |
| } |
| ], |
| "stateNoBiasIo": [ |
| { |
| "name": "input", |
| "arg": "inputT", |
| "semantic": "input", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$inputScalar" |
| }, |
| { |
| "name": "weight", |
| "arg": "weightT", |
| "semantic": "weight", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$inputScalar" |
| }, |
| { |
| "name": "past_state", |
| "arg": "pastStateT", |
| "semantic": "past_state", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$inputScalar" |
| }, |
| { |
| "name": "output", |
| "arg": "outputT", |
| "semantic": "output", |
| "buffer": { "type": "storage" }, |
| "elementType": "$outputScalar" |
| }, |
| { |
| "name": "present_state", |
| "arg": "presentStateT", |
| "semantic": "present_state", |
| "buffer": { "type": "storage" }, |
| "elementType": "$outputScalar" |
| } |
| ], |
| "stateBiasIo": [ |
| { |
| "name": "input", |
| "arg": "inputT", |
| "semantic": "input", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$inputScalar" |
| }, |
| { |
| "name": "weight", |
| "arg": "weightT", |
| "semantic": "weight", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$inputScalar" |
| }, |
| { |
| "name": "bias", |
| "arg": "biasT", |
| "semantic": "bias", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$inputScalar" |
| }, |
| { |
| "name": "past_state", |
| "arg": "pastStateT", |
| "semantic": "past_state", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$inputScalar" |
| }, |
| { |
| "name": "output", |
| "arg": "outputT", |
| "semantic": "output", |
| "buffer": { "type": "storage" }, |
| "elementType": "$outputScalar" |
| }, |
| { |
| "name": "present_state", |
| "arg": "presentStateT", |
| "semantic": "present_state", |
| "buffer": { "type": "storage" }, |
| "elementType": "$outputScalar" |
| } |
| ], |
| "zeroTiled": [ |
| { |
| "name": "input", |
| "arg": "inputT", |
| "semantic": "input", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$inputScalar" |
| }, |
| { |
| "name": "weight", |
| "arg": "weightT", |
| "semantic": "weight", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$inputScalar" |
| }, |
| { |
| "name": "output", |
| "arg": "outputT", |
| "semantic": "output", |
| "buffer": { "type": "storage" }, |
| "elementType": "$outputScalar" |
| }, |
| { |
| "name": "present_state", |
| "arg": "presentStateT", |
| "semantic": "present_state", |
| "buffer": { "type": "storage" }, |
| "elementType": "$outputScalar" |
| }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [ |
| { "name": "batchSize", "type": "u32", "value": "dim(shapes.inputT, 0)" }, |
| { "name": "channels", "type": "u32", "value": "dim(shapes.inputT, 1)" }, |
| { "name": "length", "type": "u32", "value": "dim(shapes.inputT, 2)" }, |
| { "name": "stateWindow", "type": "u32", "value": "stateWindow" }, |
| { "name": "stateSlotStride", "type": "u32", "value": "stateSlotStride" } |
| ] |
| } |
| } |
| ], |
| "biasNoStateTiled": [ |
| { |
| "name": "input", |
| "arg": "inputT", |
| "semantic": "input", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$inputScalar" |
| }, |
| { |
| "name": "weight", |
| "arg": "weightT", |
| "semantic": "weight", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$inputScalar" |
| }, |
| { |
| "name": "bias", |
| "arg": "biasT", |
| "semantic": "bias", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$inputScalar" |
| }, |
| { |
| "name": "output", |
| "arg": "outputT", |
| "semantic": "output", |
| "buffer": { "type": "storage" }, |
| "elementType": "$outputScalar" |
| }, |
| { |
| "name": "present_state", |
| "arg": "presentStateT", |
| "semantic": "present_state", |
| "buffer": { "type": "storage" }, |
| "elementType": "$outputScalar" |
| }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [ |
| { "name": "batchSize", "type": "u32", "value": "dim(shapes.inputT, 0)" }, |
| { "name": "channels", "type": "u32", "value": "dim(shapes.inputT, 1)" }, |
| { "name": "length", "type": "u32", "value": "dim(shapes.inputT, 2)" }, |
| { "name": "stateWindow", "type": "u32", "value": "stateWindow" }, |
| { "name": "stateSlotStride", "type": "u32", "value": "stateSlotStride" } |
| ] |
| } |
| } |
| ], |
| "stateNoBiasTiled": [ |
| { |
| "name": "input", |
| "arg": "inputT", |
| "semantic": "input", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$inputScalar" |
| }, |
| { |
| "name": "weight", |
| "arg": "weightT", |
| "semantic": "weight", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$inputScalar" |
| }, |
| { |
| "name": "past_state", |
| "arg": "pastStateT", |
| "semantic": "past_state", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$inputScalar" |
| }, |
| { |
| "name": "output", |
| "arg": "outputT", |
| "semantic": "output", |
| "buffer": { "type": "storage" }, |
| "elementType": "$outputScalar" |
| }, |
| { |
| "name": "present_state", |
| "arg": "presentStateT", |
| "semantic": "present_state", |
| "buffer": { "type": "storage" }, |
| "elementType": "$outputScalar" |
| }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [ |
| { "name": "batchSize", "type": "u32", "value": "dim(shapes.inputT, 0)" }, |
| { "name": "channels", "type": "u32", "value": "dim(shapes.inputT, 1)" }, |
| { "name": "length", "type": "u32", "value": "dim(shapes.inputT, 2)" }, |
| { "name": "stateWindow", "type": "u32", "value": "stateWindow" }, |
| { "name": "stateSlotStride", "type": "u32", "value": "stateSlotStride" } |
| ] |
| } |
| } |
| ], |
| "stateBiasTiled": [ |
| { |
| "name": "input", |
| "arg": "inputT", |
| "semantic": "input", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$inputScalar" |
| }, |
| { |
| "name": "weight", |
| "arg": "weightT", |
| "semantic": "weight", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$inputScalar" |
| }, |
| { |
| "name": "bias", |
| "arg": "biasT", |
| "semantic": "bias", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$inputScalar" |
| }, |
| { |
| "name": "past_state", |
| "arg": "pastStateT", |
| "semantic": "past_state", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$inputScalar" |
| }, |
| { |
| "name": "output", |
| "arg": "outputT", |
| "semantic": "output", |
| "buffer": { "type": "storage" }, |
| "elementType": "$outputScalar" |
| }, |
| { |
| "name": "present_state", |
| "arg": "presentStateT", |
| "semantic": "present_state", |
| "buffer": { "type": "storage" }, |
| "elementType": "$outputScalar" |
| }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [ |
| { "name": "batchSize", "type": "u32", "value": "dim(shapes.inputT, 0)" }, |
| { "name": "channels", "type": "u32", "value": "dim(shapes.inputT, 1)" }, |
| { "name": "length", "type": "u32", "value": "dim(shapes.inputT, 2)" }, |
| { "name": "stateWindow", "type": "u32", "value": "stateWindow" }, |
| { "name": "stateSlotStride", "type": "u32", "value": "stateSlotStride" } |
| ] |
| } |
| } |
| ], |
| "biasNoStateVec4": [ |
| { |
| "name": "input", |
| "arg": "inputT", |
| "semantic": "input", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$inputVec4" |
| }, |
| { |
| "name": "weight", |
| "arg": "weightT", |
| "semantic": "weight", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$weightElem" |
| }, |
| { |
| "name": "bias", |
| "arg": "biasT", |
| "semantic": "bias", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$inputScalar" |
| }, |
| { |
| "name": "output", |
| "arg": "outputT", |
| "semantic": "output", |
| "buffer": { "type": "storage" }, |
| "elementType": "$outputVec4" |
| }, |
| { |
| "name": "present_state", |
| "arg": "presentStateT", |
| "semantic": "present_state", |
| "buffer": { "type": "storage" }, |
| "elementType": "$outputScalar" |
| }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [ |
| { "name": "batchSize", "type": "u32", "value": "dim(shapes.inputT, 0)" }, |
| { "name": "channels", "type": "u32", "value": "dim(shapes.inputT, 1)" }, |
| { "name": "length", "type": "u32", "value": "dim(shapes.inputT, 2)" }, |
| { "name": "stateWindow", "type": "u32", "value": "stateWindow" }, |
| { "name": "stateSlotStride", "type": "u32", "value": "stateSlotStride" } |
| ] |
| } |
| } |
| ], |
| "stateNoBiasVec4": [ |
| { |
| "name": "input", |
| "arg": "inputT", |
| "semantic": "input", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$inputVec4" |
| }, |
| { |
| "name": "weight", |
| "arg": "weightT", |
| "semantic": "weight", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$weightElem" |
| }, |
| { |
| "name": "past_state", |
| "arg": "pastStateT", |
| "semantic": "past_state", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$inputScalar" |
| }, |
| { |
| "name": "output", |
| "arg": "outputT", |
| "semantic": "output", |
| "buffer": { "type": "storage" }, |
| "elementType": "$outputVec4" |
| }, |
| { |
| "name": "present_state", |
| "arg": "presentStateT", |
| "semantic": "present_state", |
| "buffer": { "type": "storage" }, |
| "elementType": "$outputScalar" |
| }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [ |
| { "name": "batchSize", "type": "u32", "value": "dim(shapes.inputT, 0)" }, |
| { "name": "channels", "type": "u32", "value": "dim(shapes.inputT, 1)" }, |
| { "name": "length", "type": "u32", "value": "dim(shapes.inputT, 2)" }, |
| { "name": "stateWindow", "type": "u32", "value": "stateWindow" }, |
| { "name": "stateSlotStride", "type": "u32", "value": "stateSlotStride" } |
| ] |
| } |
| } |
| ], |
| "stateBiasVec4": [ |
| { |
| "name": "input", |
| "arg": "inputT", |
| "semantic": "input", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$inputVec4" |
| }, |
| { |
| "name": "weight", |
| "arg": "weightT", |
| "semantic": "weight", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$weightElem" |
| }, |
| { |
| "name": "bias", |
| "arg": "biasT", |
| "semantic": "bias", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$inputScalar" |
| }, |
| { |
| "name": "past_state", |
| "arg": "pastStateT", |
| "semantic": "past_state", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$inputScalar" |
| }, |
| { |
| "name": "output", |
| "arg": "outputT", |
| "semantic": "output", |
| "buffer": { "type": "storage" }, |
| "elementType": "$outputVec4" |
| }, |
| { |
| "name": "present_state", |
| "arg": "presentStateT", |
| "semantic": "present_state", |
| "buffer": { "type": "storage" }, |
| "elementType": "$outputScalar" |
| }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [ |
| { "name": "batchSize", "type": "u32", "value": "dim(shapes.inputT, 0)" }, |
| { "name": "channels", "type": "u32", "value": "dim(shapes.inputT, 1)" }, |
| { "name": "length", "type": "u32", "value": "dim(shapes.inputT, 2)" }, |
| { "name": "stateWindow", "type": "u32", "value": "stateWindow" }, |
| { "name": "stateSlotStride", "type": "u32", "value": "stateSlotStride" } |
| ] |
| } |
| } |
| ] |
| }, |
| "variants": [ |
| { |
| "id": "zero_state_vec4", |
| "priority": 20, |
| "when": ["zeroStateContract", "kernelSize >= 2", "kernelSize <= 4", "dim(shapes.inputT, 2) >= 4", "dim(shapes.inputT, 2) % 4 == 0"], |
| "constants": { |
| "useSilu": "attrs.activation == \"silu\" or attrs.activation == \"swish\"", |
| "workgroupSize": 256, |
| "hasStateWindow": "windowed", |
| "usesF16": "tensorDtypes.inputT == \"float16\"", |
| "inputScalar": "dtypes.T", |
| "outputScalar": "dtypes.T", |
| "inputVec4": "\"vec4<f16>\" if tensorDtypes.inputT == \"float16\" else \"vec4<f32>\"", |
| "outputVec4": "\"vec4<f16>\" if tensorDtypes.inputT == \"float16\" else \"vec4<f32>\"", |
| "hasBias": false, |
| "hasState": false, |
| "kernelSize": "kernelSize", |
| "kernelSizePadded": "kernelSizePadded", |
| "weightElem": "(\"vec4<f16>\" if tensorDtypes.inputT == \"float16\" else \"vec4<f32>\") if kernelSize == 4 else dtypes.T" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "CausalConvWithState.Vec4", |
| "source": { |
| "shader": "causal-conv-with-state-vec4.wgsl.jinja", |
| "inputs": { "materializeConvBeforeActivation": false } |
| }, |
| "bindings": "zeroVec4", |
| "dispatch": { |
| "threads": "dim(shapes.outputT, 0) * dim(shapes.outputT, 1) * (dim(shapes.outputT, 2) / 4)", |
| "workgroupSize": "constants.workgroupSize" |
| } |
| } |
| ] |
| }, |
| { |
| "id": "state_bias_vec4", |
| "priority": 20, |
| "when": ["stateBiasContract", "kernelSize >= 2", "kernelSize <= 4", "dim(shapes.inputT, 2) >= 4", "dim(shapes.inputT, 2) % 4 == 0"], |
| "constants": { |
| "useSilu": "attrs.activation == \"silu\" or attrs.activation == \"swish\"", |
| "workgroupSize": 256, |
| "hasStateWindow": "windowed", |
| "usesF16": "tensorDtypes.inputT == \"float16\"", |
| "inputScalar": "dtypes.T", |
| "outputScalar": "dtypes.T", |
| "inputVec4": "\"vec4<f16>\" if tensorDtypes.inputT == \"float16\" else \"vec4<f32>\"", |
| "outputVec4": "\"vec4<f16>\" if tensorDtypes.inputT == \"float16\" else \"vec4<f32>\"", |
| "hasBias": true, |
| "hasState": true, |
| "kernelSize": "kernelSize", |
| "kernelSizePadded": "kernelSizePadded", |
| "weightElem": "(\"vec4<f16>\" if tensorDtypes.inputT == \"float16\" else \"vec4<f32>\") if kernelSize == 4 else dtypes.T" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "CausalConvWithState.Vec4", |
| "source": { |
| "shader": "causal-conv-with-state-vec4.wgsl.jinja", |
| "inputs": { "materializeConvBeforeActivation": false } |
| }, |
| "bindings": "stateBiasVec4", |
| "dispatch": { |
| "threads": "dim(shapes.outputT, 0) * dim(shapes.outputT, 1) * (dim(shapes.outputT, 2) / 4)", |
| "workgroupSize": "constants.workgroupSize" |
| } |
| } |
| ] |
| }, |
| { |
| "id": "bias_no_state_vec4", |
| "priority": 20, |
| "when": ["biasNoStateContract", "kernelSize >= 2", "kernelSize <= 4", "dim(shapes.inputT, 2) >= 4", "dim(shapes.inputT, 2) % 4 == 0"], |
| "constants": { |
| "useSilu": "attrs.activation == \"silu\" or attrs.activation == \"swish\"", |
| "workgroupSize": 256, |
| "hasStateWindow": "windowed", |
| "usesF16": "tensorDtypes.inputT == \"float16\"", |
| "inputScalar": "dtypes.T", |
| "outputScalar": "dtypes.T", |
| "inputVec4": "\"vec4<f16>\" if tensorDtypes.inputT == \"float16\" else \"vec4<f32>\"", |
| "outputVec4": "\"vec4<f16>\" if tensorDtypes.inputT == \"float16\" else \"vec4<f32>\"", |
| "hasBias": true, |
| "hasState": false, |
| "kernelSize": "kernelSize", |
| "kernelSizePadded": "kernelSizePadded", |
| "weightElem": "(\"vec4<f16>\" if tensorDtypes.inputT == \"float16\" else \"vec4<f32>\") if kernelSize == 4 else dtypes.T" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "CausalConvWithState.Vec4", |
| "source": { |
| "shader": "causal-conv-with-state-vec4.wgsl.jinja", |
| "inputs": { "materializeConvBeforeActivation": false } |
| }, |
| "bindings": "biasNoStateVec4", |
| "dispatch": { |
| "threads": "dim(shapes.outputT, 0) * dim(shapes.outputT, 1) * (dim(shapes.outputT, 2) / 4)", |
| "workgroupSize": "constants.workgroupSize" |
| } |
| } |
| ] |
| }, |
| { |
| "id": "state_no_bias_vec4", |
| "priority": 20, |
| "when": ["stateNoBiasContract", "kernelSize >= 2", "kernelSize <= 4", "dim(shapes.inputT, 2) >= 4", "dim(shapes.inputT, 2) % 4 == 0"], |
| "constants": { |
| "useSilu": "attrs.activation == \"silu\" or attrs.activation == \"swish\"", |
| "workgroupSize": 256, |
| "hasStateWindow": "windowed", |
| "usesF16": "tensorDtypes.inputT == \"float16\"", |
| "inputScalar": "dtypes.T", |
| "outputScalar": "dtypes.T", |
| "inputVec4": "\"vec4<f16>\" if tensorDtypes.inputT == \"float16\" else \"vec4<f32>\"", |
| "outputVec4": "\"vec4<f16>\" if tensorDtypes.inputT == \"float16\" else \"vec4<f32>\"", |
| "hasBias": false, |
| "hasState": true, |
| "kernelSize": "kernelSize", |
| "kernelSizePadded": "kernelSizePadded", |
| "weightElem": "(\"vec4<f16>\" if tensorDtypes.inputT == \"float16\" else \"vec4<f32>\") if kernelSize == 4 else dtypes.T" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "CausalConvWithState.Vec4", |
| "source": { |
| "shader": "causal-conv-with-state-vec4.wgsl.jinja", |
| "inputs": { "materializeConvBeforeActivation": false } |
| }, |
| "bindings": "stateNoBiasVec4", |
| "dispatch": { |
| "threads": "dim(shapes.outputT, 0) * dim(shapes.outputT, 1) * (dim(shapes.outputT, 2) / 4)", |
| "workgroupSize": "constants.workgroupSize" |
| } |
| } |
| ] |
| }, |
| { |
| "id": "zero_state_tiled_large_kernel", |
| "priority": 10, |
| "when": ["zeroStateContract", "kernelSize >= 32", "dim(shapes.inputT, 2) >= 256", "dim(shapes.inputT, 2) % 8 == 0", "tunables.tiledWorkgroupSize >= 1", "floor(tunables.tiledWorkgroupSize) == tunables.tiledWorkgroupSize", "tunables.tiledWorkgroupSize <= device.limits.maxComputeInvocationsPerWorkgroup", "tunables.tiledWorkgroupSize <= device.limits.maxComputeWorkgroupSizeX", "(tunables.tiledWorkgroupSize * 8 + 2 * kernelSizePadded - 1) * 4 <= device.limits.maxComputeWorkgroupStorageSize"], |
| "constants": { |
| "hasBias": false, |
| "hasState": false, |
| "useSilu": "attrs.activation == \"silu\" or attrs.activation == \"swish\"", |
| "inputScalar": "dtypes.T", |
| "outputScalar": "dtypes.T", |
| "workgroupSize": "tunables.tiledWorkgroupSize", |
| "tileSize": "tunables.tiledWorkgroupSize * 8", |
| "kernelSize": "kernelSize", |
| "kernelSizePadded": "kernelSizePadded", |
| "inputTileSize": "tunables.tiledWorkgroupSize * 8 + kernelSizePadded - 1", |
| "hasStateWindow": "windowed", |
| "usesF16": "tensorDtypes.inputT == \"float16\"" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "CausalConvWithState.TiledLargeKernel", |
| "source": { |
| "shader": "causal-conv-with-state-tiled.wgsl.jinja", |
| "inputs": { "materializeConvBeforeActivation": false } |
| }, |
| "bindings": "zeroTiled", |
| "dispatch": { |
| "workgroups": "dim(shapes.outputT, 0) * dim(shapes.outputT, 1) * ceilDiv(dim(shapes.outputT, 2), constants.tileSize)" |
| } |
| } |
| ] |
| }, |
| { |
| "id": "state_bias_tiled_large_kernel", |
| "priority": 10, |
| "when": ["stateBiasContract", "kernelSize >= 32", "dim(shapes.inputT, 2) >= 256", "dim(shapes.inputT, 2) % 8 == 0", "tunables.tiledWorkgroupSize >= 1", "floor(tunables.tiledWorkgroupSize) == tunables.tiledWorkgroupSize", "tunables.tiledWorkgroupSize <= device.limits.maxComputeInvocationsPerWorkgroup", "tunables.tiledWorkgroupSize <= device.limits.maxComputeWorkgroupSizeX", "(tunables.tiledWorkgroupSize * 8 + 2 * kernelSizePadded - 1) * 4 <= device.limits.maxComputeWorkgroupStorageSize"], |
| "constants": { |
| "hasBias": true, |
| "hasState": true, |
| "useSilu": "attrs.activation == \"silu\" or attrs.activation == \"swish\"", |
| "inputScalar": "dtypes.T", |
| "outputScalar": "dtypes.T", |
| "workgroupSize": "tunables.tiledWorkgroupSize", |
| "tileSize": "tunables.tiledWorkgroupSize * 8", |
| "kernelSize": "kernelSize", |
| "kernelSizePadded": "kernelSizePadded", |
| "inputTileSize": "tunables.tiledWorkgroupSize * 8 + kernelSizePadded - 1", |
| "hasStateWindow": "windowed", |
| "usesF16": "tensorDtypes.inputT == \"float16\"" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "CausalConvWithState.TiledLargeKernel", |
| "source": { |
| "shader": "causal-conv-with-state-tiled.wgsl.jinja", |
| "inputs": { "materializeConvBeforeActivation": false } |
| }, |
| "bindings": "stateBiasTiled", |
| "dispatch": { |
| "workgroups": "dim(shapes.outputT, 0) * dim(shapes.outputT, 1) * ceilDiv(dim(shapes.outputT, 2), constants.tileSize)" |
| } |
| } |
| ] |
| }, |
| { |
| "id": "bias_no_state_tiled_large_kernel", |
| "priority": 10, |
| "when": ["biasNoStateContract", "kernelSize >= 32", "dim(shapes.inputT, 2) >= 256", "dim(shapes.inputT, 2) % 8 == 0", "tunables.tiledWorkgroupSize >= 1", "floor(tunables.tiledWorkgroupSize) == tunables.tiledWorkgroupSize", "tunables.tiledWorkgroupSize <= device.limits.maxComputeInvocationsPerWorkgroup", "tunables.tiledWorkgroupSize <= device.limits.maxComputeWorkgroupSizeX", "(tunables.tiledWorkgroupSize * 8 + 2 * kernelSizePadded - 1) * 4 <= device.limits.maxComputeWorkgroupStorageSize"], |
| "constants": { |
| "hasBias": true, |
| "hasState": false, |
| "useSilu": "attrs.activation == \"silu\" or attrs.activation == \"swish\"", |
| "inputScalar": "dtypes.T", |
| "outputScalar": "dtypes.T", |
| "workgroupSize": "tunables.tiledWorkgroupSize", |
| "tileSize": "tunables.tiledWorkgroupSize * 8", |
| "kernelSize": "kernelSize", |
| "kernelSizePadded": "kernelSizePadded", |
| "inputTileSize": "tunables.tiledWorkgroupSize * 8 + kernelSizePadded - 1", |
| "hasStateWindow": "windowed", |
| "usesF16": "tensorDtypes.inputT == \"float16\"" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "CausalConvWithState.TiledLargeKernel", |
| "source": { |
| "shader": "causal-conv-with-state-tiled.wgsl.jinja", |
| "inputs": { "materializeConvBeforeActivation": false } |
| }, |
| "bindings": "biasNoStateTiled", |
| "dispatch": { |
| "workgroups": "dim(shapes.outputT, 0) * dim(shapes.outputT, 1) * ceilDiv(dim(shapes.outputT, 2), constants.tileSize)" |
| } |
| } |
| ] |
| }, |
| { |
| "id": "state_no_bias_tiled_large_kernel", |
| "priority": 10, |
| "when": ["stateNoBiasContract", "kernelSize >= 32", "dim(shapes.inputT, 2) >= 256", "dim(shapes.inputT, 2) % 8 == 0", "tunables.tiledWorkgroupSize >= 1", "floor(tunables.tiledWorkgroupSize) == tunables.tiledWorkgroupSize", "tunables.tiledWorkgroupSize <= device.limits.maxComputeInvocationsPerWorkgroup", "tunables.tiledWorkgroupSize <= device.limits.maxComputeWorkgroupSizeX", "(tunables.tiledWorkgroupSize * 8 + 2 * kernelSizePadded - 1) * 4 <= device.limits.maxComputeWorkgroupStorageSize"], |
| "constants": { |
| "hasBias": false, |
| "hasState": true, |
| "useSilu": "attrs.activation == \"silu\" or attrs.activation == \"swish\"", |
| "inputScalar": "dtypes.T", |
| "outputScalar": "dtypes.T", |
| "workgroupSize": "tunables.tiledWorkgroupSize", |
| "tileSize": "tunables.tiledWorkgroupSize * 8", |
| "kernelSize": "kernelSize", |
| "kernelSizePadded": "kernelSizePadded", |
| "inputTileSize": "tunables.tiledWorkgroupSize * 8 + kernelSizePadded - 1", |
| "hasStateWindow": "windowed", |
| "usesF16": "tensorDtypes.inputT == \"float16\"" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "CausalConvWithState.TiledLargeKernel", |
| "source": { |
| "shader": "causal-conv-with-state-tiled.wgsl.jinja", |
| "inputs": { "materializeConvBeforeActivation": false } |
| }, |
| "bindings": "stateNoBiasTiled", |
| "dispatch": { |
| "workgroups": "dim(shapes.outputT, 0) * dim(shapes.outputT, 1) * ceilDiv(dim(shapes.outputT, 2), constants.tileSize)" |
| } |
| } |
| ] |
| }, |
| { |
| "id": "zero_state", |
| "priority": 0, |
| "when": ["zeroStateContract"], |
| "constants": { |
| "hasBias": false, |
| "hasState": false, |
| "useSilu": "attrs.activation == \"silu\" or attrs.activation == \"swish\"", |
| "inputScalar": "dtypes.T", |
| "outputScalar": "dtypes.T", |
| "workgroupSize": "tunables.workgroupSize", |
| "hasStateWindow": "windowed", |
| "usesF16": "tensorDtypes.inputT == \"float16\"" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "CausalConvWithState", |
| "source": { |
| "shader": "causal-conv-with-state.wgsl.jinja", |
| "inputs": { "materializeConvBeforeActivation": false } |
| }, |
| "bindings": "zeroScalar", |
| "dispatch": { |
| "threads": "dim(shapes.outputT, 0) * dim(shapes.outputT, 1) * max(1, dim(shapes.outputT, 2))", |
| "workgroupSize": "constants.workgroupSize" |
| } |
| } |
| ] |
| }, |
| { |
| "id": "state_bias", |
| "priority": 0, |
| "when": ["stateBiasContract"], |
| "constants": { |
| "hasBias": true, |
| "hasState": true, |
| "useSilu": "attrs.activation == \"silu\" or attrs.activation == \"swish\"", |
| "inputScalar": "dtypes.T", |
| "outputScalar": "dtypes.T", |
| "workgroupSize": "tunables.workgroupSize", |
| "hasStateWindow": "windowed", |
| "usesF16": "tensorDtypes.inputT == \"float16\"" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "CausalConvWithState", |
| "source": { |
| "shader": "causal-conv-with-state.wgsl.jinja", |
| "inputs": { "materializeConvBeforeActivation": false } |
| }, |
| "bindings": "stateBias", |
| "dispatch": { |
| "threads": "dim(shapes.outputT, 0) * dim(shapes.outputT, 1) * max(1, dim(shapes.outputT, 2))", |
| "workgroupSize": "constants.workgroupSize" |
| } |
| } |
| ] |
| }, |
| { |
| "id": "bias_no_state", |
| "priority": 0, |
| "when": ["biasNoStateContract"], |
| "constants": { |
| "hasBias": true, |
| "hasState": false, |
| "useSilu": "attrs.activation == \"silu\" or attrs.activation == \"swish\"", |
| "inputScalar": "dtypes.T", |
| "outputScalar": "dtypes.T", |
| "workgroupSize": "tunables.workgroupSize", |
| "hasStateWindow": "windowed", |
| "usesF16": "tensorDtypes.inputT == \"float16\"" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "CausalConvWithState", |
| "source": { |
| "shader": "causal-conv-with-state.wgsl.jinja", |
| "inputs": { "materializeConvBeforeActivation": false } |
| }, |
| "bindings": "biasNoState", |
| "dispatch": { |
| "threads": "dim(shapes.outputT, 0) * dim(shapes.outputT, 1) * max(1, dim(shapes.outputT, 2))", |
| "workgroupSize": "constants.workgroupSize" |
| } |
| } |
| ] |
| }, |
| { |
| "id": "state_no_bias", |
| "priority": 0, |
| "when": ["stateNoBiasContract"], |
| "constants": { |
| "hasBias": false, |
| "hasState": true, |
| "useSilu": "attrs.activation == \"silu\" or attrs.activation == \"swish\"", |
| "inputScalar": "dtypes.T", |
| "outputScalar": "dtypes.T", |
| "workgroupSize": "tunables.workgroupSize", |
| "hasStateWindow": "windowed", |
| "usesF16": "tensorDtypes.inputT == \"float16\"" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "CausalConvWithState", |
| "source": { |
| "shader": "causal-conv-with-state.wgsl.jinja", |
| "inputs": { "materializeConvBeforeActivation": false } |
| }, |
| "bindings": "stateNoBias", |
| "dispatch": { |
| "threads": "dim(shapes.outputT, 0) * dim(shapes.outputT, 1) * max(1, dim(shapes.outputT, 2))", |
| "workgroupSize": "constants.workgroupSize" |
| } |
| } |
| ] |
| } |
| ] |
| } |
|
|