[mlir] Additional folding for SelectOp

[Hardcode84]

Summary

• Fold SelectOp when both true and false args are same SSA value
• Fold some cmp + select patterns

[mehdi_amini]

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

2470

Can you document the pattern? It wasn't immediately obvious to me.

[Hardcode84]

Added example

[Hardcode84]

@mehdi_amini I don't have commit access, if you ok with this patch can you please merge it

[Hardcode84]

Rebase

[mehdi_amini]

You have commit access now right?

[Hardcode84]

Will merge tomorrow

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

static LogicalResult verify(ReturnOp op) {
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // SelectOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult SelectOp::fold(ArrayRef<Attribute> operands) {
+  auto trueVal = getTrueValue();
+  auto falseVal = getFalseValue();
+  if (trueVal == falseVal)
+    return trueVal;
+
   auto condition = getCondition();
 
   // select true, %0, %1 => %0
   if (matchPattern(condition, m_One()))
-    return getTrueValue();
+    return trueVal;
 
   // select false, %0, %1 => %1
   if (matchPattern(condition, m_Zero()))
-    return getFalseValue();
+    return falseVal;
+
+  if (auto cmp = dyn_cast_or_null<CmpIOp>(condition.getDefiningOp())) {
+    auto pred = cmp.predicate();
+    if (pred == mlir::CmpIPredicate::eq || pred == mlir::CmpIPredicate::ne) {
+      auto cmpLhs = cmp.lhs();
+      auto cmpRhs = cmp.rhs();
+
+      // %0 = cmpi eq, %arg0, %arg1
+      // %1 = select %0, %arg0, %arg1 => %arg1
+
+      // %0 = cmpi ne, %arg0, %arg1
+      // %1 = select %0, %arg0, %arg1 => %arg0
+
+      if ((cmpLhs == trueVal && cmpRhs == falseVal) ||
+          (cmpRhs == trueVal && cmpLhs == falseVal))
+        return pred == mlir::CmpIPredicate::ne ? trueVal : falseVal;
+    }
+  }
   return nullptr;
 }
 
 static void print(OpAsmPrinter &p, SelectOp op) {
   p << "select " << op.getOperands();
   p.printOptionalAttrDict(op->getAttrs());
   p << " : ";
   if (ShapedType condType = op.getCondition().getType().dyn_cast<ShapedType>())

mlir/test/Dialect/Standard/canonicalize.mlir

 }
 // CHECK-LABEL: func @subtensor_insert_output_dest_canonicalize
 //  CHECK-SAME:   %[[ARG0:[a-zA-z0-9_]+]]: tensor<2x3xi32>
 //  CHECK-SAME:   %[[ARG1:[a-zA-Z0-9_]+]]: tensor<i32>
 //       CHECK:   %[[PAD:.+]] = tensor.extract %[[ARG1]]
 //       CHECK:   %[[GENERATE:.+]] = tensor.generate
 //       CHECK:   %[[RESULT:.+]] = subtensor_insert %[[ARG0]] into %[[GENERATE]]
 //       CHECK:   return %[[RESULT]]
+
+// -----
+
+// CHECK-LABEL: @select_same_val
+//       CHECK:   return %arg1
+func @select_same_val(%arg0: i1, %arg1: i64) -> i64 {
+  %0 = select %arg0, %arg1, %arg1 : i64
+  return %0 : i64
+}
+
+// -----
+
+// CHECK-LABEL: @select_cmp_eq_select
+//       CHECK:   return %arg1
+func @select_cmp_eq_select(%arg0: i64, %arg1: i64) -> i64 {
+  %0 = cmpi eq, %arg0, %arg1 : i64
+  %1 = select %0, %arg0, %arg1 : i64
+  return %1 : i64
+}
+
+// -----
+
+// CHECK-LABEL: @select_cmp_ne_select
+//       CHECK:   return %arg0
+func @select_cmp_ne_select(%arg0: i64, %arg1: i64) -> i64 {
+  %0 = cmpi ne, %arg0, %arg1 : i64
+  %1 = select %0, %arg0, %arg1 : i64
+  return %1 : i64
+}