{ "domain": "ai.onnx", "name": "ReduceLogSumExp", "sinceVersion": 18, "description": "Computes `log(sum(exp(x)))` over the specified axes of the input tensor. The output rank matches the input when `keepdims` is 1; reduced dimensions are pruned when `keepdims` is 0. Reduction over an empty set of values yields negative infinity.", "inputs": [{ "role": "data", "dtype": "T", "description": "The input tensor to reduce." }], "outputs": [ { "role": "reduced", "dtype": "T", "rank": "ranks.data if attrs.keepdims == 1 or ((attrs.axes | length) == 0 and attrs.noop_with_empty_axes == 1) else (ranks.data - (attrs.axes | length) if (attrs.axes | length) > 0 else 0)", "description": "The reduced output tensor." } ], "attributes": { "keepdims": 1, "noop_with_empty_axes": 0, "axes": [] }, "attributeDescriptions": { "keepdims": "If 1, retains the reduced dimension with size 1 in the output; if 0, the reduced dimension is removed.", "noop_with_empty_axes": "When 1 and `axes` is empty, acts as an identity (no reduction); when 0 and `axes` is empty, reduces over all axes.", "axes": "Values of the optional ONNX `axes` tensor input, supplied through this request attribute; an empty list follows `noop_with_empty_axes`." }, "attributeConstraints": { "keepdims": { "values": [0, 1] }, "noop_with_empty_axes": { "values": [0, 1] } }, "typeConstraints": { "T": ["float32", "float16", "int32"] }, "args": { "x": { "kind": "tensor", "semantic": "data", "role": "input" }, "y": { "kind": "tensor", "semantic": "reduced", "role": "output" } }, "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.data, ranks.data - 1) if ranks.data > 0 else 1", "lastAxisCols": "dim(shapes.data, ranks.data - 1) if ranks.data > 0 else 1", "rowSerialPreferred": "lastAxisRows >= tunables.ROW_SERIAL_MIN_ROWS and lastAxisCols <= tunables.ROW_SERIAL_MAX_COLS", "axis0Rows": "dim(shapes.data, 0) if ranks.data >= 2 else 0", "axis0Cols": "dim(shapes.data, 1) if ranks.data >= 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 <= device.limits.maxComputeWorkgroupsPerDimension and ceilDiv(ceilDiv(axis0Cols, reduceWorkgroupSize), device.limits.maxComputeWorkgroupsPerDimension) <= device.limits.maxComputeWorkgroupsPerDimension and axis0SplitScratchBytes <= device.limits.maxStorageBufferBindingSize and axis0SplitScratchBytes <= device.limits.maxBufferSize", "reduceAxis": "(attrs.axes[0] + ranks.data if attrs.axes[0] < 0 else attrs.axes[0]) if ((attrs.axes | length) == 1 and isUniqueIntList(attrs.axes, 0 - ranks.data, ranks.data, 1)) else ranks.data", "axisSplitDim": "dim(shapes.data, reduceAxis) if ranks.data >= 2 and reduceAxis < ranks.data else 0", "axisSplitInner": "inner(shapes.data, reduceAxis) if ranks.data >= 2 and reduceAxis < ranks.data else 1", "axisSplitOutputs": "numel(shapes.reduced)", "axisSplitCount": "min(tunables.AXIS0_MAX_SPLITS, pow2ceil(ceilDiv(axisSplitDim, tunables.AXIS0_SPLIT_TARGET_ROWS)))", "axisSplitScratchBytes": "3 * axisSplitCount * axisSplitOutputs * 4", "axisSplitPathFits": "axisSplitCount <= device.limits.maxComputeWorkgroupsPerDimension and ceilDiv(ceilDiv(axisSplitOutputs, reduceWorkgroupSize), device.limits.maxComputeWorkgroupsPerDimension) <= device.limits.maxComputeWorkgroupsPerDimension 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.data) / tunables.VECTOR_WIDTH if numel(shapes.data) % tunables.VECTOR_WIDTH == 0 else numel(shapes.data)", "flatSplitCount": "max(1, min(tunables.FULL_REDUCE_MAX_SPLITS, ceilDiv(flatItems, reduceWorkgroupSize)))", "flatScratchBytes": "3 * flatSplitCount * dtypeBytes(\"float32\")", "flatPathFits": "treeWorkgroupOk and flatSplitCount <= device.limits.maxComputeWorkgroupsPerDimension 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.reduced) == 1 and numel(shapes.data) >= tunables.FULL_REDUCE_MIN_ELEMENTS and flatPathFits", "contiguousSuffixParallelCovered": "(dtypes.T == \"f32\" or dtypes.T == \"f16\") and f16Ok(dtypes.T) and numel(shapes.reduced) > 0 and numel(shapes.data) % numel(shapes.reduced) == 0 and numel(shapes.data) / numel(shapes.reduced) >= tunables.CONTIGUOUS_SUFFIX_MIN_COLS and ((ranks.data == 3 and hasAxis(attrs.axes, 0, 3) == false and hasAxis(attrs.axes, 1, 3) and hasAxis(attrs.axes, 2, 3) and numel(shapes.reduced) == dim(shapes.data, 0)) or (ranks.data == 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.reduced) == dim(shapes.data, 0) * dim(shapes.data, 1)) or (ranks.data == 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.reduced) == dim(shapes.data, 0)))" }, "tunables": { "WORKGROUP_SIZE": 256, "VECTOR_WIDTH": 4, "ROW_PARALLEL_MIN_COLS": 64, "SUBGROUP_MIN_COLS": 256, "SUBGROUP_SMALL_ROW_LIMIT": 32768, "AXIS0_SPLIT_MIN_ROWS": 8192, "AXIS0_SPLIT_TARGET_ROWS": 256, "AXIS0_MAX_SPLITS": 128, "AXIS0_TILE_MIN_ROWS": 64, "AXIS0_TILE_MIN_COLS": 16, "AXIS0_TILE_COLS": 16, "AXIS_SPLIT_TILE_COLS": 8, "FULL_REDUCE_MIN_ELEMENTS": 8192, "FULL_REDUCE_MAX_SPLITS": 256, "CONTIGUOUS_SUFFIX_MIN_COLS": 256, "AXES02_WORKGROUP_SIZE": 256, "ROW_SERIAL_MIN_ROWS": 8192, "ROW_SERIAL_MAX_COLS": 1024 }, "bindingSets": { "axes02": [ { "name": "x", "arg": "x", "semantic": "data", "buffer": { "type": "read-only-storage" }, "elementType": "$T" }, { "name": "y", "arg": "y", "semantic": "reduced", "buffer": { "type": "storage" }, "elementType": "$T" }, { "name": "params", "semantic": "kernel.params", "buffer": { "type": "uniform" }, "struct": { "name": "Params", "fields": [ { "name": "d0", "type": "u32", "value": "dim(shapes.data, 0)" }, { "name": "d1", "type": "u32", "value": "dim(shapes.data, 1)" }, { "name": "d2", "type": "u32", "value": "dim(shapes.data, 2)" }, { "name": "outCount", "type": "u32", "value": "numel(shapes.reduced)" } ] } } ], "suffixVec4": [ { "name": "x", "arg": "x", "semantic": "data", "buffer": { "type": "read-only-storage" }, "elementType": "$vectorScalar" }, { "name": "y", "arg": "y", "semantic": "reduced", "buffer": { "type": "storage" }, "elementType": "$T" }, { "name": "params", "semantic": "kernel.params", "buffer": { "type": "uniform" }, "struct": { "name": "Params", "fields": [ { "name": "rows", "type": "u32", "value": "numel(shapes.reduced)" }, { "name": "chunkCount", "type": "u32", "value": "numel(shapes.data) / numel(shapes.reduced) / tunables.VECTOR_WIDTH" } ] } } ], "suffixScalar": [ { "name": "x", "arg": "x", "semantic": "data", "buffer": { "type": "read-only-storage" }, "elementType": "$T" }, { "name": "y", "arg": "y", "semantic": "reduced", "buffer": { "type": "storage" }, "elementType": "$T" }, { "name": "params", "semantic": "kernel.params", "buffer": { "type": "uniform" }, "struct": { "name": "Params", "fields": [ { "name": "rows", "type": "u32", "value": "numel(shapes.reduced)" }, { "name": "cols", "type": "u32", "value": "numel(shapes.data) / numel(shapes.reduced)" } ] } } ], "elementwise": [ { "name": "x", "arg": "x", "semantic": "data", "buffer": { "type": "read-only-storage" }, "elementType": "$T" }, { "name": "y", "arg": "y", "semantic": "reduced", "buffer": { "type": "storage" }, "elementType": "$T" }, { "name": "params", "semantic": "kernel.params", "buffer": { "type": "uniform" }, "struct": { "name": "Params", "fields": [{ "name": "count", "type": "u32", "value": "numel(shapes.reduced)" }] } } ], "lastAxisVec4": [ { "name": "x", "arg": "x", "semantic": "data", "buffer": { "type": "read-only-storage" }, "elementType": "$vectorScalar" }, { "name": "y", "arg": "y", "semantic": "reduced", "buffer": { "type": "storage" }, "elementType": "$T" }, { "name": "params", "semantic": "kernel.params", "buffer": { "type": "uniform" }, "struct": { "name": "Params", "fields": [ { "name": "rows", "type": "u32", "value": "rows(shapes.data, ranks.data - 1)" }, { "name": "chunkCount", "type": "u32", "value": "dim(shapes.data, ranks.data - 1) / tunables.VECTOR_WIDTH" } ] } } ], "lastAxisScalar": [ { "name": "x", "arg": "x", "semantic": "data", "buffer": { "type": "read-only-storage" }, "elementType": "$T" }, { "name": "y", "arg": "y", "semantic": "reduced", "buffer": { "type": "storage" }, "elementType": "$T" }, { "name": "params", "semantic": "kernel.params", "buffer": { "type": "uniform" }, "struct": { "name": "Params", "fields": [ { "name": "rows", "type": "u32", "value": "rows(shapes.data, ranks.data - 1)" }, { "name": "cols", "type": "u32", "value": "dim(shapes.data, ranks.data - 1)" } ] } } ], "lastAxisScalarSubgroup": [ { "name": "x", "arg": "x", "semantic": "data", "buffer": { "type": "read-only-storage" }, "elementType": "$T" }, { "name": "y", "arg": "y", "semantic": "reduced", "buffer": { "type": "storage" }, "elementType": "$T" }, { "name": "params", "semantic": "kernel.params", "buffer": { "type": "uniform" }, "struct": { "name": "Params", "fields": [ { "name": "rows", "type": "u32", "value": "rows(shapes.data, ranks.data - 1)" }, { "name": "chunkCount", "type": "u32", "value": "dim(shapes.data, ranks.data - 1)" } ] } } ], "scalar": [ { "name": "x", "arg": "x", "semantic": "data", "buffer": { "type": "read-only-storage" }, "elementType": "$T" }, { "name": "y", "arg": "y", "semantic": "reduced", "buffer": { "type": "storage" }, "elementType": "$T" }, { "name": "params", "semantic": "kernel.params", "buffer": { "type": "uniform" }, "struct": { "name": "Params", "fields": [ { "name": "rows", "type": "u32", "value": "1" }, { "name": "cols", "type": "u32", "value": "1" }, { "name": "outCount", "type": "u32", "value": "1" } ] } } ], "rank1Axis0": [ { "name": "x", "arg": "x", "semantic": "data", "buffer": { "type": "read-only-storage" }, "elementType": "$T" }, { "name": "y", "arg": "y", "semantic": "reduced", "buffer": { "type": "storage" }, "elementType": "$T" }, { "name": "params", "semantic": "kernel.params", "buffer": { "type": "uniform" }, "struct": { "name": "Params", "fields": [ { "name": "rows", "type": "u32", "value": "dim(shapes.data, 0)" }, { "name": "cols", "type": "u32", "value": "1" }, { "name": "outCount", "type": "u32", "value": "numel(shapes.reduced)" } ] } } ], "rank2Serial": [ { "name": "x", "arg": "x", "semantic": "data", "buffer": { "type": "read-only-storage" }, "elementType": "$T" }, { "name": "y", "arg": "y", "semantic": "reduced", "buffer": { "type": "storage" }, "elementType": "$T" }, { "name": "params", "semantic": "kernel.params", "buffer": { "type": "uniform" }, "struct": { "name": "Params", "fields": [ { "name": "rows", "type": "u32", "value": "dim(shapes.data, 0)" }, { "name": "cols", "type": "u32", "value": "dim(shapes.data, 1)" }, { "name": "outCount", "type": "u32", "value": "numel(shapes.reduced)" } ] } } ], "rank2SerialAxis1": [ { "name": "x", "arg": "x", "semantic": "data", "buffer": { "type": "read-only-storage" }, "elementType": "$T" }, { "name": "y", "arg": "y", "semantic": "reduced", "buffer": { "type": "storage" }, "elementType": "$T" }, { "name": "params", "semantic": "kernel.params", "buffer": { "type": "uniform" }, "struct": { "name": "Params", "fields": [ { "name": "cols", "type": "u32", "value": "dim(shapes.data, 1)" }, { "name": "outCount", "type": "u32", "value": "numel(shapes.reduced)" } ] } } ], "axis0Parallel": [ { "name": "x", "arg": "x", "semantic": "data", "buffer": { "type": "read-only-storage" }, "elementType": "$T" }, { "name": "y", "arg": "y", "semantic": "reduced", "buffer": { "type": "storage" }, "elementType": "$T" }, { "name": "params", "semantic": "kernel.params", "buffer": { "type": "uniform" }, "struct": { "name": "Params", "fields": [ { "name": "rows", "type": "u32", "value": "dim(shapes.data, 0)" }, { "name": "cols", "type": "u32", "value": "dim(shapes.data, 1)" } ] } } ], "fullReduceSerial": [ { "name": "x", "arg": "x", "semantic": "data", "buffer": { "type": "read-only-storage" }, "elementType": "$T" }, { "name": "y", "arg": "y", "semantic": "reduced", "buffer": { "type": "storage" }, "elementType": "$T" }, { "name": "params", "semantic": "kernel.params", "buffer": { "type": "uniform" }, "struct": { "name": "Params", "fields": [ { "name": "rows", "type": "u32", "value": "numel(shapes.data)" }, { "name": "cols", "type": "u32", "value": "1" }, { "name": "outCount", "type": "u32", "value": "numel(shapes.reduced)" } ] } } ], "axisSplitReduce": [ { "name": "x", "arg": "x", "semantic": "data", "buffer": { "type": "read-only-storage" }, "elementType": "$T" }, { "name": "partials", "semantic": "partials", "buffer": { "type": "storage" }, "elementType": "$partialElement" }, { "name": "params", "semantic": "kernel.params", "buffer": { "type": "uniform" }, "struct": { "name": "Params", "fields": [ { "name": "axisDim", "type": "u32", "value": "axisSplitDim" }, { "name": "inner", "type": "u32", "value": "axisSplitInner" }, { "name": "outputs", "type": "u32", "value": "axisSplitOutputs" } ] } } ], "axisSplitCombine": [ { "name": "partials", "semantic": "partials", "buffer": { "type": "read-only-storage" }, "elementType": "$partialElement" }, { "name": "y", "arg": "y", "semantic": "reduced", "buffer": { "type": "storage" }, "elementType": "$T" }, { "name": "params", "semantic": "kernel.params", "buffer": { "type": "uniform" }, "struct": { "name": "Params", "fields": [{ "name": "cols", "type": "u32", "value": "axisSplitOutputs" }] } } ], "axis0SplitReduce": [ { "name": "x", "arg": "x", "semantic": "data", "buffer": { "type": "read-only-storage" }, "elementType": "$T" }, { "name": "partials", "semantic": "partials", "buffer": { "type": "storage" }, "elementType": "$partialElement" }, { "name": "params", "semantic": "kernel.params", "buffer": { "type": "uniform" }, "struct": { "name": "Params", "fields": [ { "name": "rows", "type": "u32", "value": "dim(shapes.data, 0)" }, { "name": "cols", "type": "u32", "value": "dim(shapes.data, 1)" } ] } } ], "axis0SplitCombine": [ { "name": "partials", "semantic": "partials", "buffer": { "type": "read-only-storage" }, "elementType": "$partialElement" }, { "name": "y", "arg": "y", "semantic": "reduced", "buffer": { "type": "storage" }, "elementType": "$T" }, { "name": "params", "semantic": "kernel.params", "buffer": { "type": "uniform" }, "struct": { "name": "Params", "fields": [{ "name": "cols", "type": "u32", "value": "dim(shapes.data, 1)" }] } } ], "rankNAxis": [ { "name": "x", "arg": "x", "semantic": "data", "buffer": { "type": "read-only-storage" }, "elementType": "$T" }, { "name": "y", "arg": "y", "semantic": "reduced", "buffer": { "type": "storage" }, "elementType": "$T" }, { "name": "params", "semantic": "kernel.params", "buffer": { "type": "uniform" }, "struct": { "name": "Params", "fields": [ { "name": "axisDim", "type": "u32", "value": "axisSplitDim" }, { "name": "outCount", "type": "u32", "value": "numel(shapes.reduced)" } ] } } ], "flatPartialF32": [ { "name": "x", "arg": "x", "semantic": "data", "buffer": { "type": "read-only-storage" }, "elementType": "$flatScalar" }, { "name": "partials", "semantic": "partials", "buffer": { "type": "storage" }, "elementType": "f32" }, { "name": "params", "semantic": "kernel.params", "buffer": { "type": "uniform" }, "struct": { "name": "Params", "fields": [{ "name": "count", "type": "u32", "value": "flatItems" }] } } ], "flatCombineF32": [ { "name": "partials", "semantic": "partials", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, { "name": "y", "arg": "y", "semantic": "reduced", "buffer": { "type": "storage" }, "elementType": "$T" }, { "name": "params", "semantic": "kernel.params", "buffer": { "type": "uniform" }, "struct": { "name": "Params", "fields": [{ "name": "cols", "type": "u32", "value": "1" }] } } ], "multiAxis": [ { "name": "x", "arg": "x", "semantic": "data", "buffer": { "type": "read-only-storage" }, "elementType": "$T" }, { "name": "y", "arg": "y", "semantic": "reduced", "buffer": { "type": "storage" }, "elementType": "$T" }, { "name": "params", "semantic": "kernel.params", "buffer": { "type": "uniform" }, "struct": { "name": "Params", "fields": [{ "name": "outCount", "type": "u32", "value": "numel(shapes.reduced)" }] } } ] }, "variants": [ { "id": "contiguous_suffix_subgroup_vec4", "priority": 30, "requires": { "features": ["subgroups"] }, "when": ["device.wgslLanguageFeatures.has(\"subgroup_id\")", "not flatParallelCovered", "contiguousSuffixParallelCovered", "(numel(shapes.data) / numel(shapes.reduced)) % tunables.VECTOR_WIDTH == 0"], "constants": { "scalar": "dtypes.T", "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"", "workgroupSize": "min(reduceWorkgroupSize, max(subgroupWorkgroupFloor, pow2ceil(ceilDiv(numel(shapes.data) / numel(shapes.reduced), tunables.VECTOR_WIDTH))))" }, "passes": [ { "id": "main", "name": "ReduceLogSumExp.ContiguousSuffixSubgroupVec4", "source": { "shader": "reduce-row-subgroup.wgsl.jinja", "inputs": { "op": "\"logsumexp\"", "vec4": true, "castF32": "dtypes.T == \"f16\"", "usesF16": "dtypes.T == \"f16\"" } }, "subgroupCollectivesWidth": "portable", "bindings": "suffixVec4", "dispatch": { "workgroups": "numel(shapes.reduced)" } } ] }, { "id": "contiguous_suffix_tree_vec4", "priority": 22, "when": ["not flatParallelCovered", "contiguousSuffixParallelCovered", "(numel(shapes.data) / numel(shapes.reduced)) % tunables.VECTOR_WIDTH == 0", "treeWorkgroupOk"], "constants": { "scalar": "dtypes.T", "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"", "workgroupSize": "min(reduceWorkgroupSize, pow2ceil(ceilDiv(numel(shapes.data) / numel(shapes.reduced), tunables.VECTOR_WIDTH)))" }, "passes": [ { "id": "main", "name": "ReduceLogSumExp.ContiguousSuffixTreeVec4", "source": { "shader": "reduce-row-tree.wgsl.jinja", "inputs": { "op": "\"logsumexp\"", "vec4": true, "castF32": "dtypes.T == \"f16\"", "usesF16": "dtypes.T == \"f16\"" } }, "bindings": "suffixVec4", "dispatch": { "workgroups": "numel(shapes.reduced)" } } ] }, { "id": "contiguous_suffix_tree", "priority": 21, "when": ["not flatParallelCovered", "contiguousSuffixParallelCovered", "treeWorkgroupOk"], "constants": { "workgroupSize": "min(reduceWorkgroupSize, pow2ceil(numel(shapes.data) / numel(shapes.reduced)))", "scalar": "dtypes.T" }, "passes": [ { "id": "main", "name": "ReduceLogSumExp.ContiguousSuffixTree", "source": { "shader": "reduce-row-tree.wgsl.jinja", "inputs": { "op": "\"logsumexp\"", "castF32": "dtypes.T == \"f16\"", "usesF16": "dtypes.T == \"f16\"" } }, "bindings": "suffixScalar", "dispatch": { "workgroups": "numel(shapes.reduced)" } } ] }, { "id": "multi_axis_rank3", "priority": 8, "when": ["not flatParallelCovered", "not contiguousSuffixParallelCovered", "f16Ok(dtypes.T)", "ranks.data == 3", "(attrs.keepdims == 1 and ranks.reduced == 3 and (dim(shapes.reduced, 0) == 1 if hasAxis(attrs.axes, 0, 3) else dim(shapes.reduced, 0) == dim(shapes.data, 0)) and (dim(shapes.reduced, 1) == 1 if hasAxis(attrs.axes, 1, 3) else dim(shapes.reduced, 1) == dim(shapes.data, 1)) and (dim(shapes.reduced, 2) == 1 if hasAxis(attrs.axes, 2, 3) else dim(shapes.reduced, 2) == dim(shapes.data, 2))) or (attrs.keepdims == 0 and ranks.reduced == 1)"], "passes": [ { "id": "main", "name": "ReduceLogSumExp.MultiAxisRank3", "source": { "shader": "reduce-serial-axis.wgsl.jinja", "inputs": { "op": "\"logsumexp\"", "indexing": "\"multiaxis\"", "rank": 3, "reduce": ["hasAxis(attrs.axes, 0, 3)", "hasAxis(attrs.axes, 1, 3)", "hasAxis(attrs.axes, 2, 3)"], "dataShape": "shapes.data", "outputShape": "shapes.reduced", "outputRank": "ranks.reduced", "keepDims": "attrs.keepdims != 0", "intMode": "dtypes.T == \"i32\"", "castF32": "dtypes.T == \"f16\"", "usesF16": "dtypes.T == \"f16\"" } }, "bindings": "multiAxis", "dispatch": { "threads": "numel(shapes.reduced)", "workgroupSize": "reduceWorkgroupSize" } } ], "constants": { "scalar": "dtypes.T" } }, { "id": "multi_axis_rank4", "priority": 8, "when": ["not flatParallelCovered", "not contiguousSuffixParallelCovered", "f16Ok(dtypes.T)", "ranks.data == 4", "attrs.noop_with_empty_axes == 0", "numel(shapes.reduced) == (1 if hasAxis(attrs.axes, 0, 4) else dim(shapes.data, 0)) * (1 if hasAxis(attrs.axes, 1, 4) else dim(shapes.data, 1)) * (1 if hasAxis(attrs.axes, 2, 4) else dim(shapes.data, 2)) * (1 if hasAxis(attrs.axes, 3, 4) else dim(shapes.data, 3))", "((attrs.keepdims == 1 and ranks.reduced == 4) or (attrs.keepdims == 0 and ranks.reduced < 4))"], "passes": [ { "id": "main", "name": "ReduceLogSumExp.MultiAxisRank4", "source": { "shader": "reduce-serial-axis.wgsl.jinja", "inputs": { "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.data", "outputShape": "shapes.reduced", "outputRank": "ranks.reduced", "keepDims": "attrs.keepdims != 0", "intMode": "dtypes.T == \"i32\"", "castF32": "dtypes.T == \"f16\"", "usesF16": "dtypes.T == \"f16\"" } }, "bindings": "multiAxis", "dispatch": { "threads": "numel(shapes.reduced)", "workgroupSize": "reduceWorkgroupSize" } } ], "constants": { "scalar": "dtypes.T" } }, { "id": "int32_rank3_axes02_keepdims", "priority": 30, "when": ["dtypes.T == \"i32\"", "ranks.data == 3", "attrs.keepdims == 1", "hasAxis(attrs.axes, 0, 3)", "hasAxis(attrs.axes, 2, 3)", "hasAxis(attrs.axes, 1, 3) == false", "dim(shapes.data, 0) > 0", "dim(shapes.data, 2) > 0", "ranks.reduced == 3", "dim(shapes.reduced, 0) == 1", "dim(shapes.reduced, 1) == dim(shapes.data, 1)", "dim(shapes.reduced, 2) == 1"], "passes": [ { "id": "main", "name": "ReduceLogSumExp.Int32Rank3Axes02Keepdims", "bindings": "axes02", "dispatch": { "threads": "numel(shapes.reduced)", "workgroupSize": "axes02WorkgroupSize" }, "source": { "shader": "reduce-i32-axes02.wgsl.jinja", "inputs": { "op": "\"logsumexp\"", "workgroupSize": "axes02WorkgroupSize" } } } ], "derive": { "axes02WorkgroupSize": "min(tunables.AXES02_WORKGROUP_SIZE, device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)" } }, { "id": "noop_empty_axes", "priority": 40, "when": ["dtypes.T == \"f32\"", "attrs.noop_with_empty_axes == 1", "(attrs.axes | length) == 0", "sameShape(shapes.data, shapes.reduced)"], "passes": [ { "id": "main", "name": "ReduceLogSumExp.NoopEmptyAxes", "source": { "shader": "reduce-noop-empty-axes.wgsl.jinja", "inputs": { "op": "\"identity\"" } }, "bindings": "elementwise", "dispatch": { "threads": "numel(shapes.reduced)", "workgroupSize": "reduceWorkgroupSize" } } ] }, { "id": "tree_last_axis_vec4", "priority": 23, "demoteWhen": ["rowSerialPreferred"], "when": ["not flatParallelCovered", "(dtypes.T == \"f32\" or dtypes.T == \"f16\")", "f16Ok(dtypes.T)", "ranks.data >= 1", "reduceAxis == ranks.data - 1", "numel(shapes.reduced) == rows(shapes.data, ranks.data - 1)", "attrs.noop_with_empty_axes == 0", "lastAxisCols >= tunables.ROW_PARALLEL_MIN_COLS", "lastAxisCols % tunables.VECTOR_WIDTH == 0", "treeWorkgroupOk"], "constants": { "scalar": "dtypes.T", "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"", "workgroupSize": "min(reduceWorkgroupSize, pow2ceil(ceilDiv(lastAxisCols, tunables.VECTOR_WIDTH)))" }, "passes": [ { "id": "main", "name": "ReduceLogSumExp.TreeRowVec4", "source": { "shader": "reduce-row-tree.wgsl.jinja", "inputs": { "op": "\"logsumexp\"", "vec4": true, "castF32": "dtypes.T == \"f16\"", "usesF16": "dtypes.T == \"f16\"" } }, "bindings": "lastAxisVec4", "dispatch": { "workgroups": "lastAxisRows" } } ] }, { "id": "rank0_scalar", "priority": 40, "constants": { "axis": 0, "scalar": "dtypes.T" }, "when": ["f16Ok(dtypes.T)", "ranks.data == 0", "ranks.reduced == 0"], "passes": [ { "id": "main", "name": "ReduceLogSumExp.Rank0Scalar", "source": { "shader": "reduce-serial-axis.wgsl.jinja", "inputs": { "op": "\"logsumexp\"", "indexing": "\"axis2d\"", "intMode": "dtypes.T == \"i32\"", "castF32": "dtypes.T == \"f16\"", "usesF16": "dtypes.T == \"f16\"", "logicalBool": "tensorDtypes.data == \"bool\"" } }, "bindings": "scalar", "dispatch": { "x": 1 } } ] }, { "id": "rank1_axis0", "constants": { "axis": 0, "scalar": "dtypes.T" }, "when": ["not flatParallelCovered", "f16Ok(dtypes.T)", "ranks.data == 1", "reduceAxis == 0", "((attrs.keepdims == 0 and ranks.reduced == 0) or (attrs.keepdims == 1 and ranks.reduced == 1 and dim(shapes.reduced, 0) == 1))"], "passes": [ { "id": "main", "name": "ReduceLogSumExp.Rank1Axis0", "source": { "shader": "reduce-serial-axis.wgsl.jinja", "inputs": { "op": "\"logsumexp\"", "indexing": "\"axis2d\"", "intMode": "dtypes.T == \"i32\"", "castF32": "dtypes.T == \"f16\"", "usesF16": "dtypes.T == \"f16\"", "logicalBool": "tensorDtypes.data == \"bool\"" } }, "bindings": "rank1Axis0", "dispatch": { "threads": "numel(shapes.reduced)", "workgroupSize": "reduceWorkgroupSize" } } ] }, { "id": "axis1_parallel", "priority": 20, "demoteWhen": ["rowSerialPreferred"], "when": ["not flatParallelCovered", "(dtypes.T == \"f32\" or dtypes.T == \"f16\")", "f16Ok(dtypes.T)", "ranks.data >= 2", "reduceAxis == ranks.data - 1", "numel(shapes.reduced) == rows(shapes.data, ranks.data - 1)", "lastAxisCols >= tunables.ROW_PARALLEL_MIN_COLS", "treeWorkgroupOk"], "constants": { "workgroupSize": "min(reduceWorkgroupSize, pow2ceil(dim(shapes.data, ranks.data - 1)))" }, "passes": [ { "id": "main", "name": "ReduceLogSumExp.Axis1Parallel", "source": { "shader": "reduce-row-tree.wgsl.jinja", "inputs": { "op": "\"logsumexp\"", "castF32": "dtypes.T == \"f16\"", "usesF16": "dtypes.T == \"f16\"" } }, "bindings": "lastAxisScalar", "dispatch": { "workgroups": "rows(shapes.data, ranks.data - 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.data >= 2", "reduceAxis < ranks.data - 1", "not (ranks.data == 2 and reduceAxis == 0)", "axisSplitDim >= tunables.AXIS0_SPLIT_MIN_ROWS", "axisSplitOutputs >= 1", "axisSplitOutputs <= 4096", "axisSplitOutputs == rows(shapes.data, reduceAxis)", "axisSplitPathFits"], "derive": { "splitCount": "axisSplitCount" }, "constants": { "partialElement": "\"f32\"", "workgroupSize": "reduceWorkgroupSize", "split": "splitCount" }, "intermediates": [{ "id": "partials", "dtype": "float32", "shape": "[3 * splitCount * axisSplitOutputs]" }], "passes": [ { "id": "split_reduce", "name": "ReduceLogSumExp.AxisSplitReduce", "source": { "shader": "reduce-axis-split-reduce.wgsl.jinja", "inputs": { "op": "\"logsumexp\"", "workgroupSize": "reduceWorkgroupSize", "split": "splitCount", "castF32": "dtypes.T == \"f16\"", "usesF16": "dtypes.T == \"f16\"" } }, "bindings": "axisSplitReduce", "dispatch": { "threads": "axisSplitOutputs", "workgroupSize": "reduceWorkgroupSize", "y": "splitCount" } }, { "id": "combine", "name": "ReduceLogSumExp.AxisSplitCombine", "source": { "shader": "reduce-axis0-splitk-combine.wgsl.jinja", "inputs": { "op": "\"logsumexp\"", "workgroupSize": "reduceWorkgroupSize", "split": "splitCount", "outputF16": "dtypes.T == \"f16\"" } }, "bindings": "axisSplitCombine", "dispatch": { "threads": "axisSplitOutputs", "workgroupSize": "reduceWorkgroupSize" } } ] }, { "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.data >= 2", "reduceAxis < ranks.data - 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.data, reduceAxis)", "axisSplitPathFits"], "derive": { "splitCount": "axisSplitCount" }, "constants": { "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", "source": { "shader": "reduce-axis0-tilecols.wgsl.jinja", "inputs": { "op": "\"logsumexp\"", "workgroupSize": "reduceWorkgroupSize", "split": "splitCount", "tileCols": "tunables.AXIS_SPLIT_TILE_COLS", "castF32": "dtypes.T == \"f16\"", "usesF16": "dtypes.T == \"f16\"" } }, "bindings": "axisSplitReduce", "dispatch": { "workgroups": "ceilDiv((axisSplitOutputs), (constants.tileCols))", "y": "splitCount" } }, { "id": "combine", "name": "ReduceLogSumExp.AxisSplitCombine", "source": { "shader": "reduce-axis0-splitk-combine.wgsl.jinja", "inputs": { "op": "\"logsumexp\"", "workgroupSize": "reduceWorkgroupSize", "split": "splitCount", "outputF16": "dtypes.T == \"f16\"" } }, "bindings": "axisSplitCombine", "dispatch": { "threads": "axisSplitOutputs", "workgroupSize": "reduceWorkgroupSize" } } ] }, { "id": "axis0_splitk", "priority": 22, "when": ["not flatParallelCovered", "(dtypes.T == \"f32\" or dtypes.T == \"f16\")", "f16Ok(dtypes.T)", "ranks.data == 2", "reduceAxis == 0", "axis0Rows >= tunables.AXIS0_SPLIT_MIN_ROWS", "dim(shapes.data, 1) > 0", "((attrs.keepdims == 0 and ranks.reduced == 1 and dim(shapes.reduced, 0) == dim(shapes.data, 1)) or (attrs.keepdims == 1 and ranks.reduced == 2 and dim(shapes.reduced, 0) == 1 and dim(shapes.reduced, 1) == dim(shapes.data, 1)))", "axis0SplitPathFits"], "derive": { "splitCount": "axis0SplitCount" }, "constants": { "partialElement": "\"f32\"", "workgroupSize": "reduceWorkgroupSize", "split": "splitCount" }, "intermediates": [{ "id": "partials", "dtype": "float32", "shape": "[3 * splitCount * dim(shapes.data, 1)]" }], "passes": [ { "id": "split_reduce", "name": "ReduceLogSumExp.Axis0SplitKReduce", "source": { "shader": "reduce-axis0-splitk-reduce.wgsl.jinja", "inputs": { "op": "\"logsumexp\"", "workgroupSize": "reduceWorkgroupSize", "split": "splitCount", "castF32": "dtypes.T == \"f16\"", "usesF16": "dtypes.T == \"f16\"" } }, "bindings": "axis0SplitReduce", "dispatch": { "threads": "dim(shapes.data, 1)", "workgroupSize": "reduceWorkgroupSize", "y": "splitCount" } }, { "id": "combine", "name": "ReduceLogSumExp.Axis0SplitKCombine", "source": { "shader": "reduce-axis0-splitk-combine.wgsl.jinja", "inputs": { "op": "\"logsumexp\"", "workgroupSize": "reduceWorkgroupSize", "split": "splitCount", "outputF16": "dtypes.T == \"f16\"" } }, "bindings": "axis0SplitCombine", "dispatch": { "threads": "dim(shapes.data, 1)", "workgroupSize": "reduceWorkgroupSize" } } ] }, { "id": "axis0_tilecols", "priority": 20, "when": ["not flatParallelCovered", "(dtypes.T == \"f32\" or dtypes.T == \"f16\")", "f16Ok(dtypes.T)", "ranks.data == 2", "reduceAxis == 0", "axis0Rows >= tunables.AXIS0_TILE_MIN_ROWS", "axis0Cols >= tunables.AXIS0_TILE_MIN_COLS", "((attrs.keepdims == 0 and ranks.reduced == 1 and dim(shapes.reduced, 0) == dim(shapes.data, 1)) or (attrs.keepdims == 1 and ranks.reduced == 2 and dim(shapes.reduced, 0) == 1 and dim(shapes.reduced, 1) == dim(shapes.data, 1)))", "axis0TilePathFits"], "constants": { "workgroupSize": "reduceWorkgroupSize", "tileCols": "tunables.AXIS0_TILE_COLS" }, "passes": [ { "id": "main", "name": "ReduceLogSumExp.Axis0TileCols", "source": { "shader": "reduce-axis0-tilecols.wgsl.jinja", "inputs": { "op": "\"logsumexp\"", "castF32": "dtypes.T == \"f16\"", "usesF16": "dtypes.T == \"f16\"" } }, "bindings": "axis0Parallel", "dispatch": { "workgroups": "ceilDiv((dim(shapes.data, 1)), (constants.tileCols))" } } ] }, { "id": "all_axes_flat", "priority": 31, "constants": { "scalar": "dtypes.T", "workgroupSize": "reduceWorkgroupSize", "flatScalar": "\"vec4<\" ~ dtypes.T ~ \">\" if numel(shapes.data) % tunables.VECTOR_WIDTH == 0 else dtypes.T", "split": "flatSplitCount" }, "when": ["flatParallelCovered"], "intermediates": [{ "id": "partials", "dtype": "float32", "shape": "[3 * flatSplitCount]" }], "passes": [ { "id": "flat_partial", "name": "ReduceLogSumExp.AllAxesFlatPartial", "source": { "shader": "reduce-flat-partial-logsumexp.wgsl.jinja", "inputs": { "vec4": "numel(shapes.data) % tunables.VECTOR_WIDTH == 0", "castF32": "dtypes.T == \"f16\"", "usesF16": "dtypes.T == \"f16\"" } }, "bindings": "flatPartialF32", "dispatch": { "x": "flatSplitCount" } }, { "id": "combine", "name": "ReduceLogSumExp.AllAxesFlatCombine", "source": { "shader": "reduce-flat-combine-logsumexp.wgsl.jinja", "inputs": { "outputF16": "dtypes.T == \"f16\"" } }, "bindings": "flatCombineF32", "dispatch": { "x": 1 } } ] }, { "id": "rankn_single_axis_generic", "priority": 12, "supersededBy": ["axis_split_tiled_narrow", "axis_split", "subgroup_last_axis_vec4", "subgroup_last_axis", "tree_last_axis_vec4"], "when": ["not flatParallelCovered", "f16Ok(dtypes.T)", "ranks.data >= 3", "attrs.noop_with_empty_axes == 0", "reduceAxis < ranks.data", "numel(shapes.reduced) == rows(shapes.data, reduceAxis)", "((attrs.keepdims == 0 and ranks.reduced == ranks.data - 1) or (attrs.keepdims == 1 and ranks.reduced == ranks.data and dim(shapes.reduced, reduceAxis) == 1))"], "passes": [ { "id": "main", "name": "ReduceLogSumExp.RankNSingleAxisGeneric", "source": { "shader": "reduce-serial-axis.wgsl.jinja", "inputs": { "op": "\"logsumexp\"", "indexing": "\"rankn\"", "rank": "ranks.data", "axis": "reduceAxis", "dataShape": "shapes.data", "outputShape": "shapes.reduced", "outputRank": "ranks.reduced", "keepDims": "attrs.keepdims != 0", "intMode": "dtypes.T == \"i32\"", "castF32": "dtypes.T == \"f16\"", "usesF16": "dtypes.T == \"f16\"" } }, "bindings": "rankNAxis", "dispatch": { "threads": "numel(shapes.reduced)", "workgroupSize": "reduceWorkgroupSize" } } ], "constants": { "scalar": "dtypes.T" } }, { "id": "subgroup_last_axis_vec4", "priority": 25, "requires": { "features": ["subgroups"] }, "when": ["device.wgslLanguageFeatures.has(\"subgroup_id\")", "not flatParallelCovered", "(dtypes.T == \"f32\" or dtypes.T == \"f16\")", "f16Ok(dtypes.T)", "ranks.data >= 1", "reduceAxis == ranks.data - 1", "numel(shapes.reduced) == rows(shapes.data, ranks.data - 1)", "dim(shapes.data, ranks.data - 1) >= 4", "dim(shapes.data, ranks.data - 1) % tunables.VECTOR_WIDTH == 0", "(lastAxisCols >= tunables.SUBGROUP_MIN_COLS or lastAxisRows < tunables.SUBGROUP_SMALL_ROW_LIMIT)", "not rowSerialPreferred"], "constants": { "scalar": "dtypes.T", "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"", "workgroupSize": "min(reduceWorkgroupSize, max(subgroupWorkgroupFloor, pow2ceil(ceilDiv(lastAxisCols, tunables.VECTOR_WIDTH))))" }, "passes": [ { "id": "main", "name": "ReduceLogSumExp.SubgroupRowVec4", "source": { "shader": "reduce-row-subgroup.wgsl.jinja", "inputs": { "op": "\"logsumexp\"", "vec4": true, "castF32": "dtypes.T == \"f16\"", "usesF16": "dtypes.T == \"f16\"" } }, "subgroupCollectivesWidth": "portable", "bindings": "lastAxisVec4", "dispatch": { "workgroups": "rows(shapes.data, ranks.data - 1)" } } ] }, { "id": "subgroup_last_axis", "priority": 24, "requires": { "features": ["subgroups"] }, "when": ["device.wgslLanguageFeatures.has(\"subgroup_id\")", "not flatParallelCovered", "(dtypes.T == \"f32\" or dtypes.T == \"f16\")", "f16Ok(dtypes.T)", "ranks.data >= 1", "reduceAxis == ranks.data - 1", "numel(shapes.reduced) == rows(shapes.data, ranks.data - 1)", "dim(shapes.data, ranks.data - 1) > 0", "dim(shapes.data, ranks.data - 1) % tunables.VECTOR_WIDTH != 0", "(lastAxisCols >= tunables.SUBGROUP_MIN_COLS or lastAxisRows < tunables.SUBGROUP_SMALL_ROW_LIMIT)", "not rowSerialPreferred"], "constants": { "scalar": "dtypes.T", "workgroupSize": "min(reduceWorkgroupSize, max(subgroupWorkgroupFloor, pow2ceil(lastAxisCols)))" }, "passes": [ { "id": "main", "name": "ReduceLogSumExp.SubgroupRow", "source": { "shader": "reduce-row-subgroup.wgsl.jinja", "inputs": { "op": "\"logsumexp\"", "vec4": false, "castF32": "dtypes.T == \"f16\"", "usesF16": "dtypes.T == \"f16\"" } }, "subgroupCollectivesWidth": "portable", "bindings": "lastAxisScalarSubgroup", "dispatch": { "workgroups": "rows(shapes.data, ranks.data - 1)" } } ] }, { "id": "axis0", "priority": 0, "supersededBy": ["axis_split_tiled_narrow", "axis0_splitk", "axis0_tilecols"], "constants": { "axis": 0, "scalar": "dtypes.T" }, "when": ["not flatParallelCovered", "f16Ok(dtypes.T)", "ranks.data == 2", "reduceAxis == 0", "((attrs.keepdims == 0 and ranks.reduced == 1 and dim(shapes.reduced, 0) == dim(shapes.data, 1)) or (attrs.keepdims == 1 and ranks.reduced == 2 and dim(shapes.reduced, 0) == 1 and dim(shapes.reduced, 1) == dim(shapes.data, 1)))"], "passes": [ { "id": "main", "name": "axis0", "source": { "shader": "reduce-serial-axis.wgsl.jinja", "inputs": { "op": "\"logsumexp\"", "indexing": "\"axis2d\"", "intMode": "dtypes.T == \"i32\"", "castF32": "dtypes.T == \"f16\"", "usesF16": "dtypes.T == \"f16\"" } }, "bindings": "rank2Serial", "constants": { "axis": 0 }, "dispatch": { "threads": "numel(shapes.reduced)", "workgroupSize": "reduceWorkgroupSize" } } ] }, { "id": "axis1", "priority": 0, "constants": { "axis": 1, "scalar": "dtypes.T" }, "when": ["not flatParallelCovered", "f16Ok(dtypes.T)", "ranks.data == 2", "reduceAxis == 1", "((attrs.keepdims == 0 and ranks.reduced == 1 and dim(shapes.reduced, 0) == dim(shapes.data, 0)) or (attrs.keepdims == 1 and ranks.reduced == 2 and dim(shapes.reduced, 0) == dim(shapes.data, 0) and dim(shapes.reduced, 1) == 1))"], "passes": [ { "id": "main", "name": "axis1", "source": { "shader": "reduce-serial-axis.wgsl.jinja", "inputs": { "op": "\"logsumexp\"", "indexing": "\"axis2d\"", "intMode": "dtypes.T == \"i32\"", "castF32": "dtypes.T == \"f16\"", "usesF16": "dtypes.T == \"f16\"" } }, "bindings": "rank2SerialAxis1", "constants": { "axis": 1 }, "dispatch": { "threads": "numel(shapes.reduced)", "workgroupSize": "reduceWorkgroupSize" } } ] }, { "id": "all_axes_keepdims", "priority": 30, "constants": { "axis": 0, "scalar": "dtypes.T" }, "when": ["not flatParallelCovered", "f16Ok(dtypes.T)", "ranks.data >= 3", "attrs.keepdims == 1", "ranks.reduced == ranks.data", "numel(shapes.reduced) == 1"], "passes": [ { "id": "main", "name": "ReduceLogSumExp.Rank3AllAxesKeepdims", "source": { "shader": "reduce-serial-axis.wgsl.jinja", "inputs": { "op": "\"logsumexp\"", "indexing": "\"axis2d\"", "intMode": "dtypes.T == \"i32\"", "castF32": "dtypes.T == \"f16\"", "usesF16": "dtypes.T == \"f16\"" } }, "bindings": "fullReduceSerial", "dispatch": { "threads": "numel(shapes.reduced)", "workgroupSize": "reduceWorkgroupSize" } } ] }, { "id": "all_axes_no_keepdims", "priority": 30, "constants": { "axis": 0, "scalar": "dtypes.T" }, "when": ["not flatParallelCovered", "f16Ok(dtypes.T)", "ranks.data >= 3", "attrs.keepdims == 0", "attrs.noop_with_empty_axes == 0", "ranks.reduced == 0"], "passes": [ { "id": "main", "name": "ReduceLogSumExp.Rank3AllAxesNoKeepdims", "source": { "shader": "reduce-serial-axis.wgsl.jinja", "inputs": { "op": "\"logsumexp\"", "indexing": "\"axis2d\"", "intMode": "dtypes.T == \"i32\"", "castF32": "dtypes.T == \"f16\"", "usesF16": "dtypes.T == \"f16\"" } }, "bindings": "fullReduceSerial", "dispatch": { "threads": "numel(shapes.reduced)", "workgroupSize": "reduceWorkgroupSize" } } ] } ] }