[InstCombine] Make folding (X >s -1) ? C1 : C2 --> ((X >>s 31) & (C2 - C1)) + C1 support splat vectors

[craig.topper]

Summary

This reuses decomposeBitTestICmp to decode the compare.

I'm assuming that pointer icmps will never match decomposeBitTestICmp because we shouldn't be doing relational comparisons on points, I think? So I removed the IntegerType check, but I could be wrong.

I've deferred the equality check on the type of compare vs the select until after the decomposeBitTestICmp call because I've considering putting an m_Trunc test inside of decomposeBitTestICmp so X won't always just be the CmpLHS.

[craig.topper]

Added back the type checks and added an explicit truncate test. We don't want to look through a truncate here. Even if decomposeBitTestICmp starts looking through truncates, we need to prevent is for this usage.

[spatel]

LGTM.

lib/Transforms/InstCombine/InstCombineSelect.cpp

 //===- InstCombineSelect.cpp ----------------------------------------------===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements the visitSelect function.
 //
 //===----------------------------------------------------------------------===//
 
 #include "InstCombineInternal.h"
+#include "llvm/Analysis/CmpInstAnalysis.h"
 #include "llvm/Analysis/ConstantFolding.h"
 #include "llvm/Analysis/InstructionSimplify.h"
 #include "llvm/Analysis/ValueTracking.h"
 #include "llvm/IR/MDBuilder.h"
 #include "llvm/IR/PatternMatch.h"
 #include "llvm/Support/KnownBits.h"
 using namespace llvm;
 using namespace PatternMatch;
Instruction *InstCombiner::foldSelectInstWithICmp(SelectInst &SI,
   Value *CmpLHS = ICI->getOperand(0);
   Value *CmpRHS = ICI->getOperand(1);
 
   // Transform (X >s -1) ? C1 : C2 --> ((X >>s 31) & (C2 - C1)) + C1
   // and       (X <s  0) ? C2 : C1 --> ((X >>s 31) & (C2 - C1)) + C1
   // FIXME: Type and constness constraints could be lifted, but we have to
   //        watch code size carefully. We should consider xor instead of
   //        sub/add when we decide to do that.
-  if (IntegerType *Ty = dyn_cast<IntegerType>(CmpLHS->getType())) {
-    if (TrueVal->getType() == Ty) {
-      if (ConstantInt *Cmp = dyn_cast<ConstantInt>(CmpRHS)) {
-        ConstantInt *C1 = nullptr, *C2 = nullptr;
-        if (Pred == ICmpInst::ICMP_SGT && Cmp->isMinusOne()) {
-          C1 = dyn_cast<ConstantInt>(TrueVal);
-          C2 = dyn_cast<ConstantInt>(FalseVal);
-        } else if (Pred == ICmpInst::ICMP_SLT && Cmp->isZero()) {
-          C1 = dyn_cast<ConstantInt>(FalseVal);
-          C2 = dyn_cast<ConstantInt>(TrueVal);
-        }
-        if (C1 && C2) {
+  if (CmpLHS->getType()->isIntOrIntVectorTy() &&
+      CmpLHS->getType() == TrueVal->getType()) {
+    const APInt *C1, *C2;
+    if (match(TrueVal, m_APInt(C1)) && match(FalseVal, m_APInt(C2))) {
+      ICmpInst::Predicate Pred = ICI->getPredicate();
+      Value *X;
+      APInt Mask;
+      if (decomposeBitTestICmp(CmpLHS, CmpRHS, Pred, X, Mask)) {
+        if (Mask.isSignMask()) {
+          assert(X == CmpLHS && "Expected to use the compare input directly");
+          assert(ICmpInst::isEquality(Pred) && "Expected equality predicate");
+
+          if (Pred == ICmpInst::ICMP_NE)
+            std::swap(C1, C2);
+
           // This shift results in either -1 or 0.
-          Value *AShr = Builder.CreateAShr(CmpLHS, Ty->getBitWidth() - 1);
+          Value *AShr = Builder.CreateAShr(X, Mask.getBitWidth() - 1);
 
           // Check if we can express the operation with a single or.
-          if (C2->isMinusOne())
-            return replaceInstUsesWith(SI, Builder.CreateOr(AShr, C1));
+          if (C2->isAllOnesValue())
+            return replaceInstUsesWith(SI, Builder.CreateOr(AShr, *C1));
 
-          Value *And = Builder.CreateAnd(AShr, C2->getValue() - C1->getValue());
-          return replaceInstUsesWith(SI, Builder.CreateAdd(And, C1));
+          Value *And = Builder.CreateAnd(AShr, *C2 - *C1);
+          return replaceInstUsesWith(SI, Builder.CreateAdd(And,
+                                        ConstantInt::get(And->getType(), *C1)));
         }
       }
     }
   }
 
   // NOTE: if we wanted to, this is where to detect integer MIN/MAX
 
   if (CmpRHS != CmpLHS && isa<Constant>(CmpRHS)) {

test/Transforms/InstCombine/select.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt < %s -instcombine -S | FileCheck %s
 
 ; PR1822
 
 target datalayout = "e-p:64:64-p1:16:16-p2:32:32:32-p3:64:64:64"
 
 define i32 @test1(i32 %A, i32 %B) {
         %C = select i1 false, i32 %A, i32 %B
define i32 @test35(i32 %x) {
   ret i32 %cond
 ; CHECK-LABEL: @test35(
 ; CHECK: ashr i32 %x, 31
 ; CHECK: and i32 {{.*}}, 40
 ; CHECK: add nuw nsw i32 {{.*}}, 60
 ; CHECK: ret
 }
 
+define <2 x i32> @test35vec(<2 x i32> %x) {
+; CHECK-LABEL: @test35vec(
+; CHECK-NEXT:    [[TMP1:%.*]] = ashr <2 x i32> [[X:%.*]], <i32 31, i32 31>
+; CHECK-NEXT:    [[TMP2:%.*]] = and <2 x i32> [[TMP1]], <i32 40, i32 40>
+; CHECK-NEXT:    [[TMP3:%.*]] = add nuw nsw <2 x i32> [[TMP2]], <i32 60, i32 60>
+; CHECK-NEXT:    ret <2 x i32> [[TMP3]]
+;
+  %cmp = icmp sge <2 x i32> %x, <i32 0, i32 0>
+  %cond = select <2 x i1> %cmp, <2 x i32> <i32 60, i32 60>, <2 x i32> <i32 100, i32 100>
+  ret <2 x i32> %cond
+}
+
+; Make sure we can still perform this optimization with a truncate present
+define i32 @test35_with_trunc(i64 %x) {
+; CHECK-LABEL: @test35_with_trunc(
+; CHECK-NEXT:    [[X1:%.*]] = trunc i64 [[X:%.*]] to i32
+; CHECK-NEXT:    [[TMP1:%.*]] = ashr i32 [[X1]], 31
+; CHECK-NEXT:    [[TMP2:%.*]] = and i32 [[TMP1]], 40
+; CHECK-NEXT:    [[TMP3:%.*]] = add nuw nsw i32 [[TMP2]], 60
+; CHECK-NEXT:    ret i32 [[TMP3]]
+;
+  %x1 = trunc i64 %x to i32
+  %cmp = icmp sge i32 %x1, 0
+  %cond = select i1 %cmp, i32 60, i32 100
+  ret i32 %cond
+}
+
 define i32 @test36(i32 %x) {
   %cmp = icmp slt i32 %x, 0
   %cond = select i1 %cmp, i32 60, i32 100
   ret i32 %cond
 ; CHECK-LABEL: @test36(
 ; CHECK: ashr i32 %x, 31
 ; CHECK: and i32 {{.*}}, -40
 ; CHECK: add nsw i32 {{.*}}, 100
 ; CHECK: ret
 }
 
+define <2 x i32> @test36vec(<2 x i32> %x) {
+; CHECK-LABEL: @test36vec(
+; CHECK-NEXT:    [[TMP1:%.*]] = ashr <2 x i32> [[X:%.*]], <i32 31, i32 31>
+; CHECK-NEXT:    [[TMP2:%.*]] = and <2 x i32> [[TMP1]], <i32 -40, i32 -40>
+; CHECK-NEXT:    [[TMP3:%.*]] = add nsw <2 x i32> [[TMP2]], <i32 100, i32 100>
+; CHECK-NEXT:    ret <2 x i32> [[TMP3]]
+;
+  %cmp = icmp slt <2 x i32> %x, <i32 0, i32 0>
+  %cond = select <2 x i1> %cmp, <2 x i32> <i32 60, i32 60>, <2 x i32> <i32 100, i32 100>
+  ret <2 x i32> %cond
+}
+
 define i32 @test37(i32 %x) {
   %cmp = icmp sgt i32 %x, -1
   %cond = select i1 %cmp, i32 1, i32 -1
   ret i32 %cond
 ; CHECK-LABEL: @test37(
 ; CHECK: ashr i32 %x, 31
 ; CHECK: or i32 {{.*}}, 1
 ; CHECK: ret
 }
 
+define <2 x i32> @test37vec(<2 x i32> %x) {
+; CHECK-LABEL: @test37vec(
+; CHECK-NEXT:    [[TMP1:%.*]] = ashr <2 x i32> [[X:%.*]], <i32 31, i32 31>
+; CHECK-NEXT:    [[TMP2:%.*]] = or <2 x i32> [[TMP1]], <i32 1, i32 1>
+; CHECK-NEXT:    ret <2 x i32> [[TMP2]]
+;
+  %cmp = icmp sgt <2 x i32> %x, <i32 -1, i32 -1>
+  %cond = select <2 x i1> %cmp, <2 x i32> <i32 1, i32 1>, <2 x i32> <i32 -1, i32 -1>
+  ret <2 x i32> %cond
+}
+
 define i1 @test38(i1 %cond) {
   %zero = alloca i32
   %one = alloca i32
   %ptr = select i1 %cond, i32* %zero, i32* %one
   %isnull = icmp eq i32* %ptr, null
   ret i1 %isnull
 ; CHECK-LABEL: @test38(
 ; CHECK: ret i1 false