{
  "domain": "com.microsoft",
  "name": "FusedMatMul",
  "sinceVersion": 1,
  "description": "Matrix product of two N-dimensional tensors `A` and `B`, following NumPy-style matrix-multiplication broadcasting. Supports optional transposition of either operand's last two dimensions, optional batch-dimension transposition, and a scalar `alpha` multiplier. Float32 and float16 are supported; double and bfloat16 are not.",
  "inputs": [
    { "role": "A", "dtype": "T", "description": "N-dimensional matrix A." },
    { "role": "B", "dtype": "T", "description": "N-dimensional matrix B." }
  ],
  "outputs": [
    {
      "role": "Y",
      "dtype": "T",
      "rank": "max(ranks.A, ranks.B) - (1 if ranks.A == 1 or ranks.B == 1 else 0)",
      "shape": "matmulShape(logicalAShape, logicalBShape)",
      "description": "Matrix-multiplication result whose shape follows NumPy-style rules after applying the requested batch and matrix transpositions."
    }
  ],
  "attributes": { "alpha": 1, "transA": 0, "transB": 0, "transBatchA": 0, "transBatchB": 0 },
  "attributeDescriptions": {
    "alpha": "Scalar multiplier applied to the product of the input tensors.",
    "transA": "When non-zero, transposes `A` on its last two dimensions before multiplication.",
    "transB": "When non-zero, transposes `B` on its last two dimensions before multiplication.",
    "transBatchA": "When non-zero, transposes `A` on its first dimension and batch dimensions (dim-1 to dim-rank-2) before multiplication.",
    "transBatchB": "When non-zero, transposes `B` on its first dimension and batch dimensions (dim-1 to dim-rank-2) before multiplication."
  },
  "typeConstraints": { "T": ["float32", "float16"] },
  "args": {
    "A": { "kind": "tensor", "semantic": "A", "role": "input" },
    "B": { "kind": "tensor", "semantic": "B", "role": "input" },
    "Y": { "kind": "tensor", "semantic": "Y", "role": "output" }
  },
  "tunables": {
    "TILED_REG_MIN_WORKGROUPS": 64,
    "GEMV_TARGET_BLOCKS": 512,
    "SUBGROUP_MATRIX_MIN_M": 2,
    "SUBGROUP_MATRIX_SPLITK_TARGET_WGS": 512,
    "SUBGROUP_MATRIX_SPLITK_MIN_K": 1024,
    "SUBGROUP_MATRIX_SPLITK_MAX_TILES": 128
  },
  "bindingSets": {
    "matrix": [
      { "name": "a", "arg": "A", "semantic": "A", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" },
      { "name": "b", "arg": "B", "semantic": "B", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" },
      { "name": "y", "arg": "Y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$scalar" }
    ],
    "sgmatSplitPartial": [
      { "name": "a", "arg": "A", "semantic": "A", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" },
      { "name": "b", "arg": "B", "semantic": "B", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" },
      { "name": "partials", "semantic": "partials", "buffer": { "type": "storage" }, "elementType": "f32" }
    ],
    "sgmatSplitCombine": [
      { "name": "partials", "semantic": "partials", "buffer": { "type": "read-only-storage" }, "elementType": "f32" },
      { "name": "y", "arg": "Y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$scalar" },
      {
        "name": "params",
        "semantic": "kernel.params",
        "buffer": { "type": "uniform" },
        "struct": { "name": "Params", "fields": [{ "name": "cols", "type": "u32", "value": "numel(shapes.Y)" }] }
      }
    ],
    "gemvF32": [
      { "name": "a", "arg": "A", "semantic": "A", "buffer": { "type": "read-only-storage" }, "elementType": "f32" },
      {
        "name": "b",
        "arg": "B",
        "semantic": "B",
        "buffer": { "type": "read-only-storage" },
        "elementType": "vec4<f32>"
      },
      { "name": "c", "arg": "Y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "vec4<f32>" },
      {
        "name": "params",
        "semantic": "kernel.params",
        "buffer": { "type": "uniform" },
        "struct": {
          "name": "Params",
          "fields": [
            { "name": "K", "type": "u32", "value": "dim(shapes.A, 1)" },
            { "name": "N4", "type": "u32", "value": "dim(shapes.B, 1) / 4" }
          ]
        }
      }
    ]
  },
  "derive": {
    "gemvWorkgroups": "ceilDiv(dim(shapes.B, 1), 128)",
    "gemvSliceCap": "min(32, device.limits.maxComputeWorkgroupSizeY, floor(device.limits.maxComputeInvocationsPerWorkgroup / 32), floor(device.limits.maxComputeWorkgroupStorageSize / 512))",
    "gemvSlices": "max(1, min(gemvSliceCap, max(8, pow2ceil(ceilDiv(tunables.GEMV_TARGET_BLOCKS, gemvWorkgroups)))))",
    "batchMovedAShape": "moveAxis(shapes.A, 0, -2) if attrs.transBatchA != 0 else shapes.A",
    "batchMovedBShape": "moveAxis(shapes.B, 0, -2) if attrs.transBatchB != 0 else shapes.B",
    "logicalAShape": "moveAxis(batchMovedAShape, -1, -2) if attrs.transA != 0 and ranks.A > 1 else batchMovedAShape",
    "logicalBShape": "moveAxis(batchMovedBShape, -1, -2) if attrs.transB != 0 and ranks.B > 1 else batchMovedBShape",
    "transBatchContract": "(attrs.transBatchA == 0 and attrs.transBatchB == 0) or (ranks.A == ranks.B and ranks.A >= 3)",
    "wave32Adapter": "has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize == 32 and device.adapterInfo.subgroupMaxSize == 32",
    "canPinSubgroupSize32": "device.features.has(\"subgroups\") and device.features.has(\"subgroup-size-control\") and has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize <= 32 and device.adapterInfo.subgroupMaxSize >= 32",
    "pinSubgroupSize32": "canPinSubgroupSize32 and not wave32Adapter",
    "wave32Effective": "wave32Adapter or pinSubgroupSize32",
    "deviceWorkgroupCap": "min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)",
    "fusedSgmatRank2Ok": "ranks.A == 2 and ranks.B == 2 and ranks.Y == 2 and attrs.transA == 0 and attrs.transB == 0 and attrs.transBatchA == 0 and attrs.transBatchB == 0 and dim(shapes.A, 1) == dim(shapes.B, 0) and dim(shapes.Y, 0) == dim(shapes.A, 0) and dim(shapes.Y, 1) == dim(shapes.B, 1)",
    "subgroupMatrixResourcesFit": "128 <= deviceWorkgroupCap and ((32 * 32 + 64 * 32) * dtypeBytes(dtypes.T) + 4 * 4 * 64 * 4) <= device.limits.maxComputeWorkgroupStorageSize",
    "sgmatSplitKDepth": "dim(shapes.A, ranks.A - 1)",
    "sgmatOutTiles": "ceilDiv(dim(shapes.A, 0), 32) * ceilDiv(dim(shapes.B, 1), 64) if fusedSgmatRank2Ok else 1",
    "sgmatSplitKWant": "ceilDiv(tunables.SUBGROUP_MATRIX_SPLITK_TARGET_WGS, sgmatOutTiles)",
    "sgmatSplitK32Ok": "sgmatSplitKDepth % 1024 == 0",
    "sgmatSplitK16Ok": "sgmatSplitKDepth % 512 == 0",
    "sgmatSplitK8Ok": "sgmatSplitKDepth % 256 == 0",
    "sgmatSplitK4Ok": "sgmatSplitKDepth % 128 == 0",
    "sgmatSplitK2Ok": "sgmatSplitKDepth % 64 == 0",
    "sgmatSplitK": "32 if (sgmatSplitKWant > 16 and sgmatSplitK32Ok) else (16 if (sgmatSplitKWant > 8 and sgmatSplitK16Ok) else (8 if (sgmatSplitKWant > 4 and sgmatSplitK8Ok) else (4 if (sgmatSplitKWant > 2 and sgmatSplitK4Ok) else (2 if sgmatSplitK2Ok else 1))))"
  },
  "variants": [
    {
      "id": "f32_m1_gemv_vec4",
      "priority": 30,
      "when": ["dtypes.T == \"f32\"", "attrs.alpha == 1", "attrs.transA == 0", "attrs.transB == 0", "attrs.transBatchA == 0", "attrs.transBatchB == 0", "ranks.A == 2", "ranks.B == 2", "ranks.Y == 2", "dim(shapes.A, 0) == 1", "dim(shapes.Y, 0) == 1", "dim(shapes.A, 1) == dim(shapes.B, 0)", "dim(shapes.Y, 1) == dim(shapes.B, 1)", "dim(shapes.B, 1) > 0", "dim(shapes.B, 1) % 4 == 0", "ceil(dim(shapes.B, 1) / 128) <= device.limits.maxComputeWorkgroupsPerDimension"],
      "constants": { "T": "\"f32\"", "usesF16": false, "gemvSlices": "gemvSlices" },
      "passes": [
        {
          "id": "main",
          "name": "FusedMatMul.F32M1GemvVec4",
          "shader": "matmul-vector-matrix-vec4.wgsl.jinja",
          "bindings": "gemvF32",
          "dispatch": { "x": "ceil(dim(shapes.B, 1) / 128)" }
        }
      ]
    },
    {
      "id": "subgroup_matrix_splitk",
      "description": "Partitions the K reduction across workgroups while retaining subgroup-matrix tiles, then combines float32 partials that already carry alpha. It serves small-M rank-two products whose unsplit tile grids cannot occupy the device.",
      "priority": 12,
      "requires": {
        "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
        "subgroupMatrixConfigs": [
          { "componentType": "f16", "M": 8, "N": 8, "K": 8 },
          { "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }
        ]
      },
      "when": ["(dtypes.T == \"f16\" or dtypes.T == \"f32\") and f16Ok(dtypes.T)", "fusedSgmatRank2Ok", "dim(shapes.A, 0) >= tunables.SUBGROUP_MATRIX_MIN_M", "dim(shapes.A, 1) >= tunables.SUBGROUP_MATRIX_SPLITK_MIN_K", "dim(shapes.B, 1) % 64 == 0", "sgmatSplitK >= 2", "sgmatOutTiles < tunables.SUBGROUP_MATRIX_SPLITK_MAX_TILES", "sgmatSplitK * numel(shapes.Y) * 4 <= device.limits.maxStorageBufferBindingSize", "sgmatSplitK <= device.limits.maxComputeWorkgroupsPerDimension", "ceilDiv(dim(shapes.Y, 1), 64) <= device.limits.maxComputeWorkgroupsPerDimension", "ceilDiv(dim(shapes.Y, 0), 32) <= device.limits.maxComputeWorkgroupsPerDimension", "subgroupMatrixResourcesFit", "wave32Effective"],
      "constants": {
        "usesF16": "dtypes.T == \"f16\"",
        "fScalar": "\"f16\" if dtypes.T == \"f16\" else \"f32\"",
        "scalar": "dtypes.T",
        "hasBias": false,
        "generalAddressing": true,
        "tailSafe": false,
        "outputBuffer": "\"partials\"",
        "outScalar": "\"f32\"",
        "alpha": "attrs.alpha",
        "M": "dim(shapes.A, 0)",
        "K": "dim(shapes.A, 1)",
        "N": "dim(shapes.B, 1)",
        "batchCount": 1,
        "splitK": "sgmatSplitK",
        "kPerSplit": "dim(shapes.A, 1) / sgmatSplitK",
        "split": "sgmatSplitK",
        "workgroupSize": 256
      },
      "intermediates": [{ "id": "partials", "dtype": "float32", "shape": "[sgmatSplitK * numel(shapes.Y)]" }],
      "passes": [
        {
          "id": "partial",
          "name": "FusedMatMul.SubgroupMatrixSplitK",
          "source": {
            "shader": "matmul-subgroup-matrix-ext.wgsl.jinja",
            "inputs": {
              "aShape": ["dim(shapes.A, 0)", "dim(shapes.A, 1)"],
              "bShape": ["dim(shapes.B, 0)", "dim(shapes.B, 1)"],
              "aRank": 2,
              "bRank": 2
            }
          },
          "bindings": "sgmatSplitPartial",
          "dispatch": { "x": "ceilDiv(dim(shapes.Y, 1), 64)", "y": "ceilDiv(dim(shapes.Y, 0), 32)", "z": "sgmatSplitK" }
        },
        {
          "id": "combine",
          "name": "FusedMatMul.SubgroupMatrixSplitKCombine",
          "source": {
            "shader": "reduce-axis0-splitk-combine.wgsl.jinja",
            "inputs": { "op": "\"sum\"", "outputF16": "dtypes.T == \"f16\"", "intMode": false }
          },
          "bindings": "sgmatSplitCombine",
          "dispatch": { "threads": "numel(shapes.Y)", "workgroupSize": 256 }
        }
      ]
    },
    {
      "id": "subgroup_matrix_tail_broadcast",
      "priority": 11,
      "requires": {
        "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
        "subgroupMatrixConfigs": [{ "componentType": "f16", "M": 8, "N": 8, "K": 8 }]
      },
      "when": ["dtypes.T == \"f16\"", "f16Ok(dtypes.T)", "attrs.transA == 0", "attrs.transB == 0", "attrs.transBatchA == 0", "attrs.transBatchB == 0", "((ranks.A == 2 and ranks.B == 2 and ranks.Y == 2) or (ranks.A == 4 and ranks.B == 3 and ranks.Y == 4 and dim(shapes.Y, 0) == dim(shapes.A, 0) and (dim(shapes.A, 1) == dim(shapes.B, 0) or dim(shapes.A, 1) == 1 or dim(shapes.B, 0) == 1) and dim(shapes.Y, 1) == max(dim(shapes.A, 1), dim(shapes.B, 0))))", "dim(shapes.A, ranks.A - 1) == dim(shapes.B, ranks.B - 2)", "dim(shapes.A, ranks.A - 2) >= tunables.SUBGROUP_MATRIX_MIN_M", "dim(shapes.A, ranks.A - 1) >= 32", "dim(shapes.B, ranks.B - 1) >= 64", "dim(shapes.Y, ranks.Y - 2) == dim(shapes.A, ranks.A - 2)", "dim(shapes.Y, ranks.Y - 1) == dim(shapes.B, ranks.B - 1)", "ceil(dim(shapes.B, ranks.B - 1) / 64) <= device.limits.maxComputeWorkgroupsPerDimension", "ceil(dim(shapes.A, ranks.A - 2) / 32) <= device.limits.maxComputeWorkgroupsPerDimension", "numel(shapes.Y) / (dim(shapes.A, ranks.A - 2) * dim(shapes.B, ranks.B - 1)) <= device.limits.maxComputeWorkgroupsPerDimension", "wave32Effective"],
      "constants": {
        "hasBias": false,
        "usesF16": true,
        "fScalar": "\"f16\"",
        "outScalar": "\"f16\"",
        "scalar": "dtypes.T",
        "generalAddressing": true,
        "tailSafe": "dim(shapes.A, ranks.A - 1) % 32 != 0 or dim(shapes.B, ranks.B - 1) % 64 != 0",
        "outputBuffer": "\"y\"",
        "alpha": "attrs.alpha",
        "M": "dim(shapes.A, ranks.A - 2)",
        "K": "dim(shapes.A, ranks.A - 1)",
        "N": "dim(shapes.B, ranks.B - 1)",
        "batchCount": "numel(shapes.Y) / (dim(shapes.A, ranks.A - 2) * dim(shapes.B, ranks.B - 1))"
      },
      "passes": [
        {
          "id": "main",
          "name": "FusedMatMul.SubgroupMatrixTailBroadcast",
          "source": {
            "shader": "matmul-subgroup-matrix-ext.wgsl.jinja",
            "inputs": { "aShape": "shapes.A", "bShape": "shapes.B", "aRank": "ranks.A", "bRank": "ranks.B" }
          },
          "bindings": "matrix",
          "dispatch": { "x": "ceil(constants.N / 64)", "y": "ceil(constants.M / 32)", "z": "constants.batchCount" }
        }
      ]
    },
    {
      "id": "subgroup_matrix",
      "priority": 10,
      "requires": {
        "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
        "subgroupMatrixConfigs": [
          { "componentType": "f16", "M": 8, "N": 8, "K": 8 },
          { "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }
        ]
      },
      "when": ["f16Ok(dtypes.T)", "attrs.transBatchA == 0 or (attrs.transA == 0 and ranks.A == 3)", "attrs.transBatchB == 0", "ranks.A >= 2", "ranks.B == ranks.A", "ranks.Y == ranks.A", "(dim(shapes.A, ranks.A - 2) if attrs.transA != 0 else dim(shapes.A, ranks.A - 1)) == (dim(shapes.B, ranks.B - 1) if attrs.transB != 0 else dim(shapes.B, ranks.B - 2))", "(dim(shapes.A, ranks.A - 1) if attrs.transA != 0 else (dim(shapes.A, 0) if attrs.transBatchA != 0 else dim(shapes.A, ranks.A - 2))) >= tunables.SUBGROUP_MATRIX_MIN_M", "(dim(shapes.A, ranks.A - 2) if attrs.transA != 0 else dim(shapes.A, ranks.A - 1)) % 32 == 0", "(dim(shapes.B, ranks.B - 2) if attrs.transB != 0 else dim(shapes.B, ranks.B - 1)) % 64 == 0", "(ranks.A == 2 or (ranks.A == 3 and (dim(shapes.A, 1) if attrs.transBatchA != 0 else dim(shapes.A, 0)) == dim(shapes.B, 0) and dim(shapes.Y, 0) == dim(shapes.B, 0)) or (ranks.A == 4 and dim(shapes.A, 0) == dim(shapes.B, 0) and dim(shapes.A, 1) == dim(shapes.B, 1) and dim(shapes.Y, 0) == dim(shapes.A, 0) and dim(shapes.Y, 1) == dim(shapes.A, 1)))", "dim(shapes.Y, ranks.Y - 2) == (dim(shapes.A, ranks.A - 1) if attrs.transA != 0 else (dim(shapes.A, 0) if attrs.transBatchA != 0 else dim(shapes.A, ranks.A - 2)))", "dim(shapes.Y, ranks.Y - 1) == (dim(shapes.B, ranks.B - 2) if attrs.transB != 0 else dim(shapes.B, ranks.B - 1))", "ceil((dim(shapes.B, ranks.B - 2) if attrs.transB != 0 else dim(shapes.B, ranks.B - 1)) / 64) <= device.limits.maxComputeWorkgroupsPerDimension", "ceil((dim(shapes.A, ranks.A - 1) if attrs.transA != 0 else (dim(shapes.A, 0) if attrs.transBatchA != 0 else dim(shapes.A, ranks.A - 2))) / 32) <= device.limits.maxComputeWorkgroupsPerDimension", "numel(shapes.Y) / ((dim(shapes.A, ranks.A - 1) if attrs.transA != 0 else (dim(shapes.A, 0) if attrs.transBatchA != 0 else dim(shapes.A, ranks.A - 2))) * (dim(shapes.B, ranks.B - 2) if attrs.transB != 0 else dim(shapes.B, ranks.B - 1))) <= device.limits.maxComputeWorkgroupsPerDimension", "wave32Effective"],
      "constants": {
        "usesF16": "dtypes.T == \"f16\"",
        "fScalar": "\"f16\" if dtypes.T == \"f16\" else \"f32\"",
        "outScalar": "\"f16\" if dtypes.T == \"f16\" else \"f32\"",
        "scalar": "dtypes.T",
        "transA": "attrs.transA != 0",
        "transB": "attrs.transB != 0",
        "transBatchA": "attrs.transBatchA != 0",
        "M": "(dim(shapes.A, ranks.A - 1) if attrs.transA != 0 else (dim(shapes.A, 0) if attrs.transBatchA != 0 else dim(shapes.A, ranks.A - 2)))",
        "K": "(dim(shapes.A, ranks.A - 2) if attrs.transA != 0 else dim(shapes.A, ranks.A - 1))",
        "N": "(dim(shapes.B, ranks.B - 2) if attrs.transB != 0 else dim(shapes.B, ranks.B - 1))",
        "batchCount": "numel(shapes.Y) / ((dim(shapes.A, ranks.A - 1) if attrs.transA != 0 else (dim(shapes.A, 0) if attrs.transBatchA != 0 else dim(shapes.A, ranks.A - 2))) * (dim(shapes.B, ranks.B - 2) if attrs.transB != 0 else dim(shapes.B, ranks.B - 1)))"
      },
      "passes": [
        {
          "id": "main",
          "name": "FusedMatMul.SubgroupMatrix",
          "source": { "shader": "fused-matmul-subgroup-matrix.wgsl.jinja", "inputs": { "alpha": "attrs.alpha" } },
          "bindings": "matrix",
          "dispatch": {
            "x": "ceil(constants.N / 64)",
            "y": "ceil(constants.M / 32)",
            "z": "numel(shapes.Y) / (constants.M * constants.N)"
          }
        }
      ]
    },
    {
      "id": "broadcast_rank4_tiled_reg",
      "priority": 6,
      "when": ["f16Ok(dtypes.T)", "attrs.transA == 0", "attrs.transB == 0", "attrs.transBatchA == 0", "attrs.transBatchB == 0", "ranks.A == 4", "(ranks.B == 2 or ranks.B == 3)", "ranks.Y == 4", "dim(shapes.Y, 0) == dim(shapes.A, 0)", "(ranks.B == 2 or dim(shapes.A, 1) == dim(shapes.B, 0) or dim(shapes.A, 1) == 1 or dim(shapes.B, 0) == 1)", "dim(shapes.Y, 1) == (dim(shapes.A, 1) if ranks.B == 2 else max(dim(shapes.A, 1), dim(shapes.B, 0)))", "dim(shapes.A, 3) == dim(shapes.B, ranks.B - 2)", "dim(shapes.Y, 2) == dim(shapes.A, 2)", "dim(shapes.Y, 3) == dim(shapes.B, ranks.B - 1)", "dim(shapes.A, 2) >= 64", "dim(shapes.A, 3) >= 32", "dim(shapes.B, ranks.B - 1) >= 64", "ceil(dim(shapes.B, ranks.B - 1) / 64) <= device.limits.maxComputeWorkgroupsPerDimension", "ceil(dim(shapes.A, 2) / 64) <= device.limits.maxComputeWorkgroupsPerDimension", "numel(shapes.Y) / (dim(shapes.A, 2) * dim(shapes.B, ranks.B - 1)) <= device.limits.maxComputeWorkgroupsPerDimension"],
      "constants": { "scalar": "dtypes.T", "usesF16": "dtypes.T == \"f16\"" },
      "passes": [
        {
          "id": "main",
          "name": "FusedMatMul.BroadcastRank4TiledReg",
          "source": {
            "shader": "matmul-tiled-general-reg.wgsl.jinja",
            "inputs": {
              "aShape": "shapes.A",
              "bShape": "shapes.B",
              "alpha": "attrs.alpha",
              "aRank": "ranks.A",
              "bRank": "ranks.B",
              "transBatchA": "false"
            }
          },
          "bindings": "matrix",
          "dispatch": {
            "x": "ceil(dim(shapes.B, ranks.B - 1) / 64)",
            "y": "ceil(dim(shapes.A, 2) / 64)",
            "z": "numel(shapes.Y) / (dim(shapes.A, 2) * dim(shapes.B, ranks.B - 1))"
          }
        }
      ]
    },
    {
      "id": "plain_rank2_tiled_reg",
      "description": "Register-blocked rank-2 Y = alpha * A @ B for the tier without subgroup matrices. It is the non-transposed rung the two specialised register-blocked variants beside it never covered, so plain rank-2 shapes fell to the scalar-indexed tile.",
      "priority": 4,
      "when": ["f16Ok(dtypes.T)", "attrs.transA == 0", "attrs.transB == 0", "attrs.transBatchA == 0", "attrs.transBatchB == 0", "ranks.A == 2", "ranks.B == 2", "ranks.Y == 2", "dim(shapes.A, 1) == dim(shapes.B, 0)", "dim(shapes.Y, 0) == dim(shapes.A, 0)", "dim(shapes.Y, 1) == dim(shapes.B, 1)", "dim(shapes.A, 0) >= 64", "dim(shapes.A, 1) >= 32", "dim(shapes.B, 1) >= 64", "ceil(dim(shapes.A, 0) / 64) * ceil(dim(shapes.B, 1) / 64) >= tunables.TILED_REG_MIN_WORKGROUPS", "ceil(dim(shapes.B, 1) / 64) <= device.limits.maxComputeWorkgroupsPerDimension", "ceil(dim(shapes.A, 0) / 64) <= device.limits.maxComputeWorkgroupsPerDimension"],
      "constants": { "scalar": "dtypes.T", "usesF16": "dtypes.T == \"f16\"" },
      "passes": [
        {
          "id": "main",
          "name": "FusedMatMul.PlainRank2TiledReg",
          "source": {
            "shader": "matmul-tiled-general-reg.wgsl.jinja",
            "inputs": {
              "aShape": "shapes.A",
              "bShape": "shapes.B",
              "alpha": "attrs.alpha",
              "aRank": "ranks.A",
              "bRank": "ranks.B",
              "transBatchA": "false"
            }
          },
          "bindings": "matrix",
          "dispatch": { "x": "ceil(dim(shapes.B, 1) / 64)", "y": "ceil(dim(shapes.A, 0) / 64)", "z": 1 }
        }
      ]
    },
    {
      "id": "transbatch_a_tiled_reg",
      "priority": 5,
      "when": ["f16Ok(dtypes.T)", "attrs.transBatchA != 0", "attrs.transBatchB == 0", "attrs.transA == 0", "attrs.transB == 0", "ranks.A == 3", "ranks.B == 3", "ranks.Y == 3", "dim(shapes.A, 1) == dim(shapes.B, 0)", "dim(shapes.Y, 0) == dim(shapes.B, 0)", "dim(shapes.A, 2) == dim(shapes.B, 1)", "dim(shapes.Y, 1) == dim(shapes.A, 0)", "dim(shapes.Y, 2) == dim(shapes.B, 2)", "dim(shapes.A, 0) >= 64", "dim(shapes.A, 2) >= 32", "dim(shapes.B, 2) >= 64", "ceil(dim(shapes.B, 2) / 64) <= device.limits.maxComputeWorkgroupsPerDimension", "ceil(dim(shapes.A, 0) / 64) <= device.limits.maxComputeWorkgroupsPerDimension", "dim(shapes.Y, 0) <= device.limits.maxComputeWorkgroupsPerDimension"],
      "constants": { "scalar": "dtypes.T", "usesF16": "dtypes.T == \"f16\"" },
      "passes": [
        {
          "id": "main",
          "name": "FusedMatMul.TransBatchATiledReg",
          "source": {
            "shader": "matmul-tiled-general-reg.wgsl.jinja",
            "inputs": {
              "aShape": "shapes.A",
              "bShape": "shapes.B",
              "alpha": "attrs.alpha",
              "aRank": "ranks.A",
              "bRank": "ranks.B",
              "transBatchA": "true",
              "kTile": "4"
            }
          },
          "bindings": "matrix",
          "dispatch": {
            "x": "ceil(dim(shapes.B, 2) / 64)",
            "y": "ceil(dim(shapes.A, 0) / 64)",
            "z": "dim(shapes.Y, 0)"
          }
        }
      ]
    },
    {
      "id": "tiled",
      "priority": 0,
      "when": ["ranks.A >= 1", "ranks.B >= 1", "f16Ok(dtypes.T)", "transBatchContract", "(dim(shapes.A, 0) if ranks.A == 1 else (dim(shapes.A, ranks.A - 1) if attrs.transA == 0 else (dim(shapes.A, 0) if attrs.transBatchA != 0 else dim(shapes.A, ranks.A - 2)))) == (dim(shapes.B, 0) if ranks.B == 1 else (dim(shapes.B, ranks.B - 1) if attrs.transB != 0 else (dim(shapes.B, 0) if attrs.transBatchB != 0 else dim(shapes.B, ranks.B - 2))))", "ceil((1 if ranks.A == 1 else (dim(shapes.A, ranks.A - 1) if attrs.transA != 0 else (dim(shapes.A, 0) if attrs.transBatchA != 0 else dim(shapes.A, ranks.A - 2)))) / 16) <= device.limits.maxComputeWorkgroupsPerDimension", "ceil((1 if ranks.B == 1 else (dim(shapes.B, ranks.B - 1) if attrs.transB == 0 else (dim(shapes.B, 0) if attrs.transBatchB != 0 else dim(shapes.B, ranks.B - 2)))) / 16) <= device.limits.maxComputeWorkgroupsPerDimension", "numel(shapes.Y) / max(1, (1 if ranks.A == 1 else (dim(shapes.A, ranks.A - 1) if attrs.transA != 0 else (dim(shapes.A, 0) if attrs.transBatchA != 0 else dim(shapes.A, ranks.A - 2)))) * (1 if ranks.B == 1 else (dim(shapes.B, ranks.B - 1) if attrs.transB == 0 else (dim(shapes.B, 0) if attrs.transBatchB != 0 else dim(shapes.B, ranks.B - 2))))) <= device.limits.maxComputeWorkgroupsPerDimension"],
      "constants": { "scalar": "dtypes.T", "usesF16": "dtypes.T == \"f16\"" },
      "passes": [
        {
          "id": "main",
          "name": "FusedMatMul.Tiled",
          "source": {
            "shader": "matmul-tiled-general.wgsl.jinja",
            "inputs": {
              "aShape": "shapes.A",
              "bShape": "shapes.B",
              "transA": "attrs.transA != 0",
              "transB": "attrs.transB != 0",
              "alpha": "attrs.alpha",
              "aRank": "ranks.A",
              "bRank": "ranks.B",
              "transBatchA": "attrs.transBatchA != 0",
              "transBatchB": "attrs.transBatchB != 0"
            }
          },
          "bindings": "matrix",
          "dispatch": {
            "x": "ceil((1 if ranks.B == 1 else (dim(shapes.B, ranks.B - 1) if attrs.transB == 0 else (dim(shapes.B, 0) if attrs.transBatchB != 0 else dim(shapes.B, ranks.B - 2)))) / 32)",
            "y": "ceil((1 if ranks.A == 1 else (dim(shapes.A, ranks.A - 1) if attrs.transA != 0 else (dim(shapes.A, 0) if attrs.transBatchA != 0 else dim(shapes.A, ranks.A - 2)))) / 32)",
            "z": "numel(shapes.Y) / max(1, (1 if ranks.A == 1 else (dim(shapes.A, ranks.A - 1) if attrs.transA != 0 else (dim(shapes.A, 0) if attrs.transBatchA != 0 else dim(shapes.A, ranks.A - 2)))) * (1 if ranks.B == 1 else (dim(shapes.B, ranks.B - 1) if attrs.transB == 0 else (dim(shapes.B, 0) if attrs.transBatchB != 0 else dim(shapes.B, ranks.B - 2)))))"
          }
        }
      ]
    }
  ]
}