{
  "domain": "com.microsoft",
  "name": "NGramHashMapping",
  "sinceVersion": 1,
  "inputs": {
    "inputIdsT": { "onnx": "input_ids", "dtype": "M", "rank": 2 },
    "multipliersT": { "onnx": "multipliers", "dtype": "M", "rank": 1 },
    "vocabSizesT": { "onnx": "vocab_sizes", "dtype": "M", "rank": 1 },
    "pastIdsT": { "onnx": "past_ids", "dtype": "M", "rank": 2, "optional": true }
  },
  "outputs": {
    "hashIdsT": {
      "onnx": "hash_ids",
      "dtype": "M",
      "rank": 3,
      "shape": "[dim(shapes.inputIdsT, 0), dim(shapes.inputIdsT, 1), (attrs.max_ngram_size - 1) * attrs.n_head_per_ngram]"
    },
    "presentIdsT": {
      "onnx": "present_ids",
      "dtype": "M",
      "rank": 2,
      "shape": "[dim(shapes.inputIdsT, 0), attrs.max_ngram_size - 1]"
    }
  },
  "attributes": { "max_ngram_size": {}, "n_head_per_ngram": {}, "pad_id": {} },
  "attributeConstraints": {
    "max_ngram_size": { "required": true },
    "n_head_per_ngram": { "required": true },
    "pad_id": { "required": true }
  },
  "typeConstraints": { "M": ["int32"] },
  "tunables": { "WORKGROUP_SIZE": { "default": 256 }, "ITEMS_PER_LANE": { "default": 4 } },
  "derive": {
    "foldedDispatchCapacity": "min(device.limits.maxComputeWorkgroupsPerDimension, 65535) * min(device.limits.maxComputeWorkgroupsPerDimension, 65535)",
    "maxNgramSize": "attrs.max_ngram_size",
    "headsPerNgram": "attrs.n_head_per_ngram",
    "stateLength": "maxNgramSize - 1",
    "numHeads": "stateLength * headsPerNgram",
    "stagedTableScalars": "maxNgramSize + numHeads",
    "padId": "attrs.pad_id",
    "idsRankOk": "ranks.inputIdsT == 2",
    "batchSize": "dim(shapes.inputIdsT, 0) if idsRankOk else 0",
    "seqLength": "dim(shapes.inputIdsT, 1) if idsRankOk else 0",
    "spanLength": "seqLength + stateLength",
    "sharedTileSlots": "max(1, min(tunables.WORKGROUP_SIZE, spanLength, floor(device.limits.maxComputeWorkgroupStorageSize / (max(1, stateLength) * 4))))",
    "sharedSpanTiles": "ceilDiv(spanLength, sharedTileSlots)",
    "sharedDispatchOk": "sharedSpanTiles <= foldedDispatchCapacity",
    "sharedMemoryOk": "sharedTileSlots * stateLength * 4 <= device.limits.maxComputeWorkgroupStorageSize",
    "tileSlots": "tunables.WORKGROUP_SIZE * tunables.ITEMS_PER_LANE",
    "spanTiles": "ceilDiv(spanLength, max(1, tileSlots))",
    "idScalar": "dtypes.M",
    "attrsOk": "maxNgramSize >= 2 and maxNgramSize <= 16 and headsPerNgram >= 1 and numHeads <= 64 and padId >= -2147483648 and padId <= 2147483647",
    "dtypeOk": "tensorDtypes.inputIdsT == \"int32\" and tensorDtypes.multipliersT == \"int32\" and tensorDtypes.vocabSizesT == \"int32\" and tensorDtypes.hashIdsT == \"int32\" and tensorDtypes.presentIdsT == \"int32\"",
    "shapeOk": "idsRankOk and ranks.multipliersT == 1 and dim(shapes.multipliersT, 0) == maxNgramSize and ranks.vocabSizesT == 1 and dim(shapes.vocabSizesT, 0) == numHeads and ranks.hashIdsT == 3 and sameShape(shapes.hashIdsT, [batchSize, seqLength, numHeads]) and ranks.presentIdsT == 2 and sameShape(shapes.presentIdsT, [batchSize, stateLength])",
    "dispatchOk": "batchSize <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and spanTiles <= foldedDispatchCapacity",
    "workgroupOk": "tunables.WORKGROUP_SIZE > 0 and tunables.WORKGROUP_SIZE <= device.limits.maxComputeInvocationsPerWorkgroup and tunables.WORKGROUP_SIZE <= device.limits.maxComputeWorkgroupSizeX and tunables.ITEMS_PER_LANE >= 1 and floor(tunables.ITEMS_PER_LANE) == tunables.ITEMS_PER_LANE",
    "baseContract": "attrsOk and dtypeOk and shapeOk and dispatchOk and workgroupOk",
    "noPastContract": "baseContract and not present.pastIdsT",
    "pastContract": "baseContract and present.pastIdsT and tensorDtypes.pastIdsT == \"int32\" and ranks.pastIdsT == 2 and sameShape(shapes.pastIdsT, [batchSize, stateLength])"
  },
  "bindings": {
    "input_ids": { "arg": "inputIdsT", "elementType": "$idScalar" },
    "multipliers": { "arg": "multipliersT", "elementType": "$idScalar", "length": "$maxNgramSize" },
    "vocab_sizes": { "arg": "vocabSizesT", "elementType": "$idScalar", "length": "$numHeads" },
    "past_ids": { "arg": "pastIdsT", "elementType": "$idScalar" },
    "hash_ids": { "arg": "hashIdsT", "elementType": "$idScalar" },
    "present_ids": { "arg": "presentIdsT", "elementType": "$idScalar" },
    "params": { "struct": [{ "name": "seqLength", "type": "u32", "value": "seqLength" }] }
  },
  "variants": [
    {
      "id": "past",
      "when": ["pastContract"],
      "derive": { "hasPast": true },
      "passes": [
        {
          "id": "main",
          "name": "NGramHashMapping.Past",
          "shader": "ngram-hash-mapping.wgsl.jinja",
          "bindings": ["input_ids", "multipliers", "vocab_sizes", "past_ids", "hash_ids", "present_ids", "params"],
          "dispatch": {
            "x": "min(spanTiles, DISPATCH_FOLD_WIDTH)",
            "y": "batchSize",
            "z": "ceilDiv(spanTiles, DISPATCH_FOLD_WIDTH)"
          }
        }
      ]
    },
    {
      "id": "past_shared_prefix",
      "when": ["pastContract", "sharedDispatchOk", "sharedMemoryOk"],
      "derive": { "hasPast": true },
      "passes": [
        {
          "id": "main",
          "name": "NGramHashMapping.Past.SharedPrefix",
          "shader": "ngram-hash-mapping-shared-prefix.wgsl.jinja",
          "bindings": ["input_ids", "multipliers", "vocab_sizes", "past_ids", "hash_ids", "present_ids", "params"],
          "dispatch": {
            "x": "min(sharedSpanTiles, DISPATCH_FOLD_WIDTH)",
            "y": "batchSize",
            "z": "ceilDiv(sharedSpanTiles, DISPATCH_FOLD_WIDTH)"
          }
        }
      ],
      "priority": 10,
      "demoteWhen": ["stagedTableScalars < 59"]
    },
    {
      "id": "fresh",
      "when": ["noPastContract"],
      "derive": { "hasPast": false },
      "passes": [
        {
          "id": "main",
          "name": "NGramHashMapping.Fresh",
          "shader": "ngram-hash-mapping.wgsl.jinja",
          "bindings": ["input_ids", "multipliers", "vocab_sizes", "hash_ids", "present_ids", "params"],
          "dispatch": {
            "x": "min(spanTiles, DISPATCH_FOLD_WIDTH)",
            "y": "batchSize",
            "z": "ceilDiv(spanTiles, DISPATCH_FOLD_WIDTH)"
          }
        }
      ]
    },
    {
      "id": "fresh_shared_prefix",
      "when": ["noPastContract", "sharedDispatchOk", "sharedMemoryOk"],
      "derive": { "hasPast": false },
      "passes": [
        {
          "id": "main",
          "name": "NGramHashMapping.Fresh.SharedPrefix",
          "shader": "ngram-hash-mapping-shared-prefix.wgsl.jinja",
          "bindings": ["input_ids", "multipliers", "vocab_sizes", "hash_ids", "present_ids", "params"],
          "dispatch": {
            "x": "min(sharedSpanTiles, DISPATCH_FOLD_WIDTH)",
            "y": "batchSize",
            "z": "ceilDiv(sharedSpanTiles, DISPATCH_FOLD_WIDTH)"
          }
        }
      ],
      "priority": 10,
      "demoteWhen": ["stagedTableScalars < 59"]
    }
  ]
}