[RFC]: Move custom fusion passes to earlier in the compilation pipeline

[ProExpertProg]

Motivation.

#24604 introduces pattern matching using torch implementations of custom ops. @zou3519 pointed out is brittle and could break if Inductor transforms the graph such that the traced torch pattern for the custom op no longer matches.

Proposed Change.

We will move pattern matcher custom passes that rely on custom ops (basically all of them right now) from post_grad_custom_post_pass to post_grad_custom_pre_pass. Alternatively, we could move it to joint_graph_pass.

We will create a new pass manager object with similar semantics to the current PostGradPassManager (we'll rename them to make it clearer, perhaps PostGradPostPassManager and PostGradPrePassManager). Most current passes will move to the new manager (except for FixFunctionalizationPass, but that one will hopefully be deprecated anyway).

What we still need to figure out:

• our passes rely on some peephoole optimizations (removing noop slices, views, reshapes, etc.) - I think those currently run after post_grad_custom_pre_pass. We could run them manually, adapt our patterns, or improve pattern matcher support to match "through" certain pointwise operations.
• there is in-progress work (#164273) to allow matching mutable operations with the pattern matcher. It was going to add a hook near post_grad_custom_post_pass for going into a non-functional representation and back. We should make sure that work integrates nicely with the new approach, so perhaps it should happen near post_grad_custom_pre_pass instead.

Feedback Period.

2 weeks (1 week after PTC), 10/16-10/31

CC List.

@zou3519 @eellison @BoyuanFeng @youkaichao @mgoin @tlrmchlsmth @bnellnm @angelayi

Any Other Things.

No response

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.

[karthickai]

Currently we moved post_grad_mutable_custom_post_pass pytorch/pytorch#164273 after the post grad reinplace pass. One of the main reasons is to avoid illegal mutation pattern matching, as @eellison mentioned in the comment. Consider the illegal pattern where a clone is removed but would mutate a graph input

def pattern(x):
    y = x.clone()
    add_two_inplace(y)
    return y

def replacement(x):
    add_two_inplace(x)
    return x

If we place the mutable custom op pass before reinplace, the pattern matcher would replace the above pattern, but this is illegal because it mutates the graph input directly. After reinplace runs first, it provides only_clone_these_bases metadata indicating which tensors must be cloned for safety. This metadata is then used by decompose_auto_functionalized, which adds an extra as_strided clone op. This structural mismatch prevents the pattern matching from illegal mutation.

If we want to move post_grad_mutable_custom_post_pass near post_grad_custom_pre_pass, we need to move also the reinplace pass earlier to ensure it runs first, or we need to do differently to detect illegal mutation.