[MLIR][SCF] Combine adjacent scf.if with same condition

[mehdi_amini]

mlir/lib/Dialect/SCF/SCF.cpp

1289

The restriction on the first if isn't intrinsic to the fusion right? We could fuse the two even if they both return results, as long as the else branch of the second does not use a result of the first if I think.

1303

[wsmoses]

Address comment

mlir/lib/Dialect/SCF/SCF.cpp

1289

I think the condition you mention isn't quite sufficient, but ensuring the second if doesn't use the result of the first if should be sufficient. Adding a note that this can be extended to support that.

[mehdi_amini]

mlir/lib/Dialect/SCF/SCF.cpp

1289

Actually neither the then or the else branch should use the result indeed.
Can we implement the general fusion directly instead of implementing this partial one?

[wsmoses]

Permit yielded values in first scf.if

[mehdi_amini]

mlir/lib/Dialect/SCF/SCF.cpp

1308–1310

You have a getUsers() on Operation which allows to directly iterates on this nested loop (see inline edit)

1313

Nit: no need for braces here.

1318

(spurious braces)

and you can likely just assign() here:

mergedTypes.assign(nextIf.getResultTypes().begin(), nextIf.getResultTypes().end());

1326

rewriter.mergeBlocks will do the splicing, but at the end.
Can you merge into the prevIf and avoid manual splicing here?

1329

Nit: .front() would avoid derefencing an iterator (may just be more readable)

1338–1340

1363

Just like the other revision: there are far too much back() (and even back().back()!) here, I'd like to see some refactoring into better accessors.

1366–1368

1385

Can you use llvm::enumerate instead of a side indexing?

[wsmoses]

Refactor API usage

[ftynse]

mlir/lib/Dialect/SCF/SCF.cpp

1317

Ultra-nit: we tend to use the /*hasElse=*/false style.

1330–1331

setInsertionPointToEnd(combinedIf.thenBlock())

1334–1335

Nit: there is llvm::append_range.

[wsmoses]

Simplify splicing

[mehdi_amini]

mlir/lib/Dialect/SCF/SCF.cpp

1342

1343

Why do you still have some "splice" above instead of rewriter.mergeBlocks()?

It seems that these splice are actually adding blocks to the region instead of merging the blocks? I'm not sure I read the code right...

[wsmoses]

mlir/lib/Dialect/SCF/SCF.cpp

1342

I don't think that's correct since the splice moves the (optional) block from the nextIf to the new if, whereas your change would move instructions (rather than blocks).

1343

The splice is to conditionally add the blocks of the else region (if the block exists).

[mehdi_amini]

mlir/lib/Dialect/SCF/SCF.cpp

1343

Ah OK I see the logic, it is actually moving and preserving the block itself. And we need this because the else region can be empty but not the then one.
It didn't seem obvious to me on first read somehow, either I'm too tired to review right now, or some high-level comments could help along the way.

mlir/include/mlir/Dialect/SCF/SCFOps.td

OpBuilder getThenBodyBuilder(OpBuilder::Listener *listener = nullptr) {
       return results().empty() ? OpBuilder::atBlockTerminator(body, listener)
                                : OpBuilder::atBlockEnd(body, listener);
     }
     OpBuilder getElseBodyBuilder(OpBuilder::Listener *listener = nullptr) {
       Block* body = getBody(1);
       return results().empty() ? OpBuilder::atBlockTerminator(body, listener)
                                : OpBuilder::atBlockEnd(body, listener);
     }
+    Block* thenBlock();
+    YieldOp thenYield();
+    Block* elseBlock();
+    YieldOp elseYield();
   }];
 
   let hasCanonicalizer = 1;
 }
 
 def ParallelOp : SCF_Op<"parallel",
     [AttrSizedOperandSegments,
      DeclareOpInterfaceMethods<LoopLikeOpInterface>,

mlir/lib/Dialect/SCF/SCF.cpp

for (auto tup :
           changed = true;
         }
       }
     }
     return success(changed);
   }
 };
 
+/// Merge any consecutive scf.if's with the same condition.
+///
+///    scf.if %cond {
+///       firstCodeTrue();...
+///    } else {
+///       firstCodeFalse();...
+///    }
+///    %res = scf.if %cond {
+///       secondCodeTrue();...
+///    } else {
+///       secondCodeFalse();...
+///    }
+///
+///  becomes
+///    %res = scf.if %cmp {
+///       firstCodeTrue();...
+///       secondCodeTrue();...
+///    } else {
+///       firstCodeFalse();...
+///       secondCodeFalse();...
+///    }
+struct CombineIfs : public OpRewritePattern<IfOp> {
+  using OpRewritePattern<IfOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(IfOp nextIf,
+                                PatternRewriter &rewriter) const override {
+    Block *parent = nextIf->getBlock();
+    if (nextIf == &parent->front())
+      return failure();
+
+    auto prevIf = dyn_cast<IfOp>(nextIf->getPrevNode());
+    if (!prevIf)
+      return failure();
+
+    if (nextIf.condition() != prevIf.condition())
+      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<Type> mergedTypes(prevIf.getResultTypes());
+    llvm::append_range(mergedTypes, nextIf.getResultTypes());
+
+    IfOp combinedIf = rewriter.create<IfOp>(
+        nextIf.getLoc(), mergedTypes, nextIf.condition(), /*hasElse=*/false);
+    rewriter.eraseBlock(&combinedIf.thenRegion().back());
+
+    YieldOp thenYield = prevIf.thenYield();
+    YieldOp thenYield2 = nextIf.thenYield();
+
+    combinedIf.thenRegion().getBlocks().splice(
+        combinedIf.thenRegion().getBlocks().begin(),
+        prevIf.thenRegion().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.elseRegion().getBlocks().splice(
+        combinedIf.elseRegion().getBlocks().begin(),
+        prevIf.elseRegion().getBlocks());
+
+    if (!nextIf.elseRegion().empty()) {
+      if (combinedIf.elseRegion().empty()) {
+        combinedIf.elseRegion().getBlocks().splice(
+            combinedIf.elseRegion().getBlocks().begin(),
+            nextIf.elseRegion().getBlocks());
+      } else {
+        YieldOp elseYield = combinedIf.elseYield();
+        YieldOp elseYield2 = nextIf.elseYield();
+        rewriter.mergeBlocks(nextIf.elseBlock(), combinedIf.elseBlock());
+
+        rewriter.setInsertionPointToEnd(combinedIf.elseBlock());
+
+        SmallVector<Value> mergedElseYields(elseYield.getOperands());
+        llvm::append_range(mergedElseYields, elseYield2.getOperands());
+
+        rewriter.create<YieldOp>(elseYield2.getLoc(), mergedElseYields);
+        rewriter.eraseOp(elseYield);
+        rewriter.eraseOp(elseYield2);
+      }
+    }
+
+    SmallVector<Value> prevValues;
+    SmallVector<Value> nextValues;
+    for (auto pair : llvm::enumerate(combinedIf.getResults())) {
+      if (pair.index() < prevIf.getNumResults())
+        prevValues.push_back(pair.value());
+      else
+        nextValues.push_back(pair.value());
+    }
+    rewriter.replaceOp(prevIf, prevValues);
+    rewriter.replaceOp(nextIf, nextValues);
+    return success();
+  }
+};
+
 } // namespace
 
 void IfOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                        MLIRContext *context) {
-  results
-      .add<RemoveUnusedResults, RemoveStaticCondition, ConvertTrivialIfToSelect,
-           ConditionPropagation, ReplaceIfYieldWithConditionOrValue>(context);
+  results.add<RemoveUnusedResults, RemoveStaticCondition,
+              ConvertTrivialIfToSelect, ConditionPropagation,
+              ReplaceIfYieldWithConditionOrValue, CombineIfs>(context);
 }
 
+Block *IfOp::thenBlock() { return &thenRegion().back(); }
+YieldOp IfOp::thenYield() { return cast<YieldOp>(&thenBlock()->back()); }
+Block *IfOp::elseBlock() { return &elseRegion().back(); }
+YieldOp IfOp::elseYield() { return cast<YieldOp>(&elseBlock()->back()); }
+
 //===----------------------------------------------------------------------===//
 // ParallelOp
 //===----------------------------------------------------------------------===//
 
 void ParallelOp::build(
     OpBuilder &builder, OperationState &result, ValueRange lowerBounds,
     ValueRange upperBounds, ValueRange steps, ValueRange initVals,
     function_ref<void(OpBuilder &, Location, ValueRange, ValueRange)>

mlir/test/Dialect/SCF/canonicalize.mlir

 // CHECK-NEXT:         %{{.+}} = scf.while : () -> i1 {
 // CHECK-NEXT:           %[[cmp:.+]] = "test.condition"() : () -> i1
 // CHECK-NEXT:           scf.condition(%[[cmp]]) %[[cmp]] : i1
 // CHECK-NEXT:         } do {
 // CHECK-NEXT:         ^bb0(%arg0: i1):  // no predecessors
 // CHECK-NEXT:           "test.use"(%[[true]]) : (i1) -> ()
 // CHECK-NEXT:           scf.yield
 // CHECK-NEXT:         }
+
+// -----
+
+// CHECK-LABEL: @combineIfs
+func @combineIfs(%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"() : () -> i32
+    scf.yield %v : i32
+  } else {
+    %v2 = "test.secondCodeFalse"() : () -> 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"() : () -> i32
+// CHECK-NEXT:       scf.yield %[[tval0]], %[[tval]] : i32, i32
+// CHECK-NEXT:     } else {
+// CHECK-NEXT:       %[[fval0:.+]] = "test.firstCodeFalse"() : () -> i32
+// CHECK-NEXT:       %[[fval:.+]] = "test.secondCodeFalse"() : () -> i32
+// CHECK-NEXT:       scf.yield %[[fval0]], %[[fval]] : i32, i32
+// CHECK-NEXT:     }
+// CHECK-NEXT:     return %[[res]]#0, %[[res]]#1 : i32, i32
+
+
+// CHECK-LABEL: @combineIfs2
+func @combineIfs2(%arg0 : i1, %arg2: i64) -> i32 {
+  scf.if %arg0 {
+    "test.firstCodeTrue"() : () -> ()
+    scf.yield
+  }
+  %res = scf.if %arg0 -> i32 {
+    %v = "test.secondCodeTrue"() : () -> i32
+    scf.yield %v : i32
+  } else {
+    %v2 = "test.secondCodeFalse"() : () -> i32
+    scf.yield %v2 : i32
+  }
+  return %res : i32
+}
+// CHECK-NEXT:     %[[res:.+]] = scf.if %arg0 -> (i32) {
+// CHECK-NEXT:       "test.firstCodeTrue"() : () -> ()
+// CHECK-NEXT:       %[[tval:.+]] = "test.secondCodeTrue"() : () -> i32
+// CHECK-NEXT:       scf.yield %[[tval]] : i32
+// CHECK-NEXT:     } else {
+// CHECK-NEXT:       %[[fval:.+]] = "test.secondCodeFalse"() : () -> i32
+// CHECK-NEXT:       scf.yield %[[fval]] : i32
+// CHECK-NEXT:     }
+// CHECK-NEXT:     return %[[res]] : i32
+
+
+// CHECK-LABEL: @combineIfs3
+func @combineIfs3(%arg0 : i1, %arg2: i64) -> i32 {
+  %res = scf.if %arg0 -> i32 {
+    %v = "test.firstCodeTrue"() : () -> i32
+    scf.yield %v : i32
+  } else {
+    %v2 = "test.firstCodeFalse"() : () -> i32
+    scf.yield %v2 : i32
+  }
+  scf.if %arg0 {
+    "test.secondCodeTrue"() : () -> ()
+    scf.yield
+  }
+  return %res : i32
+}
+// CHECK-NEXT:     %[[res:.+]] = scf.if %arg0 -> (i32) {
+// CHECK-NEXT:       %[[tval:.+]] = "test.firstCodeTrue"() : () -> i32
+// CHECK-NEXT:       "test.secondCodeTrue"() : () -> ()
+// CHECK-NEXT:       scf.yield %[[tval]] : i32
+// CHECK-NEXT:     } else {
+// CHECK-NEXT:       %[[fval:.+]] = "test.firstCodeFalse"() : () -> i32
+// CHECK-NEXT:       scf.yield %[[fval]] : i32
+// CHECK-NEXT:     }
+// CHECK-NEXT:     return %[[res]] : i32
+
+// CHECK-LABEL: @combineIfs4
+func @combineIfs4(%arg0 : i1, %arg2: i64) {
+  scf.if %arg0 {
+    "test.firstCodeTrue"() : () -> ()
+    scf.yield
+  }
+  scf.if %arg0 {
+    "test.secondCodeTrue"() : () -> ()
+    scf.yield
+  }
+  return
+}
+
+// CHECK-NEXT:     scf.if %arg0 {
+// CHECK-NEXT:       "test.firstCodeTrue"() : () -> ()
+// CHECK-NEXT:       "test.secondCodeTrue"() : () -> ()
+// CHECK-NEXT:     }