diff --git a/mlir/lib/Dialect/SCF/SCF.cpp b/mlir/lib/Dialect/SCF/SCF.cpp --- a/mlir/lib/Dialect/SCF/SCF.cpp +++ b/mlir/lib/Dialect/SCF/SCF.cpp @@ -1106,12 +1106,101 @@ return success(); } }; + +struct ConditionPropagation : public OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(IfOp op, + PatternRewriter &rewriter) const override { + if (op.condition().getDefiningOp()) { + return failure(); + } + bool Changed = false; + mlir::Type i1Ty = rewriter.getI1Type(); + for (OpOperand &use : + llvm::make_early_inc_range(op.condition().getUses())) { + if (op.thenRegion().isAncestor(use.getOwner()->getParentRegion())) { + Changed = true; + rewriter.updateRootInPlace(use.getOwner(), [&]() { + use.set(rewriter.create( + op.getLoc(), i1Ty, rewriter.getIntegerAttr(i1Ty, 1))); + }); + } else if (op.elseRegion().isAncestor( + use.getOwner()->getParentRegion())) { + Changed = true; + rewriter.updateRootInPlace(use.getOwner(), [&]() { + use.set(rewriter.create( + op.getLoc(), i1Ty, rewriter.getIntegerAttr(i1Ty, 0))); + }); + } + } + if (Changed) + return success(); + else + return failure(); + } +}; + +struct ReplaceIfYieldWithCondition : public OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(IfOp op, + PatternRewriter &rewriter) const override { + // Replace the operation if only a subset of its results have uses. + if (op.getNumResults() == 0) + return failure(); + + auto trueYield = cast(op.thenRegion().back().getTerminator()); + auto falseYield = + cast(op.elseRegion().back().getTerminator()); + + rewriter.setInsertionPoint(op->getBlock(), + op.getOperation()->getIterator()); + bool Changed = false; + mlir::Type i1Ty = rewriter.getI1Type(); + for (auto Tup : + llvm::zip(trueYield.results(), falseYield.results(), op.results())) { + if (!std::get<0>(Tup).getType().isInteger(1)) + continue; + if (auto TrueYield = std::get<0>(Tup).getDefiningOp()) { + if (auto FalseYield = std::get<1>(Tup).getDefiningOp()) { + bool TrueVal = TrueYield.getValue().cast().getValue(); + bool FalseVal = FalseYield.getValue().cast().getValue(); + if (!TrueVal && FalseVal) { + for (OpOperand &use : + llvm::make_early_inc_range(std::get<2>(Tup).getUses())) { + Changed = true; + rewriter.updateRootInPlace(use.getOwner(), [&]() { + auto notv = rewriter.create( + op.getLoc(), op.condition(), + rewriter.create( + op.getLoc(), i1Ty, rewriter.getIntegerAttr(i1Ty, 1))); + use.set(notv); + }); + } + } + if (TrueVal && !FalseVal) { + for (OpOperand &use : + llvm::make_early_inc_range(std::get<2>(Tup).getUses())) { + Changed = true; + rewriter.updateRootInPlace(use.getOwner(), + [&]() { use.set(op.condition()); }); + } + } + } + } + } + return Changed ? success() : failure(); + } +}; + } // namespace void IfOp::getCanonicalizationPatterns(RewritePatternSet &results, MLIRContext *context) { - results.add(context); + results + .add(context); } //===----------------------------------------------------------------------===// diff --git a/mlir/test/Dialect/SCF/canonicalize.mlir b/mlir/test/Dialect/SCF/canonicalize.mlir --- a/mlir/test/Dialect/SCF/canonicalize.mlir +++ b/mlir/test/Dialect/SCF/canonicalize.mlir @@ -610,3 +610,90 @@ %res = subtensor_insert %2 into %t1[0, 0] [32, 1024] [1, 1] : tensor<32x1024xf32> into tensor<1024x1024xf32> return %res : tensor<1024x1024xf32> } + + + +// CHECK-LABEL: @cond_prop +func @cond_prop(%arg0 : i1) -> index { + %c1 = constant 1 : index + %c2 = constant 2 : index + %c3 = constant 3 : index + %c4 = constant 4 : index + %res = scf.if %arg0 -> index { + %res1 = scf.if %arg0 -> index { + %v1 = "test.get_some_value"() : () -> i32 + scf.yield %c1 : index + } else { + %v2 = "test.get_some_value"() : () -> i32 + scf.yield %c2 : index + } + scf.yield %res1 : index + } else { + %res2 = scf.if %arg0 -> index { + %v3 = "test.get_some_value"() : () -> i32 + scf.yield %c3 : index + } else { + %v4 = "test.get_some_value"() : () -> i32 + scf.yield %c4 : index + } + scf.yield %res2 : index + } + return %res : index +} +// CHECK-NEXT: %[[c1:.+]] = constant 1 : index +// CHECK-NEXT: %[[c4:.+]] = constant 4 : index +// CHECK-NEXT: %[[if:.+]] = scf.if %arg0 -> (index) { +// CHECK-NEXT: %{{.+}} = "test.get_some_value"() : () -> i32 +// CHECK-NEXT: scf.yield %[[c1]] : index +// CHECK-NEXT: } else { +// CHECK-NEXT: %{{.+}} = "test.get_some_value"() : () -> i32 +// CHECK-NEXT: scf.yield %[[c4]] : index +// CHECK-NEXT: } +// CHECK-NEXT: return %[[if]] : index +// CHECK-NEXT:} + +// CHECK-LABEL: @replace_if_with_cond1 +func @replace_if_with_cond1(%arg0 : i1) -> (i32, i1) { + %true = constant true + %false = constant false + %res:2 = scf.if %arg0 -> (i32, i1) { + %v = "test.get_some_value"() : () -> i32 + scf.yield %v, %true : i32, i1 + } else { + %v2 = "test.get_some_value"() : () -> i32 + scf.yield %v2, %false : i32, i1 + } + return %res#0, %res#1 : i32, i1 +} +// CHECK-NEXT: %[[if:.+]] = scf.if %arg0 -> (i32) { +// CHECK-NEXT: %[[sv1:.+]] = "test.get_some_value"() : () -> i32 +// CHECK-NEXT: scf.yield %[[sv1]] : i32 +// CHECK-NEXT: } else { +// CHECK-NEXT: %[[sv2:.+]] = "test.get_some_value"() : () -> i32 +// CHECK-NEXT: scf.yield %[[sv2]] : i32 +// CHECK-NEXT: } +// CHECK-NEXT: return %[[if]], %arg0 : i32, i1 + +// CHECK-LABEL: @replace_if_with_cond2 +func @replace_if_with_cond2(%arg0 : i1) -> (i32, i1) { + %true = constant true + %false = constant false + %res:2 = scf.if %arg0 -> (i32, i1) { + %v = "test.get_some_value"() : () -> i32 + scf.yield %v, %false : i32, i1 + } else { + %v2 = "test.get_some_value"() : () -> i32 + scf.yield %v2, %true : i32, i1 + } + return %res#0, %res#1 : i32, i1 +} +// CHECK-NEXT: %true = constant true +// CHECK-NEXT: %[[toret:.+]] = xor %arg0, %true : i1 +// CHECK-NEXT: %[[if:.+]] = scf.if %arg0 -> (i32) { +// CHECK-NEXT: %[[sv1:.+]] = "test.get_some_value"() : () -> i32 +// CHECK-NEXT: scf.yield %[[sv1]] : i32 +// CHECK-NEXT: } else { +// CHECK-NEXT: %[[sv2:.+]] = "test.get_some_value"() : () -> i32 +// CHECK-NEXT: scf.yield %[[sv2]] : i32 +// CHECK-NEXT: } +// CHECK-NEXT: return %[[if]], %[[toret]] : i32, i1 \ No newline at end of file diff --git a/mlir/test/Dialect/SCF/ops.mlir b/mlir/test/Dialect/SCF/ops.mlir --- a/mlir/test/Dialect/SCF/ops.mlir +++ b/mlir/test/Dialect/SCF/ops.mlir @@ -278,4 +278,4 @@ scf.yield } return -} +} \ No newline at end of file diff --git a/mlir/test/Transforms/canonicalize.mlir b/mlir/test/Transforms/canonicalize.mlir --- a/mlir/test/Transforms/canonicalize.mlir +++ b/mlir/test/Transforms/canonicalize.mlir @@ -1177,11 +1177,12 @@ // ----- // CHECK-LABEL: func @clone_nested_region -func @clone_nested_region(%arg0: index, %arg1: index) -> memref { +func @clone_nested_region(%arg0: index, %arg1: index, %arg2: index) -> memref { + %cmp = cmpi eq, %arg0, %arg1 : index %0 = cmpi eq, %arg0, %arg1 : index %1 = memref.alloc(%arg0, %arg0) : memref %2 = scf.if %0 -> (memref) { - %3 = scf.if %0 -> (memref) { + %3 = scf.if %cmp -> (memref) { %9 = memref.clone %1 : memref to memref scf.yield %9 : memref } else {