{
  "domain": "ai.onnx",
  "name": "LpNormalization",
  "sinceVersion": 1,
  "inputs": { "input": { "dtype": "T" } },
  "outputs": { "output": { "dtype": "T", "rank": "ranks.input", "shape": "shapes.input" } },
  "attributes": { "axis": { "default": -1 }, "p": { "default": 2 } },
  "attributeConstraints": { "p": { "values": [1, 2] } },
  "typeConstraints": { "T": ["float32", "float16"] },
  "tunables": {
    "WORKGROUP_SIZE": { "default": 256 },
    "SPLIT_MIN_DIM": { "default": 512 },
    "SPLIT_MIN_ROWS": { "default": 32 },
    "SPLIT_TARGET_DIM": { "default": 256 },
    "MAX_SPLITS": { "default": 128 }
  },
  "derive": {
    "deviceWorkgroupCap": "min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)",
    "workgroupSize": "min(tunables.WORKGROUP_SIZE, deviceWorkgroupCap)",
    "hasSubgroupId": "device.features.has(\"subgroups\") and device.wgslLanguageFeatures.has(\"subgroup_id\")",
    "scalar": "dtypes.T"
  },
  "when": ["ranks.output == ranks.input", "numel(shapes.input) == numel(shapes.output)", "f16Ok(dtypes.T)"],
  "bindings": {
    "input": { "buffer": "read-only-storage", "elementType": "$scalar" },
    "params": {
      "buffer": "uniform",
      "struct": [
        { "name": "rows", "type": "u32", "value": "axisRows" },
        { "name": "dim", "type": "u32", "value": "axisDim" },
        { "name": "inner", "type": "u32", "value": "axisInner" },
        { "name": "p", "type": "u32", "value": "attrs.p" }
      ]
    },
    "input_2": { "name": "input", "buffer": "read-only-storage", "elementType": "$divElem" },
    "rowNorms": { "buffer": "read-only-storage", "elementType": "f32" },
    "output": { "buffer": "storage", "elementType": "$divElem" },
    "params_3": {
      "name": "params",
      "buffer": "uniform",
      "struct": [
        { "name": "count", "type": "u32", "value": "divideCount" },
        { "name": "dim", "type": "u32", "value": "axisDim" },
        { "name": "inner", "type": "u32", "value": "axisInner" }
      ]
    }
  },
  "variants": [
    {
      "id": "axis_splitk",
      "priority": 20,
      "when": ["ranks.input >= 2", "inner(shapes.input, attrs.axis) > 1", "dim(shapes.input, attrs.axis) >= tunables.SPLIT_MIN_DIM", "numel(shapes.input) / dim(shapes.input, attrs.axis) >= tunables.SPLIT_MIN_ROWS", "min(tunables.MAX_SPLITS, pow2ceil(ceilDiv(dim(shapes.input, attrs.axis), tunables.SPLIT_TARGET_DIM))) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "numel(shapes.input) / dim(shapes.input, attrs.axis) * 4 <= device.limits.maxStorageBufferBindingSize", "numel(shapes.input) / dim(shapes.input, attrs.axis) * 4 <= device.limits.maxBufferSize", "numel(shapes.input) / dim(shapes.input, attrs.axis) * min(tunables.MAX_SPLITS, pow2ceil(ceilDiv(dim(shapes.input, attrs.axis), tunables.SPLIT_TARGET_DIM))) * 4 <= device.limits.maxStorageBufferBindingSize", "numel(shapes.input) / dim(shapes.input, attrs.axis) * min(tunables.MAX_SPLITS, pow2ceil(ceilDiv(dim(shapes.input, attrs.axis), tunables.SPLIT_TARGET_DIM))) * 4 <= device.limits.maxBufferSize"],
      "derive": {
        "axisDim": "dim(shapes.input, attrs.axis)",
        "axisInner": "inner(shapes.input, attrs.axis)",
        "axisRows": "numel(shapes.input) / axisDim",
        "split": "min(tunables.MAX_SPLITS, pow2ceil(ceilDiv(axisDim, tunables.SPLIT_TARGET_DIM)))",
        "divideVec4": "numel(shapes.output) % 4 == 0 and axisInner % 4 == 0",
        "divideCount": "numel(shapes.output) / 4 if divideVec4 else numel(shapes.output)",
        "divVec4": "divideVec4",
        "divElem": "(\"vec4<\" ~ dtypes.T ~ \">\") if divideVec4 else dtypes.T"
      },
      "intermediates": [
        { "id": "partials", "dtype": "float32", "shape": "[split * axisRows]" },
        { "id": "rowNorms", "dtype": "float32", "shape": "[axisRows]" }
      ],
      "passes": [
        {
          "id": "split_reduce",
          "name": "LpNormalization.SplitReduce",
          "shader": "lp-norm-reduce.wgsl.jinja",
          "derive": { "segmented": true },
          "bindings": ["input", { "scratch": "partials", "name": "normOutput", "elementType": "f32" }, "params"],
          "dispatch": {
            "x": "min(ceilDiv(axisRows, workgroupSize), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))",
            "y": "split"
          }
        },
        {
          "id": "combine",
          "name": "LpNormalization.SplitCombine",
          "shader": "lp-norm-split-combine.wgsl.jinja",
          "derive": {},
          "bindings": [
            { "name": "partials", "buffer": "read-only-storage", "elementType": "f32" },
            { "name": "rowNorms", "buffer": "storage", "elementType": "f32" },
            {
              "name": "params",
              "struct": [
                { "name": "rows", "type": "u32", "value": "axisRows" },
                { "name": "p", "type": "u32", "value": "attrs.p" }
              ]
            }
          ],
          "dispatch": {
            "x": "min(ceilDiv((axisRows), (workgroupSize)), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))",
            "y": 1,
            "z": 1
          }
        },
        {
          "id": "divide",
          "name": "LpNormalization.Divide",
          "shader": "lp-norm-divide.wgsl.jinja",
          "bindings": ["input_2", "rowNorms", "output", "params_3"],
          "dispatch": {
            "x": "min(ceilDiv((divideCount), (workgroupSize)), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))",
            "y": 1,
            "z": 1
          }
        }
      ]
    },
    {
      "id": "axis",
      "when": ["ranks.input >= 1", "numel(shapes.input) / dim(shapes.input, attrs.axis) * 4 <= device.limits.maxStorageBufferBindingSize", "numel(shapes.input) / dim(shapes.input, attrs.axis) * 4 <= device.limits.maxBufferSize"],
      "derive": {
        "axisDim": "dim(shapes.input, attrs.axis)",
        "axisInner": "inner(shapes.input, attrs.axis)",
        "axisRows": "numel(shapes.input) / axisDim",
        "divideCount": "numel(shapes.output)",
        "divVec4": false,
        "divElem": "dtypes.T"
      },
      "intermediates": [{ "id": "rowNorms", "dtype": "float32", "shape": "[axisRows]" }],
      "passes": [
        {
          "id": "reduce",
          "name": "LpNormalization.RowReduce",
          "shader": "lp-norm-reduce.wgsl.jinja",
          "derive": { "segmented": false },
          "bindings": ["input", { "scratch": "rowNorms", "name": "normOutput", "elementType": "f32" }, "params"],
          "dispatch": {
            "x": "min(ceilDiv((axisRows), (workgroupSize)), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))",
            "y": 1,
            "z": 1
          }
        },
        {
          "id": "divide",
          "name": "LpNormalization.Divide",
          "shader": "lp-norm-divide.wgsl.jinja",
          "bindings": ["input_2", "rowNorms", "output", "params_3"],
          "dispatch": {
            "x": "min(ceilDiv((divideCount), (workgroupSize)), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))",
            "y": 1,
            "z": 1
          }
        }
      ]
    },
    {
      "id": "last_axis_row_vec4",
      "priority": 110,
      "when": ["ranks.input >= 1", "(attrs.axis == -1 or attrs.axis == ranks.input - 1)", "dim(shapes.input, -1) % 4 == 0"],
      "derive": { "ioElement": "\"vec4<\" ~ dtypes.T ~ \">\"" },
      "passes": [
        {
          "id": "main",
          "name": "LpNormalization.LastAxisRow",
          "shader": "norm-row-stats.wgsl.jinja",
          "derive": {
            "modeSpec": "\"lp\"",
            "vec4": true,
            "scalar": "dtypes.T",
            "usesF16Spec": "dtypes.T == \"f16\"",
            "hidden": "dim(shapes.input, -1)",
            "wg": "min(workgroupSize, pow2ceil(dim(shapes.input, -1) / 4))",
            "p": "attrs.p",
            "hiddenVec": "dim(shapes.input, -1) / 4",
            "vecType": "\"vec4<\" ~ dtypes.T ~ \">\"",
            "combineSubgroups": "hasSubgroupId"
          },
          "bindings": [
            { "arg": "input", "name": "x", "elementType": "$ioElement" },
            { "arg": "output", "name": "y", "elementType": "$ioElement" },
            {
              "name": "params",
              "struct": [
                { "name": "rows", "type": "u32", "value": "numel(shapes.input) / dim(shapes.input, -1)" },
                {
                  "name": "rowStride",
                  "type": "u32",
                  "value": "max(1, min(numel(shapes.input) / dim(shapes.input, -1), min(device.limits.maxComputeWorkgroupsPerDimension, 65535)))"
                }
              ]
            }
          ],
          "dispatch": {
            "x": "min(numel(shapes.input) / dim(shapes.input, -1), 65535)",
            "y": "ceilDiv(numel(shapes.input) / dim(shapes.input, -1), 65535)",
            "z": 1
          },
          "subgroupCollectivesWidth": "portable"
        }
      ]
    },
    {
      "id": "last_axis_row",
      "priority": 100,
      "when": ["ranks.input >= 1", "(attrs.axis == -1 or attrs.axis == ranks.input - 1)", "true"],
      "derive": { "ioElement": "dtypes.T" },
      "passes": [
        {
          "id": "main",
          "name": "LpNormalization.LastAxisRow",
          "shader": "norm-row-stats.wgsl.jinja",
          "derive": {
            "modeSpec": "\"lp\"",
            "vec4": false,
            "scalar": "dtypes.T",
            "usesF16Spec": "dtypes.T == \"f16\"",
            "hidden": "dim(shapes.input, -1)",
            "wg": "min(workgroupSize, pow2ceil(dim(shapes.input, -1)))",
            "p": "attrs.p",
            "hiddenVec": 1,
            "vecType": "\"vec4<\" ~ dtypes.T ~ \">\"",
            "combineSubgroups": "hasSubgroupId"
          },
          "bindings": [
            { "arg": "input", "name": "x", "elementType": "$ioElement" },
            { "arg": "output", "name": "y", "elementType": "$ioElement" },
            {
              "name": "params",
              "struct": [
                { "name": "rows", "type": "u32", "value": "numel(shapes.input) / dim(shapes.input, -1)" },
                {
                  "name": "rowStride",
                  "type": "u32",
                  "value": "max(1, min(numel(shapes.input) / dim(shapes.input, -1), min(device.limits.maxComputeWorkgroupsPerDimension, 65535)))"
                }
              ]
            }
          ],
          "dispatch": {
            "x": "min(numel(shapes.input) / dim(shapes.input, -1), 65535)",
            "y": "ceilDiv(numel(shapes.input) / dim(shapes.input, -1), 65535)",
            "z": 1
          },
          "subgroupCollectivesWidth": "portable"
        }
      ]
    }
  ]
}