[Mlir-commits] [mlir] [mlir][math] Propagate fast math attrs in AlgebraicSimplification (PR #166802)

[Aleksei Nurmukhametov]

nurmukhametov wrote:

> In addition to adding LIT tests, can you explain why each propagation step is correct and why it's always safe to apply the same flags to intermediate computation?

My initial thought was that it is more or less obvious and I need just to pass all flags to all intermediate operations like [here](https://github.com/Xilinx/llvm-project/blob/193995d5a21dc8b923e19d9370aa8e1f374cd940/mlir/lib/Dialect/Arith/Transforms/ExpandOps.cpp#L646) or [here](https://github.com/Xilinx/llvm-project/blob/193995d5a21dc8b923e19d9370aa8e1f374cd940/mlir/lib/Conversion/ComplexToStandard/ComplexToStandard.cpp#L204). However, after thorough consideration, I am not quite sure that this is the case. As there are quite a few separate flags and transformations, I'll try to go through them one by one:

1. `pow(x, 2.0)` -> `x * x`
No intermediate computation here. The propagation is correct because both operations have the same input/output domain and behavior across it.
3. `pow(x, 3.0)` -> `x * x * x`
There are intermediate computations here, so it looks like `reassoc` can be placed only on the last multiplication to prevent situations like: `a * (x * x * x)` → `(a * x * x) * x`. Similar things may happen with `arcp`/`contract` flags when intermediate operations could be involved in transformations with outer-context computations. Values covered with `ninf`/`nsz`/`nan` are not produced by intermediate operations, so there is no reason to clear those flags for them. The propagation is correct due to the same domain and behavior.
4. `pow(x, -1.0)` -> `1.0 / x`
No intermediate computation here. The propagation is correct due to the same domain and behavior.
5. `pow(x, 0.5)` -> `sqrt(x)`
No intermediate computation here. These do not have the same behavior for all inputs (e.g., -inf and -0.0), and this simplification should not happen for these values. The propagation is correct if the simplification is applicable.
6. `pow(x, -0.5)` -> `rsqrt(x)`
No intermediate computation here. Same as 3 and 4.
7. `pow(x, 0.75)` -> `sqrt(v = sqrt(x)) * v`
There are intermediate computations here, so it is similar to 2. The same considerations about propagation as in case 4 are applicable here.

For 1, 3, and 4, the same propagation is applied [here](https://github.com/llvm/llvm-project/blob/f22d5884b2847c575b83b067c1c6b69f224f0854/llvm/lib/Transforms/Utils/SimplifyLibCalls.cpp#L2365).

I dropped `reassoc`/`contract`/`arcp` flags for intermediate operations and preserved them for the final operation after this analysis.

https://github.com/llvm/llvm-project/pull/166802