[InstCombine] Fix scalable-vector bitwise select matching

[frasercrmck]

Summary

D113035 enhanced the matching of bitwise selects from vector types. This
change unfortunately introduced crashes as it tries to cast scalable
vector types to integers.

[spatel]

Do we have a test for scalable vectors where the remaining part of the transform is valid?

[spatel]

In D124997#3493929, @spatel wrote:

Do we have a test for scalable vectors where the remaining part of the transform is valid?

Ah, the 1st test covers that. But then couldn't we just make a one-line change to avoid the problematic part?

if (auto *VecTy = dyn_cast<FixedVectorType>(Cond->getType())) {

[frasercrmck]

In D124997#3493954, @spatel wrote:

In D124997#3493929, @spatel wrote:

Do we have a test for scalable vectors where the remaining part of the transform is valid?

Ah, the 1st test covers that. But then couldn't we just make a one-line change to avoid the problematic part?

if (auto *VecTy = dyn_cast<FixedVectorType>(Cond->getType())) {

Aye, that's what the 1st test is for.

I was under the impression that if Cond->getType() is any vector type we must either be able to do the bitcast or bail. So wouldn't we just need the other checks I've added in an else?

[spatel]

In D124997#3493988, @frasercrmck wrote:

I was under the impression that if Cond->getType() is any vector type we must either be able to do the bitcast or bail. So wouldn't we just need the other checks I've added in an else?

We can still have a vector bitcast where the cast is to a type that matches the number of elements in the vector condition. That's the 2nd test IIUC, so if we just do the dyn_cast, it becomes:

define <vscale x 1 x i64> @vec_of_casted_bools_scalable(<vscale x 1 x i64> %a, <vscale x 1 x i64> %b, <vscale x 8 x i1> %cond) {
  %1 = bitcast <vscale x 1 x i64> %a to <vscale x 8 x i8>
  %2 = bitcast <vscale x 1 x i64> %b to <vscale x 8 x i8>
  %3 = select <vscale x 8 x i1> %cond, <vscale x 8 x i8> %1, <vscale x 8 x i8> %2
  %4 = bitcast <vscale x 8 x i8> %3 to <vscale x 1 x i64>
  ret <vscale x 1 x i64> %4
}

[frasercrmck]

address comments

[frasercrmck]

In D124997#3494017, @spatel wrote:

We can still have a vector bitcast where the cast is to a type that matches the number of elements in the vector condition. That's the 2nd test IIUC, so if we just do the dyn_cast, it becomes:

define <vscale x 1 x i64> @vec_of_casted_bools_scalable(<vscale x 1 x i64> %a, <vscale x 1 x i64> %b, <vscale x 8 x i1> %cond) {
  %1 = bitcast <vscale x 1 x i64> %a to <vscale x 8 x i8>
  %2 = bitcast <vscale x 1 x i64> %b to <vscale x 8 x i8>
  %3 = select <vscale x 8 x i1> %cond, <vscale x 8 x i8> %1, <vscale x 8 x i8> %2
  %4 = bitcast <vscale x 8 x i8> %3 to <vscale x 1 x i64>
  ret <vscale x 1 x i64> %4
}

Oh, right you are! Sorry, I had the wrong impression of how this worked. Thank you.

[spatel]

In D124997#3494120, @frasercrmck wrote:

Oh, right you are! Sorry, I had the wrong impression of how this worked. Thank you.

The code comments could be improved. There are many different potential patterns within this block, and scalable vectors just make it harder to keep it all straight. :)
But after mucking around in here for a long time, I don't think the fix is sufficient.
I added a test with 7bad1d281c798929a to try to uncover another potential bug.
But I can't find a case currently where we would go wrong because code before here (ComputeNumSignBits) prevents matching a scalable vector with the right combination of bitcasts to trigger that bug.

A better fix should do something like this:

// If this is a vector, we may need to cast to match the condition's length.
Type *SelTy = A->getType();
if (auto *VecTy = dyn_cast<VectorType>(Cond->getType())) {
  // For a fixed or scalable vector get N from <{vscale x} N x iM>
  unsigned Elts = VecTy->getElementCount().getKnownMinValue();
  // For a fixed or scalable vector, get the value N x iM; for a scalar this is just M.
  unsigned SelEltSize = SelTy->getPrimitiveSizeInBits().getKnownMinSize();
  Type *EltTy = Builder.getIntNTy(SelEltSize / Elts);
  SelTy = VectorType::get(EltTy, VecTy->getElementCount());
}

I had a hard time making sense of the type and size APIs, so if anyone knows that better, please correct/improve if possible.

[frasercrmck]

fix and enable bitcasts for scalable vectors

[frasercrmck]

In D124997#3494745, @spatel wrote:

The code comments could be improved. There are many different potential patterns within this block, and scalable vectors just make it harder to keep it all straight. :)
But after mucking around in here for a long time, I don't think the fix is sufficient.
I added a test with 7bad1d281c798929a to try to uncover another potential bug.
But I can't find a case currently where we would go wrong because code before here (ComputeNumSignBits) prevents matching a scalable vector with the right combination of bitcasts to trigger that bug.

A better fix should do something like this:

// If this is a vector, we may need to cast to match the condition's length.
Type *SelTy = A->getType();
if (auto *VecTy = dyn_cast<VectorType>(Cond->getType())) {
  // For a fixed or scalable vector get N from <{vscale x} N x iM>
  unsigned Elts = VecTy->getElementCount().getKnownMinValue();
  // For a fixed or scalable vector, get the value N x iM; for a scalar this is just M.
  unsigned SelEltSize = SelTy->getPrimitiveSizeInBits().getKnownMinSize();
  Type *EltTy = Builder.getIntNTy(SelEltSize / Elts);
  SelTy = VectorType::get(EltTy, VecTy->getElementCount());
}

I had a hard time making sense of the type and size APIs, so if anyone knows that better, please correct/improve if possible.

Thanks for digging in. I also found it very difficult to get scalable-vector tests which would actually trigger all possible cases we're trying to handle.

I think what you've got looks okay. I adjusted the comments a little.

[spatel]

LGTM

llvm/lib/Transforms/InstCombine/InstCombineAndOrXor.cpp

Value *InstCombinerImpl::matchSelectFromAndOr(Value *A, Value *C, Value *B,
   B = peekThroughBitcast(B, true);
   if (Value *Cond = getSelectCondition(A, B)) {
     // ((bc Cond) & C) | ((bc ~Cond) & D) --> bc (select Cond, (bc C), (bc D))
     // If this is a vector, we may need to cast to match the condition's length.
     // The bitcasts will either all exist or all not exist. The builder will
     // not create unnecessary casts if the types already match.
     Type *SelTy = A->getType();
     if (auto *VecTy = dyn_cast<VectorType>(Cond->getType())) {
+      // For a fixed or scalable vector get N from <{vscale x} N x iM>
       unsigned Elts = VecTy->getElementCount().getKnownMinValue();
-      Type *EltTy = Builder.getIntNTy(SelTy->getPrimitiveSizeInBits() / Elts);
+      // For a fixed or scalable vector, get the size in bits of N x iM; for a
+      // scalar this is just M.
+      unsigned SelEltSize = SelTy->getPrimitiveSizeInBits().getKnownMinSize();
+      Type *EltTy = Builder.getIntNTy(SelEltSize / Elts);
       SelTy = VectorType::get(EltTy, VecTy->getElementCount());
     }
     Value *BitcastC = Builder.CreateBitCast(C, SelTy);
     Value *BitcastD = Builder.CreateBitCast(D, SelTy);
     Value *Select = Builder.CreateSelect(Cond, BitcastC, BitcastD);
     return Builder.CreateBitCast(Select, OrigType);
   }
 

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

 ;
   %not = xor <4 x i1> %c, <i1 true, i1 true, i1 true, i1 true>
   %and1 = and <4 x i1> %not, %a
   %and2 = and <4 x i1> %b, %c
   %or = or <4 x i1> %and2, %and1
   ret <4 x i1> %or
 }
 
+define <vscale x 1 x i1> @vec_of_bools_scalable(<vscale x 1 x i1> %a, <vscale x 1 x i1> %c, <vscale x 1 x i1> %d) {
+; CHECK-LABEL: @vec_of_bools_scalable(
+; CHECK-NEXT:    [[TMP1:%.*]] = select <vscale x 1 x i1> [[A:%.*]], <vscale x 1 x i1> [[C:%.*]], <vscale x 1 x i1> [[D:%.*]]
+; CHECK-NEXT:    ret <vscale x 1 x i1> [[TMP1]]
+;
+  %b = xor <vscale x 1 x i1> %a, shufflevector (<vscale x 1 x i1> insertelement (<vscale x 1 x i1> poison, i1 true, i32 0), <vscale x 1 x i1> poison, <vscale x 1 x i32> zeroinitializer)
+  %t11 = and <vscale x 1 x i1> %a, %c
+  %t12 = and <vscale x 1 x i1> %b, %d
+  %r = or <vscale x 1 x i1> %t11, %t12
+  ret <vscale x 1 x i1> %r
+}
+
 define i4 @vec_of_casted_bools(i4 %a, i4 %b, <4 x i1> %c) {
 ; CHECK-LABEL: @vec_of_casted_bools(
 ; CHECK-NEXT:    [[TMP1:%.*]] = bitcast i4 [[B:%.*]] to <4 x i1>
 ; CHECK-NEXT:    [[TMP2:%.*]] = bitcast i4 [[A:%.*]] to <4 x i1>
 ; CHECK-NEXT:    [[TMP3:%.*]] = select <4 x i1> [[C:%.*]], <4 x i1> [[TMP1]], <4 x i1> [[TMP2]]
 ; CHECK-NEXT:    [[TMP4:%.*]] = bitcast <4 x i1> [[TMP3]] to i4
 ; CHECK-NEXT:    ret i4 [[TMP4]]
 ;
   %not = xor <4 x i1> %c, <i1 true, i1 true, i1 true, i1 true>
   %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
 }
 
+define <vscale x 1 x i64> @vec_of_casted_bools_scalable(<vscale x 1 x i64> %a, <vscale x 1 x i64> %b, <vscale x 8 x i1> %cond) {
+; CHECK-LABEL: @vec_of_casted_bools_scalable(
+; CHECK-NEXT:    [[TMP1:%.*]] = bitcast <vscale x 1 x i64> [[A:%.*]] to <vscale x 8 x i8>
+; CHECK-NEXT:    [[TMP2:%.*]] = bitcast <vscale x 1 x i64> [[B:%.*]] to <vscale x 8 x i8>
+; CHECK-NEXT:    [[TMP3:%.*]] = select <vscale x 8 x i1> [[COND:%.*]], <vscale x 8 x i8> [[TMP1]], <vscale x 8 x i8> [[TMP2]]
+; CHECK-NEXT:    [[TMP4:%.*]] = bitcast <vscale x 8 x i8> [[TMP3]] to <vscale x 1 x i64>
+; CHECK-NEXT:    ret <vscale x 1 x i64> [[TMP4]]
+;
+  %scond = sext <vscale x 8 x i1> %cond to <vscale x 8 x i8>
+  %notcond = xor <vscale x 8 x i1> %cond, shufflevector (<vscale x 8 x i1> insertelement (<vscale x 8 x i1> poison, i1 true, i32 0), <vscale x 8 x i1> poison, <vscale x 8 x i32> zeroinitializer)
+  %snotcond = sext <vscale x 8 x i1> %notcond to <vscale x 8 x i8>
+  %bc1 = bitcast <vscale x 8 x i8> %scond to <vscale x 1 x i64>
+  %bc2 = bitcast <vscale x 8 x i8> %snotcond to <vscale x 1 x i64>
+  %and1 = and <vscale x 1 x i64> %a, %bc1
+  %and2 = and <vscale x 1 x i64> %bc2, %b
+  %or = or <vscale x 1 x i64> %and1, %and2
+  ret <vscale x 1 x i64> %or
+}
+
 ; Inverted 'and' constants mean this is a select which is canonicalized to a shuffle.
 
 define <4 x i32> @vec_sel_consts(<4 x i32> %a, <4 x i32> %b) {
 ; CHECK-LABEL: @vec_sel_consts(
 ; CHECK-NEXT:    [[TMP1:%.*]] = shufflevector <4 x i32> [[A:%.*]], <4 x i32> [[B:%.*]], <4 x i32> <i32 0, i32 5, i32 6, i32 3>
 ; CHECK-NEXT:    ret <4 x i32> [[TMP1]]
 ;
   %and1 = and <4 x i32> %a, <i32 -1, i32 0, i32 0, i32 -1>