[MLIR] Simplify CmpIOp when operands are identical

[gtbercea]

Summary

This patch simplifies the cmpi Standard dialect operation when its operands are identical and it replaces it with a constant representing the outcome of the comparison.

[rriddle]

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

883

This looks like it should be part of the fold method instead.

892

Can you use proper operand accessors instead of this? (.lhs()/.rhs())

mlir/test/Dialect/Standard/canonicalize.mlir

14

Please avoid the use of SSA value names inside of check lines, use captures instead.

[gtbercea]

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

883

I tried that but it led to many lit test failures. I think the problem is that the input to fold are attributes whereas here I have Value operands.

[rriddle]

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

883

The input attributes to fold are operands that are known to be constant, i.e. you can still use the operands directly from the operation. If you are getting lit failures, that is due to something else. (e.g. tests being broken by better folding that changed the expected output)

[gtbercea]

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

883

Good point, I completely ignored the fact that I could get the operands as usual even when inside the fold function! It all works now. :)

[rriddle]

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

860

nit: Can you use BoolAttr::get here instead? BoolAttr is a utility wrapper around i1 IntegerAttr.

mlir/test/Dialect/Standard/canonicalize.mlir

8

Can you trim out all of the unnecessary operations here? You should be able to test these cases with just the cmpi and function arguments.

25

Can you make sure to add tests for each of the other cases that you are folding?

[rriddle]

mlir/test/Dialect/Standard/canonicalize.mlir

3

Do you need the select in any of these? Why not return the comparison result directly?

[rriddle]

Thanks. LGTM after resolving the last few comments.

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

849

This isn't possible, the operands to an operation are always valid. You only have to check null when accessing attributes from the operands parameter, as a null Attribute means that the operand isn't constant.

mlir/test/Dialect/Standard/canonicalize.mlir

7

nit: The general style is to have this before not after.

8

nit: These CHECK-NOTS can be simplified to just CHECK-NOT: cmpi

[AlexEichenberger]

LGTM

[gtbercea]

Would someone be able to commit this? Thanks a lot.

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

bool mlir::applyCmpPredicate(CmpIPredicate predicate, const APInt &lhs,
   }
   llvm_unreachable("unknown comparison predicate");
 }
 
 // Constant folding hook for comparisons.
 OpFoldResult CmpIOp::fold(ArrayRef<Attribute> operands) {
   assert(operands.size() == 2 && "cmpi takes two arguments");
 
+  auto lhsOperand = lhs();
+  auto rhsOperand = rhs();
+  if (!lhsOperand || !rhsOperand)
+    return {};
+
+  // Operands of the compare are equal.
+  if (lhsOperand == rhsOperand) {
+    // Cmp Ops which are false when operands are the same.
+    if (getPredicate() == CmpIPredicate::sgt ||
+        getPredicate() == CmpIPredicate::slt ||
+        getPredicate() == CmpIPredicate::ugt ||
+        getPredicate() == CmpIPredicate::ult ||
+        getPredicate() == CmpIPredicate::ne)
+      return BoolAttr::get(false, getContext());
+
+    // Cmp Ops which are true when operands are the same.
+    if (getPredicate() == CmpIPredicate::eq ||
+        getPredicate() == CmpIPredicate::sle ||
+        getPredicate() == CmpIPredicate::ule ||
+        getPredicate() == CmpIPredicate::sge ||
+        getPredicate() == CmpIPredicate::uge)
+      return BoolAttr::get(true, getContext());
+  }
+
   auto lhs = operands.front().dyn_cast_or_null<IntegerAttr>();
   auto rhs = operands.back().dyn_cast_or_null<IntegerAttr>();
   if (!lhs || !rhs)
     return {};
 
   auto val = applyCmpPredicate(getPredicate(), lhs.getValue(), rhs.getValue());
   return IntegerAttr::get(IntegerType::get(1, getContext()), APInt(1, val));
 }
 
 //===----------------------------------------------------------------------===//
 // CmpFOp
 //===----------------------------------------------------------------------===//
 
 static void buildCmpFOp(OpBuilder &build, OperationState &result,
                         CmpFPredicate predicate, Value lhs, Value rhs) {
   result.addOperands({lhs, rhs});
   result.types.push_back(getI1SameShape(lhs.getType()));
   result.addAttribute(CmpFOp::getPredicateAttrName(),
                       build.getI64IntegerAttr(static_cast<int64_t>(predicate)));
 }
 
 /// Compute `lhs` `pred` `rhs`, where `pred` is one of the known floating point
 /// comparison predicates.
 bool mlir::applyCmpPredicate(CmpFPredicate predicate, const APFloat &lhs,
                              const APFloat &rhs) {
   auto cmpResult = lhs.compare(rhs);
   switch (predicate) {
   case CmpFPredicate::AlwaysFalse:
     return false;
   case CmpFPredicate::OEQ:

mlir/test/Dialect/Standard/canonicalize.mlir

+// RUN: mlir-opt %s -allow-unregistered-dialect -pass-pipeline='func(canonicalize)' -split-input-file | FileCheck %s
+func @test_simplify_sgt(%arg0: index, %arg1 : index, %arg2 : index) -> index {
+  %0 = cmpi "sgt", %arg0, %arg0 : index
+  %1 = select %0, %arg1, %arg2 : index
+  return %1 : index
+
+  // CHECK-LABEL: test_simplify_sgt
+  // CHECK-NOT: cmpi "sgt", %arg0, %arg0 : index
+  // CHECK: return %arg2 : index
+}
+
+func @test_simplify_slt(%arg0: index, %arg1 : index, %arg2 : index) -> index {
+  %0 = cmpi "slt", %arg0, %arg0 : index
+  %1 = select %0, %arg1, %arg2 : index
+  return %1 : index
+
+  // CHECK-LABEL: test_simplify_slt
+  // CHECK-NOT: cmpi "slt", %arg0, %arg0 : index
+  // CHECK: return %arg2 : index
+}
+
+func @test_simplify_ugt(%arg0: index, %arg1 : index, %arg2 : index) -> index {
+  %0 = cmpi "ugt", %arg0, %arg0 : index
+  %1 = select %0, %arg1, %arg2 : index
+  return %1 : index
+
+  // CHECK-LABEL: test_simplify_ugt
+  // CHECK-NOT: cmpi "ugt", %arg0, %arg0 : index
+  // CHECK: return %arg2 : index
+}
+
+func @test_simplify_ult(%arg0: index, %arg1 : index, %arg2 : index) -> index {
+  %0 = cmpi "ult", %arg0, %arg0 : index
+  %1 = select %0, %arg1, %arg2 : index
+  return %1 : index
+
+  // CHECK-LABEL: test_simplify_ult
+  // CHECK-NOT: cmpi "ult", %arg0, %arg0 : index
+  // CHECK: return %arg2 : index
+}
+
+func @test_simplify_ne(%arg0: index, %arg1 : index, %arg2 : index) -> index {
+  %0 = cmpi "ne", %arg0, %arg0 : index
+  %1 = select %0, %arg1, %arg2 : index
+  return %1 : index
+
+  // CHECK-LABEL: test_simplify_ne
+  // CHECK-NOT: cmpi "ne", %arg0, %arg0 : index
+  // CHECK: return %arg2 : index
+}
+
+func @test_simplify_eq(%arg0: index, %arg1 : index, %arg2 : index) -> index {
+  %0 = cmpi "eq", %arg0, %arg0 : index
+  %1 = select %0, %arg1, %arg2 : index
+  return %1 : index
+
+  // CHECK-LABEL: test_simplify_eq
+  // CHECK-NOT: cmpi "eq", %arg0, %arg0 : index
+  // CHECK: return %arg1 : index
+}
+
+func @test_simplify_sle(%arg0: index, %arg1 : index, %arg2 : index) -> index {
+  %0 = cmpi "sle", %arg0, %arg0 : index
+  %1 = select %0, %arg1, %arg2 : index
+  return %1 : index
+
+  // CHECK-LABEL: test_simplify_sle
+  // CHECK-NOT: cmpi "sle", %arg0, %arg0 : index
+  // CHECK: return %arg1 : index
+}
+
+func @test_simplify_ule(%arg0: index, %arg1 : index, %arg2 : index) -> index {
+  %0 = cmpi "ule", %arg0, %arg0 : index
+  %1 = select %0, %arg1, %arg2 : index
+  return %1 : index
+
+  // CHECK-LABEL: test_simplify_ule
+  // CHECK-NOT: cmpi "ule", %arg0, %arg0 : index
+  // CHECK: return %arg1 : index
+}
+
+func @test_simplify_sge(%arg0: index, %arg1 : index, %arg2 : index) -> index {
+  %0 = cmpi "sge", %arg0, %arg0 : index
+  %1 = select %0, %arg1, %arg2 : index
+  return %1 : index
+
+  // CHECK-LABEL: test_simplify_sge
+  // CHECK-NOT: cmpi "sge", %arg0, %arg0 : index
+  // CHECK: return %arg1 : index
+}
+
+func @test_simplify_uge(%arg0: index, %arg1 : index, %arg2 : index) -> index {
+  %0 = cmpi "uge", %arg0, %arg0 : index
+  %1 = select %0, %arg1, %arg2 : index
+  return %1 : index
+
+  // CHECK-LABEL: test_simplify_uge
+  // CHECK-NOT: cmpi "uge", %arg0, %arg0 : index
+  // CHECK: return %arg1 : index
+}