Repository: llvm/llvm-project
File: mlir/lib/Dialect/ControlFlow/IR/ControlFlowOps.cpp (C++)

Code before:
///
static bool simplifyUniformBlockArgs(Block *dest, PatternRewriter &rewriter) {
  if (dest->hasNoPredecessors() ||
      llvm::hasSingleElement(dest->getPredecessors()))
    return false;

  bool changed = false;
  for (BlockArgument arg : dest->getArguments()) {
    if (arg.use_empty())
      continue;

    Value commonValue;
    bool allSame = true;
    for (Block *pred : dest->getPredecessors()) {
      auto branch = dyn_cast<BranchOpInterface>(pred->getTerminator());
      if (!branch) {
        allSame = false;
        break;
      }

      for (auto [i, succ] : llvm::enumerate(branch->getSuccessors())) {
        if (succ != dest)
          continue;
        Value val = branch.getSuccessorOperands(i)[arg.getArgNumber()];
        if (commonValue && commonValue != val) {
          allSame = false;
          break;
        }
        commonValue = val;
      }

      if (!allSame)
        break;
    }

    if (allSame && commonValue && commonValue != arg) {
      rewriter.replaceAllUsesWith(arg, commonValue);
      changed = true;
    }
  }
  return changed;
}

namespace {
/// Replaces block arguments with a uniform incoming value across all
/// predecessors, for any op implementing BranchOpInterface.
struct SimplifyUniformBlockArguments
    : public OpInterfaceRewritePattern<BranchOpInterface> {
  using OpInterfaceRewritePattern::OpInterfaceRewritePattern;
  LogicalResult matchAndRewrite(BranchOpInterface op,
                                PatternRewriter &rewriter) const override {

Review comment (suggestion):
Can you set `commonValue = Value()` here? I think then you won't need the `allSame` variable anymore. (If `commonValue` is non-null, you perform the replacement, otherwise, you don't.) Please double-check if my thinking is correct.


Code after:
///
static LogicalResult simplifyUniformBlockArgs(Block *dest,
                                              PatternRewriter &rewriter) {
  if (dest->hasNoPredecessors() ||
      llvm::hasSingleElement(dest->getPredecessors()))
    return failure();

  bool changed = false;
  for (BlockArgument arg : dest->getArguments()) {
    if (arg.use_empty())
      continue;

    Value commonValue;
    for (Block *pred : dest->getPredecessors()) {
      auto branch = dyn_cast<BranchOpInterface>(pred->getTerminator());
      if (!branch) {
        commonValue = Value();
        break;
      }

      for (auto [i, succ] : llvm::enumerate(branch->getSuccessors())) {
        if (succ != dest)
          continue;

        // Produced operands are modeled by BranchOpInterface as null Values.
        Value val = branch.getSuccessorOperands(i)[arg.getArgNumber()];
        if (commonValue && commonValue != val) {
          commonValue = Value();
          break;
        }
        commonValue = val;
      }

      if (!commonValue)
        break;
    }

    if (commonValue && commonValue != arg) {
      rewriter.replaceAllUsesWith(arg, commonValue);
      changed = true;
    }
  }
  return success(changed);
}

namespace {
/// Replaces block arguments with a uniform incoming value across all
/// predecessors, for any op implementing BranchOpInterface.
struct SimplifyUniformBlockArguments
    : public OpInterfaceRewritePattern<BranchOpInterface> {
  using OpInterfaceRewritePattern::OpInterfaceRewritePattern;