Repository navigation
Conversation
nwh
commented
Nov 10, 2026
Collaborator
- ☑️ I understand it is strictly prohibited to use AI to write PR description
- AI usage disclosure:
This change lets `mx.compile` fuse an elementwise graph into a `Reduce` operation. The input of the reduction is computed in registers and is not written to memory. Example: before the change `mx.sum(x * y, -1)`ran an elementwise kernel that wrote `x * y` to memory, then a reduce kernel that read it back. With the change the elementwise multiplication and sum run in a single kernel avoiding the memory traffic of an intermediate result. `compile_fuse` walks the tape backwards to find groups of operations that can be fused into a single kernel. A `Reduce` (sum, prod, max, min, all, any) can now be the final operation on a fused tape on Metal for dtypes that are not complex. Setting the `MLX_DISABLE_REDUCE_FUSION` environment variable disables fusion of reduction operators. The result is a `Compiled` primitive whose tape is `[elementwise ops..., Reduce]`. In `export_to_dot` its name ends in the reduce name, e.g. `CompiledMultiplySum`. `Compiled::eval_gpu` calls `compiled_reduce` if the last op on the tape is `Reduce`. The kernel is chosen for each call from the actual input shapes and strides, so one compiled function can use different kernels. If no fused kernel fits, `compiled_reduce` uses a fallback: a regular `Compiled` writes the reduce input to a temporary, and a regular `Reduce` reads it. This occurs when the reduced axes do not collapse to one dimension, the inner strides are not 0 or 1 (e.g. negative strides), the kernel needs more than 31 Metal buffers, or the grid is too large.
the unary_fused_2 and binary_fused_3 blocks now run on the CPU (StreamContext ctx(Device::cpu);). Reductions don't fuse there, so the original Reduce → Compiled structure checks are kept. I added REQUIRE(...has_primitive()) before each .primitive() call on an input. That covers both of the blocks mentioned above, plus the unary_fused_3 block, which reads a primitive from an input the same way. A failed check now ends that test case instead of crashing the whole run. I didn't make a test fail on purpose to confirm this, but REQUIRE stops a test case at the first failed check.
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.