[InstCombine] extend vector select matching for non-splat constants

[spatel]

Summary

In D21740, we discussed trying to make this a more general matcher. However, I didn't see a clean way to handle the regular m_Not cases and these non-splat vector patterns, so I've just opted for the direct approach here. If there are other potential uses of areInverseVectorBitmasks(), we could move that to a higher level.

[spatel]

Patch updated:
The previous version that I uploaded was an incorrect local draft.

[eli.friedman]

Do we care at all about the number of uses of the operands?

lib/Transforms/InstCombine/InstCombineAndOrXor.cpp

1601 ↗

(On Diff #63145)

What if one is undef?

[spatel]

In D22114#477445, @eli.friedman wrote:

Do we care at all about the number of uses of the operands?

Good question. I was assuming one-use based on my motivating examples, but since we're potentially creating an xor here and bitcasts above this, we probably do need to check that.

lib/Transforms/InstCombine/InstCombineAndOrXor.cpp

1601 ↗

(On Diff #63145)

Ah, right. Or even both elts could be undef?
Ok to make that a FIXME comment for this patch? I'll follow-up with the extra logic and test cases.

[eli.friedman]

lib/Transforms/InstCombine/InstCombineAndOrXor.cpp

1601 ↗

(On Diff #63145)

Sure.

[spatel]

Patch updated:

1. Added FIXME to handle undef elements for non-splat constant vectors.
2. Added testcase for multiple-uses of the non-splat mask vectors.
3. Ahead of this patch, limited the transform for multiple-use scenarios in:

http://reviews.llvm.org/rL274926
http://reviews.llvm.org/rL274932

[eli.friedman]

LGTM.

test/Transforms/InstCombine/logical-select.ll

374 ↗

(On Diff #63327)

Is a select actually the canonical form for this? It seems like we should prefer a shufflevector.

[spatel]

test/Transforms/InstCombine/logical-select.ll

374 ↗

(On Diff #63327)

I saw that x86 SSE or AVX didn't care, eg:

vpblendw	$60, %xmm1, %xmm0, %xmm0 # xmm0 = xmm0[0,1],xmm1[2,3,4,5],xmm0[6,7]

..so I figured one was as good as the other, but this is not true in general. AArch64 and PPC+VSX generate different code:

define <4 x i32> @foo(<4 x i32> %a, <4 x i32> %b) {
  %sel = select <4 x i1> <i1 true, i1 false, i1 false, i1 true>, <4 x i32> %a, <4 x i32> %b
  ret <4 x i32> %sel
}

define <4 x i32> @goo(<4 x i32> %a, <4 x i32> %b) {
  %shuf = shufflevector <4 x i32> %a, <4 x i32> %b, <4 x i32> <i32 0, i32 5, i32 6, i32 3>
  ret <4 x i32> %shuf
}

$ ./llc shufsel.ll -o - -mtriple=aarch64

adrp	x8, .LCPI0_0
ldr	q2, [x8, :lo12:.LCPI0_0]
bsl	v2.16b, v0.16b, v1.16b
mov		v0.16b, v2.16b
ret

ext	v1.16b, v0.16b, v1.16b, #12
ext	v0.16b, v1.16b, v0.16b, #4
ext	v1.16b, v1.16b, v1.16b, #8
ext	v0.16b, v0.16b, v1.16b, #12
ret

$ ./llc shufsel.ll -o - -mtriple=powerpc64 -mattr=vsx

addis 3, 2, .LCPI0_0@toc@ha
addi 3, 3, .LCPI0_0@toc@l
lxvw4x 0, 0, 3
xxsel 34, 35, 34, 0
blr

addis 3, 2, .LCPI1_0@toc@ha
addi 3, 3, .LCPI1_0@toc@l
lxvw4x 36, 0, 3
vperm 2, 2, 3, 4
blr

I don't know AArch at all, but bsl seems better than 4 exts. Is there a better way than either of those? PPC xxsel vs. vperm seem equivalent, but I have no knowledge of any recent PPC HW.

Does the current codegen affect what we do here in IR?

[eli.friedman]

I would lean towards making shufflevector canonical. It interacts better with existing transforms involving shuffles and vector insert/extract. I can't think of any benefit to making the select form canonical.

It's good to be aware of how specific targets handle it, but it doesn't really matter that much given that we can just fix the target to pattern-match this kind of shufflevector.

[spatel]

In D22114#482283, @eli.friedman wrote:

I would lean towards making shufflevector canonical. It interacts better with existing transforms involving shuffles and vector insert/extract. I can't think of any benefit to making the select form canonical.

I don't have a strong opinion on what we do here in IR, but IMO, a bitwise or element-level vector select is a required feature of any vector ISA - although x86 took 5+ ISA revs to get it :) . The more general vector shuffle/permute is not required.

I filed PR28530 and PR28531 to make at least ARM and PPC aware of their existing behavior and the proposal:
https://llvm.org/bugs/show_bug.cgi?id=28530
https://llvm.org/bugs/show_bug.cgi?id=28531 (note the VMX tragedy - we eliminated the logic ops in IR, and they came back in codegen!)

[eli.friedman]

The question isn't really what instructions we expect targets to have, but rather which version makes LLVM simpler overall. That said, if we assume almost every vector target has some sort of blend instruction, and therefore almost every target would need to pattern match the shuffle into a select, that would be a reason to prefer the select form, I guess.

[spatel]

In D22114#482459, @eli.friedman wrote:

The question isn't really what instructions we expect targets to have, but rather which version makes LLVM simpler overall. That said, if we assume almost every vector target has some sort of blend instruction, and therefore almost every target would need to pattern match the shuffle into a select, that would be a reason to prefer the select form, I guess.

Agree - that was the point I was trying to ineffectively make: if targets are going to prefer the select/blend anyway, we save some slight codegen churn from producing that form here in IR directly.

llvm/trunk/lib/Transforms/InstCombine/InstCombineAndOrXor.cpp

Instruction *InstCombiner::MatchBSwap(BinaryOperator &I) {
   Instruction *LastInst = Insts.pop_back_val();
   LastInst->removeFromParent();
 
   for (auto *Inst : Insts)
     Worklist.Add(Inst);
   return LastInst;
 }
 
+/// If all elements of two constant vectors are 0/-1 and inverses, return true.
+static bool areInverseVectorBitmasks(Constant *C1, Constant *C2) {
+  unsigned NumElts = C1->getType()->getVectorNumElements();
+  for (unsigned i = 0; i != NumElts; ++i) {
+    Constant *EltC1 = C1->getAggregateElement(i);
+    Constant *EltC2 = C2->getAggregateElement(i);
+    if (!EltC1 || !EltC2)
+      return false;
+
+    // One element must be all ones, and the other must be all zeros.
+    // FIXME: Allow undef elements.
+    if (!((match(EltC1, m_Zero()) && match(EltC2, m_AllOnes())) ||
+          (match(EltC2, m_Zero()) && match(EltC1, m_AllOnes()))))
+      return false;
+  }
+  return true;
+}
+
 /// We have an expression of the form (A & C) | (B & D). If A is a scalar or
 /// vector composed of all-zeros or all-ones values and is the bitwise 'not' of
 /// B, it can be used as the condition operand of a select instruction.
-static Value *getSelectCondition(Value *A, Value *B) {
+static Value *getSelectCondition(Value *A, Value *B,
+                                 InstCombiner::BuilderTy &Builder) {
   // If these are scalars or vectors of i1, A can be used directly.
   Type *Ty = A->getType();
   if (match(A, m_Not(m_Specific(B))) && Ty->getScalarType()->isIntegerTy(1))
     return A;
 
   // If A and B are sign-extended, look through the sexts to find the booleans.
   Value *Cond;
   if (match(A, m_SExt(m_Value(Cond))) &&
       Cond->getType()->getScalarType()->isIntegerTy(1) &&
       match(B, m_CombineOr(m_Not(m_SExt(m_Specific(Cond))),
                            m_SExt(m_Not(m_Specific(Cond))))))
     return Cond;
 
-  // TODO: Try more matches that only apply to non-splat constant vectors.
+  // All scalar (and most vector) possibilities should be handled now.
+  // Try more matches that only apply to non-splat constant vectors.
+  if (!Ty->isVectorTy())
+    return nullptr;
 
+  // If both operands are constants, see if the constants are inverse bitmasks.
+  Constant *AC, *BC;
+  if (match(A, m_Constant(AC)) && match(B, m_Constant(BC)) &&
+      areInverseVectorBitmasks(AC, BC))
+    return ConstantExpr::getTrunc(AC, CmpInst::makeCmpResultType(Ty));
+
+  // If both operands are xor'd with constants using the same sexted boolean
+  // operand, see if the constants are inverse bitmasks.
+  if (match(A, (m_Xor(m_SExt(m_Value(Cond)), m_Constant(AC)))) &&
+      match(B, (m_Xor(m_SExt(m_Specific(Cond)), m_Constant(BC)))) &&
+      Cond->getType()->getScalarType()->isIntegerTy(1) &&
+      areInverseVectorBitmasks(AC, BC)) {
+    AC = ConstantExpr::getTrunc(AC, CmpInst::makeCmpResultType(Ty));
+    return Builder.CreateXor(Cond, AC);
+  }
   return nullptr;
 }
 
 /// We have an expression of the form (A & C) | (B & D). Try to simplify this
 /// to "A' ? C : D", where A' is a boolean or vector of booleans.
 static Value *matchSelectFromAndOr(Value *A, Value *C, Value *B, Value *D,
                                    InstCombiner::BuilderTy &Builder) {
   // The potential condition of the select may be bitcasted. In that case, look
   // through its bitcast and the corresponding bitcast of the 'not' condition.
   Type *OrigType = A->getType();
   Value *SrcA, *SrcB;
   if (match(A, m_OneUse(m_BitCast(m_Value(SrcA)))) &&
       match(B, m_OneUse(m_BitCast(m_Value(SrcB))))) {
     A = SrcA;
     B = SrcB;
   }
 
-  if (Value *Cond = getSelectCondition(A, B)) {
+  if (Value *Cond = getSelectCondition(A, B, Builder)) {
     // ((bc Cond) & C) | ((bc ~Cond) & D) --> bc (select Cond, (bc C), (bc D))
     // The bitcasts will either all exist or all not exist. The builder will
     // not create unnecessary casts if the types already match.
     Value *BitcastC = Builder.CreateBitCast(C, A->getType());
     Value *BitcastD = Builder.CreateBitCast(D, A->getType());
     Value *Select = Builder.CreateSelect(Cond, BitcastC, BitcastD);
     return Builder.CreateBitCast(Select, OrigType);
   }

llvm/trunk/test/Transforms/InstCombine/logical-select.ll

;
   %bc1 = bitcast <4 x i1> %not to i4
   %bc2 = bitcast <4 x i1> %c to i4
   %and1 = and i4 %a, %bc1
   %and2 = and i4 %bc2, %b
   %or = or i4 %and1, %and2
   ret i4 %or
 }
 
-; FIXME: Missed conversions to select below here.
 ; Inverted 'and' constants mean this is a select.
 
 define <4 x i32> @vec_sel_consts(<4 x i32> %a, <4 x i32> %b) {
 ; CHECK-LABEL: @vec_sel_consts(
-; CHECK-NEXT:    [[AND1:%.*]] = and <4 x i32> %a, <i32 -1, i32 0, i32 0, i32 -1>
-; CHECK-NEXT:    [[AND2:%.*]] = and <4 x i32> %b, <i32 0, i32 -1, i32 -1, i32 0>
-; CHECK-NEXT:    [[OR:%.*]] = or <4 x i32> [[AND1]], [[AND2]]
-; CHECK-NEXT:    ret <4 x i32> [[OR]]
+; CHECK-NEXT:    [[TMP1:%.*]] = select <4 x i1> <i1 true, i1 false, i1 false, i1 true>, <4 x i32> %a, <4 x i32> %b
+; CHECK-NEXT:    ret <4 x i32> [[TMP1]]
 ;
   %and1 = and <4 x i32> %a, <i32 -1, i32 0, i32 0, i32 -1>
   %and2 = and <4 x i32> %b, <i32 0, i32 -1, i32 -1, i32 0>
   %or = or <4 x i32> %and1, %and2
   ret <4 x i32> %or
 }
 
 ; The select condition constant is always derived from the first operand of the 'or'.
 
 define <3 x i129> @vec_sel_consts_weird(<3 x i129> %a, <3 x i129> %b) {
 ; CHECK-LABEL: @vec_sel_consts_weird(
-; CHECK-NEXT:    [[AND1:%.*]] = and <3 x i129> %a, <i129 -1, i129 0, i129 -1>
-; CHECK-NEXT:    [[AND2:%.*]] = and <3 x i129> %b, <i129 0, i129 -1, i129 0>
-; CHECK-NEXT:    [[OR:%.*]] = or <3 x i129> [[AND2]], [[AND1]]
-; CHECK-NEXT:    ret <3 x i129> [[OR]]
+; CHECK-NEXT:    [[TMP1:%.*]] = select <3 x i1> <i1 false, i1 true, i1 false>, <3 x i129> %b, <3 x i129> %a
+; CHECK-NEXT:    ret <3 x i129> [[TMP1]]
 ;
   %and1 = and <3 x i129> %a, <i129 -1, i129 0, i129 -1>
   %and2 = and <3 x i129> %b, <i129 0, i129 -1, i129 0>
   %or = or <3 x i129> %and2, %and1
   ret <3 x i129> %or
 }
 
 ; The mask elements must be inverted for this to be a select.
;
   %or = or <4 x i32> %and1, %and2
   ret <4 x i32> %or
 }
 
 ; The inverted constants may be operands of xor instructions.
 
 define <4 x i32> @vec_sel_xor(<4 x i32> %a, <4 x i32> %b, <4 x i1> %c) {
 ; CHECK-LABEL: @vec_sel_xor(
+; CHECK-NEXT:    [[TMP1:%.*]] = xor <4 x i1> %c, <i1 false, i1 true, i1 true, i1 true>
+; CHECK-NEXT:    [[TMP2:%.*]] = select <4 x i1> [[TMP1]], <4 x i32> %a, <4 x i32> %b
+; CHECK-NEXT:    ret <4 x i32> [[TMP2]]
+;
+  %mask = sext <4 x i1> %c to <4 x i32>
+  %mask_flip1 = xor <4 x i32> %mask, <i32 -1, i32 0, i32 0, i32 0>
+  %not_mask_flip1 = xor <4 x i32> %mask, <i32 0, i32 -1, i32 -1, i32 -1>
+  %and1 = and <4 x i32> %not_mask_flip1, %a
+  %and2 = and <4 x i32> %mask_flip1, %b
+  %or = or <4 x i32> %and1, %and2
+  ret <4 x i32> %or
+}
+
+; Allow the transform even if the mask values have multiple uses because
+; there's still a net reduction of instructions from removing the and/and/or.
+
+define <4 x i32> @vec_sel_xor_multi_use(<4 x i32> %a, <4 x i32> %b, <4 x i1> %c) {
+; CHECK-LABEL: @vec_sel_xor_multi_use(
 ; CHECK-NEXT:    [[MASK:%.*]] = sext <4 x i1> %c to <4 x i32>
 ; CHECK-NEXT:    [[MASK_FLIP1:%.*]] = xor <4 x i32> [[MASK]], <i32 -1, i32 0, i32 0, i32 0>
-; CHECK-NEXT:    [[NOT_MASK_FLIP1:%.*]] = xor <4 x i32> [[MASK]], <i32 0, i32 -1, i32 -1, i32 -1>
-; CHECK-NEXT:    [[AND1:%.*]] = and <4 x i32> [[NOT_MASK_FLIP1]], %a
-; CHECK-NEXT:    [[AND2:%.*]] = and <4 x i32> [[MASK_FLIP1]], %b
-; CHECK-NEXT:    [[OR:%.*]] = or <4 x i32> [[AND1]], [[AND2]]
-; CHECK-NEXT:    ret <4 x i32> [[OR]]
+; CHECK-NEXT:    [[TMP1:%.*]] = xor <4 x i1> %c, <i1 false, i1 true, i1 true, i1 true>
+; CHECK-NEXT:    [[TMP2:%.*]] = select <4 x i1> [[TMP1]], <4 x i32> %a, <4 x i32> %b
+; CHECK-NEXT:    [[ADD:%.*]] = add <4 x i32> [[TMP2]], [[MASK_FLIP1]]
+; CHECK-NEXT:    ret <4 x i32> [[ADD]]
 ;
   %mask = sext <4 x i1> %c to <4 x i32>
   %mask_flip1 = xor <4 x i32> %mask, <i32 -1, i32 0, i32 0, i32 0>
   %not_mask_flip1 = xor <4 x i32> %mask, <i32 0, i32 -1, i32 -1, i32 -1>
   %and1 = and <4 x i32> %not_mask_flip1, %a
   %and2 = and <4 x i32> %mask_flip1, %b
   %or = or <4 x i32> %and1, %and2
-  ret <4 x i32> %or
+  %add = add <4 x i32> %or, %mask_flip1
+  ret <4 x i32> %add
 }