{
  "domain": "com.microsoft",
  "name": "VarlenCausalConvWithState",
  "sinceVersion": 1,
  "inputs": {
    "inputT": { "onnx": "input", "dtype": "T", "rank": 2 },
    "weightT": { "onnx": "weight", "dtype": "T", "rank": 3 },
    "cumulativeSequenceLengthT": { "onnx": "cumulative_sequence_length", "dtype": "M", "rank": 1, "storage": "int32" },
    "biasT": { "onnx": "bias", "dtype": "T", "rank": 1, "optional": true },
    "initialStateT": { "onnx": "initial_state", "dtype": "T", "rank": 3 },
    "captureCountT": { "onnx": "capture_count", "dtype": "M", "rank": 1, "optional": true, "storage": "int32" }
  },
  "outputs": {
    "outputT": { "onnx": "output", "dtype": "T", "rank": "ranks.inputT", "shape": "shapes.inputT" },
    "finalStateT": { "onnx": "final_state", "dtype": "T", "rank": 3, "shape": "[batchSize, channels, stateLength]" },
    "stateUpdateT": {
      "onnx": "state_update",
      "dtype": "T",
      "rank": 3,
      "optional": true,
      "shape": "[batchSize, stateUpdateCapacity, channels]"
    }
  },
  "attributes": { "activation": { "default": "none" }, "state_update_capacity": { "default": 0 } },
  "attributeConstraints": { "activation": { "values": ["none", "silu", "swish"] } },
  "typeConstraints": { "T": ["float32", "float16"], "M": ["int32"] },
  "tunables": {
    "workgroupSize": { "default": 256 },
    "streamWorkgroupSize": { "default": 64 },
    "streamChunk": { "default": 16 }
  },
  "derive": {
    "totalTokens": "dim(shapes.inputT, 0)",
    "channels": "dim(shapes.inputT, 1)",
    "batchSize": "dim(shapes.cumulativeSequenceLengthT, 0) - 1",
    "weightRankOk": "ranks.weightT == 3 and dim(shapes.weightT, 1) == 1",
    "kernelSize": "dim(shapes.weightT, 2)",
    "stateLength": "kernelSize - 1",
    "stateUpdateCapacity": "attrs.state_update_capacity",
    "capturePairOk": "present.captureCountT == (stateUpdateCapacity > 0)",
    "stateShape": "ranks.initialStateT == 3 and dim(shapes.initialStateT, 0) == batchSize and dim(shapes.initialStateT, 1) == channels and dim(shapes.initialStateT, 2) == stateLength",
    "finalStateShape": "ranks.finalStateT == 3 and dim(shapes.finalStateT, 0) == batchSize and dim(shapes.finalStateT, 1) == channels and dim(shapes.finalStateT, 2) == stateLength",
    "captureShape": "not present.captureCountT or (ranks.captureCountT == 1 and dim(shapes.captureCountT, 0) == batchSize and tensorDtypes.captureCountT == \"int32\")",
    "stateUpdateShape": "not present.stateUpdateT or (ranks.stateUpdateT == 3 and dim(shapes.stateUpdateT, 0) == batchSize and dim(shapes.stateUpdateT, 1) == stateUpdateCapacity and dim(shapes.stateUpdateT, 2) == channels and tensorDtypes.stateUpdateT == tensorDtypes.inputT)",
    "commonContract": "ranks.inputT == 2 and weightRankOk and ranks.outputT == 2 and ranks.cumulativeSequenceLengthT == 1 and channels >= 1 and kernelSize >= 1 and batchSize >= 1 and totalTokens >= batchSize and stateUpdateCapacity >= 0 and stateUpdateCapacity <= 8 and floor(stateUpdateCapacity) == stateUpdateCapacity and capturePairOk and captureShape and stateUpdateShape and (tensorDtypes.inputT == \"float32\" or tensorDtypes.inputT == \"float16\") and tensorDtypes.weightT == tensorDtypes.inputT and tensorDtypes.outputT == tensorDtypes.inputT and tensorDtypes.initialStateT == tensorDtypes.inputT and tensorDtypes.finalStateT == tensorDtypes.inputT and tensorDtypes.cumulativeSequenceLengthT == \"int32\" and f16Ok(dtypes.T) and dim(shapes.weightT, 0) == channels and dim(shapes.outputT, 0) == totalTokens and dim(shapes.outputT, 1) == channels and stateShape and finalStateShape",
    "biasOk": "ranks.biasT == 1 and tensorDtypes.biasT == tensorDtypes.inputT and dim(shapes.biasT, 0) == channels",
    "plainContract": "commonContract and not present.biasT and (not present.stateUpdateT or stateUpdateCapacity == 0)",
    "biasContract": "commonContract and present.biasT and biasOk and (not present.stateUpdateT or stateUpdateCapacity == 0)",
    "stateUpdateContract": "commonContract and not present.biasT and present.stateUpdateT and stateUpdateCapacity > 0",
    "biasStateUpdateContract": "commonContract and present.biasT and biasOk and present.stateUpdateT and stateUpdateCapacity > 0"
  },
  "bindings": {
    "input": { "arg": "inputT", "buffer": "read-only-storage", "elementType": "$inputScalar" },
    "weight": { "arg": "weightT", "buffer": "read-only-storage", "elementType": "$inputScalar" },
    "cumulative_sequence_length": {
      "arg": "cumulativeSequenceLengthT",
      "buffer": "read-only-storage",
      "elementType": "i32"
    },
    "initial_state": { "arg": "initialStateT", "buffer": "read-only-storage", "elementType": "$inputScalar" },
    "output": { "arg": "outputT", "buffer": "storage", "elementType": "$outputScalar" },
    "final_state": { "arg": "finalStateT", "buffer": "storage", "elementType": "$outputScalar" },
    "params": {
      "buffer": "uniform",
      "struct": [
        { "name": "batchSize", "type": "u32", "value": "batchSize" },
        { "name": "channels", "type": "u32", "value": "channels" },
        { "name": "totalTokens", "type": "u32", "value": "totalTokens" },
        { "name": "kernelSize", "type": "u32", "value": "kernelSize" }
      ]
    },
    "bias": { "arg": "biasT", "buffer": "read-only-storage", "elementType": "$inputScalar" },
    "capture_count": { "arg": "captureCountT", "buffer": "read-only-storage", "elementType": "i32" },
    "state_update": { "arg": "stateUpdateT", "buffer": "storage", "elementType": "$outputScalar" },
    "params_2": {
      "name": "params",
      "buffer": "uniform",
      "struct": [
        { "name": "batchSize", "type": "u32", "value": "batchSize" },
        { "name": "channels", "type": "u32", "value": "channels" },
        { "name": "totalTokens", "type": "u32", "value": "totalTokens" },
        { "name": "stateUpdateCapacity", "type": "u32", "value": "stateUpdateCapacity" }
      ]
    }
  },
  "variants": [
    {
      "id": "plain_stream",
      "priority": 20,
      "when": ["plainContract", "kernelSize >= 2", "kernelSize <= 8", "tunables.streamChunk >= 1", "floor(tunables.streamChunk) == tunables.streamChunk", "tunables.streamWorkgroupSize >= 1", "floor(tunables.streamWorkgroupSize) == tunables.streamWorkgroupSize", "tunables.streamWorkgroupSize <= device.limits.maxComputeInvocationsPerWorkgroup", "tunables.streamWorkgroupSize <= device.limits.maxComputeWorkgroupSizeX"],
      "derive": {
        "hasBias": false,
        "useSilu": "attrs.activation != \"none\"",
        "inputScalar": "dtypes.T",
        "outputScalar": "dtypes.T",
        "workgroupSize": "tunables.streamWorkgroupSize",
        "chunkSize": "min(tunables.streamChunk, ceilDiv(totalTokens, batchSize))"
      },
      "passes": [
        {
          "id": "main",
          "name": "VarlenCausalConvWithState.Stream",
          "shader": "varlen-causal-conv-stream.wgsl.jinja",
          "bindings": ["input", "weight", "cumulative_sequence_length", "initial_state", "output", "final_state", "params"],
          "dispatch": {
            "x": "min(ceilDiv((ceilDiv(totalTokens, chunkSize) * channels), (workgroupSize)), 65535)",
            "y": "ceilDiv(ceilDiv((ceilDiv(totalTokens, chunkSize) * channels), (workgroupSize)), 65535)",
            "z": 1
          }
        }
      ]
    },
    {
      "id": "bias_stream",
      "priority": 20,
      "when": ["biasContract", "kernelSize >= 2", "kernelSize <= 8", "tunables.streamChunk >= 1", "floor(tunables.streamChunk) == tunables.streamChunk", "tunables.streamWorkgroupSize >= 1", "floor(tunables.streamWorkgroupSize) == tunables.streamWorkgroupSize", "tunables.streamWorkgroupSize <= device.limits.maxComputeInvocationsPerWorkgroup", "tunables.streamWorkgroupSize <= device.limits.maxComputeWorkgroupSizeX"],
      "derive": {
        "hasBias": true,
        "useSilu": "attrs.activation != \"none\"",
        "inputScalar": "dtypes.T",
        "outputScalar": "dtypes.T",
        "workgroupSize": "tunables.streamWorkgroupSize",
        "chunkSize": "min(tunables.streamChunk, ceilDiv(totalTokens, batchSize))"
      },
      "passes": [
        {
          "id": "main",
          "name": "VarlenCausalConvWithState.Stream",
          "shader": "varlen-causal-conv-stream.wgsl.jinja",
          "bindings": ["input", "weight", "cumulative_sequence_length", "initial_state", "bias", "output", "final_state", "params"],
          "dispatch": {
            "x": "min(ceilDiv((ceilDiv(totalTokens, chunkSize) * channels), (workgroupSize)), 65535)",
            "y": "ceilDiv(ceilDiv((ceilDiv(totalTokens, chunkSize) * channels), (workgroupSize)), 65535)",
            "z": 1
          }
        }
      ]
    },
    {
      "id": "state_update_stream",
      "priority": 20,
      "when": ["stateUpdateContract", "kernelSize >= 2", "kernelSize <= 8", "tunables.streamChunk >= 1", "floor(tunables.streamChunk) == tunables.streamChunk", "tunables.streamWorkgroupSize >= 1", "floor(tunables.streamWorkgroupSize) == tunables.streamWorkgroupSize", "tunables.streamWorkgroupSize <= device.limits.maxComputeInvocationsPerWorkgroup", "tunables.streamWorkgroupSize <= device.limits.maxComputeWorkgroupSizeX"],
      "derive": {
        "hasBias": false,
        "useSilu": "attrs.activation != \"none\"",
        "inputScalar": "dtypes.T",
        "outputScalar": "dtypes.T",
        "workgroupSize": "tunables.streamWorkgroupSize",
        "chunkSize": "min(tunables.streamChunk, ceilDiv(totalTokens, batchSize))"
      },
      "passes": [
        {
          "id": "state_update",
          "name": "VarlenCausalConvWithState.StateUpdate",
          "shader": "varlen-state-update.wgsl.jinja",
          "bindings": ["input", "cumulative_sequence_length", "capture_count", "state_update", "params_2"],
          "dispatch": {
            "x": "min(ceilDiv((batchSize * stateUpdateCapacity * channels), (workgroupSize)), 65535)",
            "y": "ceilDiv(ceilDiv((batchSize * stateUpdateCapacity * channels), (workgroupSize)), 65535)",
            "z": 1
          }
        },
        {
          "id": "main",
          "name": "VarlenCausalConvWithState.Stream",
          "shader": "varlen-causal-conv-stream.wgsl.jinja",
          "bindings": ["input", "weight", "cumulative_sequence_length", "initial_state", "output", "final_state", "params"],
          "dispatch": {
            "x": "min(ceilDiv((ceilDiv(totalTokens, chunkSize) * channels), (workgroupSize)), 65535)",
            "y": "ceilDiv(ceilDiv((ceilDiv(totalTokens, chunkSize) * channels), (workgroupSize)), 65535)",
            "z": 1
          }
        }
      ]
    },
    {
      "id": "bias_state_update_stream",
      "priority": 20,
      "when": ["biasStateUpdateContract", "kernelSize >= 2", "kernelSize <= 8", "tunables.streamChunk >= 1", "floor(tunables.streamChunk) == tunables.streamChunk", "tunables.streamWorkgroupSize >= 1", "floor(tunables.streamWorkgroupSize) == tunables.streamWorkgroupSize", "tunables.streamWorkgroupSize <= device.limits.maxComputeInvocationsPerWorkgroup", "tunables.streamWorkgroupSize <= device.limits.maxComputeWorkgroupSizeX"],
      "derive": {
        "hasBias": true,
        "useSilu": "attrs.activation != \"none\"",
        "inputScalar": "dtypes.T",
        "outputScalar": "dtypes.T",
        "workgroupSize": "tunables.streamWorkgroupSize",
        "chunkSize": "min(tunables.streamChunk, ceilDiv(totalTokens, batchSize))"
      },
      "passes": [
        {
          "id": "state_update",
          "name": "VarlenCausalConvWithState.StateUpdate",
          "shader": "varlen-state-update.wgsl.jinja",
          "bindings": ["input", "cumulative_sequence_length", "capture_count", "state_update", "params_2"],
          "dispatch": {
            "x": "min(ceilDiv((batchSize * stateUpdateCapacity * channels), (workgroupSize)), 65535)",
            "y": "ceilDiv(ceilDiv((batchSize * stateUpdateCapacity * channels), (workgroupSize)), 65535)",
            "z": 1
          }
        },
        {
          "id": "main",
          "name": "VarlenCausalConvWithState.Stream",
          "shader": "varlen-causal-conv-stream.wgsl.jinja",
          "bindings": ["input", "weight", "cumulative_sequence_length", "initial_state", "bias", "output", "final_state", "params"],
          "dispatch": {
            "x": "min(ceilDiv((ceilDiv(totalTokens, chunkSize) * channels), (workgroupSize)), 65535)",
            "y": "ceilDiv(ceilDiv((ceilDiv(totalTokens, chunkSize) * channels), (workgroupSize)), 65535)",
            "z": 1
          }
        }
      ]
    },
    {
      "id": "plain",
      "priority": 0,
      "when": ["plainContract", "tunables.workgroupSize >= 1", "floor(tunables.workgroupSize) == tunables.workgroupSize", "tunables.workgroupSize <= device.limits.maxComputeInvocationsPerWorkgroup", "tunables.workgroupSize <= device.limits.maxComputeWorkgroupSizeX"],
      "derive": {
        "hasBias": false,
        "useSilu": "attrs.activation != \"none\"",
        "inputScalar": "dtypes.T",
        "outputScalar": "dtypes.T",
        "workgroupSize": "tunables.workgroupSize"
      },
      "passes": [
        {
          "id": "main",
          "name": "VarlenCausalConvWithState",
          "shader": "varlen-causal-conv.wgsl.jinja",
          "bindings": ["input", "weight", "cumulative_sequence_length", "initial_state", "output", "final_state", "params"],
          "dispatch": {
            "x": "min(ceilDiv((totalTokens * channels), (workgroupSize)), 65535)",
            "y": "ceilDiv(ceilDiv((totalTokens * channels), (workgroupSize)), 65535)",
            "z": 1
          }
        }
      ]
    },
    {
      "id": "bias",
      "priority": 0,
      "when": ["biasContract", "tunables.workgroupSize >= 1", "floor(tunables.workgroupSize) == tunables.workgroupSize", "tunables.workgroupSize <= device.limits.maxComputeInvocationsPerWorkgroup", "tunables.workgroupSize <= device.limits.maxComputeWorkgroupSizeX"],
      "derive": {
        "hasBias": true,
        "useSilu": "attrs.activation != \"none\"",
        "inputScalar": "dtypes.T",
        "outputScalar": "dtypes.T",
        "workgroupSize": "tunables.workgroupSize"
      },
      "passes": [
        {
          "id": "main",
          "name": "VarlenCausalConvWithState",
          "shader": "varlen-causal-conv.wgsl.jinja",
          "bindings": ["input", "weight", "cumulative_sequence_length", "initial_state", "bias", "output", "final_state", "params"],
          "dispatch": {
            "x": "min(ceilDiv((totalTokens * channels), (workgroupSize)), 65535)",
            "y": "ceilDiv(ceilDiv((totalTokens * channels), (workgroupSize)), 65535)",
            "z": 1
          }
        }
      ]
    },
    {
      "id": "state_update",
      "priority": 0,
      "when": ["stateUpdateContract", "tunables.workgroupSize >= 1", "floor(tunables.workgroupSize) == tunables.workgroupSize", "tunables.workgroupSize <= device.limits.maxComputeInvocationsPerWorkgroup", "tunables.workgroupSize <= device.limits.maxComputeWorkgroupSizeX"],
      "derive": {
        "hasBias": false,
        "useSilu": "attrs.activation != \"none\"",
        "inputScalar": "dtypes.T",
        "outputScalar": "dtypes.T",
        "workgroupSize": "tunables.workgroupSize"
      },
      "passes": [
        {
          "id": "state_update",
          "name": "VarlenCausalConvWithState.StateUpdate",
          "shader": "varlen-state-update.wgsl.jinja",
          "bindings": ["input", "cumulative_sequence_length", "capture_count", "state_update", "params_2"],
          "dispatch": {
            "x": "min(ceilDiv((batchSize * stateUpdateCapacity * channels), (workgroupSize)), 65535)",
            "y": "ceilDiv(ceilDiv((batchSize * stateUpdateCapacity * channels), (workgroupSize)), 65535)",
            "z": 1
          }
        },
        {
          "id": "main",
          "name": "VarlenCausalConvWithState",
          "shader": "varlen-causal-conv.wgsl.jinja",
          "bindings": ["input", "weight", "cumulative_sequence_length", "initial_state", "output", "final_state", "params"],
          "dispatch": {
            "x": "min(ceilDiv((totalTokens * channels), (workgroupSize)), 65535)",
            "y": "ceilDiv(ceilDiv((totalTokens * channels), (workgroupSize)), 65535)",
            "z": 1
          }
        }
      ]
    },
    {
      "id": "bias_state_update",
      "priority": 0,
      "when": ["biasStateUpdateContract", "tunables.workgroupSize >= 1", "floor(tunables.workgroupSize) == tunables.workgroupSize", "tunables.workgroupSize <= device.limits.maxComputeInvocationsPerWorkgroup", "tunables.workgroupSize <= device.limits.maxComputeWorkgroupSizeX"],
      "derive": {
        "hasBias": true,
        "useSilu": "attrs.activation != \"none\"",
        "inputScalar": "dtypes.T",
        "outputScalar": "dtypes.T",
        "workgroupSize": "tunables.workgroupSize"
      },
      "passes": [
        {
          "id": "state_update",
          "name": "VarlenCausalConvWithState.StateUpdate",
          "shader": "varlen-state-update.wgsl.jinja",
          "bindings": ["input", "cumulative_sequence_length", "capture_count", "state_update", "params_2"],
          "dispatch": {
            "x": "min(ceilDiv((batchSize * stateUpdateCapacity * channels), (workgroupSize)), 65535)",
            "y": "ceilDiv(ceilDiv((batchSize * stateUpdateCapacity * channels), (workgroupSize)), 65535)",
            "z": 1
          }
        },
        {
          "id": "main",
          "name": "VarlenCausalConvWithState",
          "shader": "varlen-causal-conv.wgsl.jinja",
          "bindings": ["input", "weight", "cumulative_sequence_length", "initial_state", "bias", "output", "final_state", "params"],
          "dispatch": {
            "x": "min(ceilDiv((totalTokens * channels), (workgroupSize)), 65535)",
            "y": "ceilDiv(ceilDiv((totalTokens * channels), (workgroupSize)), 65535)",
            "z": 1
          }
        }
      ]
    }
  ]
}