ConvertConcat: don't rebuild the concat-to-batch input (non-terminating rewrite) - #2988
Open
AsafManela wants to merge 1 commit into
Open
ConvertConcat: don't rebuild the concat-to-batch input (non-terminating rewrite)#2988AsafManela wants to merge 1 commit into
AsafManela wants to merge 1 commit into
Conversation
`convert(concat(reshape_i(x_i)))`, where each operand is a reshape-like op
inserting the concatenation dimension, is the shape `ConcatInsertDimToBatch`
produces out of `concat(reshape_i(convert(x_i)))`. Pushing the convert back
into the operands there re-creates that pattern's input:
convert_concat -> concat(convert(reshape(x_i)))
elementwise_reshape_like -> concat(reshape(convert(x_i)))
concat_insert_dim_elementwise -> convert(concat(reshape(x_i))) (start)
so the greedy rewriter never reaches a fixed point, and each trip through the
batching rewrite leaves another `enzymexla_unbatched_ConcatInsertDimToBatch_*`
wrapper function in the module. On a six-line function that loop runs for over
ten minutes with the module still growing (~1.3 GB -> 1.7 GB RSS in three
minutes); Reactant.jl's `:all` pipeline hit it on layered bf16 graphs, where
rotary-style `concat(reshape(...))` feeds a bf16 cast.
Decline the rewrite in that shape. Nothing is gained by it there: the converts
are batched straight back together.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
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.
Fixes #2983.
convert(concat(reshape_i(x_i))), where every operand is a reshape-like op inserting theconcatenation dimension, is the shape
ConcatInsertDimToBatchproduces out ofconcat(reshape_i(convert(x_i))). Pushing the convert back into the operands therere-creates that pattern's input:
so the greedy driver never reaches a fixed point, and each trip through the batching rewrite
leaves another
enzymexla_unbatched_ConcatInsertDimToBatch_*wrapper function behind -- themodule grows without bound. On the six-line function in the issue that loop runs for over
ten minutes with RSS still climbing (1.33 GB -> 1.70 GB in three minutes). Any graph with
fp32 master weights and bf16 compute hits it: a rotary embedding's
concat(reshape(...))feeding a bf16 cast is enough.
This declines
ConvertConcatin that shape. Nothing is lost by it there -- the converts getbatched straight back together. It is one of three possible cut points; say the word if
you'd rather break the loop in
ElementwiseReshapeLikeor in the batching rewrite and I'llredo it.
reorder_elementwise_and_shape_opiselementwise_reshape_like's counterpart in thetranspose-propagation group, so the transpose side loops the same way and is fixed by the
same guard.
Lit test:
test/lit_tests/convert_concat_insert_dim_cycle.mlir(before this change thattest does not fail, it hangs).
Reactant.jl gets a matching change so its default
:allpipeline stops feeding all threepatterns into one greedy run -- EnzymeAD/Reactant.jl#3225 -- which is also what users on an
older jll need.
Testing
The C++ is not built or run here -- no bazel in this environment. The reproducer,
timings and memory numbers come from running the shipped patterns (
Reactant_jllv0.0.405,same compiled code) through Reactant's pass-pipeline entry point. The lit test is written
against the current sources but has not been executed; please let CI have it, and tell me if
the CHECK lines need adjusting.
🤖 Generated with Claude Code
https://claude.ai/code/session_012k6KZhhCNSCvZFKzErHaHH