[MLIR] Add SCF.if Condition Canonicalizations

[wsmoses]

Summary

Add two canoncalizations for scf.if.

1. A canonicalization that allows users of a condition within an if to assume the condition is true if in the true region, etc.
2. A canonicalization that removes yielded statements that are equivalent to the condition or its negation

[nicolasvasilache]

mlir/lib/Dialect/SCF/SCF.cpp

1110

doc with (pseudo)-IR examples plz

1144

doc with (pseudo)-IR examples plz

1159

can you please fix casing and other clang-tidy stuff?
Filtering out the noise will make it easier to read.

1165

That's deep but we're not in NN land :)
Can you please turn into early continue and reduce nesting?

[wsmoses]

Fix casing and add description as documentation

[wsmoses]

Fix capitalization of Tup

[mehdi_amini]

mlir/lib/Dialect/SCF/SCF.cpp

1130

Nit: remove trivial braces.

1138–1139

Can you try to create a single constant op in the pattern instead of a new one for each new use?

1153

You can replace the if/else with : return success(changed);

1180

I don't understand this comment here

[wsmoses]

Reduce redundant constantop creation and provide better comments

[kumasento]

mlir/lib/Dialect/SCF/SCF.cpp

1171

What would happen if the input is like:

%false = constant false
%res:2 = scf.if %cmp {
  yield something(), %false
} else {
  yield something2(), %false
}

return %res#0, %res#1

I think this case can be canonicalized, but the current implementation doesn't support it?

1206

nit: I would prefer std::tie instead of std::get especially in the case that you need to access every field in the tuple several times. The code can be more readble IMHO.

e.g.

Value trueResult, falseResult, opResult;
std::tie(trueResult, falseResult, opResult) = tup;

mlir/test/Dialect/SCF/canonicalize.mlir

705

nit: Add a newline at the end of the file (you can configure this in your editor, e.g., if you use VSCode: https://stackoverflow.com/questions/44704968/visual-studio-code-insert-new-line-at-the-end-of-files). It is just for not introducing unnecessary changes.

mlir/test/Transforms/canonicalize.mlir

1201–1202

Why we need %arg2 here?

[mehdi_amini]

mlir/lib/Dialect/SCF/SCF.cpp

1171

Seems like a different case that should be handle by a more generic canonicalizer which does not even require i1 : if the same SSA value is yielded in both regions it can be directly RAUW.

If it isn't handled elsewhere already, it could be added as a separate canonicalization, or by patching the loop below:

for (auto tup :  llvm::zip(trueYield.results(), falseYield.results(), op.results())) {
  if (std::get<0>(tup) == std::get<1>(tup)) {
     std::get<1>(tup).replaceAllUsesWith(std::get<2>(tup));
     changed = true;
     continue;
  }
  ....
}

[wsmoses]

Add scf.if common yield elimination

[wsmoses]

mlir/lib/Dialect/SCF/SCF.cpp

1171

I would've thought this "if invariant " motion would already be handled, but indeed it's not. It's easy enough to add here so I'll bring it in. And yes I agree with medhi that this should be done in a more general way (in fact that just checks if the the operands are the same).

mlir/test/Transforms/canonicalize.mlir

1201–1202

Without modification this test would be optimized to have the nested if eliminated [since we know %0 is true]. Adding the extra arg was an easy way to preserve (what I believe to be) the intended behavior of the test. If you have a better suggestion, lmk!

[ftynse]

LGTM with feedback from others addressed.

mlir/lib/Dialect/SCF/SCF.cpp

1122

Nit: top-level comments in MLIR use three slashes ///

1134

Nit: drop mlir::.

1136–1137

Nit: please use proper capitalization and punctuation (trailing periods) in comments.

1264–1268

This magic with updateRootInPlace isn't strictly necessary in canonicalizer patterns because the canonicalizer driver doesn't undo them anyway. It's probably nice to have for future-proofness, but you are already using direct RAUW in this pattern, so it's not safe anway. Maybe factor this loop out into a helper function and use it everywhere + propose this helper to be promoted to PatternRewriter interface? Otherwise, I'm fine with using RAUW everywhere in canonicalizer patterns.

mlir/test/Dialect/SCF/canonicalize.mlir

649

Do we care about lines being next to each other? I would rather use CHECK and make the test less constrained on unnecessary details.

727

Add newlines plz.

[wsmoses]

Address Alex's comments [and finally add newline]

mlir/lib/Dialect/SCF/SCF.cpp

for (auto it : llvm::enumerate(llvm::zip(thenYieldArgs, elseYieldArgs))) {
         results[it.index()] =
             rewriter.create<SelectOp>(op.getLoc(), cond, trueVal, falseVal);
     }
 
     rewriter.replaceOp(op, results);
     return success();
   }
 };
+
+// Allow the true region of an if to assume the condition is true
+// and vice versa. For example:
+//
+//   scf.if %cmp {
+//      print(%cmp)
+//   }
+//
+//  becomes
+//
+//   scf.if %cmp {
+//      print(true)
+//   }
+//
+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 (op.condition().getDefiningOp<ConstantOp>())
+      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.condition().getUses())) {
+      if (op.thenRegion().isAncestor(use.getOwner()->getParentRegion())) {
+        changed = true;
+
+        if (!constantTrue)
+          constantTrue = rewriter.create<mlir::ConstantOp>(
+              op.getLoc(), i1Ty, rewriter.getIntegerAttr(i1Ty, 1));
+
+        rewriter.updateRootInPlace(use.getOwner(),
+                                   [&]() { use.set(constantTrue); });
+      } else if (op.elseRegion().isAncestor(
+                     use.getOwner()->getParentRegion())) {
+        changed = true;
+
+        if (!constantFalse)
+          constantFalse = rewriter.create<mlir::ConstantOp>(
+              op.getLoc(), i1Ty, rewriter.getIntegerAttr(i1Ty, 0));
+
+        rewriter.updateRootInPlace(use.getOwner(),
+                                   [&]() { use.set(constantFalse); });
+      }
+    }
+
+    return success(changed);
+  }
+};
+
+/// Remove any statements from an if that are equivalent to the condition
+/// or its negation. For example:
+///
+///    %res:2 = scf.if %cmp {
+///       yield something(), true
+///    } else {
+///       yield something2(), false
+///    }
+///    print(%res#1)
+///
+///  becomes
+///    %res = scf.if %cmp {
+///       yield something()
+///    } else {
+///       yield something2()
+///    }
+///    print(%cmp)
+///
+/// Additionally if both branches yield the same value, replace all uses
+/// of the result with the yielded value
+///
+///    %res:2 = scf.if %cmp {
+///       yield something(), %arg1
+///    } else {
+///       yield something2(), %arg1
+///    }
+///    print(%res#1)
+///
+///  becomes
+///    %res = scf.if %cmp {
+///       yield something()
+///    } else {
+///       yield something2()
+///    }
+//    print(%arg1)
+struct ReplaceIfYieldWithConditionOrValue : public OpRewritePattern<IfOp> {
+  using OpRewritePattern<IfOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(IfOp op,
+                                PatternRewriter &rewriter) const override {
+    // Early exit if there are no results that could be replaced.
+    if (op.getNumResults() == 0)
+      return failure();
+
+    auto trueYield = cast<scf::YieldOp>(op.thenRegion().back().getTerminator());
+    auto falseYield =
+        cast<scf::YieldOp>(op.elseRegion().back().getTerminator());
+
+    rewriter.setInsertionPoint(op->getBlock(),
+                               op.getOperation()->getIterator());
+    bool changed = false;
+    Type i1Ty = rewriter.getI1Type();
+    for (auto tup :
+         llvm::zip(trueYield.results(), falseYield.results(), op.results())) {
+      Value trueResult, falseResult, opResult;
+      std::tie(trueResult, falseResult, opResult) = tup;
+
+      if (trueResult == falseResult) {
+        if (!opResult.use_empty()) {
+          opResult.replaceAllUsesWith(trueResult);
+          changed = true;
+        }
+        continue;
+      }
+
+      auto trueYield = trueResult.getDefiningOp<ConstantOp>();
+      if (!trueYield)
+        continue;
+
+      if (!trueYield.getType().isInteger(1))
+        continue;
+
+      auto falseYield = falseResult.getDefiningOp<ConstantOp>();
+      if (!falseYield)
+        continue;
+
+      bool trueVal = trueYield.getValue().cast<BoolAttr>().getValue();
+      bool falseVal = falseYield.getValue().cast<BoolAttr>().getValue();
+      if (!trueVal && falseVal) {
+        if (!opResult.use_empty()) {
+          Value notCond = rewriter.create<XOrOp>(
+              op.getLoc(), op.condition(),
+              rewriter.create<mlir::ConstantOp>(
+                  op.getLoc(), i1Ty, rewriter.getIntegerAttr(i1Ty, 1)));
+          opResult.replaceAllUsesWith(notCond);
+          changed = true;
+        }
+      }
+      if (trueVal && !falseVal) {
+        if (!opResult.use_empty()) {
+          opResult.replaceAllUsesWith(op.condition());
+          changed = true;
+        }
+      }
+    }
+    return success(changed);
+  }
+};
+
 } // namespace
 
 void IfOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                        MLIRContext *context) {
-  results.add<RemoveUnusedResults, RemoveStaticCondition,
-              ConvertTrivialIfToSelect>(context);
+  results
+      .add<RemoveUnusedResults, RemoveStaticCondition, ConvertTrivialIfToSelect,
+           ConditionPropagation, ReplaceIfYieldWithConditionOrValue>(context);
 }
 
 //===----------------------------------------------------------------------===//
 // ParallelOp
 //===----------------------------------------------------------------------===//
 
 void ParallelOp::build(
     OpBuilder &builder, OperationState &result, ValueRange lowerBounds,

mlir/test/Dialect/SCF/canonicalize.mlir

 // CHECK:               return [[V0]], [[V1]]
 
 // -----
 
 func private @side_effect()
 func @one_unused(%cond: i1) -> (index) {
   %c0 = constant 0 : index
   %c1 = constant 1 : index
+  %c2 = constant 2 : index
+  %c3 = constant 3 : index
   %0, %1 = scf.if %cond -> (index, index) {
     call @side_effect() : () -> ()
     scf.yield %c0, %c1 : index, index
   } else {
-    scf.yield %c0, %c1 : index, index
+    scf.yield %c2, %c3 : index, index
   }
   return %1 : index
 }
 
 // CHECK-LABEL:   func @one_unused
 // CHECK:           [[C0:%.*]] = constant 1 : index
+// CHECK:           [[C3:%.*]] = constant 3 : index
 // CHECK:           [[V0:%.*]] = scf.if %{{.*}} -> (index) {
 // CHECK:             call @side_effect() : () -> ()
 // CHECK:             scf.yield [[C0]] : index
 // CHECK:           } else
-// CHECK:             scf.yield [[C0]] : index
+// CHECK:             scf.yield [[C3]] : index
 // CHECK:           }
 // CHECK:           return [[V0]] : index
 
 // -----
 
 func private @side_effect()
 func @nested_unused(%cond1: i1, %cond2: i1) -> (index) {
   %c0 = constant 0 : index
   %c1 = constant 1 : index
+  %c2 = constant 2 : index
+  %c3 = constant 3 : index
   %0, %1 = scf.if %cond1 -> (index, index) {
     %2, %3 = scf.if %cond2 -> (index, index) {
       call @side_effect() : () -> ()
       scf.yield %c0, %c1 : index, index
     } else {
-      scf.yield %c0, %c1 : index, index
+      scf.yield %c2, %c3 : index, index
     }
     scf.yield %2, %3 : index, index
   } else {
     scf.yield %c0, %c1 : index, index
   }
   return %1 : index
 }
 
 // CHECK-LABEL:   func @nested_unused
 // CHECK:           [[C0:%.*]] = constant 1 : index
+// CHECK:           [[C3:%.*]] = constant 3 : index
 // CHECK:           [[V0:%.*]] = scf.if {{.*}} -> (index) {
 // CHECK:             [[V1:%.*]] = scf.if {{.*}} -> (index) {
 // CHECK:               call @side_effect() : () -> ()
 // CHECK:               scf.yield [[C0]] : index
 // CHECK:             } else
-// CHECK:               scf.yield [[C0]] : index
+// CHECK:               scf.yield [[C3]] : index
 // CHECK:             }
 // CHECK:             scf.yield [[V1]] : index
 // CHECK:           } else
 // CHECK:             scf.yield [[C0]] : index
 // CHECK:           }
 // CHECK:           return [[V0]] : index
 
 // -----
//       CHECK:   scf.yield %[[UNCAST]] : tensor<32x1024xf32>
   }
 //   CHECK-NOT: tensor.cast
 //       CHECK: %[[RES:.*]] = subtensor_insert %[[FOR_RES]] into %[[T1]][0, 0] [32, 1024] [1, 1] : tensor<32x1024xf32> into tensor<1024x1024xf32>
 //       CHECK: return %[[RES]] : tensor<1024x1024xf32>
   %2 = tensor.cast %1 : tensor<?x?xf32> to tensor<32x1024xf32>
   %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-DAG:  %[[c1:.+]] = constant 1 : index
+// CHECK-DAG:  %[[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
+
+
+// CHECK-LABEL: @replace_if_with_cond3
+func @replace_if_with_cond3(%arg0 : i1, %arg2: i64) -> (i32, i64) {
+  %res:2 = scf.if %arg0 -> (i32, i64) {
+    %v = "test.get_some_value"() : () -> i32
+    scf.yield %v, %arg2 : i32, i64
+  } else {
+    %v2 = "test.get_some_value"() : () -> i32
+    scf.yield %v2, %arg2 : i32, i64
+  }
+  return %res#0, %res#1 : i32, i64
+}
+// 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]], %arg1 : i32, i64

mlir/test/Transforms/canonicalize.mlir

 // CHECK-NEXT: %[[ALLOC2:.*]] = memref.alloc
 // CHECK-NEXT: scf.yield %[[ALLOC2]]
 // CHECK: linalg.copy(%[[ALLOC1]]
 // CHECK-NEXT: memref.dealloc %[[ALLOC1]]
 
 // -----
 
 // CHECK-LABEL: func @clone_nested_region
-func @clone_nested_region(%arg0: index, %arg1: index) -> memref<?x?xf32> {
+func @clone_nested_region(%arg0: index, %arg1: index, %arg2: index) -> memref<?x?xf32> {
+  %cmp = cmpi eq, %arg0, %arg1 : index
   %0 = cmpi eq, %arg0, %arg1 : index
   %1 = memref.alloc(%arg0, %arg0) : memref<?x?xf32>
   %2 = scf.if %0 -> (memref<?x?xf32>) {
-    %3 = scf.if %0 -> (memref<?x?xf32>) {
+    %3 = scf.if %cmp -> (memref<?x?xf32>) {
       %9 = memref.clone %1 : memref<?x?xf32> to memref<?x?xf32>
       scf.yield %9 : memref<?x?xf32>
     } else {
       %7 = memref.alloc(%arg0, %arg1) : memref<?x?xf32>
       %10 = memref.clone %7 : memref<?x?xf32> to memref<?x?xf32>
       memref.dealloc %7 : memref<?x?xf32>
       scf.yield %10 : memref<?x?xf32>
     }