[Mlir-commits] [mlir] [mlir][reducer] Separate Reduction Steps in `findOptimal` and `applyPatterns` (PR #190560)

[llvmlistbot at llvm.org]

aidint wrote:

@linuxlonelyeagle Thank you so much for your comment! I'll try to clarify what I tried to do with more details.

**First Motivation**: There is currently a bug in `mlir-reduce`. `mlir-reduce` will cause a segfault if:
1. the initial module contains a trivially dead operation (let's call it `op1`)
2. `op1` is the operation that makes the module interesting for the test script

The reason this segfault happens is that the initial reduction node (`root`) gets updated multiple times in the `findOptimal` function. Once [here](https://github.com/llvm/llvm-project/blob/996efeaaa20247ac6bcdf5de9bf533c94c8a360f/mlir/lib/Reducer/ReductionTreePass.cpp#L104), and once inside the loop ([here](https://github.com/llvm/llvm-project/blob/996efeaaa20247ac6bcdf5de9bf533c94c8a360f/mlir/lib/Reducer/ReductionTreePass.cpp#L115)) after calling `applyPatterns` on its region. This is how this happens:

1. `root` node is interesting (it has the initial module)
2. we create iterator on `root` 
https://github.com/llvm/llvm-project/blob/996efeaaa20247ac6bcdf5de9bf533c94c8a360f/mlir/lib/Reducer/ReductionTreePass.cpp#L107
3. we start iterating, and choose `root` as our first reduction node in the loop
4. we call `applyPatterns` on its region, but with empty patterns.
5. every operation is in range for `root`, therefore `applyPatterns` will call `applyOpPatternsGreedily` on every single operation inside `root`'s region.
6. once it gets to `op1`  (our trivially dead but problematic operation), it will remove it from the region. That's because `applyOpPatternsGreedily` will run the following code on its worklist:
https://github.com/llvm/llvm-project/blob/996efeaaa20247ac6bcdf5de9bf533c94c8a360f/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp#L483-L486
7. Once `op1` is erased from the region, the module becomes uninteresting, and the `root` will be updated:
https://github.com/llvm/llvm-project/blob/996efeaaa20247ac6bcdf5de9bf533c94c8a360f/mlir/lib/Reducer/ReductionTreePass.cpp#L115
8. This will both release the pointer and erase `root.module`. `root.module` becomes a `nullptr`.
https://github.com/llvm/llvm-project/blob/996efeaaa20247ac6bcdf5de9bf533c94c8a360f/mlir/lib/Reducer/ReductionNode.cpp#L114
9. Once we hit `++iter` in the loop, `root` tries to generate new variants. New variants try to clone their parent's module
https://github.com/llvm/llvm-project/blob/996efeaaa20247ac6bcdf5de9bf533c94c8a360f/mlir/lib/Reducer/ReductionNode.cpp#L41
10. `root.module` is `nullptr`, so it will cause a segfault.

**Implementaion**
We can resolve this bug either
(1) by keeping `root` untouched, and start iterating from the first variation. 
This changes the logic of mlir-reduce. Currently `applyPatterns` is first applied on the whole region (`root`) and after reducing that, new variations are generated. This will have performance consequences so I didn't go with this option.

(2) or by having `root` untouched as the root of the reduction tree, but creating a clone of it as its first variation, and start iterating from this clone.
This will not change the current logic of mlir-reduce, but will keep `root` untouched.
That's why I introduced `addVariant` method on `ReductionNode` to be able to inject this clone as a variation to `root`.

**Second Motivation**: Even with resolving the bug, our initial module can't be reduced, because `applyPatterns` will call `applyOpPatternsGreedily` on its operations and therefore no variation will be interesting.

This is because, as I mentioned in the PR's description, "`findOptimal` and `applyPatterns` were each handling multiple types of reduction simultaneously".  Take `applyPatterns` for instance, it has multiple responsibilities given different arguments:

1. with `patterns = {}` and `eraseOpNotInRange = true` it acts not as a "pattern application" function, but more as an "operation elimination" one.
2. with `patterns != {}` and `eraseOpNotInRange = false` it skips operation elimination, but applies patterns.

To keep both of these responsibilities, the function is calling `applyOpPatternsGreedily`, probably assuming that it doesn't have a side effect if `patterns = {}`. This is why we can't reduce the said module.

**Implementation**
Separating these two responsibilities, will both separate the concerns and also let us reduce the modules with problematic trivially dead operation without hurting the logic.

As both elimination and pattern application can be considered a function call on reduction node's region inside `findOptimal`'s iteration loop, we can also reuse `findOptimal` by adding a template parameter to it that accepts both elimination and pattern application functions as template arguments (What I called `ApplyFn`).

In this structure, we can also integerate `eraseAllOpsInRegion` as another template argument. It uses the same tree traversal, but it will not generate new variants and it will finish on first iteration, logically equivalent to what's happening right now.

So, to summarize, the PR is trying to resolve an issue for "modules with problematic trivially dead operations", but the changes also reflect a refactor, without changing the logic of mlir-reduce. I believe any other solution could either end up in more functions or even duplication of code.

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