[MLIR][SCF] Allow combining subsequent if statements that yield & negated condition

[wsmoses]

Summary

This patch extends the existing if combining canonicalization to also handle the case where a value returned by the first if is used within the body of the second if.

This patch also extends if combining to support if's whose conditions are logical negations of each other.

[wsmoses]

Add support for negated conditions

[ftynse]

LGTM with comments addressed.

mlir/lib/Dialect/SCF/SCF.cpp

1543–1558

Something like rewriter.replaceOpWithinBlock or rewriter.replaceOpWithIf should work here.

mlir/test/Dialect/SCF/canonicalize.mlir

1140

Please don't match exact value names (%arg0), here and below.

[wsmoses]

mlir/lib/Dialect/SCF/SCF.cpp

1543–1558

It's not clear to me this is possible as the values yielded by the if may be block arguments and thus not an op. I don't believe those methods work to replace a value (as opposed to an operation).

[wsmoses]

Fix test

[ftynse]

mlir/lib/Dialect/SCF/SCF.cpp

1543–1558

This code is updating the uses of results of prevIf, not the uses of values that are yielded. So I expect something like

rewriter.replaceOpWithinBlock(prevIf, prevIf.thenYield().getOperands(), nextThen);
rewriter.replaceOpWithinBlock(prevIf, prevElseYielded, nextElse);

to work.

mlir/lib/Dialect/SCF/SCF.cpp

LogicalResult matchAndRewrite(IfOp nextIf,
 
     auto prevIf = dyn_cast<IfOp>(nextIf->getPrevNode());
     if (!prevIf)
       return failure();
 
     if (nextIf.getCondition() != prevIf.getCondition())
       return failure();
 
-    // Don't permit merging if a result of the first if is used
-    // within the second.
-    if (llvm::any_of(prevIf->getUsers(),
-                     [&](Operation *user) { return nextIf->isAncestor(user); }))
-      return failure();
+    SmallVector<Value> prevElseYielded;
+    if (!prevIf.getElseRegion().empty())
+      prevElseYielded = prevIf.elseYield().getOperands();
+    // Replace all uses of return values of op within nextIf with the
+    // corresponding yields
+    for (auto it : llvm::zip(prevIf.getResults(),
+                             prevIf.thenYield().getOperands(), prevElseYielded))
+      for (OpOperand &use :
+           llvm::make_early_inc_range(std::get<0>(it).getUses())) {
+        if (nextIf.getThenRegion().isAncestor(
+                use.getOwner()->getParentRegion())) {
+          rewriter.startRootUpdate(use.getOwner());
+          use.set(std::get<1>(it));
+          rewriter.finalizeRootUpdate(use.getOwner());
+        } else if (nextIf.getElseRegion().isAncestor(
+                       use.getOwner()->getParentRegion())) {
+          rewriter.startRootUpdate(use.getOwner());
+          use.set(std::get<2>(it));
+          rewriter.finalizeRootUpdate(use.getOwner());
+        }
+      }
 
     SmallVector<Type> mergedTypes(prevIf.getResultTypes());
     llvm::append_range(mergedTypes, nextIf.getResultTypes());
 
     IfOp combinedIf = rewriter.create<IfOp>(
-        nextIf.getLoc(), mergedTypes, nextIf.getCondition(), /*hasElse=*/false);
+        nextIf.getLoc(), mergedTypes, prevIf.getCondition(), /*hasElse=*/false);
     rewriter.eraseBlock(&combinedIf.getThenRegion().back());
 
     YieldOp thenYield = prevIf.thenYield();
     YieldOp thenYield2 = nextIf.thenYield();
 
     combinedIf.getThenRegion().getBlocks().splice(
         combinedIf.getThenRegion().getBlocks().begin(),
         prevIf.getThenRegion().getBlocks());
 
     rewriter.mergeBlocks(nextIf.thenBlock(), combinedIf.thenBlock());
     rewriter.setInsertionPointToEnd(combinedIf.thenBlock());
 
     SmallVector<Value> mergedYields(thenYield.getOperands());
     llvm::append_range(mergedYields, thenYield2.getOperands());
     rewriter.create<YieldOp>(thenYield2.getLoc(), mergedYields);
     rewriter.eraseOp(thenYield);
     rewriter.eraseOp(thenYield2);
 
     combinedIf.getElseRegion().getBlocks().splice(
         combinedIf.getElseRegion().getBlocks().begin(),
         prevIf.getElseRegion().getBlocks());
 
     if (!nextIf.getElseRegion().empty()) {
       if (combinedIf.getElseRegion().empty()) {

mlir/test/Dialect/SCF/canonicalize.mlir

func @combineIfs4(%arg0 : i1, %arg2: i64) {
   return
 }
 
 // CHECK-NEXT:     scf.if %arg0 {
 // CHECK-NEXT:       "test.firstCodeTrue"() : () -> ()
 // CHECK-NEXT:       "test.secondCodeTrue"() : () -> ()
 // CHECK-NEXT:     }
 
+// CHECK-LABEL: @combineIfsUsed
+func @combineIfsUsed(%arg0 : i1, %arg2: i64) -> (i32, i32) {
+  %res = scf.if %arg0 -> i32 {
+    %v = "test.firstCodeTrue"() : () -> i32
+    scf.yield %v : i32
+  } else {
+    %v2 = "test.firstCodeFalse"() : () -> i32
+    scf.yield %v2 : i32
+  }
+  %res2 = scf.if %arg0 -> i32 {
+    %v = "test.secondCodeTrue"(%res) : (i32) -> i32
+    scf.yield %v : i32
+  } else {
+    %v2 = "test.secondCodeFalse"(%res) : (i32) -> i32
+    scf.yield %v2 : i32
+  }
+  return %res, %res2 : i32, i32
+}
+// CHECK-NEXT:     %[[res:.+]]:2 = scf.if %arg0 -> (i32, i32) {
+// CHECK-NEXT:       %[[tval0:.+]] = "test.firstCodeTrue"() : () -> i32
+// CHECK-NEXT:       %[[tval:.+]] = "test.secondCodeTrue"(%[[tval0]]) : (i32) -> i32
+// CHECK-NEXT:       scf.yield %[[tval0]], %[[tval]] : i32, i32
+// CHECK-NEXT:     } else {
+// CHECK-NEXT:       %[[fval0:.+]] = "test.firstCodeFalse"() : () -> i32
+// CHECK-NEXT:       %[[fval:.+]] = "test.secondCodeFalse"(%[[fval0]]) : (i32) -> i32
+// CHECK-NEXT:       scf.yield %[[fval0]], %[[fval]] : i32, i32
+// CHECK-NEXT:     }
+// CHECK-NEXT:     return %[[res]]#0, %[[res]]#1 : i32, i32
+
 // -----
 
 // CHECK-LABEL: func @propagate_into_execute_region
 func @propagate_into_execute_region() {
   %cond = arith.constant 0 : i1
   affine.for %i = 0 to 100 {
     "test.foo"() : () -> ()
     %v = scf.execute_region -> i64 {