ConditionPropagation can be slow

[wsmoses]

x/ref:

struct ConditionPropagation : public OpRewritePattern<IfOp> {
  using OpRewritePattern<IfOp>::OpRewritePattern;

  LogicalResult matchAndRewrite(IfOp op,
                                PatternRewriter &rewriter) const override {
    // Early exit if the condition is constant since replacing a constant
    // in the body with another constant isn't a simplification.
    if (matchPattern(op.getCondition(), m_Constant()))
      return failure();

    bool changed = false;
    mlir::Type i1Ty = rewriter.getI1Type();

    // These variables serve to prevent creating duplicate constants
    // and hold constant true or false values.
    Value constantTrue = nullptr;
    Value constantFalse = nullptr;

    for (OpOperand &use :
         llvm::make_early_inc_range(op.getCondition().getUses())) {
      if (op.getThenRegion().isAncestor(use.getOwner()->getParentRegion())) {
        changed = true;

        if (!constantTrue)
          constantTrue = arith::ConstantOp::create(
              rewriter, op.getLoc(), i1Ty, rewriter.getIntegerAttr(i1Ty, 1));

        rewriter.modifyOpInPlace(use.getOwner(),
                                 [&]() { use.set(constantTrue); });
      } else if (op.getElseRegion().isAncestor(
                     use.getOwner()->getParentRegion())) {
        changed = true;

        if (!constantFalse)
          constantFalse = arith::ConstantOp::create(
              rewriter, op.getLoc(), i1Ty, rewriter.getIntegerAttr(i1Ty, 0));

        rewriter.modifyOpInPlace(use.getOwner(),
                                 [&]() { use.set(constantFalse); });
      }
    }

    return success(changed);
  }
};

If the condition has many uses (not just within this if statement), the query of isAncestor can be quite expensive (and in practice is quite expensive when cmopiling large c++ codebases with polygeist/reactant).

We can create a cache where if we've previously proven a region to be within then/else/neither, we don't need to re-query.

cc @ftynse

[llvmbot]

@llvm/issue-subscribers-mlir

Author: William Moses (wsmoses)

x/ref: ```

struct ConditionPropagation : public OpRewritePattern<IfOp> {
using OpRewritePattern<IfOp>::OpRewritePattern;

LogicalResult matchAndRewrite(IfOp op,
PatternRewriter &rewriter) const override {
// Early exit if the condition is constant since replacing a constant
// in the body with another constant isn't a simplification.
if (matchPattern(op.getCondition(), m_Constant()))
return failure();

bool changed = false;
mlir::Type i1Ty = rewriter.getI1Type();

// These variables serve to prevent creating duplicate constants
// and hold constant true or false values.
Value constantTrue = nullptr;
Value constantFalse = nullptr;

for (OpOperand &use :
     llvm::make_early_inc_range(op.getCondition().getUses())) {
  if (op.getThenRegion().isAncestor(use.getOwner()->getParentRegion())) {
    changed = true;

    if (!constantTrue)
      constantTrue = arith::ConstantOp::create(
          rewriter, op.getLoc(), i1Ty, rewriter.getIntegerAttr(i1Ty, 1));

    rewriter.modifyOpInPlace(use.getOwner(),
                             [&]() { use.set(constantTrue); });
  } else if (op.getElseRegion().isAncestor(
                 use.getOwner()->getParentRegion())) {
    changed = true;

    if (!constantFalse)
      constantFalse = arith::ConstantOp::create(
          rewriter, op.getLoc(), i1Ty, rewriter.getIntegerAttr(i1Ty, 0));

    rewriter.modifyOpInPlace(use.getOwner(),
                             [&]() { use.set(constantFalse); });
  }
}

return success(changed);

}
};

If the condition has many uses (not just within this if statement), the query of isAncestor can be quite expensive (and in practice is quite expensive when cmopiling large c++ codebases with polygeist/reactant).

We can create a cache where if we've previously proven a region to be within then/else/neither, we don't need to re-query.

cc @<!-- -->ftynse 
</details>