Adding traditional element-wise operations to stablehlo-complex-math-expander

[ddilbazTT]

Request description

I am currently working on adding complex data type and operation support to an MLIR-based Compiler. I want to add support for simple element-wise operations in complex data type in a similar manner to stablehlo-complex-math-expander. However, I want to understand what is the best way to go about adding support for operations like add, multiply, divide.

For example, I want to add expansion for multiplication such that:

t1 * t2 = complex(re(t1) * re(t2) - im(t1) * im(t2), re(t1) * im(t2) + re(t2) * im(t1))

Could you please inform me if I should add support for these operations in pearu/functional_algorithms or through a different way? It seems like stablehlo-complex-math-expander is targeting an approximation for operations with complex tensors to improve accuracy. However, the goal with adding support for operations like add, multiplication, divide is to keep arithmetic in real domain. Therefore, I wanted to ensure the best place to locate expansion for such ops.

@pearu @mrguenther @GleasonK I would appreciate if you could give me advice on how to go about this. I can open a PR on adding support for operations of my interest once I know where to add them to.

[GleasonK]

If they're relatively simple, you can add them manually without functional_algorithms. That library was very useful for generating patterns for larger decomps that were difficult to reason about in MLIR patterns. The multiplication example you use here seems a bit more simple and can likely be manually added to StablehloComplexMathExpanderPatterns.td (the file notes that its generated, but it hasn't/doesn't change much so I think its fine to edit directly there).

[ddilbazTT]

@GleasonK Sounds good, thanks for confirming that I can edit StablehloComplexMathExpanderPatterns.td directly. I will follow-up on this with a PR!

[pearu]

I agree that implementing simple operations such as multiplication of complex numbers within stablehlo does not require functional_algorithms. Doing that via functional_algorithms would make sense if one needs results with extended precision but that does not seem to be the case here.

[GleasonK]

Actually that's a good point - could you hide these behind a flag.. I need to think more on the name.

It sounds like we have a distinction here between:
(1) decomposing because hardware doesn't support complex numbers (new patterns).
(2) decomposing via a special algorithm because we need more precision (current state of the pass).

Maybe something like a "decompose-all" flag or option "mode=[all | precision]" that defaults to precision?

[ddilbazTT]

Hi @GleasonK I put up a draft PR on how to add basic arithmetic operations to complex math expander: #2929 I would appreciate your input before I continue with it.