{ "domain": "ai.onnx", "name": "ReduceLogSumExp", "sinceVersion": 18, "inputs": { "x": { "onnx": "data", "dtype": "T" } }, "outputs": { "y": { "onnx": "reduced", "dtype": "T", "rank": "ranks.x if attrs.keepdims == 1 or ((attrs.axes | length) == 0 and attrs.noop_with_empty_axes == 1) else (ranks.x - (attrs.axes | length) if (attrs.axes | length) > 0 else 0)" } }, "attributes": { "keepdims": { "default": 1 }, "noop_with_empty_axes": { "default": 0 }, "axes": { "default": [] } }, "attributeConstraints": { "keepdims": { "values": [0, 1] }, "noop_with_empty_axes": { "values": [0, 1] } }, "typeConstraints": { "T": ["float32", "float16", "int32"] }, "tunables": { "WORKGROUP_SIZE": { "default": 256 }, "VECTOR_WIDTH": { "default": 4 }, "ROW_PARALLEL_MIN_COLS": { "default": 64 }, "SUBGROUP_MIN_COLS": { "default": 256 }, "SUBGROUP_SMALL_ROW_LIMIT": { "default": 32768 }, "AXIS0_SPLIT_MIN_ROWS": { "default": 8192 }, "AXIS0_SPLIT_TARGET_ROWS": { "default": 256 }, "AXIS0_MAX_SPLITS": { "default": 128 }, "AXIS0_TILE_MIN_ROWS": { "default": 64 }, "AXIS0_TILE_MIN_COLS": { "default": 16 }, "AXIS0_TILE_COLS": { "default": 16 }, "AXIS_SPLIT_TILE_COLS": { "default": 8 }, "FULL_REDUCE_MIN_ELEMENTS": { "default": 8192 }, "FULL_REDUCE_MAX_SPLITS": { "default": 256 }, "CONTIGUOUS_SUFFIX_MIN_COLS": { "default": 256 }, "AXES02_WORKGROUP_SIZE": { "default": 256 }, "ROW_SERIAL_MIN_ROWS": { "default": 8192 }, "ROW_SERIAL_MAX_COLS": { "default": 1024 } }, "derive": { "deviceWorkgroupCap": "min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)", "reduceWorkgroupSize": "min(tunables.WORKGROUP_SIZE, deviceWorkgroupCap)", "treeWorkgroupOk": "reduceWorkgroupSize > 0 and pow2ceil(reduceWorkgroupSize) == reduceWorkgroupSize and reduceWorkgroupSize * dtypeBytes(\"float32\") <= device.limits.maxComputeWorkgroupStorageSize", "subgroupWorkgroupFloor": "min(reduceWorkgroupSize, max(1, device.adapterInfo.subgroupMaxSize))", "lastAxisRows": "rows(shapes.x, ranks.x - 1) if ranks.x > 0 else 1", "lastAxisCols": "dim(shapes.x, ranks.x - 1) if ranks.x > 0 else 1", "rowSerialPreferred": "lastAxisRows >= tunables.ROW_SERIAL_MIN_ROWS and lastAxisCols <= tunables.ROW_SERIAL_MAX_COLS", "axis0Rows": "dim(shapes.x, 0) if ranks.x >= 2 else 0", "axis0Cols": "dim(shapes.x, 1) if ranks.x >= 2 else 0", "axis0SplitCount": "min(tunables.AXIS0_MAX_SPLITS, pow2ceil(ceilDiv(axis0Rows, tunables.AXIS0_SPLIT_TARGET_ROWS)))", "axis0SplitScratchBytes": "3 * axis0SplitCount * axis0Cols * dtypeBytes(\"float32\")", "axis0SplitPathFits": "axis0SplitCount <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and ceilDiv(ceilDiv(axis0Cols, reduceWorkgroupSize), min(device.limits.maxComputeWorkgroupsPerDimension, 65535)) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and axis0SplitScratchBytes <= device.limits.maxStorageBufferBindingSize and axis0SplitScratchBytes <= device.limits.maxBufferSize", "reduceAxis": "(attrs.axes[0] + ranks.x if attrs.axes[0] < 0 else attrs.axes[0]) if ((attrs.axes | length) == 1 and isUniqueIntList(attrs.axes, 0 - ranks.x, ranks.x, 1)) else ranks.x", "axisSplitDim": "dim(shapes.x, reduceAxis) if ranks.x >= 2 and reduceAxis < ranks.x else 0", "axisSplitInner": "inner(shapes.x, reduceAxis) if ranks.x >= 2 and reduceAxis < ranks.x else 1", "axisSplitOutputs": "numel(shapes.y)", "axisSplitCount": "min(tunables.AXIS0_MAX_SPLITS, pow2ceil(ceilDiv(axisSplitDim, tunables.AXIS0_SPLIT_TARGET_ROWS)))", "axisSplitScratchBytes": "3 * axisSplitCount * axisSplitOutputs * 4", "axisSplitPathFits": "axisSplitCount <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and ceilDiv(ceilDiv(axisSplitOutputs, reduceWorkgroupSize), min(device.limits.maxComputeWorkgroupsPerDimension, 65535)) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and axisSplitScratchBytes <= device.limits.maxStorageBufferBindingSize and axisSplitScratchBytes <= device.limits.maxBufferSize", "axis0TilePathFits": "treeWorkgroupOk and tunables.AXIS0_TILE_COLS > 0 and tunables.AXIS0_TILE_COLS <= reduceWorkgroupSize and reduceWorkgroupSize % tunables.AXIS0_TILE_COLS == 0", "flatItems": "numel(shapes.x) / tunables.VECTOR_WIDTH if numel(shapes.x) % tunables.VECTOR_WIDTH == 0 else numel(shapes.x)", "flatSplitCount": "max(1, min(tunables.FULL_REDUCE_MAX_SPLITS, ceilDiv(flatItems, reduceWorkgroupSize)))", "flatScratchBytes": "3 * flatSplitCount * dtypeBytes(\"float32\")", "flatPathFits": "treeWorkgroupOk and flatSplitCount <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and flatScratchBytes <= device.limits.maxStorageBufferBindingSize and flatScratchBytes <= device.limits.maxBufferSize", "flatParallelCovered": "(dtypes.T == \"f32\" or dtypes.T == \"f16\") and f16Ok(dtypes.T) and numel(shapes.y) == 1 and numel(shapes.x) >= tunables.FULL_REDUCE_MIN_ELEMENTS and flatPathFits", "contiguousSuffixParallelCovered": "(dtypes.T == \"f32\" or dtypes.T == \"f16\") and f16Ok(dtypes.T) and numel(shapes.y) > 0 and numel(shapes.x) % numel(shapes.y) == 0 and numel(shapes.x) / numel(shapes.y) >= tunables.CONTIGUOUS_SUFFIX_MIN_COLS and ((ranks.x == 3 and hasAxis(attrs.axes, 0, 3) == false and hasAxis(attrs.axes, 1, 3) and hasAxis(attrs.axes, 2, 3) and numel(shapes.y) == dim(shapes.x, 0)) or (ranks.x == 4 and hasAxis(attrs.axes, 0, 4) == false and hasAxis(attrs.axes, 1, 4) == false and hasAxis(attrs.axes, 2, 4) and hasAxis(attrs.axes, 3, 4) and numel(shapes.y) == dim(shapes.x, 0) * dim(shapes.x, 1)) or (ranks.x == 4 and hasAxis(attrs.axes, 0, 4) == false and hasAxis(attrs.axes, 1, 4) and hasAxis(attrs.axes, 2, 4) and hasAxis(attrs.axes, 3, 4) and numel(shapes.y) == dim(shapes.x, 0)))" }, "bindings": { "x": { "buffer": "read-only-storage", "elementType": "$vectorScalar" }, "y": { "buffer": "storage", "elementType": "$T" }, "params": { "buffer": "uniform", "struct": [ { "name": "rows", "type": "u32", "value": "numel(shapes.y)" }, { "name": "chunkCount", "type": "u32", "value": "numel(shapes.x) / numel(shapes.y) / tunables.VECTOR_WIDTH" } ] }, "x_2": { "name": "x", "buffer": "read-only-storage", "elementType": "$T" }, "params_2": { "name": "params", "buffer": "uniform", "struct": [ { "name": "rows", "type": "u32", "value": "numel(shapes.y)" }, { "name": "cols", "type": "u32", "value": "numel(shapes.x) / numel(shapes.y)" } ] }, "params_3": { "name": "params", "buffer": "uniform", "struct": [{ "name": "outCount", "type": "u32", "value": "numel(shapes.y)" }] }, "params_5": { "name": "params", "buffer": "uniform", "struct": [{ "name": "count", "type": "u32", "value": "numel(shapes.y)" }] }, "params_6": { "name": "params", "buffer": "uniform", "struct": [ { "name": "rows", "type": "u32", "value": "rows(shapes.x, ranks.x - 1)" }, { "name": "chunkCount", "type": "u32", "value": "dim(shapes.x, ranks.x - 1) / tunables.VECTOR_WIDTH" } ] }, "params_7": { "name": "params", "buffer": "uniform", "struct": [ { "name": "rows", "type": "u32", "value": "1" }, { "name": "cols", "type": "u32", "value": "1" }, { "name": "outCount", "type": "u32", "value": "1" } ] }, "params_8": { "name": "params", "buffer": "uniform", "struct": [ { "name": "rows", "type": "u32", "value": "dim(shapes.x, 0)" }, { "name": "cols", "type": "u32", "value": "1" }, { "name": "outCount", "type": "u32", "value": "numel(shapes.y)" } ] }, "params_9": { "name": "params", "buffer": "uniform", "struct": [ { "name": "rows", "type": "u32", "value": "rows(shapes.x, ranks.x - 1)" }, { "name": "cols", "type": "u32", "value": "dim(shapes.x, ranks.x - 1)" } ] }, "partials": { "buffer": "storage", "elementType": "$partialElement" }, "params_10": { "name": "params", "buffer": "uniform", "struct": [ { "name": "axisDim", "type": "u32", "value": "axisSplitDim" }, { "name": "inner", "type": "u32", "value": "axisSplitInner" }, { "name": "outputs", "type": "u32", "value": "axisSplitOutputs" } ] }, "partials_2": { "name": "partials", "buffer": "read-only-storage", "elementType": "$partialElement" }, "params_11": { "name": "params", "buffer": "uniform", "struct": [{ "name": "cols", "type": "u32", "value": "axisSplitOutputs" }] }, "params_12": { "name": "params", "buffer": "uniform", "struct": [ { "name": "rows", "type": "u32", "value": "dim(shapes.x, 0)" }, { "name": "cols", "type": "u32", "value": "dim(shapes.x, 1)" } ] }, "params_13": { "name": "params", "buffer": "uniform", "struct": [{ "name": "cols", "type": "u32", "value": "dim(shapes.x, 1)" }] }, "params_16": { "name": "params", "buffer": "uniform", "struct": [ { "name": "axisDim", "type": "u32", "value": "axisSplitDim" }, { "name": "outCount", "type": "u32", "value": "numel(shapes.y)" } ] }, "params_17": { "name": "params", "buffer": "uniform", "struct": [ { "name": "rows", "type": "u32", "value": "rows(shapes.x, ranks.x - 1)" }, { "name": "chunkCount", "type": "u32", "value": "dim(shapes.x, ranks.x - 1)" } ] }, "params_18": { "name": "params", "buffer": "uniform", "struct": [ { "name": "rows", "type": "u32", "value": "dim(shapes.x, 0)" }, { "name": "cols", "type": "u32", "value": "dim(shapes.x, 1)" }, { "name": "outCount", "type": "u32", "value": "numel(shapes.y)" } ] }, "params_19": { "name": "params", "buffer": "uniform", "struct": [ { "name": "cols", "type": "u32", "value": "dim(shapes.x, 1)" }, { "name": "outCount", "type": "u32", "value": "numel(shapes.y)" } ] } }, "variants": [ { "id": "contiguous_suffix_subgroup_vec4", "priority": 30, "when": ["device.wgslLanguageFeatures.has(\"subgroup_id\")", "not flatParallelCovered", "contiguousSuffixParallelCovered", "(numel(shapes.x) / numel(shapes.y)) % tunables.VECTOR_WIDTH == 0"], "requires": { "features": ["subgroups"] }, "derive": { "scalar": "dtypes.T", "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"", "workgroupSize": "min(reduceWorkgroupSize, max(subgroupWorkgroupFloor, pow2ceil(ceilDiv(numel(shapes.x) / numel(shapes.y), tunables.VECTOR_WIDTH))))" }, "passes": [ { "id": "main", "name": "ReduceLogSumExp.ContiguousSuffixSubgroupVec4", "shader": "reduce-row-subgroup.wgsl.jinja", "derive": { "op": "\"logsumexp\"", "vec4": true, "castF32": "dtypes.T == \"f16\"", "usesF16Spec": "dtypes.T == \"f16\"" }, "bindings": ["x", "y", "params"], "dispatch": { "x": "min(numel(shapes.y), 65535)", "y": "ceilDiv(numel(shapes.y), 65535)", "z": 1 }, "subgroupCollectivesWidth": "portable" } ] }, { "id": "contiguous_suffix_tree_vec4", "priority": 22, "when": ["not flatParallelCovered", "contiguousSuffixParallelCovered", "(numel(shapes.x) / numel(shapes.y)) % tunables.VECTOR_WIDTH == 0", "treeWorkgroupOk"], "derive": { "scalar": "dtypes.T", "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"", "workgroupSize": "min(reduceWorkgroupSize, pow2ceil(ceilDiv(numel(shapes.x) / numel(shapes.y), tunables.VECTOR_WIDTH)))" }, "passes": [ { "id": "main", "name": "ReduceLogSumExp.ContiguousSuffixTreeVec4", "shader": "reduce-row-tree.wgsl.jinja", "derive": { "op": "\"logsumexp\"", "vec4": true, "castF32": "dtypes.T == \"f16\"", "usesF16Spec": "dtypes.T == \"f16\"" }, "bindings": ["x", "y", "params"], "dispatch": { "x": "min(numel(shapes.y), 65535)", "y": "ceilDiv(numel(shapes.y), 65535)", "z": 1 } } ] }, { "id": "contiguous_suffix_tree", "priority": 21, "when": ["not flatParallelCovered", "contiguousSuffixParallelCovered", "treeWorkgroupOk"], "derive": { "workgroupSize": "min(reduceWorkgroupSize, pow2ceil(numel(shapes.x) / numel(shapes.y)))", "scalar": "dtypes.T" }, "passes": [ { "id": "main", "name": "ReduceLogSumExp.ContiguousSuffixTree", "shader": "reduce-row-tree.wgsl.jinja", "derive": { "op": "\"logsumexp\"", "castF32": "dtypes.T == \"f16\"", "usesF16Spec": "dtypes.T == \"f16\"" }, "bindings": ["x_2", "y", "params_2"], "dispatch": { "x": "min(numel(shapes.y), 65535)", "y": "ceilDiv(numel(shapes.y), 65535)", "z": 1 } } ] }, { "id": "multi_axis_rank3", "priority": 8, "when": ["not flatParallelCovered", "not contiguousSuffixParallelCovered", "f16Ok(dtypes.T)", "ranks.x == 3", "(attrs.keepdims == 1 and ranks.y == 3 and (dim(shapes.y, 0) == 1 if hasAxis(attrs.axes, 0, 3) else dim(shapes.y, 0) == dim(shapes.x, 0)) and (dim(shapes.y, 1) == 1 if hasAxis(attrs.axes, 1, 3) else dim(shapes.y, 1) == dim(shapes.x, 1)) and (dim(shapes.y, 2) == 1 if hasAxis(attrs.axes, 2, 3) else dim(shapes.y, 2) == dim(shapes.x, 2))) or (attrs.keepdims == 0 and ranks.y == 1)"], "derive": { "scalar": "dtypes.T", "reduceWorkgroupSize": "reduceWorkgroupSize" }, "passes": [ { "id": "main", "name": "ReduceLogSumExp.MultiAxisRank3", "shader": "reduce-serial-axis.wgsl.jinja", "derive": { "op": "\"logsumexp\"", "indexing": "\"multiaxis\"", "rank": 3, "reduce": ["hasAxis(attrs.axes, 0, 3)", "hasAxis(attrs.axes, 1, 3)", "hasAxis(attrs.axes, 2, 3)"], "dataShape": "shapes.x", "outputShape": "shapes.y", "outputRank": "ranks.y", "keepDims": "attrs.keepdims != 0", "intMode": "dtypes.T == \"i32\"", "castF32": "dtypes.T == \"f16\"", "usesF16Spec": "dtypes.T == \"f16\"" }, "bindings": ["x_2", "y", "params_3"], "dispatch": { "x": "min(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)", "y": "ceilDiv(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)", "z": 1 } } ] }, { "id": "multi_axis_rank4", "priority": 8, "when": ["not flatParallelCovered", "not contiguousSuffixParallelCovered", "f16Ok(dtypes.T)", "ranks.x == 4", "attrs.noop_with_empty_axes == 0", "numel(shapes.y) == (1 if hasAxis(attrs.axes, 0, 4) else dim(shapes.x, 0)) * (1 if hasAxis(attrs.axes, 1, 4) else dim(shapes.x, 1)) * (1 if hasAxis(attrs.axes, 2, 4) else dim(shapes.x, 2)) * (1 if hasAxis(attrs.axes, 3, 4) else dim(shapes.x, 3))", "((attrs.keepdims == 1 and ranks.y == 4) or (attrs.keepdims == 0 and ranks.y < 4))"], "derive": { "scalar": "dtypes.T", "reduceWorkgroupSize": "reduceWorkgroupSize" }, "passes": [ { "id": "main", "name": "ReduceLogSumExp.MultiAxisRank4", "shader": "reduce-serial-axis.wgsl.jinja", "derive": { "op": "\"logsumexp\"", "indexing": "\"multiaxis\"", "rank": 4, "reduce": ["hasAxis(attrs.axes, 0, 4)", "hasAxis(attrs.axes, 1, 4)", "hasAxis(attrs.axes, 2, 4)", "hasAxis(attrs.axes, 3, 4)"], "dataShape": "shapes.x", "outputShape": "shapes.y", "outputRank": "ranks.y", "keepDims": "attrs.keepdims != 0", "intMode": "dtypes.T == \"i32\"", "castF32": "dtypes.T == \"f16\"", "usesF16Spec": "dtypes.T == \"f16\"" }, "bindings": ["x_2", "y", "params_3"], "dispatch": { "x": "min(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)", "y": "ceilDiv(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)", "z": 1 } } ] }, { "id": "int32_rank3_axes02_keepdims", "priority": 30, "when": ["dtypes.T == \"i32\"", "ranks.x == 3", "attrs.keepdims == 1", "hasAxis(attrs.axes, 0, 3)", "hasAxis(attrs.axes, 2, 3)", "hasAxis(attrs.axes, 1, 3) == false", "dim(shapes.x, 0) > 0", "dim(shapes.x, 2) > 0", "ranks.y == 3", "dim(shapes.y, 0) == 1", "dim(shapes.y, 1) == dim(shapes.x, 1)", "dim(shapes.y, 2) == 1"], "derive": { "axes02WorkgroupSize": "min(tunables.AXES02_WORKGROUP_SIZE, device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)" }, "passes": [ { "id": "main", "name": "ReduceLogSumExp.Int32Rank3Axes02Keepdims", "shader": "reduce-i32-axes02.wgsl.jinja", "derive": { "op": "\"logsumexp\"", "workgroupSizeSpec": "axes02WorkgroupSize" }, "bindings": [ "x_2", "y", { "name": "params", "struct": [ { "name": "d0", "type": "u32", "value": "dim(shapes.x, 0)" }, { "name": "d1", "type": "u32", "value": "dim(shapes.x, 1)" }, { "name": "d2", "type": "u32", "value": "dim(shapes.x, 2)" }, { "name": "outCount", "type": "u32", "value": "numel(shapes.y)" } ] } ], "dispatch": { "x": "min(ceilDiv((numel(shapes.y)), (axes02WorkgroupSize)), 65535)", "y": "ceilDiv(ceilDiv((numel(shapes.y)), (axes02WorkgroupSize)), 65535)", "z": 1 } } ] }, { "id": "noop_empty_axes", "priority": 40, "when": ["dtypes.T == \"f32\"", "attrs.noop_with_empty_axes == 1", "(attrs.axes | length) == 0", "sameShape(shapes.x, shapes.y)"], "derive": { "reduceWorkgroupSize": "reduceWorkgroupSize" }, "passes": [ { "id": "main", "name": "ReduceLogSumExp.NoopEmptyAxes", "shader": "reduce-noop-empty-axes.wgsl.jinja", "derive": { "op": "\"identity\"" }, "bindings": ["x_2", "y", "params_5"], "dispatch": { "x": "min(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)", "y": "ceilDiv(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)", "z": 1 } } ] }, { "id": "subgroup_rows_last_axis_vec4", "priority": 26, "when": ["not flatParallelCovered", "(dtypes.T == \"f32\" or dtypes.T == \"f16\")", "f16Ok(dtypes.T)", "ranks.x >= 1", "reduceAxis == ranks.x - 1", "numel(shapes.y) == rows(shapes.x, ranks.x - 1)", "attrs.noop_with_empty_axes == 0", "lastAxisCols % tunables.VECTOR_WIDTH == 0", "lastAxisCols >= tunables.VECTOR_WIDTH", "has(device.adapterInfo, \"subgroupMinSize\")", "has(device.adapterInfo, \"subgroupMaxSize\")", "device.adapterInfo.subgroupMinSize >= 16", "ceilDiv(lastAxisCols / tunables.VECTOR_WIDTH, device.adapterInfo.subgroupMinSize) <= 8", "lastAxisRows >= 64", "treeWorkgroupOk", "device.adapterInfo.subgroupMaxSize <= reduceWorkgroupSize", "(not rowSerialPreferred or lastAxisCols >= tunables.SUBGROUP_MIN_COLS)"], "requires": { "features": ["subgroups"] }, "derive": { "scalar": "dtypes.T", "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"", "workgroupSize": "reduceWorkgroupSize", "vecsPerLane": "ceilDiv(lastAxisCols / tunables.VECTOR_WIDTH, device.adapterInfo.subgroupMinSize)" }, "passes": [ { "id": "main", "name": "ReduceLogSumExp.SubgroupRowsVec4", "shader": "reduce-row-subgroup-rows.wgsl.jinja", "derive": { "op": "\"logsumexp\"", "castF32": "dtypes.T == \"f16\"", "usesF16Spec": "dtypes.T == \"f16\"" }, "bindings": ["x", "y", "params_6"], "dispatch": { "x": "min(ceilDiv(lastAxisRows, reduceWorkgroupSize / device.adapterInfo.subgroupMaxSize), 65535)", "y": "ceilDiv(ceilDiv(lastAxisRows, reduceWorkgroupSize / device.adapterInfo.subgroupMaxSize), 65535)", "z": 1 }, "subgroupCollectivesWidth": "portable" } ] }, { "id": "tree_last_axis_vec4", "priority": 23, "when": ["not flatParallelCovered", "(dtypes.T == \"f32\" or dtypes.T == \"f16\")", "f16Ok(dtypes.T)", "ranks.x >= 1", "reduceAxis == ranks.x - 1", "numel(shapes.y) == rows(shapes.x, ranks.x - 1)", "attrs.noop_with_empty_axes == 0", "lastAxisCols >= tunables.ROW_PARALLEL_MIN_COLS", "lastAxisCols % tunables.VECTOR_WIDTH == 0", "treeWorkgroupOk"], "demoteWhen": ["rowSerialPreferred"], "derive": { "scalar": "dtypes.T", "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"", "workgroupSize": "min(reduceWorkgroupSize, pow2ceil(ceilDiv(lastAxisCols, tunables.VECTOR_WIDTH)))" }, "passes": [ { "id": "main", "name": "ReduceLogSumExp.TreeRowVec4", "shader": "reduce-row-tree.wgsl.jinja", "derive": { "op": "\"logsumexp\"", "vec4": true, "castF32": "dtypes.T == \"f16\"", "usesF16Spec": "dtypes.T == \"f16\"" }, "bindings": ["x", "y", "params_6"], "dispatch": { "x": "min(lastAxisRows, 65535)", "y": "ceilDiv(lastAxisRows, 65535)", "z": 1 } } ] }, { "id": "rank0_scalar", "priority": 40, "when": ["f16Ok(dtypes.T)", "ranks.x == 0", "ranks.y == 0"], "derive": { "axis": 0, "scalar": "dtypes.T", "reduceWorkgroupSize": "reduceWorkgroupSize" }, "passes": [ { "id": "main", "name": "ReduceLogSumExp.Rank0Scalar", "shader": "reduce-serial-axis.wgsl.jinja", "derive": { "op": "\"logsumexp\"", "indexing": "\"axis2d\"", "intMode": "dtypes.T == \"i32\"", "castF32": "dtypes.T == \"f16\"", "usesF16Spec": "dtypes.T == \"f16\"", "logicalBool": "tensorDtypes.x == \"bool\"" }, "bindings": ["x_2", "y", "params_7"], "dispatch": { "x": 1 } } ] }, { "id": "rank1_axis0", "when": ["not flatParallelCovered", "f16Ok(dtypes.T)", "ranks.x == 1", "reduceAxis == 0", "((attrs.keepdims == 0 and ranks.y == 0) or (attrs.keepdims == 1 and ranks.y == 1 and dim(shapes.y, 0) == 1))"], "derive": { "axis": 0, "scalar": "dtypes.T", "reduceWorkgroupSize": "reduceWorkgroupSize" }, "passes": [ { "id": "main", "name": "ReduceLogSumExp.Rank1Axis0", "shader": "reduce-serial-axis.wgsl.jinja", "derive": { "op": "\"logsumexp\"", "indexing": "\"axis2d\"", "intMode": "dtypes.T == \"i32\"", "castF32": "dtypes.T == \"f16\"", "usesF16Spec": "dtypes.T == \"f16\"", "logicalBool": "tensorDtypes.x == \"bool\"" }, "bindings": ["x_2", "y", "params_8"], "dispatch": { "x": "min(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)", "y": "ceilDiv(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)", "z": 1 } } ] }, { "id": "axis1_parallel", "priority": 20, "when": ["not flatParallelCovered", "(dtypes.T == \"f32\" or dtypes.T == \"f16\")", "f16Ok(dtypes.T)", "ranks.x >= 2", "reduceAxis == ranks.x - 1", "numel(shapes.y) == rows(shapes.x, ranks.x - 1)", "lastAxisCols >= tunables.ROW_PARALLEL_MIN_COLS", "treeWorkgroupOk"], "demoteWhen": ["rowSerialPreferred"], "derive": { "workgroupSize": "min(reduceWorkgroupSize, pow2ceil(dim(shapes.x, ranks.x - 1)))" }, "passes": [ { "id": "main", "name": "ReduceLogSumExp.Axis1Parallel", "shader": "reduce-row-tree.wgsl.jinja", "derive": { "op": "\"logsumexp\"", "castF32": "dtypes.T == \"f16\"", "usesF16Spec": "dtypes.T == \"f16\"" }, "bindings": ["x_2", "y", "params_9"], "dispatch": { "x": "min(rows(shapes.x, ranks.x - 1), 65535)", "y": "ceilDiv(rows(shapes.x, ranks.x - 1), 65535)", "z": 1 } } ] }, { "id": "axis_split", "priority": 24, "when": ["not flatParallelCovered", "(dtypes.T == \"f32\" or dtypes.T == \"f16\")", "f16Ok(dtypes.T)", "attrs.noop_with_empty_axes == 0", "ranks.x >= 2", "reduceAxis < ranks.x - 1", "not (ranks.x == 2 and reduceAxis == 0)", "axisSplitDim >= tunables.AXIS0_SPLIT_MIN_ROWS", "axisSplitOutputs >= 1", "axisSplitOutputs <= 4096", "axisSplitOutputs == rows(shapes.x, reduceAxis)", "axisSplitPathFits"], "derive": { "splitCount": "axisSplitCount", "partialElement": "\"f32\"", "workgroupSize": "reduceWorkgroupSize", "split": "splitCount" }, "intermediates": [{ "id": "partials", "dtype": "float32", "shape": "[3 * splitCount * axisSplitOutputs]" }], "passes": [ { "id": "split_reduce", "name": "ReduceLogSumExp.AxisSplitReduce", "shader": "reduce-axis-split-reduce.wgsl.jinja", "derive": { "op": "\"logsumexp\"", "splitSpec": "splitCount", "castF32": "dtypes.T == \"f16\"", "usesF16Spec": "dtypes.T == \"f16\"" }, "bindings": ["x_2", "partials", "params_10"], "dispatch": { "x": "min(ceilDiv((axisSplitOutputs), (reduceWorkgroupSize)), DISPATCH_FOLD_WIDTH)", "y": "splitCount", "z": "ceilDiv(ceilDiv((axisSplitOutputs), (reduceWorkgroupSize)), DISPATCH_FOLD_WIDTH)" } }, { "id": "combine", "name": "ReduceLogSumExp.AxisSplitCombine", "shader": "reduce-axis0-splitk-combine.wgsl.jinja", "derive": { "op": "\"logsumexp\"", "splitSpec": "splitCount", "outputF16": "dtypes.T == \"f16\"" }, "bindings": ["partials_2", "y", "params_11"], "dispatch": { "x": "min(ceilDiv((axisSplitOutputs), (reduceWorkgroupSize)), 65535)", "y": "ceilDiv(ceilDiv((axisSplitOutputs), (reduceWorkgroupSize)), 65535)", "z": 1 } } ] }, { "id": "axis_split_tiled_narrow", "priority": 25, "when": ["not flatParallelCovered", "(dtypes.T == \"f32\" or dtypes.T == \"f16\")", "f16Ok(dtypes.T)", "attrs.noop_with_empty_axes == 0", "ranks.x >= 2", "reduceAxis < ranks.x - 1", "axisSplitDim >= tunables.AXIS0_SPLIT_MIN_ROWS", "axisSplitOutputs >= 1", "axisSplitOutputs <= 2 * tunables.AXIS_SPLIT_TILE_COLS", "reduceWorkgroupSize % tunables.AXIS_SPLIT_TILE_COLS == 0", "axisSplitOutputs == rows(shapes.x, reduceAxis)", "axisSplitPathFits"], "derive": { "splitCount": "axisSplitCount", "partialElement": "\"f32\"", "scalar": "dtypes.T", "workgroupSize": "reduceWorkgroupSize", "split": "splitCount", "tileCols": "tunables.AXIS_SPLIT_TILE_COLS" }, "intermediates": [{ "id": "partials", "dtype": "float32", "shape": "[3 * splitCount * axisSplitOutputs]" }], "passes": [ { "id": "split_reduce", "name": "ReduceLogSumExp.AxisSplitTiledReduce", "shader": "reduce-axis0-tilecols.wgsl.jinja", "derive": { "op": "\"logsumexp\"", "splitSpec": "splitCount", "tileCols": "tunables.AXIS_SPLIT_TILE_COLS", "castF32": "dtypes.T == \"f16\"", "usesF16Spec": "dtypes.T == \"f16\"" }, "bindings": ["x_2", "partials", "params_10"], "dispatch": { "x": "min(ceilDiv((axisSplitOutputs), (tileCols)), DISPATCH_FOLD_WIDTH)", "y": "splitCount", "z": "ceilDiv(ceilDiv((axisSplitOutputs), (tileCols)), DISPATCH_FOLD_WIDTH)" } }, { "id": "combine", "name": "ReduceLogSumExp.AxisSplitCombine", "shader": "reduce-axis0-splitk-combine.wgsl.jinja", "derive": { "op": "\"logsumexp\"", "splitSpec": "splitCount", "outputF16": "dtypes.T == \"f16\"" }, "bindings": ["partials_2", "y", "params_11"], "dispatch": { "x": "min(ceilDiv((axisSplitOutputs), (reduceWorkgroupSize)), 65535)", "y": "ceilDiv(ceilDiv((axisSplitOutputs), (reduceWorkgroupSize)), 65535)", "z": 1 } } ] }, { "id": "axis0_splitk", "priority": 22, "when": ["not flatParallelCovered", "(dtypes.T == \"f32\" or dtypes.T == \"f16\")", "f16Ok(dtypes.T)", "ranks.x == 2", "reduceAxis == 0", "axis0Rows >= tunables.AXIS0_SPLIT_MIN_ROWS", "dim(shapes.x, 1) > 0", "((attrs.keepdims == 0 and ranks.y == 1 and dim(shapes.y, 0) == dim(shapes.x, 1)) or (attrs.keepdims == 1 and ranks.y == 2 and dim(shapes.y, 0) == 1 and dim(shapes.y, 1) == dim(shapes.x, 1)))", "axis0SplitPathFits"], "derive": { "splitCount": "axis0SplitCount", "partialElement": "\"f32\"", "workgroupSize": "reduceWorkgroupSize", "split": "splitCount" }, "intermediates": [{ "id": "partials", "dtype": "float32", "shape": "[3 * splitCount * dim(shapes.x, 1)]" }], "passes": [ { "id": "split_reduce", "name": "ReduceLogSumExp.Axis0SplitKReduce", "shader": "reduce-axis0-splitk-reduce.wgsl.jinja", "derive": { "op": "\"logsumexp\"", "splitSpec": "splitCount", "castF32": "dtypes.T == \"f16\"", "usesF16Spec": "dtypes.T == \"f16\"" }, "bindings": ["x_2", "partials", "params_12"], "dispatch": { "x": "min(ceilDiv((dim(shapes.x, 1)), (reduceWorkgroupSize)), DISPATCH_FOLD_WIDTH)", "y": "splitCount", "z": "ceilDiv(ceilDiv((dim(shapes.x, 1)), (reduceWorkgroupSize)), DISPATCH_FOLD_WIDTH)" } }, { "id": "combine", "name": "ReduceLogSumExp.Axis0SplitKCombine", "shader": "reduce-axis0-splitk-combine.wgsl.jinja", "derive": { "op": "\"logsumexp\"", "splitSpec": "splitCount", "outputF16": "dtypes.T == \"f16\"" }, "bindings": ["partials_2", "y", "params_13"], "dispatch": { "x": "min(ceilDiv((dim(shapes.x, 1)), (reduceWorkgroupSize)), 65535)", "y": "ceilDiv(ceilDiv((dim(shapes.x, 1)), (reduceWorkgroupSize)), 65535)", "z": 1 } } ] }, { "id": "axis0_tilecols", "priority": 20, "when": ["not flatParallelCovered", "(dtypes.T == \"f32\" or dtypes.T == \"f16\")", "f16Ok(dtypes.T)", "ranks.x == 2", "reduceAxis == 0", "axis0Rows >= tunables.AXIS0_TILE_MIN_ROWS", "axis0Cols >= tunables.AXIS0_TILE_MIN_COLS", "((attrs.keepdims == 0 and ranks.y == 1 and dim(shapes.y, 0) == dim(shapes.x, 1)) or (attrs.keepdims == 1 and ranks.y == 2 and dim(shapes.y, 0) == 1 and dim(shapes.y, 1) == dim(shapes.x, 1)))", "axis0TilePathFits"], "derive": { "workgroupSize": "reduceWorkgroupSize", "tileCols": "tunables.AXIS0_TILE_COLS" }, "passes": [ { "id": "main", "name": "ReduceLogSumExp.Axis0TileCols", "shader": "reduce-axis0-tilecols.wgsl.jinja", "derive": { "op": "\"logsumexp\"", "castF32": "dtypes.T == \"f16\"", "usesF16Spec": "dtypes.T == \"f16\"" }, "bindings": ["x_2", "y", "params_12"], "dispatch": { "x": "min(ceilDiv((dim(shapes.x, 1)), (tileCols)), 65535)", "y": "ceilDiv(ceilDiv((dim(shapes.x, 1)), (tileCols)), 65535)", "z": 1 } } ] }, { "id": "all_axes_flat", "priority": 31, "when": ["flatParallelCovered"], "derive": { "scalar": "dtypes.T", "workgroupSize": "reduceWorkgroupSize", "flatScalar": "\"vec4<\" ~ dtypes.T ~ \">\" if numel(shapes.x) % tunables.VECTOR_WIDTH == 0 else dtypes.T", "split": "flatSplitCount" }, "intermediates": [{ "id": "partials", "dtype": "float32", "shape": "[3 * flatSplitCount]" }], "passes": [ { "id": "flat_partial", "name": "ReduceLogSumExp.AllAxesFlatPartial", "shader": "reduce-flat-partial-logsumexp.wgsl.jinja", "derive": { "vec4": "numel(shapes.x) % tunables.VECTOR_WIDTH == 0", "castF32": "dtypes.T == \"f16\"", "usesF16Spec": "dtypes.T == \"f16\"" }, "bindings": [ { "arg": "x", "elementType": "$flatScalar" }, { "name": "partials", "buffer": "storage", "elementType": "f32" }, { "name": "params", "struct": [{ "name": "count", "type": "u32", "value": "flatItems" }] } ], "dispatch": { "x": "flatSplitCount" } }, { "id": "combine", "name": "ReduceLogSumExp.AllAxesFlatCombine", "shader": "reduce-flat-combine-logsumexp.wgsl.jinja", "derive": { "outputF16": "dtypes.T == \"f16\"" }, "bindings": [ { "name": "partials", "buffer": "read-only-storage", "elementType": "f32" }, "y", { "name": "params", "struct": [{ "name": "cols", "type": "u32", "value": "1" }] } ], "dispatch": { "x": 1 } } ] }, { "id": "rankn_single_axis_generic", "priority": 12, "when": ["not flatParallelCovered", "f16Ok(dtypes.T)", "ranks.x >= 3", "attrs.noop_with_empty_axes == 0", "reduceAxis < ranks.x", "numel(shapes.y) == rows(shapes.x, reduceAxis)", "((attrs.keepdims == 0 and ranks.y == ranks.x - 1) or (attrs.keepdims == 1 and ranks.y == ranks.x and dim(shapes.y, reduceAxis) == 1))"], "supersededBy": ["axis_split_tiled_narrow", "axis_split", "subgroup_last_axis_vec4", "subgroup_last_axis", "tree_last_axis_vec4"], "derive": { "scalar": "dtypes.T", "reduceWorkgroupSize": "reduceWorkgroupSize" }, "passes": [ { "id": "main", "name": "ReduceLogSumExp.RankNSingleAxisGeneric", "shader": "reduce-serial-axis.wgsl.jinja", "derive": { "op": "\"logsumexp\"", "indexing": "\"rankn\"", "rank": "ranks.x", "axisSpec": "reduceAxis", "dataShape": "shapes.x", "outputShape": "shapes.y", "outputRank": "ranks.y", "keepDims": "attrs.keepdims != 0", "intMode": "dtypes.T == \"i32\"", "castF32": "dtypes.T == \"f16\"", "usesF16Spec": "dtypes.T == \"f16\"" }, "bindings": ["x_2", "y", "params_16"], "dispatch": { "x": "min(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)", "y": "ceilDiv(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)", "z": 1 } } ] }, { "id": "subgroup_last_axis_vec4", "priority": 25, "when": ["device.wgslLanguageFeatures.has(\"subgroup_id\")", "not flatParallelCovered", "(dtypes.T == \"f32\" or dtypes.T == \"f16\")", "f16Ok(dtypes.T)", "ranks.x >= 1", "reduceAxis == ranks.x - 1", "numel(shapes.y) == rows(shapes.x, ranks.x - 1)", "dim(shapes.x, ranks.x - 1) >= 4", "dim(shapes.x, ranks.x - 1) % tunables.VECTOR_WIDTH == 0", "(lastAxisCols >= tunables.SUBGROUP_MIN_COLS or lastAxisRows < tunables.SUBGROUP_SMALL_ROW_LIMIT)", "not rowSerialPreferred"], "requires": { "features": ["subgroups"] }, "derive": { "scalar": "dtypes.T", "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"", "workgroupSize": "min(reduceWorkgroupSize, max(subgroupWorkgroupFloor, pow2ceil(ceilDiv(lastAxisCols, tunables.VECTOR_WIDTH))))" }, "passes": [ { "id": "main", "name": "ReduceLogSumExp.SubgroupRowVec4", "shader": "reduce-row-subgroup.wgsl.jinja", "derive": { "op": "\"logsumexp\"", "vec4": true, "castF32": "dtypes.T == \"f16\"", "usesF16Spec": "dtypes.T == \"f16\"" }, "bindings": ["x", "y", "params_6"], "dispatch": { "x": "min(rows(shapes.x, ranks.x - 1), 65535)", "y": "ceilDiv(rows(shapes.x, ranks.x - 1), 65535)", "z": 1 }, "subgroupCollectivesWidth": "portable" } ] }, { "id": "subgroup_last_axis", "priority": 24, "when": ["device.wgslLanguageFeatures.has(\"subgroup_id\")", "not flatParallelCovered", "(dtypes.T == \"f32\" or dtypes.T == \"f16\")", "f16Ok(dtypes.T)", "ranks.x >= 1", "reduceAxis == ranks.x - 1", "numel(shapes.y) == rows(shapes.x, ranks.x - 1)", "dim(shapes.x, ranks.x - 1) > 0", "dim(shapes.x, ranks.x - 1) % tunables.VECTOR_WIDTH != 0", "(lastAxisCols >= tunables.SUBGROUP_MIN_COLS or lastAxisRows < tunables.SUBGROUP_SMALL_ROW_LIMIT)", "not rowSerialPreferred"], "requires": { "features": ["subgroups"] }, "derive": { "scalar": "dtypes.T", "workgroupSize": "min(reduceWorkgroupSize, max(subgroupWorkgroupFloor, pow2ceil(lastAxisCols)))" }, "passes": [ { "id": "main", "name": "ReduceLogSumExp.SubgroupRow", "shader": "reduce-row-subgroup.wgsl.jinja", "derive": { "op": "\"logsumexp\"", "vec4": false, "castF32": "dtypes.T == \"f16\"", "usesF16Spec": "dtypes.T == \"f16\"" }, "bindings": ["x_2", "y", "params_17"], "dispatch": { "x": "min(rows(shapes.x, ranks.x - 1), 65535)", "y": "ceilDiv(rows(shapes.x, ranks.x - 1), 65535)", "z": 1 }, "subgroupCollectivesWidth": "portable" } ] }, { "id": "axis0", "priority": 0, "when": ["not flatParallelCovered", "f16Ok(dtypes.T)", "ranks.x == 2", "reduceAxis == 0", "((attrs.keepdims == 0 and ranks.y == 1 and dim(shapes.y, 0) == dim(shapes.x, 1)) or (attrs.keepdims == 1 and ranks.y == 2 and dim(shapes.y, 0) == 1 and dim(shapes.y, 1) == dim(shapes.x, 1)))"], "supersededBy": ["axis_split_tiled_narrow", "axis0_splitk", "axis0_tilecols"], "derive": { "axis": 0, "scalar": "dtypes.T" }, "passes": [ { "id": "main", "name": "axis0", "shader": "reduce-serial-axis.wgsl.jinja", "derive": { "axis": 0, "op": "\"logsumexp\"", "indexing": "\"axis2d\"", "intMode": "dtypes.T == \"i32\"", "castF32": "dtypes.T == \"f16\"", "usesF16Spec": "dtypes.T == \"f16\"" }, "bindings": ["x_2", "y", "params_18"], "dispatch": { "x": "min(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)", "y": "ceilDiv(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)", "z": 1 } } ] }, { "id": "axis1", "priority": 0, "when": ["not flatParallelCovered", "f16Ok(dtypes.T)", "ranks.x == 2", "reduceAxis == 1", "((attrs.keepdims == 0 and ranks.y == 1 and dim(shapes.y, 0) == dim(shapes.x, 0)) or (attrs.keepdims == 1 and ranks.y == 2 and dim(shapes.y, 0) == dim(shapes.x, 0) and dim(shapes.y, 1) == 1))"], "derive": { "axis": 1, "scalar": "dtypes.T" }, "passes": [ { "id": "main", "name": "axis1", "shader": "reduce-serial-axis.wgsl.jinja", "derive": { "axis": 1, "op": "\"logsumexp\"", "indexing": "\"axis2d\"", "intMode": "dtypes.T == \"i32\"", "castF32": "dtypes.T == \"f16\"", "usesF16Spec": "dtypes.T == \"f16\"" }, "bindings": ["x_2", "y", "params_19"], "dispatch": { "x": "min(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)", "y": "ceilDiv(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)", "z": 1 } } ] }, { "id": "all_axes_keepdims", "priority": 30, "when": ["not flatParallelCovered", "f16Ok(dtypes.T)", "ranks.x >= 3", "attrs.keepdims == 1", "ranks.y == ranks.x", "numel(shapes.y) == 1"], "derive": { "axis": 0, "scalar": "dtypes.T" }, "passes": [ { "id": "main", "name": "ReduceLogSumExp.Rank3AllAxesKeepdims", "shader": "reduce-serial-axis.wgsl.jinja", "derive": { "op": "\"logsumexp\"", "indexing": "\"axis2d\"", "intMode": "dtypes.T == \"i32\"", "castF32": "dtypes.T == \"f16\"", "usesF16Spec": "dtypes.T == \"f16\"" }, "bindings": [ "x_2", "y", { "name": "params", "struct": [ { "name": "rows", "type": "u32", "value": "numel(shapes.x)" }, { "name": "cols", "type": "u32", "value": "1" }, { "name": "outCount", "type": "u32", "value": "numel(shapes.y)" } ] } ], "dispatch": { "x": "min(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)", "y": "ceilDiv(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)", "z": 1 } } ] }, { "id": "all_axes_no_keepdims", "priority": 30, "when": ["not flatParallelCovered", "f16Ok(dtypes.T)", "ranks.x >= 3", "attrs.keepdims == 0", "attrs.noop_with_empty_axes == 0", "ranks.y == 0"], "derive": { "axis": 0, "scalar": "dtypes.T" }, "passes": [ { "id": "main", "name": "ReduceLogSumExp.Rank3AllAxesNoKeepdims", "shader": "reduce-serial-axis.wgsl.jinja", "derive": { "op": "\"logsumexp\"", "indexing": "\"axis2d\"", "intMode": "dtypes.T == \"i32\"", "castF32": "dtypes.T == \"f16\"", "usesF16Spec": "dtypes.T == \"f16\"" }, "bindings": [ "x_2", "y", { "name": "params", "struct": [ { "name": "rows", "type": "u32", "value": "numel(shapes.x)" }, { "name": "cols", "type": "u32", "value": "1" }, { "name": "outCount", "type": "u32", "value": "numel(shapes.y)" } ] } ], "dispatch": { "x": "min(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)", "y": "ceilDiv(ceilDiv((numel(shapes.y)), (reduceWorkgroupSize)), 65535)", "z": 1 } } ] } ] }