How to disable folding for particular ops in applyPatternsAndFoldGreedily?

[apivovarov]

LowerComplex mhlo pass uses applyPatternsAndFoldGreedily.
The doc says

/// Also performs folding and simple dead-code elimination before attempting to
/// match any of the provided patterns.

How to disable folding for particular ops?
In particular I need to preserve mhlo.copy op

The only workaround I found so far is to patch hlo_ops.cc and return {} in CopyOp::fold() instead of getOperand().

I would appreciate any help and ideas

[jreiffers]

mhlo.copy is semantically a no-op. You can:

1. Create your own copy op that isn't.
2. Try to change the semantics of mhlo.copy.

If you really have to do this, I recommend the former option. But I also recommend taking a good look at the semantics you're trying to get here. If your op also changes the tensor, this seems fine (and maybe option two could be appropriate). If you're really trying to get an identical copy of the value, and you don't want it to be cleaned up by any passes, I suspect you might run into many issues if you use MLIR infrastructure.

[apivovarov]

hlo has CopyInsertion <https://github.com/openxla/xla/blob/main/xla/service/copy_insertion.h#L46> pass If I apply CopyInsertion pass in hlo, then convert hlo to mhlo and run applyPatternsAndFoldGreedily (used in LowerComplex <https://github.com/openxla/xla/blob/main/xla/mlir_hlo/mhlo/transforms/lower_complex/lower_complex.cc#L64> for example) - it removes copy ops inserted by hlo CopyInsertion

…

On Tue, May 28, 2024 at 4:31 AM Johannes Reifferscheid < ***@***.***> wrote: mhlo.copy is semantically a no-op. You can: 1. Create your own copy op that isn't. 2. Try to change the semantics of mhlo.copy. If you really have to do this, I recommend the former option. But I also recommend taking a good look at the semantics you're trying to get here. If your op also changes the tensor, this seems fine (and maybe option two could be appropriate). If you're really trying to get an identical copy of the value, and you don't want it to be cleaned up by any passes, I suspect you might run into many issues if you use MLIR infrastructure. — Reply to this email directly, view it on GitHub <#13042 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AABVHZ7FKBU4DZVL2QUHIVDZERTI7AVCNFSM6AAAAABIG5AH5OVHI2DSMVQWIX3LMV43SRDJONRXK43TNFXW4Q3PNVWWK3TUHM4TKNZZHE3TA> . You are receiving this because you authored the thread.Message ID: ***@***.***>

[jreiffers]

I see. The problem is that there are invariants that are guaranteed by CopyInsertion, but don't match the semantics of MHLO. You're probably the first one to mix MLIR/HLO passes in this particular way.

I suggest checking for the presence of attributes in the folding function and not folding if there are any (regardless of what they are). This would also prevent folding of copies with this attribute here, which is probably correct.

If you don't want to touch MLIR, the other option is my option 1 above: rewrite all the mhlo.copy ops to my.copy, run the pass, rewrite them back, go back to HLO. But that seems strictly worse.

[apivovarov]

Is these any guidelines / recommendations regarding proper usage of hlo and mhlo passes in a typical DL compiler frontend? is it beneficial to use both hlo and mhlo passes? is there a recommended order? what passes should run first? hlo or mhlo? Will appreciate any chance information regarding this. Thank you

…

On Tue, May 28, 2024 at 10:49 PM Johannes Reifferscheid < ***@***.***> wrote: I see. The problem is that there are invariants that are guaranteed by CopyInsertion, but don't match the semantics of MHLO. You're probably the first one to mix MLIR/HLO passes in this particular way. I suggest checking for the presence of attributes in the folding function and not folding if there are any (regardless of what they are). This would also prevent folding of copies with this attribute here <https://github.com/openxla/xla/blob/main/xla/mlir_hlo/mhlo/IR/hlo_ops_attrs.td#L192>, which is probably correct. If you don't want to touch MLIR, the other option is my option 1 above: rewrite all the mhlo.copy ops to my.copy, run the pass, rewrite them back, go back to HLO. But that seems strictly worse. — Reply to this email directly, view it on GitHub <#13042 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AABVHZ3QNHEHBBGJX6HP67TZEVT5FAVCNFSM6AAAAABIG5AH5OVHI2DSMVQWIX3LMV43SRDJONRXK43TNFXW4Q3PNVWWK3TUHM4TKOBZGAYDO> . You are receiving this because you authored the thread.Message ID: ***@***.***>

[jreiffers]

If you use MHLO to LLVM lowerings:

• Once you're on MHLO, avoid going back to HLO.
• Therefore, HLO passes must run first.
• It can be beneficial to use both HLO and MHLO passes, since we don't have MHLO equivalents for everything, in particular the expansion passes.

If you don't, you'll have to figure out the right order to do things yourself, I don't think anyone has done it that way. You'll definitely want to avoid running late passes (like copy insertion) before the MHLO parts. Something like this should work:

1. expansion, simplification etc. HLO passes
2. conversion to MHLO
3. MHLO passes
4. conversion back to HLO
5. copy insertion, buffer assignment, scheduling, etc.