| { |
| "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" } |
| } |
| ] |
| } |
| ] |
| } |
|
|
