[CIR][Transform] Add simplify transformation for select op (#816)

As mentioned at
#809 (comment) , this
PR adds more simplify transformations for select op:

- `cir.select if %0 then x else x` -> `x`
- `cir.select if %0 then #true else #false` -> `%0`
- `cir.select if %0 then #false else #true` -> `cir.unary not %0`

Author: Lancern

clang/lib/CIR/Dialect/IR/CIRDialect.cpp

@@ -1400,11 +1400,20 @@ void TernaryOp::build(OpBuilder &builder, OperationState &result, Value cond,
 
 OpFoldResult SelectOp::fold(FoldAdaptor adaptor) {
   auto condition = adaptor.getCondition();
-  if (!condition)
-    return nullptr;
+  if (condition) {
+    auto conditionValue = mlir::cast<mlir::cir::BoolAttr>(condition).getValue();
+    return conditionValue ? getTrueValue() : getFalseValue();
+  }
 
-  auto conditionValue = mlir::cast<mlir::cir::BoolAttr>(condition).getValue();
-  return conditionValue ? getTrueValue() : getFalseValue();
+  // cir.select if %0 then x else x -> x
+  auto trueValue = adaptor.getTrueValue();
+  auto falseValue = adaptor.getFalseValue();
+  if (trueValue && trueValue == falseValue)
+    return trueValue;
+  if (getTrueValue() == getFalseValue())
+    return getTrueValue();
+
+  return nullptr;
 }
 
 //===----------------------------------------------------------------------===//

clang/lib/CIR/Dialect/Transforms/CIRSimplify.cpp

@@ -107,6 +107,45 @@ struct RemoveTrivialTry : public OpRewritePattern<TryOp> {
   }
 };
 
+struct SimplifySelect : public OpRewritePattern<SelectOp> {
+  using OpRewritePattern<SelectOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(SelectOp op,
+                                PatternRewriter &rewriter) const final {
+    mlir::Operation *trueValueOp = op.getTrueValue().getDefiningOp();
+    mlir::Operation *falseValueOp = op.getFalseValue().getDefiningOp();
+    auto trueValueConstOp =
+        mlir::dyn_cast_if_present<mlir::cir::ConstantOp>(trueValueOp);
+    auto falseValueConstOp =
+        mlir::dyn_cast_if_present<mlir::cir::ConstantOp>(falseValueOp);
+    if (!trueValueConstOp || !falseValueConstOp)
+      return mlir::failure();
+
+    auto trueValue =
+        mlir::dyn_cast<mlir::cir::BoolAttr>(trueValueConstOp.getValue());
+    auto falseValue =
+        mlir::dyn_cast<mlir::cir::BoolAttr>(falseValueConstOp.getValue());
+    if (!trueValue || !falseValue)
+      return mlir::failure();
+
+    // cir.select if %0 then #true else #false -> %0
+    if (trueValue.getValue() && !falseValue.getValue()) {
+      rewriter.replaceAllUsesWith(op, op.getCondition());
+      rewriter.eraseOp(op);
+      return mlir::success();
+    }
+
+    // cir.seleft if %0 then #false else #true -> cir.unary not %0
+    if (!trueValue.getValue() && falseValue.getValue()) {
+      rewriter.replaceOpWithNewOp<mlir::cir::UnaryOp>(
+          op, mlir::cir::UnaryOpKind::Not, op.getCondition());
+      return mlir::success();
+    }
+
+    return mlir::failure();
+  }
+};
+
 //===----------------------------------------------------------------------===//
 // CIRSimplifyPass
 //===----------------------------------------------------------------------===//
@@ -131,7 +170,8 @@ void populateMergeCleanupPatterns(RewritePatternSet &patterns) {
     RemoveRedundantBranches,
     RemoveEmptyScope,
     RemoveEmptySwitch,
-    RemoveTrivialTry
+    RemoveTrivialTry,
+    SimplifySelect
   >(patterns.getContext());
   // clang-format on
 }

clang/test/CIR/Transforms/select.cir

@@ -1,4 +1,4 @@
-// RUN: cir-opt --canonicalize -o %t.cir %s
+// RUN: cir-opt -cir-simplify -o %t.cir %s
 // RUN: FileCheck --input-file=%t.cir %s
 
 !s32i = !cir.int<s, 32>
@@ -23,4 +23,38 @@ module {
   //      CHECK: cir.func @fold_false(%[[ARG0:.+]]: !s32i, %[[ARG1:.+]]: !s32i) -> !s32i {
   // CHECK-NEXT:   cir.return %[[ARG1]] : !s32i
   // CHECK-NEXT: }
+
+  cir.func @fold_to_const(%arg0 : !cir.bool) -> !s32i {
+    %0 = cir.const #cir.int<42> : !s32i
+    %1 = cir.select if %arg0 then %0 else %0 : (!cir.bool, !s32i, !s32i) -> !s32i
+    cir.return %1 : !s32i
+  }
+
+  //      CHECK: cir.func @fold_to_const(%{{.+}}: !cir.bool) -> !s32i {
+  // CHECK-NEXT:   %[[#A:]] = cir.const #cir.int<42> : !s32i
+  // CHECK-NEXT:   cir.return %[[#A]] : !s32i
+  // CHECK-NEXT: }
+
+  cir.func @simplify_1(%arg0 : !cir.bool) -> !cir.bool {
+    %0 = cir.const #cir.bool<true> : !cir.bool
+    %1 = cir.const #cir.bool<false> : !cir.bool
+    %2 = cir.select if %arg0 then %0 else %1 : (!cir.bool, !cir.bool, !cir.bool) -> !cir.bool
+    cir.return %2 : !cir.bool
+  }
+
+  //      CHECK: cir.func @simplify_1(%[[ARG0:.+]]: !cir.bool) -> !cir.bool {
+  // CHECK-NEXT:   cir.return %[[ARG0]] : !cir.bool
+  // CHECK-NEXT: }
+
+  cir.func @simplify_2(%arg0 : !cir.bool) -> !cir.bool {
+    %0 = cir.const #cir.bool<false> : !cir.bool
+    %1 = cir.const #cir.bool<true> : !cir.bool
+    %2 = cir.select if %arg0 then %0 else %1 : (!cir.bool, !cir.bool, !cir.bool) -> !cir.bool
+    cir.return %2 : !cir.bool
+  }
+
+  //      CHECK: cir.func @simplify_2(%[[ARG0:.+]]: !cir.bool) -> !cir.bool {
+  // CHECK-NEXT:   %[[#A:]] = cir.unary(not, %[[ARG0]]) : !cir.bool, !cir.bool
+  // CHECK-NEXT:   cir.return %[[#A]] : !cir.bool
+  // CHECK-NEXT: }
 }