[MLIR][SCF] Fix nested if merging bug
authorWilliam S. Moses <gh@wsmoses.com>
Mon, 21 Mar 2022 00:32:36 +0000 (20:32 -0400)
committerWilliam S. Moses <gh@wsmoses.com>
Mon, 21 Mar 2022 15:42:26 +0000 (11:42 -0400)
commit195de3dd6c86f01956f2d1f87b2b7dd25f8c0aed
tree4d864f644796b4d028f66f62df6b3bd9518140e4
parent6761dd7d0a7f4882e2f79cf10052652622a6d284
[MLIR][SCF] Fix nested if merging bug

The current nested if merging has a bug. Specifically, consider the following code:

```
    %r = scf.if %arg3 -> (i32) {
      scf.if %arg1 {
        "test.op"() : () -> ()
      }
      scf.yield %arg0 : i32
    } else {
      scf.yield %arg2 : i32
    }
```

When the above gets merged, it will become:
```
    %r = scf.if %arg3 && %arg1-> (i32) {
      "test.op"() : () -> ()
      scf.yield %arg0 : i32
    } else {
      scf.yield %arg2 : i32
    }
```

However, this means that when only %arg3 is true, we will incorrectly return %arg2 instead
of %arg0. This change updates the behavior of the pass to only enable nested if merging where
the outer yield contains only values from the inner if, or values defined outside of the if.

In the case of the latter, they can turned into a select of only the outer if condition, thus
maintaining correctness.

Reviewed By: ftynse

Differential Revision: https://reviews.llvm.org/D122108
mlir/lib/Dialect/SCF/SCF.cpp
mlir/test/Dialect/SCF/canonicalize.mlir