[MLIR][SCF] Fix nested if merging bug

[wsmoses]

Summary

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.

[ftynse]

mlir/lib/Dialect/SCF/SCF.cpp

1791–1807

Nit: continue here and drop the else + unindent?

1799

Did you mean can be defined outside the "if"?

mlir/lib/Dialect/SCF/SCF.cpp

LogicalResult matchAndRewrite(IfOp op,
     if (nestedIf.elseBlock() && !llvm::hasSingleElement(*nestedIf.elseBlock()))
       return failure();
 
     SmallVector<Value> thenYield(op.thenYield().getOperands());
     SmallVector<Value> elseYield;
     if (op.elseBlock())
       llvm::append_range(elseYield, op.elseYield().getOperands());
 
+    // A list of indices for which we should upgrade the value yielded
+    // in the else to a select.
+    SmallVector<unsigned> elseYieldsToUpgradeToSelect;
+
     // If the outer scf.if yields a value produced by the inner scf.if,
     // only permit combining if the value yielded when the condition
     // is false in the outer scf.if is the same value yielded when the
     // inner scf.if condition is false.
     // Note that the array access to elseYield will not go out of bounds
     // since it must have the same length as thenYield, since they both
     // come from the same scf.if.
     for (auto tup : llvm::enumerate(thenYield)) {
       if (tup.value().getDefiningOp() == nestedIf) {
         auto nestedIdx = tup.value().cast<OpResult>().getResultNumber();
         if (nestedIf.elseYield().getOperand(nestedIdx) !=
             elseYield[tup.index()]) {
           return failure();
         }
         // If the correctness test passes, we will yield
         // corresponding value from the inner scf.if
         thenYield[tup.index()] = nestedIf.thenYield().getOperand(nestedIdx);
+        continue;
+      }
+
+      // Otherwise, we need to ensure the else block of the combined
+      // condition still returns the same value when the outer condition is
+      // true and the inner condition is false. This can be accomplished if
+      // the then value is defined outside the outer scf.if and we replace the
+      // value with a select that considers just the outer condition. Since
+      // the else region contains just the yield, its yielded value is
+      // defined outside the scf.if, by definition.
+
+      // If the then value is defined within the scf.if, bail.
+      if (tup.value().getParentRegion() == &op.getThenRegion()) {
+        return failure();
+      } else {
+        elseYieldsToUpgradeToSelect.push_back(tup.index());
       }
     }
 
     Location loc = op.getLoc();
     Value newCondition = rewriter.create<arith::AndIOp>(
         loc, op.getCondition(), nestedIf.getCondition());
     auto newIf = rewriter.create<IfOp>(loc, op.getResultTypes(), newCondition);
+
+    SmallVector<Value> results;
+    llvm::append_range(results, newIf.getResults());
+    rewriter.setInsertionPoint(newIf);
+
+    for (auto idx : elseYieldsToUpgradeToSelect)
+      results[idx] = rewriter.create<arith::SelectOp>(
+          op.getLoc(), op.getCondition(), thenYield[idx], elseYield[idx]);
+
     Block *newIfBlock = newIf.thenBlock();
     if (newIfBlock)
       rewriter.eraseOp(newIfBlock->getTerminator());
     else
       newIfBlock = rewriter.createBlock(&newIf.getThenRegion());
     rewriter.mergeBlocks(nestedIf.thenBlock(), newIfBlock);
     rewriter.setInsertionPointToEnd(newIf.thenBlock());
     rewriter.replaceOpWithNewOp<YieldOp>(newIf.thenYield(), thenYield);
     if (!elseYield.empty()) {
       rewriter.createBlock(&newIf.getElseRegion());
       rewriter.setInsertionPointToEnd(newIf.elseBlock());
       rewriter.create<YieldOp>(loc, elseYield);
     }
-    rewriter.replaceOp(op, newIf.getResults());
+    rewriter.replaceOp(op, results);
     return success();
   }
 };
 
 } // namespace
 
 void IfOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                        MLIRContext *context) {

mlir/test/Dialect/SCF/canonicalize.mlir

 }
 
 // CHECK-LABEL: @merge_yielding_nested_if_nv2
 // CHECK-SAME: (%[[ARG0:.*]]: i1, %[[ARG1:.*]]: i1)
 func @merge_yielding_nested_if_nv2(%arg0: i1, %arg1: i1) -> i32 {
 // CHECK: %[[PRE0:.*]] = "test.op"() : () -> i32
 // CHECK: %[[PRE1:.*]] = "test.op1"() : () -> i32
 // CHECK: %[[COND:.*]] = arith.andi %[[ARG0]], %[[ARG1]]
-// CHECK: %[[RES:.*]] = arith.select %[[COND]], %[[PRE0]], %[[PRE1]]
+// CHECK: %[[RES:.*]] = arith.select %[[ARG0]], %[[PRE0]], %[[PRE1]]
 // CHECK: scf.if %[[COND]] 
 // CHECK:   "test.run"() : () -> ()
 // CHECK: }
 // CHECK: return %[[RES]]
   %0 = "test.op"() : () -> (i32)
   %1 = "test.op1"() : () -> (i32)
   %r = scf.if %arg0 -> i32 {
     scf.if %arg1 {
%a:2 = scf.if %arg1 -> (i32, f32) {
       scf.yield %0, %1 : i32, f32
     }
     scf.yield %0, %a#1, %a#0, %3 : i32, f32, i32, i8
   } else {
     scf.yield %0, %1, %2, %3 : i32, f32, i32, i8
   }
   return %r#0, %r#1, %r#2, %r#3 : i32, f32, i32, i8
 }
+
 // -----
 
 // CHECK-LABEL:   func @if_condition_swap
 // CHECK-NEXT:     %{{.*}} = scf.if %arg0 -> (index) {
 // CHECK-NEXT:       %[[i1:.+]] = "test.origFalse"() : () -> index
 // CHECK-NEXT:       scf.yield %[[i1]] : index
 // CHECK-NEXT:     } else {
 // CHECK-NEXT:       %[[i2:.+]] = "test.origTrue"() : () -> index