[DAG] Conservatively check the VSELECT Operation Action in tryToFoldExtendSelectLoad

[xiangzhangllvm]

Summary

tryToFoldExtendSelectLoad can be called in DAG combine2 which has already finish legalization
So we should conservatively check the Operation Action for new VSELECT.

[LuoYuanke]

llvm/lib/Target/X86/X86ISelDAGToDAG.cpp

1043

vXi16?

llvm/test/CodeGen/X86/vselect-post-combine.ll

5

Can we simplify the variable name?

[RKSimon]

llvm/test/CodeGen/X86/vselect-post-combine.ll

27

You might be able to further simplify this test?

[LuoYuanke]

llvm/test/CodeGen/X86/select-ext.ll

97 ↗

(On Diff #507607)

It seems generate more code. Probably you need check "Level <= AfterLegalizeVectorOps" in tryToFoldExtendSelectLoad(), so that there is more combine opportunity.

[LuoYuanke]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

12246

Check "Level <= AfterLegalizeVectorOps" instead?

[LuoYuanke]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

12246

Maybe changed as below.

if ((N0->getOpcode() == ISD::VSELECT && Level > AfterLegalizeVectorOps &&
     TLI.getOperationAction(ISD::VSELECT, VT) != TargetLowering::Legal
     ))
  return SDValue();

[LuoYuanke]

llvm/lib/Target/X86/X86ISelDAGToDAG.cpp

1043

On the other hand, why can't use VPBLENDVB ymm1, ymm2, ymm3/m256 (https://www.felixcloutier.com/x86/pblendvb) to vselect v16i16? We can bitcast v16i16 to v32i8.

[pengfei]

llvm/test/CodeGen/X86/vselect-post-combine.ll

3

This problem is not limited to Windows. Linux has the some problem if passing a ptr and load it to <32 x i8>.

[xiangzhangllvm]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

12246

Yes, I also had a quick think about it, just need to pass Level here.

llvm/lib/Target/X86/X86ISelDAGToDAG.cpp

1043

Yes, but not only bitcast, it also need to bring new instructions to transfer value of ymm1.

llvm/test/CodeGen/X86/vselect-post-combine.ll

3

Yes, I tended to rm windows then I found it is not triggered, duo to the previous optimization difference.

5

Sure

27

Let me have a try, in fact I spend a lot time to simplify the raw test to current status : )

[pengfei]

llvm/test/CodeGen/X86/vselect-post-combine.ll

3

The difference first comes from ABI, so you can remove it by

define ptr @test_mul(ptr %0) {
%vecinit.i.i.i.i.i92 = load <32 x i8>, ptr %0

[xiangzhangllvm]

llvm/test/CodeGen/X86/vselect-post-combine.ll

3

Sorry, there is a delay for me to see this comment, I have update it : ), the test options is mainly to trigger the combine, now we can trigger it. Can we go on now?

[LuoYuanke]

llvm/test/CodeGen/X86/vselect-post-combine.ll

19

Could you simplify %shuffle.i.i.i6.i.i as well?

[LuoYuanke]

LGMT with NIT comments.

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

 
 /// Fold
 ///   (sext (select c, load x, load y)) -> (select c, sextload x, sextload y)
 ///   (zext (select c, load x, load y)) -> (select c, zextload x, zextload y)
 ///   (aext (select c, load x, load y)) -> (select c, extload x, extload y)
 /// This function is called by the DAGCombiner when visiting sext/zext/aext
 /// dag nodes (see for example method DAGCombiner::visitSIGN_EXTEND).
 static SDValue tryToFoldExtendSelectLoad(SDNode *N, const TargetLowering &TLI,
-                                         SelectionDAG &DAG) {
+                                         SelectionDAG &DAG,
+                                         CombineLevel Level) {
   unsigned Opcode = N->getOpcode();
   SDValue N0 = N->getOperand(0);
   EVT VT = N->getValueType(0);
   SDLoc DL(N);
 
   assert((Opcode == ISD::SIGN_EXTEND || Opcode == ISD::ZERO_EXTEND ||
           Opcode == ISD::ANY_EXTEND) &&
          "Expected EXTEND dag node in input!");
 
   if (!(N0->getOpcode() == ISD::SELECT || N0->getOpcode() == ISD::VSELECT) ||
       !N0.hasOneUse())
     return SDValue();
 
   SDValue Op1 = N0->getOperand(1);
   SDValue Op2 = N0->getOperand(2);
   if (!isCompatibleLoad(Op1, Opcode) || !isCompatibleLoad(Op2, Opcode))
     return SDValue();
 
   auto ExtLoadOpcode = ISD::EXTLOAD;
   if (Opcode == ISD::SIGN_EXTEND)
     ExtLoadOpcode = ISD::SEXTLOAD;
   else if (Opcode == ISD::ZERO_EXTEND)
     ExtLoadOpcode = ISD::ZEXTLOAD;
 
+  // Illegal VSELECT may ISel fail if happen after legalization (DAG
+  // Combine2), so we should conservatively check the OperationAction.
   LoadSDNode *Load1 = cast<LoadSDNode>(Op1);
   LoadSDNode *Load2 = cast<LoadSDNode>(Op2);
   if (!TLI.isLoadExtLegal(ExtLoadOpcode, VT, Load1->getMemoryVT()) ||
-      !TLI.isLoadExtLegal(ExtLoadOpcode, VT, Load2->getMemoryVT()))
+      !TLI.isLoadExtLegal(ExtLoadOpcode, VT, Load2->getMemoryVT()) ||
+      (N0->getOpcode() == ISD::VSELECT && Level >= AfterLegalizeTypes &&
+       TLI.getOperationAction(ISD::VSELECT, VT) != TargetLowering::Legal))
     return SDValue();
 
   SDValue Ext1 = DAG.getNode(Opcode, DL, VT, Op1);
   SDValue Ext2 = DAG.getNode(Opcode, DL, VT, Op2);
   return DAG.getSelect(DL, VT, N0->getOperand(0), Ext1, Ext2);
 }
 
 /// Try to fold a sext/zext/aext dag node into a ConstantSDNode or
if (SDValue NewXor = visitXOR(N0.getNode())) {
       // Return a new sext with the new xor.
       return DAG.getNode(ISD::SIGN_EXTEND, DL, VT, NewXor);
     }
 
     SDValue Zext = DAG.getNode(ISD::ZERO_EXTEND, DL, VT, N0.getOperand(0));
     return DAG.getNode(ISD::ADD, DL, VT, Zext, DAG.getAllOnesConstant(DL, VT));
   }
 
-  if (SDValue Res = tryToFoldExtendSelectLoad(N, TLI, DAG))
+  if (SDValue Res = tryToFoldExtendSelectLoad(N, TLI, DAG, Level))
     return Res;
 
   return SDValue();
 }
 
 // isTruncateOf - If N is a truncate of some other value, return true, record
 // the value being truncated in Op and which of Op's bits are zero/one in Known.
 // This function computes KnownBits to avoid a duplicated call to
if (SDValue NewVSel = matchVSelectOpSizesWithSetCC(N))
     return NewVSel;
 
   if (SDValue NewCtPop = widenCtPop(N, DAG))
     return NewCtPop;
 
   if (SDValue V = widenAbs(N, DAG))
     return V;
 
-  if (SDValue Res = tryToFoldExtendSelectLoad(N, TLI, DAG))
+  if (SDValue Res = tryToFoldExtendSelectLoad(N, TLI, DAG, Level))
     return Res;
 
   return SDValue();
 }
 
 SDValue DAGCombiner::visitANY_EXTEND(SDNode *N) {
   SDValue N0 = N->getOperand(0);
   EVT VT = N->getValueType(0);
if (SDValue SCC = SimplifySelectCC(
             DAG.getConstant(0, DL, VT),
             cast<CondCodeSDNode>(N0.getOperand(2))->get(), true))
       return SCC;
   }
 
   if (SDValue NewCtPop = widenCtPop(N, DAG))
     return NewCtPop;
 
-  if (SDValue Res = tryToFoldExtendSelectLoad(N, TLI, DAG))
+  if (SDValue Res = tryToFoldExtendSelectLoad(N, TLI, DAG, Level))
     return Res;
 
   return SDValue();
 }
 
 SDValue DAGCombiner::visitAssertExt(SDNode *N) {
   unsigned Opcode = N->getOpcode();
   SDValue N0 = N->getOperand(0);

llvm/lib/Target/X86/X86ISelDAGToDAG.cpp

case X86ISD::VBROADCAST_LOAD: {
       break;
     }
     case ISD::VSELECT: {
       // Replace VSELECT with non-mask conditions with with BLENDV.
       if (N->getOperand(0).getValueType().getVectorElementType() == MVT::i1)
         break;
 
       assert(Subtarget->hasSSE41() && "Expected SSE4.1 support!");
+      assert(N->getValueType(0).getVectorElementType() != MVT::i16 &&
+             "We can't replace VSELECT with BLENDV in vXi16!");
       SDValue Blendv =
           CurDAG->getNode(X86ISD::BLENDV, SDLoc(N), N->getValueType(0),
                           N->getOperand(0), N->getOperand(1), N->getOperand(2));
       --I;
       CurDAG->ReplaceAllUsesWith(N, Blendv.getNode());
       ++I;
       MadeChange = true;
       continue;

llvm/test/CodeGen/X86/vselect-post-combine.ll

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc < %s -mtriple=x86_64-unknown-unknown -mattr=+avx2 | FileCheck %s --check-prefix=AVX2
+
+define ptr @test_mul(ptr %addr) {
+; AVX2-LABEL: test_mul:
+; AVX2:       # %bb.0: # %entry
+; AVX2-NEXT:    vmovdqa {{.*#+}} xmm0 = [255,0,0,0]
+; AVX2-NEXT:    vmovdqa {{.*#+}} xmm1 = [1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1]
+; AVX2-NEXT:    vpblendvb %xmm0, (%rdi), %xmm1, %xmm0
+; AVX2-NEXT:    vpmovzxbw {{.*#+}} ymm0 = xmm0[0],zero,xmm0[1],zero,xmm0[2],zero,xmm0[3],zero,xmm0[4],zero,xmm0[5],zero,xmm0[6],zero,xmm0[7],zero,xmm0[8],zero,xmm0[9],zero,xmm0[10],zero,xmm0[11],zero,xmm0[12],zero,xmm0[13],zero,xmm0[14],zero,xmm0[15],zero
+; AVX2-NEXT:    vmovdqu %ymm0, 0
+; AVX2-NEXT:    xorl %eax, %eax
+; AVX2-NEXT:    vzeroupper
+; AVX2-NEXT:    retq
+entry:
+  %vec0 = load <32 x i8>, ptr %addr
+  %vec1 = shufflevector <32 x i8> %vec0, <32 x i8> <i8 1, i8 1, i8 1, i8 1, i8 1, i8 1, i8 1, i8 1, i8 1, i8 1, i8 1, i8 1, i8 1, i8 1, i8 1, i8 1, i8 1, i8 1, i8 1, i8 1, i8 1, i8 1, i8 1, i8 1, i8 1, i8 1, i8 1, i8 1, i8 1, i8 1, i8 1, i8 1>, <32 x i32> <i32 0, i32 33, i32 34, i32 35, i32 36, i32 37, i32 38, i32 39, i32 40, i32 41, i32 42, i32 43, i32 44, i32 45, i32 46, i32 47, i32 48, i32 49, i32 50, i32 51, i32 52, i32 53, i32 54, i32 55, i32 56, i32 57, i32 58, i32 59, i32 60, i32 61, i32 62, i32 63>
+  %0 = bitcast <32 x i8> %vec1 to <4 x i64>
+  %shuffle = shufflevector <4 x i64> %0, <4 x i64> zeroinitializer, <2 x i32> <i32 0, i32 1>
+  %1 = bitcast <2 x i64> %shuffle to <16 x i8>
+  %conv = zext <16 x i8> %1 to <16 x i16>
+  store <16 x i16> %conv, ptr null, align 1
+  ret ptr null
+}