[X86][DAGISel] Don't widen shuffle element with AVX512

[LuoYuanke]

Summary

Currently the X86 shuffle lowering would widen the element type for
shuffle if the mask element value is adjacent. For below example

  %t2 = add nsw <16 x i32> %t0, %t1
  %t3 = sub nsw <16 x i32> %t0, %t1
  %t4 = shufflevector <16 x i32> %t2, <16 x i32> %t3,
                      <16 x i32> <i32 16, i32 17, i32 2, i32 3, i32 4,
                       i32 5, i32 6, i32 7, i32 8, i32 9, i32 10,
                       i32 11, i32 12, i32 13, i32 14, i32 15>

  ret <16 x i32> %t4

Compiler would transform the shuffle to
  %t4 = shufflevector <8 x i64> %t2, <8 x i64> %t3,
                      <8 x i64> <i32 8, i32 1, i32 2, i32 3, i32 4,
                                 i32 5, i32 6, i32 7>
This may lose the oppotunity to let ISel select mask instruction when
avx512 is enabled.

This patch is to prevent the tranform when avx512 feature is enabled.

[LuoYuanke]

Rebase.

[RKSimon]

llvm/lib/Target/X86/X86ISelLowering.cpp

44372

This doesn't look right - shouldn't it be something like:

APInt Mask = APIntOps::ScaleBitMask(ConstCond->getAPIntValue(), NumElts * 2) ?

[LuoYuanke]

llvm/lib/Target/X86/X86ISelLowering.cpp

44372

Thanks, Simon. ScaleBitMask perfectly fit this coputation.

[LuoYuanke]

Address Simon's comments.

[RKSimon]

I think we can probably generalize this very easily to any legal widening, e.g. to handle: https://gcc.godbolt.org/z/Pjz5qfYT7 (v32i16 -> v16i32)

[RKSimon]

llvm/lib/Target/X86/X86ISelLowering.cpp

44346

Check if N->getOpcode() == ISD::VSELECT

44350

auto *ConstCond = dyn_cast<ConstantSDNode>(Cond.getOperand(0));
if (!ConstCond)
  return SDValue();

[LuoYuanke]

llvm/lib/Target/X86/X86ISelLowering.cpp

44346

Is it possible that CondVT is not vXi1 for ISD::VSELECT? I ask the question becasue the comments for ISD::VSELECT says "targets may change the condition type".

/// At first, the VSELECT condition is of vXi1 type. Later, targets may
/// change the condition type in order to match the VSELECT node using a
/// pattern. The condition follows the BooleanContent format of the target.

[LuoYuanke]

Address Simon's comments.

[LuoYuanke]

In D129537#3659947, @RKSimon wrote:

I think we can probably generalize this very easily to any legal widening, e.g. to handle: https://gcc.godbolt.org/z/Pjz5qfYT7 (v32i16 -> v16i32)

Good suggestion. Let me take a look at it.

[LuoYuanke]

Address Simon's address to generalize the blend/select combine.

[LuoYuanke]

llvm/test/CodeGen/X86/avx512-shuffles/shuffle-blend.ll

60

For this case not sure if it is worse than left side code (previous code).

[RKSimon]

I've added some additional test coverage to shuffle-blend.ll - please can you rebase?

llvm/test/CodeGen/X86/avx512-shuffles/shuffle-blend.ll

60

Add a 128-bit vector limit?

[RKSimon]

Please can you update the patch title/summary?

[LuoYuanke]

Rebase and update test case.

[LuoYuanke]

llvm/test/CodeGen/X86/avx512-shuffles/shuffle-blend.ll

60

I'll add 128-bit vector limit. However in this case it is 128-bit vector.

[LuoYuanke]

Limit the vector bit width >=128 and add test cases.

[RKSimon]

llvm/test/CodeGen/X86/avx512-shuffles/shuffle-blend.ll

2

CHECK,AVX512BW,X86-AVX512BW
CHECK,AVX512BW,X64-AVX512BW

[LuoYuanke]

Address Simon's comments.

[LuoYuanke]

llvm/test/CodeGen/X86/avx512-shuffles/shuffle-blend.ll

68

retl can be merged to retq by ret{{[l|q]}}. Not sure why utils/update_llc_test_checks.py doesn't merge.

[RKSimon]

llvm/test/CodeGen/X86/avx512-shuffles/shuffle-blend.ll

75

regression? you might need to improve the 128-bit limit logic to account for vXi16 specifically

[LuoYuanke]

llvm/test/CodeGen/X86/avx512-shuffles/shuffle-blend.ll

75

There is PBLENDW, but there is no PBLENDB instruction, so it is better to widen to v2Xi16 from vXi8. Besides there is more instruction for 16-bit element (e.g., movsh). I'll investigate more on this issue.

[LuoYuanke]

Specially handling for vXi8, because vXi16 can be applied PBLENDW while vXi8 can't.

[RKSimon]

llvm/test/CodeGen/X86/avx512-shuffles/shuffle-blend.ll

68

"kmovd %eax" vs "kmovq %rax"

74

pre-commit these additional tests

[LuoYuanke]

llvm/test/CodeGen/X86/avx512-shuffles/shuffle-blend.ll

68

Got it. :)

74

Sure. I'll do it.

[LuoYuanke]

Rebase

[RKSimon]

llvm/lib/Target/X86/X86ISelLowering.cpp

19326

This list is going to get longer, and we're likely to miss patterns that only fold to target nodes later on - I'm wondering whether we could consider accepting any TLI.isBinOp() case here?

[RKSimon]

LGTM - with one minor comment for future work

llvm/lib/Target/X86/X86ISelLowering.cpp

19326

Please can you add a TODO about maybe converting this to TLI.isBinOp()?

[LuoYuanke]

Address Simon's comments.

[LuoYuanke]

llvm/test/CodeGen/X86/haddsub-undef.ll

1053 ↗

(On Diff #447318)

This looks a regression, I'll take a look at it.

[LuoYuanke]

llvm/lib/Target/X86/X86ISelLowering.cpp

19326

Thank Simon for the suggestion. It seems there is regression on some cases, I'll take a look at the regression.

[RKSimon]

Yeah - that was what I saw as well - if you want to get this in for 15.x I'd recommend going back to the old switch statement - then investigate the binop general case later (if you solve it soon enough you request a merge)

[LuoYuanke]

Revert to previous version and add TODO for checking TLI.isBinOp().

[LuoYuanke]

In D129537#3676257, @RKSimon wrote:

Yeah - that was what I saw as well - if you want to get this in for 15.x I'd recommend going back to the old switch statement - then investigate the binop general case later (if you solve it soon enough you request a merge)

Thank Simon for the review. All the comments are valuable. Let me land the old switch statement version first and then investigate the general binop.

[fhahn]

This patch unfortunately causes crashes when building llvm-test-suite optimizing for AVX512.

Reproducer for llc:

target datalayout = "e-m:o-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
target triple = "x86_64-apple-macosx"

define i32 @test(<32 x i32> %0) #0 {
entry:
  %1 = mul <32 x i32> %0, <i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1, i32 1>
  %2 = tail call i32 @llvm.vector.reduce.add.v32i32(<32 x i32> %1)
  ret i32 %2
}

; Function Attrs: nocallback nofree nosync nounwind readnone willreturn
declare i32 @llvm.vector.reduce.add.v32i32(<32 x i32>) #1

attributes #0 = { "min-legal-vector-width"="0" "target-cpu"="skylake-avx512" }
attributes #1 = { nocallback nofree nosync nounwind readnone willreturn }

I've reverted the patch in the meantime to get current main back into a good state.

llvm/lib/Target/X86/X86ISelLowering.cpp

     for (int &M : NewMask)
       if (M >= NumElements)
         M = -1;
     return DAG.getVectorShuffle(VT, DL, V1, V2, NewMask);
   }
 
   // Check for illegal shuffle mask element index values.
   int MaskUpperLimit = OrigMask.size() * (V2IsUndef ? 1 : 2);
   (void)MaskUpperLimit;
   assert(llvm::all_of(OrigMask,
                       [&](int M) { return -1 <= M && M < MaskUpperLimit; }) &&
          "Out of bounds shuffle index");
 
   // We actually see shuffles that are entirely re-arrangements of a set of
   // zero inputs. This mostly happens while decomposing complex shuffles into
   // simple ones. Directly lower these as a buildvector of zeros.
   APInt KnownUndef, KnownZero;

   // sub accomplishes the negation of the replacement pattern.
   if (V == Y)
     std::swap(SubOp1, SubOp2);
 
   SDValue Res = DAG.getNode(ISD::SUB, DL, MaskVT, SubOp1, SubOp2);
   return DAG.getBitcast(VT, Res);
 }
 
+static SDValue combineSelectVxi64(SDNode *N, SelectionDAG &DAG,
+                                  const X86Subtarget &Subtarget) {
+  SDLoc DL(N);
+  SDValue Cond = N->getOperand(0);
+  SDValue LHS = N->getOperand(1);
+  SDValue RHS = N->getOperand(2);
+  EVT VT = LHS.getValueType();
+  EVT CondVT = Cond.getValueType();
+  // Combine
+  // select vXi1 bitcast (int cond),
+  //         (<vXi64> bitcast <v2Xi32> a),
+  //         (<vXi64> bitcast <v2Xi32> b)
+  // to
+  // select <v2Xi1> cond, <v2Xi32> a, <v2Xi32> b
+  // to create opportunity for mask instructions with AVX512 instructions.
+  if (!Subtarget.hasAVX512())
+    return SDValue();
+
+  if (!CondVT.isVector() || CondVT.getVectorElementType() != MVT::i1)
+    return SDValue();
+  if (Cond.getOpcode() != ISD::BITCAST)
+    return SDValue();
+  if (!dyn_cast<ConstantSDNode>(Cond.getOperand(0)))
+    return SDValue();
+
+  if (VT.getVectorElementType() != MVT::i64)
+    return SDValue();
+
+  if (LHS.getOpcode() != ISD::BITCAST ||
+      LHS.getOperand(0).getValueType().getVectorElementType() != MVT::i32)
+    return SDValue();
+  if (RHS.getOpcode() != ISD::BITCAST ||
+      RHS.getOperand(0).getValueType().getVectorElementType() != MVT::i32)
+    return SDValue();
+
+  if (!Cond.hasOneUse() || !LHS.hasOneUse() || !RHS.hasOneUse())
+    return SDValue();
+
+  int NumElts = VT.getVectorNumElements();
+  EVT ExpandCondVT = EVT::getVectorVT(*DAG.getContext(), MVT::i1, NumElts * 2);
+  EVT ExpandVT = EVT::getVectorVT(*DAG.getContext(), MVT::i32, NumElts * 2);
+
+  ConstantSDNode *ConstCond = cast<ConstantSDNode>(Cond.getOperand(0));
+  uint64_t Mask = ConstCond->getZExtValue();
+  Mask = (Mask << 1) | Mask;
+  SDValue MaskVal = DAG.getConstant(
+      Mask, DL, EVT::getIntegerVT(*DAG.getContext(), NumElts * 2));
+  SDValue NewCond = DAG.getBitcast(ExpandCondVT, MaskVal);
+  return DAG.getBitcast(VT,
+                        DAG.getSelect(DL, ExpandVT, NewCond, LHS.getOperand(0),
+                                      RHS.getOperand(0)));
+}
+
 /// Do target-specific dag combines on SELECT and VSELECT nodes.
 static SDValue combineSelect(SDNode *N, SelectionDAG &DAG,
                              TargetLowering::DAGCombinerInfo &DCI,
                              const X86Subtarget &Subtarget) {
   SDLoc DL(N);
   SDValue Cond = N->getOperand(0);
   SDValue LHS = N->getOperand(1);
   SDValue RHS = N->getOperand(2);

         break;
       }
     }
 
     if (Opcode)
       return DAG.getNode(Opcode, DL, N->getValueType(0), LHS, RHS);
   }
 
+  if (SDValue V = combineSelectVxi64(N, DAG, Subtarget))
+    return V;
+
   // Some mask scalar intrinsics rely on checking if only one bit is set
   // and implement it in C code like this:
   // A[0] = (U & 1) ? A[0] : W[0];
   // This creates some redundant instructions that break pattern matching.
   // fold (select (setcc (and (X, 1), 0, seteq), Y, Z)) -> select(and(X, 1),Z,Y)
   if (Subtarget.hasAVX512() && N->getOpcode() == ISD::SELECT &&
       Cond.getOpcode() == ISD::SETCC && (VT == MVT::f32 || VT == MVT::f64)) {
     ISD::CondCode CC = cast<CondCodeSDNode>(Cond.getOperand(2))->get();

llvm/test/CodeGen/X86/avx512-shuffles/shuffle-blend.ll

 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
 ; RUN: llc -mtriple=x86_64-unknown-linux-gnu -mattr=+avx512f,+avx512vl,+avx512bw %s -o - | FileCheck %s
 
 define <16 x i32> @shuffle_v8i64(<16 x i32> %t0, <16 x i32> %t1) {
 ; CHECK-LABEL: shuffle_v8i64:
 ; CHECK:       # %bb.0: # %entry
 ; CHECK-NEXT:    vpaddd %zmm1, %zmm0, %zmm2
-; CHECK-NEXT:    vpsubd %zmm1, %zmm0, %zmm0
-; CHECK-NEXT:    movb $-86, %al
+; CHECK-NEXT:    movw $510, %ax # imm = 0x1FE
 ; CHECK-NEXT:    kmovd %eax, %k1
-; CHECK-NEXT:    vmovdqa64 %zmm0, %zmm2 {%k1}
+; CHECK-NEXT:    vpsubd %zmm1, %zmm0, %zmm2 {%k1}
 ; CHECK-NEXT:    vmovdqa64 %zmm2, %zmm0
 ; CHECK-NEXT:    retq
 entry:
   %t2 = add nsw <16 x i32> %t0, %t1
   %t3 = sub nsw <16 x i32> %t0, %t1
   %t4 = shufflevector <16 x i32> %t2, <16 x i32> %t3, <16 x i32> <i32 0, i32 1, i32 18, i32 19, i32 4, i32 5, i32 22, i32 23, i32 8, i32 9, i32 26, i32 27, i32 12, i32 13, i32 30, i32 31>
   ret <16 x i32> %t4
 }

 ; CHECK-NEXT:    vpsubd %xmm1, %xmm0, %xmm0
 ; CHECK-NEXT:    vpblendd {{.*#+}} xmm0 = xmm2[0],xmm0[1],xmm2[2,3]
 ; CHECK-NEXT:    retq
 entry:
   %t2 = add nsw <2 x i32> %t0, %t1
   %t3 = sub nsw <2 x i32> %t0, %t1
   %t4 = shufflevector <2 x i32> %t2, <2 x i32> %t3, <2 x i32> <i32 0, i32 3>
   ret <2 x i32> %t4
 }