Skip to content

ConvertConcat: don't rebuild the concat-to-batch input (non-terminating rewrite) - #2988

Open
AsafManela wants to merge 1 commit into
EnzymeAD:mainfrom
AsafManela:fix-convertconcat-batch-cycle
Open

ConvertConcat: don't rebuild the concat-to-batch input (non-terminating rewrite)#2988
AsafManela wants to merge 1 commit into
EnzymeAD:mainfrom
AsafManela:fix-convertconcat-batch-cycle

Conversation

@AsafManela

Copy link
Copy Markdown

Fixes #2983.

convert(concat(reshape_i(x_i))), where every 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:

start:                            convert(concat(reshape(a), reshape(b)))
convert_concat                 -> concat(convert(reshape(a)), convert(reshape(b)))
elementwise_reshape_like       -> concat(reshape(convert(a)), reshape(convert(b)))
concat_insert_dim_elementwise  -> convert(concat(reshape(a), reshape(b)))   == start

so the greedy driver never reaches a fixed point, and each trip through the batching rewrite
leaves another enzymexla_unbatched_ConcatInsertDimToBatch_* wrapper function behind -- the
module 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 ConvertConcat in that shape. Nothing is lost by it there -- the converts get
batched straight back together. It is one of three possible cut points; say the word if
you'd rather break the loop in ElementwiseReshapeLike or in the batching rewrite and I'll
redo it.

reorder_elementwise_and_shape_op is elementwise_reshape_like's counterpart in the
transpose-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 that
test does not fail, it hangs).

Reactant.jl gets a matching change so its default :all pipeline stops feeding all three
patterns 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_jll v0.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

`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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

convert_concat / elementwise_reshape_like / concat_insert_dim_elementwise never terminate on convert(concat(reshape))

1 participant