[InstCombine] extend (select X, C1, C2 --> ext X) to vectors

[spatel]

Summary

The code change is hopefully straightforward: replace dyn_casts with m_APInt, and we get transforms for splat vectors.

But these transforms raise some questions:

1. In the cases where we need a 'not', we're increasing the instruction count. Is this justified because ext/not is always assumed cheaper/more canonical than a select?

2. In the not+zext case, notice that the zext gets moved ahead of the xor. This is because visitZext() has: // zext (xor i1 X, true) to i32 --> xor (zext i1 X to i32), 1

There's no code comment for the motivation. Assuming there is good reason to break the m_Not pattern should sext+xor do the same?

3. If we're ok with increasing the instruction count for #1, is the scalar select of vectors example also a good transform? In the worst case, we'd need 4 instructions in place of the select: xor, zext, insertelement, shufflevector.

[eli.friedman]

Generally looks fine.

lib/Transforms/InstCombine/InstCombineSelect.cpp

968

Extra newline.

971

Won't this crash if SI.getType() is <4 x i1>?

[eli.friedman]

Re: your questions about canonicalization:

1. The canonical representation of a calculation is whatever instcombine declares it to be. We generally try to favor arithmetic over selects because selects tend to be more expensive, but there's a limit to where that's worthwhile.
2. Not sure exactly why we prefer to move xor out of a zext rather than in... apparently the transform was added in r21713 to cover this testcase:

define i32 @test22(i1 %X) {
; CHECK-LABEL: @test22(
; CHECK-NEXT: %1 = zext i1 %X to i32
; CHECK-NEXT: ret i32 %1
	%Y = xor i1 %X, true		; <i1> [#uses=1]
	%Z = zext i1 %Y to i32		; <i32> [#uses=1]
	%Q = xor i32 %Z, 1		; <i32> [#uses=1]
	ret i32 %Q
}

Obviously, that doesn't really motivate the canonical representation either way. That said, moving the xor in rather than out in general raises awkward questions about what to do in cases like "(zext i8 x to i32) ^ 257". And changing the canonical representation probably involves some work to figure out what other transforms depend on the canonical representation of zext+xor.

[spatel]

Patch updated:
Added assert for i1 types (vectors should be handled after rL274465).

[spatel]

In D21899#473188, @eli.friedman wrote:

Re: your questions about canonicalization:

Obviously, that doesn't really motivate the canonical representation either way. That said, moving the xor in rather than out in general raises awkward questions about what to do in cases like "(zext i8 x to i32) ^ 257". And changing the canonical representation probably involves some work to figure out what other transforms depend on the canonical representation of zext+xor.

Thanks for the answers! I don't think we need to worry about your ^257 example because we're only looking at i1 types for this transform.

lib/Transforms/InstCombine/InstCombineSelect.cpp

971

Yes - scalar i1 selects were all handled above, but vectors were not. That should be fixed with rL274465, so now we can assert that no i1 types make it down here.

[eli.friedman]

LGTM.

lib/Transforms/InstCombine/InstCombineSelect.cpp

if (SI.getType()->getScalarType()->isIntegerTy(1) &&
     // select a, ~a, b -> (~a) & b
     // select a, b, ~a -> (~a) | b
     if (match(TrueVal, m_Not(m_Specific(CondVal))))
       return BinaryOperator::CreateAnd(TrueVal, FalseVal);
     if (match(FalseVal, m_Not(m_Specific(CondVal))))
       return BinaryOperator::CreateOr(TrueVal, FalseVal);
   }
 
-  // Selecting between two integer constants?
-  if (ConstantInt *TrueValC = dyn_cast<ConstantInt>(TrueVal))
-    if (ConstantInt *FalseValC = dyn_cast<ConstantInt>(FalseVal)) {
-      // select C, 1, 0 -> zext C to int
-      if (FalseValC->isZero() && TrueValC->getValue() == 1)
-        return new ZExtInst(CondVal, SI.getType());
+  // Selecting between two integer or vector splat integer constants?
+  //
+  // Note that we don't handle a scalar select of vectors:
+  // select i1 %c, <2 x i8> <1, 1>, <2 x i8> <0, 0>
+  // because that may need 3 instructions to splat the condition value:
+  // extend, insertelement, shufflevector.
+  const APInt *TrueC, *FalseC;
+  if (match(TrueVal, m_APInt(TrueC)) && match(FalseVal, m_APInt(FalseC)) &&
+      (CondVal->getType()->isVectorTy() == SI.getType()->isVectorTy())) {
+    assert(!SI.getType()->getScalarType()->isIntegerTy(1) &&
+           "Bool operands should already be handled");
+
+    if (*FalseC == 0 && *TrueC == 1)
 
-      // select C, -1, 0 -> sext C to int
-      if (FalseValC->isZero() && TrueValC->isAllOnesValue())
-        return new SExtInst(CondVal, SI.getType());
+    if (*FalseC == 0 && TrueC->isAllOnesValue())
 
-      // select C, 0, 1 -> zext !C to int
-      if (TrueValC->isZero() && FalseValC->getValue() == 1) {
-        Value *NotCond = Builder->CreateNot(CondVal, "not."+CondVal->getName());
-        return new ZExtInst(NotCond, SI.getType());
-      }
+    if (*TrueC == 0 && *FalseC == 1) {
+      Value *NotCond = Builder->CreateNot(CondVal, "not." + CondVal->getName());
 
-      // select C, 0, -1 -> sext !C to int
-      if (TrueValC->isZero() && FalseValC->isAllOnesValue()) {
-        Value *NotCond = Builder->CreateNot(CondVal, "not."+CondVal->getName());
-        return new SExtInst(NotCond, SI.getType());
-      }
+    if (*TrueC == 0 && FalseC->isAllOnesValue()) {
+      Value *NotCond = Builder->CreateNot(CondVal, "not." + CondVal->getName());
+  }
 
+  if (ConstantInt *TrueValC = dyn_cast<ConstantInt>(TrueVal))
+    if (ConstantInt *FalseValC = dyn_cast<ConstantInt>(FalseVal))
       if (Value *V = foldSelectICmpAnd(SI, TrueValC, FalseValC, Builder))
         return replaceInstUsesWith(SI, V);
-    }
 
   // See if we are selecting two values based on a comparison of the two values.
   if (FCmpInst *FCI = dyn_cast<FCmpInst>(CondVal)) {
     if (FCI->getOperand(0) == TrueVal && FCI->getOperand(1) == FalseVal) {
       // Transform (X == Y) ? X : Y  -> Y
       if (FCI->getPredicate() == FCmpInst::FCMP_OEQ) {
         // This is not safe in general for floating point:
         // consider X== -0, Y== +0.

test/Transforms/InstCombine/apint-select.ll

 ; CHECK-NEXT:    [[NOT_C:%.*]] = xor i1 %C, true
 ; CHECK-NEXT:    [[V:%.*]] = sext i1 [[NOT_C]] to i999
 ; CHECK-NEXT:    ret i999 [[V]]
 ;
   %V = select i1 %C, i999 0, i999 -1
   ret i999 %V
 }
 
-; FIXME: Vector selects of vector splat constants match APInt too.
+; Vector selects of vector splat constants match APInt too.
 
 define <2 x i41> @zext_vec(<2 x i1> %C) {
 ; CHECK-LABEL: @zext_vec(
-; CHECK-NEXT:    [[V:%.*]] = select <2 x i1> %C, <2 x i41> <i41 1, i41 1>, <2 x i41> zeroinitializer
+; CHECK-NEXT:    [[V:%.*]] = zext <2 x i1> %C to <2 x i41>
 ; CHECK-NEXT:    ret <2 x i41> [[V]]
 ;
   %V = select <2 x i1> %C, <2 x i41> <i41 1, i41 1>, <2 x i41> <i41 0, i41 0>
   ret <2 x i41> %V
 }
 
 define <2 x i32> @sext_vec(<2 x i1> %C) {
 ; CHECK-LABEL: @sext_vec(
-; CHECK-NEXT:    [[V:%.*]] = select <2 x i1> %C, <2 x i32> <i32 -1, i32 -1>, <2 x i32> zeroinitializer
+; CHECK-NEXT:    [[V:%.*]] = sext <2 x i1> %C to <2 x i32>
 ; CHECK-NEXT:    ret <2 x i32> [[V]]
 ;
   %V = select <2 x i1> %C, <2 x i32> <i32 -1, i32 -1>, <2 x i32> <i32 0, i32 0>
   ret <2 x i32> %V
 }
 
 define <2 x i999> @not_zext_vec(<2 x i1> %C) {
 ; CHECK-LABEL: @not_zext_vec(
-; CHECK-NEXT:    [[V:%.*]] = select <2 x i1> %C, <2 x i999> zeroinitializer, <2 x i999> <i999 1, i999 1>
+; CHECK-NEXT:    [[TMP1:%.*]] = zext <2 x i1> %C to <2 x i999>
+; CHECK-NEXT:    [[V:%.*]] = xor <2 x i999> [[TMP1]], <i999 1, i999 1>
 ; CHECK-NEXT:    ret <2 x i999> [[V]]
 ;
   %V = select <2 x i1> %C, <2 x i999> <i999 0, i999 0>, <2 x i999> <i999 1, i999 1>
   ret <2 x i999> %V
 }
 
 define <2 x i64> @not_sext_vec(<2 x i1> %C) {
 ; CHECK-LABEL: @not_sext_vec(
-; CHECK-NEXT:    [[V:%.*]] = select <2 x i1> %C, <2 x i64> zeroinitializer, <2 x i64> <i64 -1, i64 -1>
+; CHECK-NEXT:    [[NOT_C:%.*]] = xor <2 x i1> %C, <i1 true, i1 true>
+; CHECK-NEXT:    [[V:%.*]] = sext <2 x i1> [[NOT_C]] to <2 x i64>
 ; CHECK-NEXT:    ret <2 x i64> [[V]]
 ;
   %V = select <2 x i1> %C, <2 x i64> <i64 0, i64 0>, <2 x i64> <i64 -1, i64 -1>
   ret <2 x i64> %V
 }
 
+; But don't touch this - we would need 3 instructions to extend and splat the scalar select condition.
+
+define <2 x i32> @scalar_select_of_vectors(i1 %c) {
+; CHECK-LABEL: @scalar_select_of_vectors(
+; CHECK-NEXT:    [[V:%.*]] = select i1 %c, <2 x i32> <i32 1, i32 1>, <2 x i32> zeroinitializer
+; CHECK-NEXT:    ret <2 x i32> [[V]]
+;
+  %V = select i1 %c, <2 x i32> <i32 1, i32 1>, <2 x i32> zeroinitializer
+  ret <2 x i32> %V
+}
+
 ;; (x <s 0) ? -1 : 0 -> ashr x, 31
 
 define i41 @test3(i41 %X) {
 ; CHECK-LABEL: @test3(
 ; CHECK-NEXT:    [[X_LOBIT:%.*]] = ashr i41 %X, 40
 ; CHECK-NEXT:    ret i41 [[X_LOBIT]]
 ;
   %t = icmp slt i41 %X, 0