[instcombine] recognize three way comparison idioms

[reames]

Summary

Many languages have a three way comparison idiom where comparing two values produces not a boolean, but a tri-state value. Typical values (e.g. as used in the lcmp/fcmp bytecodes from Java) are -1 for less than, 0 for equality, and +1 for greater than.

We actually do a great job already of converting three way comparisons into binary comparisons when the result produced has one a single use. Unfortunately, such values can have more than one use, and in that case, our existing optimizations break down.

The patch adds a peephole which converts a three-way compare + test idiom into a binary comparison on the original inputs. It focused on replacing the test on the result of the three way compare and does nothing about removing the three way compare itself. That's left to other optimizations (which do actually kick in commonly.)

I specifically only matching a single select idiom at the moment. We don't need to worry about phis since we canonicalize selects into phis. We should handle (or canonicalize) a larger family of select idioms, but I wanted to start small and iterate. Similarly, the current patch only adds support for signed compare idioms. Future work could add unsigned idioms, and possibly floating point idioms. I did try to generalize over the actual values returns by the three way compare so that the transform might kick on more than just the idiomatic three way comparison.

[reames]

ping

[sanjoy]

lib/Transforms/InstCombine/InstCombineCompares.cpp

2173

Start with uppercase.

2175

I'd put this in InstCombineInternal.h, that'll make it more likely that this gets used in other parts of IC.

2183

Can you please replace -1, 0 and 1 with Less, Equal, and Greater in the comments, so that is is more in tune with the code.

2191

Minor, but I'd put the check on PredA right after you matched it, and the check on PredB after you matched it; since they're logically part of matching operation.

2210

I'd avoid the explicit combinatorial logic, but instead have:

Value *Cond = getConstant(false);
if (TrueWhenLessThan)
  Cond = CreateOr(Cond, CreateICmp(SLT, OrigLHS, OrigRHS))
if (TrueWhenEqual)
  Cond = CreateOr(Cond, CreateICmp(EQ, OrigLHS, OrigRHS))
if (TrueWhenGreater)
  Cond = CreateOr(Cond, CreateICmp(SGT, OrigLHS, OrigRHS))

And leave the rest of instcombine to combine, e.g., a s< b || a == b
into a s<= b.

2238

s/obvious/obviously/ or drop "obvious"

[anna]

This change is now resurrected as https://reviews.llvm.org/D34278

lib/Transforms/InstCombine/InstCombineCompares.cpp

if (!ICI.isEquality()) {
       //       C2+1 is a power of 2
       if (ICI.getPredicate() == ICmpInst::ICMP_UGT && LHSI->hasOneUse() &&
           (RHSV + 1).isPowerOf2() && (LHSV & RHSV) == 0)
         return new ICmpInst(ICmpInst::ICMP_NE,
                             Builder->CreateAnd(LHSI->getOperand(0), ~RHSV),
                             ConstantExpr::getNeg(LHSC));
     }
     break;
+  case Instruction::Select:
+    // match a select chain which produces one of three values based on whether
+    // the LHS is less than, equal to, or greater than RHS respectively.
+    auto matchThreeWaySignedCompare = [](SelectInst *SI, Value *&LHS,
+                                         Value*& RHS, ConstantInt *&Less,
+                                         ConstantInt *&Equal,
+                                         ConstantInt *&Greater) {
+      // TODO: Generalize this to work with other comparison idioms or ensure
+      // they get canonicalized into this form.
+
+      // select i1 (a == b), i32 0, i32 (select i1 (a < b), i32 -1, i32 1),
+      // except -1, 0, 1 are placeholders for any three constants
+      ICmpInst::Predicate PredA, PredB;
+      if (match(SI->getTrueValue(), m_ConstantInt(Equal)) &&
+          match(SI->getCondition(), m_ICmp(PredA, m_Value(LHS),
+                                           m_Value(RHS))) &&
+          match(SI->getFalseValue(),
+                m_Select(m_ICmp(PredB, m_Specific(LHS), m_Specific(RHS)),
+                         m_ConstantInt(Less), m_ConstantInt(Greater)))) {
+        if (PredA == ICmpInst::ICMP_EQ &&
+            PredB == ICmpInst::ICMP_SLT)
+          return true;
+      }
+      return false;
+    };
+
+    // If we're testing a constant value against the result of a three way
+    // comparison, the result can be expressed directly in terms of the
+    // original values being compared.  Note: We could possibly be more
+    // aggressive here and remove the hasOneUse test.  The original select is
+    // really likely to simplify or sink when we remove a test of the result.
+    Value *OrigLHS, *OrigRHS;
+    ConstantInt *C1, *C2, *C3;
+    if (ICI.hasOneUse() &&
+        matchThreeWaySignedCompare(cast<SelectInst>(LHSI), OrigLHS, OrigRHS,
+                                   C1, C2, C3)) {
+      assert(C1 && C2 && C3 && RHS);
+
+      bool TrueWhenLessThan =
+        ConstantExpr::getCompare(ICI.getPredicate(), C1, RHS)->isAllOnesValue();
+      bool TrueWhenEqual =
+        ConstantExpr::getCompare(ICI.getPredicate(), C2, RHS)->isAllOnesValue();
+      bool TrueWhenGreaterThan =
+        ConstantExpr::getCompare(ICI.getPredicate(), C3, RHS)->isAllOnesValue();
+
+      // Eight combinations ...
+      if (TrueWhenLessThan && TrueWhenEqual && TrueWhenGreaterThan)
+        return replaceInstUsesWith(ICI, Builder->getTrue());
+      else if (TrueWhenLessThan && TrueWhenEqual && !TrueWhenGreaterThan)
+        return new ICmpInst(ICmpInst::ICMP_SLE, OrigLHS, OrigRHS);
+      else if (TrueWhenLessThan && !TrueWhenEqual && TrueWhenGreaterThan)
+        return new ICmpInst(ICmpInst::ICMP_NE, OrigLHS, OrigRHS);
+      else if (TrueWhenLessThan && !TrueWhenEqual && !TrueWhenGreaterThan)
+        return new ICmpInst(ICmpInst::ICMP_SLT, OrigLHS, OrigRHS);
+      else if (!TrueWhenLessThan && TrueWhenEqual && TrueWhenGreaterThan)
+        return new ICmpInst(ICmpInst::ICMP_SGE, OrigLHS, OrigRHS);
+      else if (!TrueWhenLessThan && TrueWhenEqual && !TrueWhenGreaterThan)
+        return new ICmpInst(ICmpInst::ICMP_EQ, OrigLHS, OrigRHS);
+      else if (!TrueWhenLessThan && !TrueWhenEqual && TrueWhenGreaterThan)
+        return new ICmpInst(ICmpInst::ICMP_SGT, OrigLHS, OrigRHS);
+      else if (!TrueWhenLessThan && !TrueWhenEqual && !TrueWhenGreaterThan)
+        return replaceInstUsesWith(ICI, Builder->getFalse());
+      else
+        llvm_unreachable("missed a combination");
+    }
+
+    // TODO: This could obvious be extended to handle three way unsigned
+    // compare idioms as well.  Also, fcmp idioms?
+    break;
   }
 
   // Simplify icmp_eq and icmp_ne instructions with integer constant RHS.
   if (ICI.isEquality()) {
     bool isICMP_NE = ICI.getPredicate() == ICmpInst::ICMP_NE;
 
     // If the first operand is (add|sub|and|or|xor|rem) with a constant, and
     // the second operand is a constant, simplify a bit.

test/Transforms/InstCombine/compare-3way.ll

+; RUN: opt -S -instcombine < %s | FileCheck %s
+
+declare void @use(i32)
+
+; These 18 excersice all combinations of signed comparison
+; for each of the three values produced by your typical 
+; 3way compare function (-1, 0, 1)
+
+define void @test_low_sgt(i64 %a, i64 %b) {
+; CHECK-LABEL: @test_low_sgt
+; CHECK:  %cmp = icmp slt i64 %a, %b
+; CHECK: br i1 %cmp, label %normal, label %unreached
+  %eq = icmp eq i64 %a, %b
+  %slt = icmp slt i64 %a, %b
+  %. = select i1 %slt, i32 -1, i32 1
+  %result = select i1 %eq, i32 0, i32 %.
+  %cmp = icmp sgt i32 %result, -1
+  br i1 %cmp, label %unreached, label %normal
+normal:
+  ret void
+unreached:
+  call void @use(i32 %result)
+  ret void
+}
+
+define void @test_low_slt(i64 %a, i64 %b) {
+; CHECK-LABEL: @test_low_slt
+; CHECK: br i1 false, label %unreached, label %normal
+  %eq = icmp eq i64 %a, %b
+  %slt = icmp slt i64 %a, %b
+  %. = select i1 %slt, i32 -1, i32 1
+  %result = select i1 %eq, i32 0, i32 %.
+  %cmp = icmp slt i32 %result, -1
+  br i1 %cmp, label %unreached, label %normal
+normal:
+  ret void
+unreached:
+  call void @use(i32 %result)
+  ret void
+}
+
+define void @test_low_sge(i64 %a, i64 %b) {
+; CHECK-LABEL: @test_low_sge
+; CHECK: br i1 true, label %unreached, label %normal
+  %eq = icmp eq i64 %a, %b
+  %slt = icmp slt i64 %a, %b
+  %. = select i1 %slt, i32 -1, i32 1
+  %result = select i1 %eq, i32 0, i32 %.
+  %cmp = icmp sge i32 %result, -1
+  br i1 %cmp, label %unreached, label %normal
+normal:
+  ret void
+unreached:
+  call void @use(i32 %result)
+  ret void
+}
+
+define void @test_low_sle(i64 %a, i64 %b) {
+; CHECK-LABEL: @test_low_sle
+; CHECK: %cmp = icmp slt i64 %a, %b
+; CHECK: br i1 %cmp, label %unreached, label %normal
+  %eq = icmp eq i64 %a, %b
+  %slt = icmp slt i64 %a, %b
+  %. = select i1 %slt, i32 -1, i32 1
+  %result = select i1 %eq, i32 0, i32 %.
+  %cmp = icmp sle i32 %result, -1
+  br i1 %cmp, label %unreached, label %normal
+normal:
+  ret void
+unreached:
+  call void @use(i32 %result)
+  ret void
+}
+
+define void @test_low_ne(i64 %a, i64 %b) {
+; CHECK-LABEL: @test_low_ne
+; CHECK:  %cmp = icmp slt i64 %a, %b
+; CHECK: br i1 %cmp, label %normal, label %unreached
+  %eq = icmp eq i64 %a, %b
+  %slt = icmp slt i64 %a, %b
+  %. = select i1 %slt, i32 -1, i32 1
+  %result = select i1 %eq, i32 0, i32 %.
+  %cmp = icmp ne i32 %result, -1
+  br i1 %cmp, label %unreached, label %normal
+normal:
+  ret void
+unreached:
+  call void @use(i32 %result)
+  ret void
+}
+
+define void @test_low_eq(i64 %a, i64 %b) {
+; CHECK-LABEL: @test_low_eq
+; CHECK:  %cmp = icmp slt i64 %a, %b
+; CHECK: br i1 %cmp, label %unreached, label %normal
+  %eq = icmp eq i64 %a, %b
+  %slt = icmp slt i64 %a, %b
+  %. = select i1 %slt, i32 -1, i32 1
+  %result = select i1 %eq, i32 0, i32 %.
+  %cmp = icmp eq i32 %result, -1
+  br i1 %cmp, label %unreached, label %normal
+normal:
+  ret void
+unreached:
+  call void @use(i32 %result)
+  ret void
+}
+
+define void @test_mid_sgt(i64 %a, i64 %b) {
+; CHECK-LABEL: @test_mid_sgt
+; CHECK:  %cmp = icmp sgt i64 %a, %b
+; CHECK: br i1 %cmp, label %unreached, label %normal
+  %eq = icmp eq i64 %a, %b
+  %slt = icmp slt i64 %a, %b
+  %. = select i1 %slt, i32 -1, i32 1
+  %result = select i1 %eq, i32 0, i32 %.
+  %cmp = icmp sgt i32 %result, 0
+  br i1 %cmp, label %unreached, label %normal
+normal:
+  ret void
+unreached:
+  call void @use(i32 %result)
+  ret void
+}
+
+define void @test_mid_slt(i64 %a, i64 %b) {
+; CHECK-LABEL: @test_mid_slt
+; CHECK:  %cmp = icmp slt i64 %a, %b
+; CHECK: br i1 %cmp, label %unreached, label %normal
+  %eq = icmp eq i64 %a, %b
+  %slt = icmp slt i64 %a, %b
+  %. = select i1 %slt, i32 -1, i32 1
+  %result = select i1 %eq, i32 0, i32 %.
+  %cmp = icmp slt i32 %result, 0
+  br i1 %cmp, label %unreached, label %normal
+normal:
+  ret void
+unreached:
+  call void @use(i32 %result)
+  ret void
+}
+
+define void @test_mid_sge(i64 %a, i64 %b) {
+; CHECK-LABEL: @test_mid_sge
+; CHECK: %cmp = icmp slt i64 %a, %b
+; CHECK: br i1 %cmp, label %normal, label %unreached
+  %eq = icmp eq i64 %a, %b
+  %slt = icmp slt i64 %a, %b
+  %. = select i1 %slt, i32 -1, i32 1
+  %result = select i1 %eq, i32 0, i32 %.
+  %cmp = icmp sge i32 %result, 0
+  br i1 %cmp, label %unreached, label %normal
+normal:
+  ret void
+unreached:
+  call void @use(i32 %result)
+  ret void
+}
+
+define void @test_mid_sle(i64 %a, i64 %b) {
+; CHECK-LABEL: @test_mid_sle
+; CHECK: %cmp = icmp sgt i64 %a, %b
+; CHECK: br i1 %cmp, label %normal, label %unreached
+  %eq = icmp eq i64 %a, %b
+  %slt = icmp slt i64 %a, %b
+  %. = select i1 %slt, i32 -1, i32 1
+  %result = select i1 %eq, i32 0, i32 %.
+  %cmp = icmp sle i32 %result, 0
+  br i1 %cmp, label %unreached, label %normal
+normal:
+  ret void
+unreached:
+  call void @use(i32 %result)
+  ret void
+}
+
+define void @test_mid_ne(i64 %a, i64 %b) {
+; CHECK-LABEL: @test_mid_ne
+; CHECK:  %cmp = icmp eq i64 %a, %b
+; CHECK: br i1 %cmp, label %normal, label %unreached
+  %eq = icmp eq i64 %a, %b
+  %slt = icmp slt i64 %a, %b
+  %. = select i1 %slt, i32 -1, i32 1
+  %result = select i1 %eq, i32 0, i32 %.
+  %cmp = icmp ne i32 %result, 0
+  br i1 %cmp, label %unreached, label %normal
+normal:
+  ret void
+unreached:
+  call void @use(i32 %result)
+  ret void
+}
+
+define void @test_mid_eq(i64 %a, i64 %b) {
+; CHECK-LABEL: @test_mid_eq
+; CHECK:  icmp eq i64 %a, %b
+; CHECK: br i1 %eq, label %unreached, label %normal
+  %eq = icmp eq i64 %a, %b
+  %slt = icmp slt i64 %a, %b
+  %. = select i1 %slt, i32 -1, i32 1
+  %result = select i1 %eq, i32 0, i32 %.
+  %cmp = icmp eq i32 %result, 0
+  br i1 %cmp, label %unreached, label %normal
+normal:
+  ret void
+unreached:
+  call void @use(i32 %result)
+  ret void
+}
+
+define void @test_high_sgt(i64 %a, i64 %b) {
+; CHECK-LABEL: @test_high_sgt
+; CHECK: br i1 false, label %unreached, label %normal
+  %eq = icmp eq i64 %a, %b
+  %slt = icmp slt i64 %a, %b
+  %. = select i1 %slt, i32 -1, i32 1
+  %result = select i1 %eq, i32 0, i32 %.
+  %cmp = icmp sgt i32 %result, 1
+  br i1 %cmp, label %unreached, label %normal
+normal:
+  ret void
+unreached:
+  call void @use(i32 %result)
+  ret void
+}
+
+define void @test_high_slt(i64 %a, i64 %b) {
+; CHECK-LABEL: @test_high_slt
+; CHECK:  %cmp = icmp sgt i64 %a, %b
+; CHECK: br i1 %cmp, label %normal, label %unreached
+  %eq = icmp eq i64 %a, %b
+  %slt = icmp slt i64 %a, %b
+  %. = select i1 %slt, i32 -1, i32 1
+  %result = select i1 %eq, i32 0, i32 %.
+  %cmp = icmp slt i32 %result, 1
+  br i1 %cmp, label %unreached, label %normal
+normal:
+  ret void
+unreached:
+  call void @use(i32 %result)
+  ret void
+}
+
+define void @test_high_sge(i64 %a, i64 %b) {
+; CHECK-LABEL: @test_high_sge
+; CHECK: %cmp = icmp sgt i64 %a, %b
+; CHECK: br i1 %cmp, label %unreached, label %normal
+  %eq = icmp eq i64 %a, %b
+  %slt = icmp slt i64 %a, %b
+  %. = select i1 %slt, i32 -1, i32 1
+  %result = select i1 %eq, i32 0, i32 %.
+  %cmp = icmp sge i32 %result, 1
+  br i1 %cmp, label %unreached, label %normal
+normal:
+  ret void
+unreached:
+  call void @use(i32 %result)
+  ret void
+}
+
+define void @test_high_sle(i64 %a, i64 %b) {
+; CHECK-LABEL: @test_high_sle
+; CHECK: br i1 true, label %unreached, label %normal
+  %eq = icmp eq i64 %a, %b
+  %slt = icmp slt i64 %a, %b
+  %. = select i1 %slt, i32 -1, i32 1
+  %result = select i1 %eq, i32 0, i32 %.
+  %cmp = icmp sle i32 %result, 1
+  br i1 %cmp, label %unreached, label %normal
+normal:
+  ret void
+unreached:
+  call void @use(i32 %result)
+  ret void
+}
+
+define void @test_high_ne(i64 %a, i64 %b) {
+; CHECK-LABEL: @test_high_ne
+; CHECK:  %cmp = icmp sgt i64 %a, %b
+; CHECK: br i1 %cmp, label %normal, label %unreached
+  %eq = icmp eq i64 %a, %b
+  %slt = icmp slt i64 %a, %b
+  %. = select i1 %slt, i32 -1, i32 1
+  %result = select i1 %eq, i32 0, i32 %.
+  %cmp = icmp ne i32 %result, 1
+  br i1 %cmp, label %unreached, label %normal
+normal:
+  ret void
+unreached:
+  call void @use(i32 %result)
+  ret void
+}
+
+define void @test_high_eq(i64 %a, i64 %b) {
+; CHECK-LABEL: @test_high_eq
+; CHECK:  %cmp = icmp sgt i64 %a, %b
+; CHECK: br i1 %cmp, label %unreached, label %normal
+  %eq = icmp eq i64 %a, %b
+  %slt = icmp slt i64 %a, %b
+  %. = select i1 %slt, i32 -1, i32 1
+  %result = select i1 %eq, i32 0, i32 %.
+  %cmp = icmp eq i32 %result, 1
+  br i1 %cmp, label %unreached, label %normal
+normal:
+  ret void
+unreached:
+  call void @use(i32 %result)
+  ret void
+}
+
+; These five make sure we didn't accidentally hard code one of the
+; produced values
+
+define void @non_standard_low(i64 %a, i64 %b) {
+; CHECK-LABEL: @non_standard_low
+; CHECK:  %cmp = icmp slt i64 %a, %b
+; CHECK: br i1 %cmp, label %unreached, label %normal
+  %eq = icmp eq i64 %a, %b
+  %slt = icmp slt i64 %a, %b
+  %. = select i1 %slt, i32 -3, i32 -1
+  %result = select i1 %eq, i32 -2, i32 %.
+  %cmp = icmp eq i32 %result, -3
+  br i1 %cmp, label %unreached, label %normal
+normal:
+  ret void
+unreached:
+  call void @use(i32 %result)
+  ret void
+}
+
+define void @non_standard_mid(i64 %a, i64 %b) {
+; CHECK-LABEL: @non_standard_mid
+; CHECK:  icmp eq i64 %a, %b
+; CHECK: br i1 %eq, label %unreached, label %normal
+  %eq = icmp eq i64 %a, %b
+  %slt = icmp slt i64 %a, %b
+  %. = select i1 %slt, i32 -3, i32 -1
+  %result = select i1 %eq, i32 -2, i32 %.
+  %cmp = icmp eq i32 %result, -2
+  br i1 %cmp, label %unreached, label %normal
+normal:
+  ret void
+unreached:
+  call void @use(i32 %result)
+  ret void
+}
+
+define void @non_standard_high(i64 %a, i64 %b) {
+; CHECK-LABEL: @non_standard_high
+; CHECK:  %cmp = icmp sgt i64 %a, %b
+; CHECK: br i1 %cmp, label %unreached, label %normal
+  %eq = icmp eq i64 %a, %b
+  %slt = icmp slt i64 %a, %b
+  %. = select i1 %slt, i32 -3, i32 -1
+  %result = select i1 %eq, i32 -2, i32 %.
+  %cmp = icmp eq i32 %result, -1
+  br i1 %cmp, label %unreached, label %normal
+normal:
+  ret void
+unreached:
+  call void @use(i32 %result)
+  ret void
+}
+
+define void @non_standard_bound1(i64 %a, i64 %b) {
+; CHECK-LABEL: @non_standard_bound1
+; CHECK:  br i1 false, label %unreached, label %normal
+  %eq = icmp eq i64 %a, %b
+  %slt = icmp slt i64 %a, %b
+  %. = select i1 %slt, i32 -3, i32 -1
+  %result = select i1 %eq, i32 -2, i32 %.
+  %cmp = icmp eq i32 %result, -20
+  br i1 %cmp, label %unreached, label %normal
+normal:
+  ret void
+unreached:
+  call void @use(i32 %result)
+  ret void
+}
+
+define void @non_standard_bound2(i64 %a, i64 %b) {
+; CHECK-LABEL: @non_standard_bound2
+; CHECK:  br i1 false, label %unreached, label %normal
+  %eq = icmp eq i64 %a, %b
+  %slt = icmp slt i64 %a, %b
+  %. = select i1 %slt, i32 -3, i32 -1
+  %result = select i1 %eq, i32 -2, i32 %.
+  %cmp = icmp eq i32 %result, 0
+  br i1 %cmp, label %unreached, label %normal
+normal:
+  ret void
+unreached:
+  call void @use(i32 %result)
+  ret void
+}