[Mlir-commits] [flang] [mlir] [mlir][CSE] Introduce hoist-pure-ops logic to CSE pass (PR #180556)

[lonely eagle]

linuxlonelyeagle wrote:

Thank you for looking into this. Let me describe the issue you encountered above. I have logged the changes locally after each call to the `hoistPureOp` function. I’ve included some comments, which I hope will help clarify the logic. Note: `replaceAllUsesWith`  is not used here.
```
[cse CSE.cpp:248 1] module {
[cse CSE.cpp:248 1]   func.func @cse_multiple_regions(%arg0: i1, %arg1: i32) -> (i32, i32) {
[cse CSE.cpp:248 1]     %0 = "test.producer"() : () -> i32
[cse CSE.cpp:248 1]     %1 = scf.if %arg0 -> (i32) {
                                     //
                                     // hoist the arith.add op.
                                    //
[cse CSE.cpp:248 1]       %3 = arith.addi %0, %0 : i32
[cse CSE.cpp:248 1]       %4 = arith.addi %0, %0 : i32
[cse CSE.cpp:248 1]       %5 = scf.if %arg0 -> (i32) {
[cse CSE.cpp:248 1]         scf.yield %3 : i32
[cse CSE.cpp:248 1]       } else {
[cse CSE.cpp:248 1]         scf.yield %4 : i32
[cse CSE.cpp:248 1]       }
[cse CSE.cpp:248 1]       scf.yield %5 : i32
[cse CSE.cpp:248 1]     } else {
[cse CSE.cpp:248 1]       scf.yield %arg1 : i32
[cse CSE.cpp:248 1]     }
[cse CSE.cpp:248 1]     %2 = scf.if %arg0 -> (i32) {
[cse CSE.cpp:248 1]       %3 = scf.if %arg0 -> (i32) {
[cse CSE.cpp:248 1]         %4 = arith.addi %0, %0 : i32
[cse CSE.cpp:248 1]         scf.yield %4 : i32
[cse CSE.cpp:248 1]       } else {
[cse CSE.cpp:248 1]         %4 = arith.addi %0, %0 : i32
[cse CSE.cpp:248 1]         scf.yield %4 : i32
[cse CSE.cpp:248 1]       }
[cse CSE.cpp:248 1]       scf.yield %3 : i32
[cse CSE.cpp:248 1]     } else {
[cse CSE.cpp:248 1]       scf.yield %arg1 : i32
[cse CSE.cpp:248 1]     }
[cse CSE.cpp:248 1]     return %1, %2 : i32, i32
[cse CSE.cpp:248 1]   }
[cse CSE.cpp:248 1] }
```

```
[cse CSE.cpp:248 1] moduleOp:
[cse CSE.cpp:248 1] module {
[cse CSE.cpp:248 1]   func.func @cse_multiple_regions(%arg0: i1, %arg1: i32) -> (i32, i32) {
[cse CSE.cpp:248 1]     %0 = "test.producer"() : () -> i32

// we hoist all add op, during the second hoisting of %1, 
// the operation %2 used to be left behind in the outer scf.if, 
// which prevented the scf.if from being simplified or removed. 
// Now, %2 is hoisted along with %1.

[cse CSE.cpp:248 1]     %1 = arith.addi %0, %0 : i32
[cse CSE.cpp:248 1]     %2 = arith.addi %0, %0 : i32
[cse CSE.cpp:248 1]     %3 = arith.addi %0, %0 : i32
[cse CSE.cpp:248 1]     %4 = scf.if %arg0 -> (i32) {
[cse CSE.cpp:248 1]       %6 = scf.if %arg0 -> (i32) {
[cse CSE.cpp:248 1]         scf.yield %1 : i32
[cse CSE.cpp:248 1]       } else {
[cse CSE.cpp:248 1]         scf.yield %1 : i32 //  use replaceAllUse.
[cse CSE.cpp:248 1]       }
[cse CSE.cpp:248 1]       scf.yield %6 : i32
[cse CSE.cpp:248 1]     } else {
[cse CSE.cpp:248 1]       scf.yield %arg1 : i32
[cse CSE.cpp:248 1]     }
[cse CSE.cpp:248 1]     %5 = scf.if %arg0 -> (i32) {
[cse CSE.cpp:248 1]       %6 = scf.if %arg0 -> (i32) {
[cse CSE.cpp:248 1]         scf.yield %3 : i32
[cse CSE.cpp:248 1]       } else {
[cse CSE.cpp:248 1]         %7 = arith.addi %0, %0 : i32
[cse CSE.cpp:248 1]         scf.yield %7 : i32
[cse CSE.cpp:248 1]       }
[cse CSE.cpp:248 1]       scf.yield %6 : i32
[cse CSE.cpp:248 1]     } else {
[cse CSE.cpp:248 1]       scf.yield %arg1 : i32
[cse CSE.cpp:248 1]     }
[cse CSE.cpp:248 1]     return %4, %5 : i32, i32
[cse CSE.cpp:248 1]   }
[cse CSE.cpp:248 1] }

```

```
[cse CSE.cpp:248 1] module {
[cse CSE.cpp:248 1]   func.func @cse_multiple_regions(%arg0: i1, %arg1: i32) -> (i32, i32) {
[cse CSE.cpp:248 1]     %0 = "test.producer"() : () -> i32
[cse CSE.cpp:248 1]     %1 = arith.addi %0, %0 : i32

// We hoist the scf.if, the reason for the invalid IR generated above is
//  that the use of %1 within the scf.if region was not taken into account. 
// Since %1 is now defined before the scf.if, this issue no longer exists.

[cse CSE.cpp:248 1]     %2 = scf.if %arg0 -> (i32) {
[cse CSE.cpp:248 1]       scf.yield %1 : i32
[cse CSE.cpp:248 1]     } else {
[cse CSE.cpp:248 1]       scf.yield %1 : i32
[cse CSE.cpp:248 1]     }
[cse CSE.cpp:248 1]     %3 = scf.if %arg0 -> (i32) {
[cse CSE.cpp:248 1]       scf.yield %1 : i32
[cse CSE.cpp:248 1]     } else {
[cse CSE.cpp:248 1]       scf.yield %1 : i32
[cse CSE.cpp:248 1]     }
[cse CSE.cpp:248 1]     %4 = arith.addi %0, %0 : i32
[cse CSE.cpp:248 1]     %5 = arith.addi %0, %0 : i32
[cse CSE.cpp:248 1]     %6 = arith.addi %0, %0 : i32
[cse CSE.cpp:248 1]     %7 = scf.if %arg0 -> (i32) {
[cse CSE.cpp:248 1]       scf.yield %2 : i32
[cse CSE.cpp:248 1]     } else {
[cse CSE.cpp:248 1]       scf.yield %arg1 : i32
[cse CSE.cpp:248 1]     }
[cse CSE.cpp:248 1]     %8 = scf.if %arg0 -> (i32) {
[cse CSE.cpp:248 1]       scf.yield %3 : i32
[cse CSE.cpp:248 1]     } else {
[cse CSE.cpp:248 1]       scf.yield %arg1 : i32
[cse CSE.cpp:248 1]     }
[cse CSE.cpp:248 1]     return %7, %8 : i32, i32
[cse CSE.cpp:248 1]   }
[cse CSE.cpp:248 1] }
[cse CSE.cpp:276 1] replace %3 = scf.if %arg0 -> (i32) {...} else {...} with %2 = scf.if %arg0 -> (i32) {...} else {...}
[cse CSE.cpp:279 1] add %3 = scf.if %arg0 -> (i32) {...} else {...} to opsToErase
[cse CSE.cpp:406 1] visit operation: scf.yield %2 : i32
[cse CSE.cpp:515 1] visit region #1 of %8 = scf.if %arg0 -> (i32) {...} else {...}
[cse CSE.cpp:479 1] visit block #0 of %8 = scf.if %arg0 -> (i32) {...} else {...}
[cse CSE.cpp:406 1] visit operation: scf.yield %arg1 : i32
[cse CSE.cpp:406 1] visit operation: %8 = scf.if %arg0 -> (i32) {...} else {...}
[cse CSE.cpp:276 1] replace %8 = scf.if %arg0 -> (i32) {...} else {...} with %7 = scf.if %arg0 -> (i32) {...} else {...}
[cse CSE.cpp:279 1] add %8 = scf.if %arg0 -> (i32) {...} else {...} to opsToErase
[cse CSE.cpp:406 1] visit operation: func.return %7, %7 : i32, i32
[cse CSE.cpp:406 1] visit operation: func.func @cse_multiple_regions(%arg0: i1, %arg1: i32) -> (i32, i32) {...}
module {
  func.func @cse_multiple_regions(%arg0: i1, %arg1: i32) -> (i32, i32) {
    %0 = "test.producer"() : () -> i32
    %1 = arith.addi %0, %0 : i32
    %2 = scf.if %arg0 -> (i32) {
      scf.yield %1 : i32
    } else {
      scf.yield %1 : i32
    }
    %3 = scf.if %arg0 -> (i32) {
      scf.yield %2 : i32
    } else {
      scf.yield %arg1 : i32
    }
    return %3, %3 : i32, i32
  }
}
```
If you have any further questions, please let me know.😉

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