{
  "domain": "com.microsoft",
  "name": "BiasSoftmax",
  "sinceVersion": 1,
  "description": "Computes `softmax(data + bias)` over the flattened suffix beginning at `axis`. The required `is_inner_broadcast` attribute selects how bias rows are reused: consecutive groups for inner broadcast or cyclic groups for outer broadcast. This specializes the `softmax(scores + additive_mask)` pattern used by transformer attention. Float16 and float32 are supported; the schema's double type is not.",
  "inputs": [
    { "role": "data", "dtype": "T", "description": "The input data tensor." },
    {
      "role": "bias",
      "dtype": "T",
      "description": "The bias (or additive mask) tensor. Its element count must be an integral number of flattened softmax rows and that row count must divide the data row count."
    }
  ],
  "outputs": [
    {
      "role": "output",
      "dtype": "T",
      "rank": "ranks.data",
      "shape": "shapes.data",
      "description": "The output tensor; same shape as data."
    }
  ],
  "attributes": { "axis": 1 },
  "attributeConstraints": { "is_inner_broadcast": { "required": true } },
  "attributeDescriptions": {
    "axis": "The axis from which softmax is applied; dimensions from `axis` onward are included in the softmax reduction.",
    "is_inner_broadcast": "When 1, bias is broadcast across dimensions from `broadcast_axis` to `axis-1`; when 0, bias is broadcast across dimensions 0 to `broadcast_axis-1`."
  },
  "typeConstraints": { "T": ["float32", "float16"] },
  "args": {
    "data": { "kind": "tensor", "semantic": "data", "role": "input" },
    "bias": { "kind": "tensor", "semantic": "bias", "role": "input" },
    "output": { "kind": "tensor", "semantic": "output", "role": "output" }
  },
  "derive": {
    "axisNorm": "attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.data",
    "batchCount": "outer(shapes.data, axisNorm)",
    "blockSize": "dim(shapes.data, axisNorm) * inner(shapes.data, axisNorm)",
    "biasBlockCount": "numel(shapes.bias) / max(1, blockSize)",
    "biasContract": "(numel(shapes.data) == 0 and numel(shapes.bias) == 0) or (blockSize > 0 and biasBlockCount > 0 and numel(shapes.bias) % blockSize == 0 and biasBlockCount <= batchCount and batchCount % biasBlockCount == 0)"
  },
  "tunables": { "WORKGROUP_SIZE": 256, "BLOCK_COLS": 2048 },
  "bindingSets": {
    "biasSoftmax": [
      {
        "name": "data",
        "arg": "data",
        "semantic": "data",
        "buffer": { "type": "read-only-storage" },
        "elementType": "$scalar"
      },
      {
        "name": "bias",
        "arg": "bias",
        "semantic": "bias",
        "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": "blockSize", "type": "u32", "value": "blockSize" },
            { "name": "batchCount", "type": "u32", "value": "batchCount" }
          ]
        }
      }
    ]
  },
  "variants": [
    {
      "id": "longrow_split",
      "priority": 40,
      "when": ["numel(shapes.data) == numel(shapes.output)", "ranks.data >= 1", "attrs.axis + ranks.data >= 0", "attrs.axis < ranks.data", "biasContract", "blockSize >= 65536", "batchCount > 0", "batchCount <= device.limits.maxComputeWorkgroupsPerDimension", "ceilDiv(blockSize, tunables.BLOCK_COLS) <= device.limits.maxComputeWorkgroupsPerDimension", "f16Ok(dtypes.T)"],
      "constants": { "scalar": "dtypes.T", "usesF16": "dtypes.T == \"f16\"", "combineSubgroups": false },
      "intermediates": [
        {
          "id": "blockMax",
          "dtype": "float32",
          "shape": "[outer(shapes.data, axisNorm) * ceilDiv(dim(shapes.data, axisNorm) * inner(shapes.data, axisNorm), tunables.BLOCK_COLS)]"
        },
        {
          "id": "blockSum",
          "dtype": "float32",
          "shape": "[outer(shapes.data, axisNorm) * ceilDiv(dim(shapes.data, axisNorm) * inner(shapes.data, axisNorm), tunables.BLOCK_COLS)]"
        },
        { "id": "rowMax", "dtype": "float32", "shape": "[outer(shapes.data, axisNorm)]" },
        { "id": "rowSum", "dtype": "float32", "shape": "[outer(shapes.data, axisNorm)]" }
      ],
      "passes": [
        {
          "id": "block_stats",
          "name": "BiasSoftmax.LongRowBlockStats",
          "source": {
            "shader": "bias-softmax-longrow-stats.wgsl.jinja",
            "inputs": {
              "stage": "\"block\"",
              "isInnerBroadcast": "attrs.is_inner_broadcast != 0",
              "biasBlockCount": "max(1, biasBlockCount)",
              "innerRepeat": "max(1, batchCount / max(1, biasBlockCount))"
            }
          },
          "bindings": [
            {
              "name": "data",
              "arg": "data",
              "semantic": "data",
              "buffer": { "type": "read-only-storage" },
              "elementType": "$scalar"
            },
            {
              "name": "bias",
              "arg": "bias",
              "semantic": "bias",
              "buffer": { "type": "read-only-storage" },
              "elementType": "$scalar"
            },
            { "name": "blockMax", "semantic": "blockMax", "buffer": { "type": "storage" }, "elementType": "f32" },
            { "name": "blockSum", "semantic": "blockSum", "buffer": { "type": "storage" }, "elementType": "f32" },
            {
              "name": "params",
              "semantic": "kernel.params",
              "buffer": { "type": "uniform" },
              "struct": {
                "name": "Params",
                "fields": [
                  { "name": "blockSize", "type": "u32", "value": "blockSize" },
                  { "name": "blocks", "type": "u32", "value": "ceilDiv(blockSize, tunables.BLOCK_COLS)" }
                ]
              }
            }
          ],
          "dispatch": { "x": "ceilDiv(blockSize, tunables.BLOCK_COLS)", "y": "batchCount" }
        },
        {
          "id": "row_stats",
          "name": "BiasSoftmax.LongRowStats",
          "source": { "shader": "bias-softmax-longrow-stats.wgsl.jinja", "inputs": { "stage": "\"row\"" } },
          "bindings": [
            {
              "name": "blockMax",
              "semantic": "blockMax",
              "buffer": { "type": "read-only-storage" },
              "elementType": "f32"
            },
            {
              "name": "blockSum",
              "semantic": "blockSum",
              "buffer": { "type": "read-only-storage" },
              "elementType": "f32"
            },
            { "name": "rowMax", "semantic": "rowMax", "buffer": { "type": "storage" }, "elementType": "f32" },
            { "name": "rowSum", "semantic": "rowSum", "buffer": { "type": "storage" }, "elementType": "f32" },
            {
              "name": "params",
              "semantic": "kernel.params",
              "buffer": { "type": "uniform" },
              "struct": {
                "name": "Params",
                "fields": [{ "name": "blocks", "type": "u32", "value": "ceilDiv(blockSize, tunables.BLOCK_COLS)" }]
              }
            }
          ],
          "dispatch": { "x": "batchCount" }
        },
        {
          "id": "normalize",
          "name": "BiasSoftmax.LongRowNormalize",
          "source": {
            "shader": "bias-softmax-longrow-normalize.wgsl.jinja",
            "inputs": {
              "isInnerBroadcast": "attrs.is_inner_broadcast != 0",
              "biasBlockCount": "max(1, biasBlockCount)",
              "innerRepeat": "max(1, batchCount / max(1, biasBlockCount))"
            }
          },
          "bindings": [
            {
              "name": "data",
              "arg": "data",
              "semantic": "data",
              "buffer": { "type": "read-only-storage" },
              "elementType": "$scalar"
            },
            {
              "name": "bias",
              "arg": "bias",
              "semantic": "bias",
              "buffer": { "type": "read-only-storage" },
              "elementType": "$scalar"
            },
            { "name": "rowMax", "semantic": "rowMax", "buffer": { "type": "read-only-storage" }, "elementType": "f32" },
            { "name": "rowSum", "semantic": "rowSum", "buffer": { "type": "read-only-storage" }, "elementType": "f32" },
            {
              "name": "output",
              "arg": "output",
              "semantic": "output",
              "buffer": { "type": "storage" },
              "elementType": "$scalar"
            },
            {
              "name": "params",
              "semantic": "kernel.params",
              "buffer": { "type": "uniform" },
              "struct": { "name": "Params", "fields": [{ "name": "blockSize", "type": "u32", "value": "blockSize" }] }
            }
          ],
          "dispatch": { "x": "ceilDiv(blockSize, tunables.BLOCK_COLS)", "y": "batchCount" }
        }
      ]
    },
    {
      "id": "packed_rows",
      "priority": 30,
      "when": ["numel(shapes.data) == numel(shapes.output)", "ranks.data >= 1", "attrs.axis + ranks.data >= 0", "attrs.axis < ranks.data", "biasContract", "blockSize > 0", "blockSize <= 8", "batchCount >= 64", "f16Ok(dtypes.T)"],
      "constants": {
        "scalar": "dtypes.T",
        "usesF16": "dtypes.T == \"f16\"",
        "combineSubgroups": false,
        "packedRows": true
      },
      "passes": [
        {
          "id": "main",
          "name": "BiasSoftmax.PackedRows",
          "source": {
            "shader": "bias-softmax.wgsl.jinja",
            "inputs": {
              "blockSize": "blockSize",
              "isInnerBroadcast": "attrs.is_inner_broadcast != 0",
              "biasBlockCount": "max(1, biasBlockCount)",
              "innerRepeat": "max(1, batchCount / max(1, biasBlockCount))"
            }
          },
          "bindings": "biasSoftmax",
          "dispatch": { "gridStride": "batchCount", "workgroupSize": "tunables.WORKGROUP_SIZE" }
        }
      ]
    },
    {
      "id": "adaptive_row",
      "priority": 10,
      "when": ["numel(shapes.data) == numel(shapes.output)", "ranks.data >= 1", "attrs.axis + ranks.data >= 0", "attrs.axis < ranks.data", "biasContract", "numel(shapes.data) >= 0", "f16Ok(dtypes.T)"],
      "constants": {
        "packedRows": false,
        "scalar": "dtypes.T",
        "usesF16": "dtypes.T == \"f16\"",
        "combineSubgroups": "device.features.has(\"subgroups\")"
      },
      "passes": [
        {
          "id": "main",
          "name": "BiasSoftmax.AdaptiveRow",
          "source": {
            "shader": "bias-softmax.wgsl.jinja",
            "inputs": {
              "blockSize": "blockSize",
              "isInnerBroadcast": "attrs.is_inner_broadcast != 0",
              "biasBlockCount": "max(1, biasBlockCount)",
              "innerRepeat": "max(1, batchCount / max(1, biasBlockCount))"
            }
          },
          "bindings": "biasSoftmax",
          "dispatch": { "workgroups": "batchCount" }
        }
      ]
    }
  ]
}