[SDAG] avoid vector extract/insert around binop

[spatel]

Summary

scalar-to-vector (binop (extractelt V, Idx), C) --> shuffle (bo V, C'), {Idx, -1, -1...}

We generally try to avoid ad-hoc vectorization in SDAG, but the motivating case from issue #39482
escapes our normal vectorization folds in IR. It seems like it should always be a win to transform this pattern in cases where we have the same vector type for input and output and the target supports the vector operation because we avoid transfers from vector to scalar and back.

In the x86 shift examples, we create the scalar-to-vector node during legalization. I'm not sure if there's a more general way to create the pattern for testing. (If so, I could add tests for other targets.)

[RKSimon]

llvm/test/CodeGen/X86/urem-vector-lkk.ll

149

can the fold be expanded to work with insert_vector_elt as well?

[spatel]

llvm/test/CodeGen/X86/urem-vector-lkk.ll

149

Yes - but I expect it'll be a closer call whether things are always profitable in those cases.

On this particular example, I think we'd need to change something else or have an x86-specific solution because we convert to target-specific ops during legalization and before there's a chance to clean up intermediate casts:

        t90: i32 = X86ISD::PEXTRW t2, TargetConstant:i8<1>
      t91: i16 = truncate t90
    t76: i16 = and t91, Constant:i16<31>
  t102: i32 = any_extend t76
t103: v8i16 = X86ISD::PINSRW t82, t102, TargetConstant:i8<1>

[RKSimon]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

23548

Do we need a isShuffleMaskLegal check?

[dmgreen]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

23548

Oh that might be a good idea, some instruction sets like MVE do not have many shuffles. It will depend on the type, it may just be a lane move which is cheap, but may require going through grp regs. I'm not sure there is a lot in it either way though, it might be slow in either case for architectures like that.

[RKSimon]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

23548

I guess we should always allow idx == 0 as that's effectively what (vt scalar_to_vector(extract_vector_elt(vt x,0))) is

[spatel]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

23548

I'll add the legality check. If a target says {0,-1, -1...} is not legal, that seems like a bug for that target since that's an identity (no-op) mask.

In the x86 shift tests that I stepped through, we always shuffle the demanded element down to element 0 before the extract, so we won't actually exercise that extra predicate, but I'll try to find some other case where it happens.

[spatel]

Patch updated:

1. Added isShuffleMaskLegal() check.
2. Added one-use check (this might be eased as a follow-up, but needs more tests - I added one test to verify, so it's currently showing no diff).
3. Added TODO comments for enhancements.

[RKSimon]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

23518

is anything stopping a vectorized integer division with undef elements occurring apart from the fact that almost no target has such an instruction?

[spatel]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

23518

Good point - I don't know how to test it, but we can't do this for opcodes that are not speculatable (because of immediate UB).

A similar problem came up here:
17dcbd816547
...so I should've remembered that! :)

[spatel]

Updated to use isSafeToSpeculativelyExecute() (added with 3aec0211187b98).
That avoids executing div/rem on unknown/undef vector elements.

[RKSimon]

LGTM - cheers!

[dyung]

@spatel, after this commit, we started seeing an assertion failure being hit in two of our internal tests. I have filed this as issue #58661, can you take a look?

[thakis]

In D136713#3890855, @dyung wrote:

@spatel, after this commit, we started seeing an assertion failure being hit in two of our internal tests. I have filed this as issue #58661, can you take a look?

Thanks for filing this! We hit the same problem, and the fix fixes it for us too.

It would've been nice if this had been reverted once the problem was found – then the probably would've gone away 7 hours earlier and I wouldn't have had to spend time to track it down as well :) If things take a while to fix, I think it's good to revert first and then fix asynchronously, to save the duplicate work of several people tracking down their compile failures to the same commit, only to discover that the problem is already known.

[thakis]

(Thanks for the fix too, of course!)

[spatel]

In D136713#3891887, @thakis wrote:

In D136713#3890855, @dyung wrote:

@spatel, after this commit, we started seeing an assertion failure being hit in two of our internal tests. I have filed this as issue #58661, can you take a look?

Thanks for filing this! We hit the same problem, and the fix fixes it for us too.

It would've been nice if this had been reverted once the problem was found – then the probably would've gone away 7 hours earlier and I wouldn't have had to spend time to track it down as well :) If things take a while to fix, I think it's good to revert first and then fix asynchronously, to save the duplicate work of several people tracking down their compile failures to the same commit, only to discover that the problem is already known.

I agree - if there's a clear connection between a patch and a crash, then please do revert. This was actually the best I could have done to post a fix given the timing; it was patched in about an hour after waking up and seeing the bug report. :)
Sorry about the bug, but thanks to all for the reports and reduced tests.

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

SDValue DAGCombiner::visitVECTOR_SHUFFLE(SDNode *N) {
 
   if (SDValue V = foldShuffleOfConcatUndefs(SVN, DAG))
     return V;
 
   return SDValue();
 }
 
 SDValue DAGCombiner::visitSCALAR_TO_VECTOR(SDNode *N) {
-  SDValue Scalar = N->getOperand(0);
   EVT VT = N->getValueType(0);
   if (!VT.isFixedLengthVector())
     return SDValue();
 
+  // Try to convert a scalar binop with an extracted vector element to a vector
+  // binop. This is intended to reduce potentially expensive register moves.
+  // TODO: Check if both operands are extracted.
+  // TODO: Generalize this, so it can be called from visitINSERT_VECTOR_ELT().
+  SDValue Scalar = N->getOperand(0);
+  unsigned Opcode = Scalar.getOpcode();
+  if (Scalar.hasOneUse() && Scalar->getNumValues() == 1 &&
+      TLI.isBinOp(Opcode) && VT.getScalarType() == Scalar.getValueType() &&
+      DAG.isSafeToSpeculativelyExecute(Opcode) && hasOperation(Opcode, VT)) {
+    // Match an extract element and get a shuffle mask equivalent.
+    SmallVector<int, 8> ShufMask(VT.getVectorNumElements(), -1);
+    auto getShuffleMaskForExtElt = [&](SDValue EE) {
+      if (EE.getOpcode() == ISD::EXTRACT_VECTOR_ELT &&
+          EE.getOperand(0).getValueType() == VT &&
+          isa<ConstantSDNode>(EE.getOperand(1))) {
+        // Mask = {ExtractIndex, undef, undef....}
+        ShufMask[0] = EE.getConstantOperandVal(1);
+        // Make sure the shuffle is legal if we are crossing lanes.
+        return TLI.isShuffleMaskLegal(ShufMask, VT);
+      }
+      return false;
+    };
+
+    // s2v (bo (extelt V, Idx), C) --> shuffle (bo V, C'), {Idx, -1, -1...}
+    if (auto *C = dyn_cast<ConstantSDNode>(Scalar.getOperand(1))) {
+      if (getShuffleMaskForExtElt(Scalar.getOperand(0))) {
+        SDLoc DL(N);
+        SDValue V = Scalar.getOperand(0).getOperand(0);
+        SDValue VecC = DAG.getConstant(C->getAPIntValue(), DL, VT);
+        SDValue VecBO = DAG.getNode(Opcode, DL, VT, V, VecC);
+        return DAG.getVectorShuffle(VT, DL, VecBO, DAG.getUNDEF(VT), ShufMask);
+      }
+    }
+    // s2v (bo C, (extelt V, Idx)) --> shuffle (bo C', V), {Idx, -1, -1...}
+    if (auto *C = dyn_cast<ConstantSDNode>(Scalar.getOperand(0))) {
+      if (getShuffleMaskForExtElt(Scalar.getOperand(1))) {
+        SDLoc DL(N);
+        SDValue V = Scalar.getOperand(1).getOperand(0);
+        SDValue VecC = DAG.getConstant(C->getAPIntValue(), DL, VT);
+        SDValue VecBO = DAG.getNode(Opcode, DL, VT, VecC, V);
+        return DAG.getVectorShuffle(VT, DL, VecBO, DAG.getUNDEF(VT), ShufMask);
+      }
+    }
+  }
+
   // Replace a SCALAR_TO_VECTOR(EXTRACT_VECTOR_ELT(V,C0)) pattern
   // with a VECTOR_SHUFFLE and possible truncate.
-  if (Scalar.getOpcode() != ISD::EXTRACT_VECTOR_ELT ||
+  if (Opcode != ISD::EXTRACT_VECTOR_ELT ||
       !Scalar.getOperand(0).getValueType().isFixedLengthVector())
     return SDValue();
 
   // If we have an implicit truncate, truncate here if it is legal.
   if (VT.getScalarType() != Scalar.getValueType() &&
       Scalar.getValueType().isScalarInteger() &&
       isTypeLegal(VT.getScalarType())) {
     SDValue Val =

llvm/test/CodeGen/X86/cvt16.ll

 ; LIBCALL-NEXT:    callq __truncsfhf2@PLT
 ; LIBCALL-NEXT:    popq %rax
 ; LIBCALL-NEXT:    .cfi_def_cfa_offset 8
 ; LIBCALL-NEXT:    jmp __extendhfsf2@PLT # TAILCALL
 ;
 ; F16C-LABEL: test3:
 ; F16C:       # %bb.0:
 ; F16C-NEXT:    vcvtps2ph $4, %xmm0, %xmm0
-; F16C-NEXT:    vmovd %xmm0, %eax
-; F16C-NEXT:    movzwl %ax, %eax
-; F16C-NEXT:    vmovd %eax, %xmm0
+; F16C-NEXT:    vpmovzxwq {{.*#+}} xmm0 = xmm0[0],zero,zero,zero,xmm0[1],zero,zero,zero
 ; F16C-NEXT:    vcvtph2ps %xmm0, %xmm0
 ; F16C-NEXT:    retq
 ;
 ; SOFTFLOAT-LABEL: test3:
 ; SOFTFLOAT:       # %bb.0:
 ; SOFTFLOAT-NEXT:    pushq %rax
 ; SOFTFLOAT-NEXT:    .cfi_def_cfa_offset 16
 ; SOFTFLOAT-NEXT:    callq __gnu_f2h_ieee@PLT

llvm/test/CodeGen/X86/half.ll

 ; BWON-F16C-NEXT:    vcvtsi2ss %edi, %xmm0, %xmm0
 ; BWON-F16C-NEXT:    vcvtps2ph $4, %xmm0, %xmm0
 ; BWON-F16C-NEXT:    vpmovzxwq {{.*#+}} xmm0 = xmm0[0],zero,zero,zero,xmm0[1],zero,zero,zero
 ; BWON-F16C-NEXT:    vcvtph2ps %xmm0, %xmm0
 ; BWON-F16C-NEXT:    vmovd %eax, %xmm1
 ; BWON-F16C-NEXT:    vcvtph2ps %xmm1, %xmm1
 ; BWON-F16C-NEXT:    vaddss %xmm0, %xmm1, %xmm0
 ; BWON-F16C-NEXT:    vcvtps2ph $4, %xmm0, %xmm0
-; BWON-F16C-NEXT:    vmovd %xmm0, %eax
-; BWON-F16C-NEXT:    movzwl %ax, %eax
-; BWON-F16C-NEXT:    vmovd %eax, %xmm0
+; BWON-F16C-NEXT:    vpmovzxwq {{.*#+}} xmm0 = xmm0[0],zero,zero,zero,xmm0[1],zero,zero,zero
 ; BWON-F16C-NEXT:    vcvtph2ps %xmm0, %xmm0
 ; BWON-F16C-NEXT:    retq
 ;
 ; CHECK-I686-LABEL: test_sitofp_fadd_i32:
 ; CHECK-I686:       # %bb.0:
 ; CHECK-I686-NEXT:    subl $60, %esp
 ; CHECK-I686-NEXT:    movl {{[0-9]+}}(%esp), %eax
 ; CHECK-I686-NEXT:    pinsrw $0, (%eax), %xmm0

llvm/test/CodeGen/X86/urem-vector-lkk.ll

; AVX-NEXT:    retq
   %3 = add <4 x i16> %1, %2
   ret <4 x i16> %3
 }
 
 ; Don't fold for divisors that are a power of two.
 define <4 x i16> @dont_fold_urem_power_of_two(<4 x i16> %x) {
 ; SSE-LABEL: dont_fold_urem_power_of_two:
 ; SSE:       # %bb.0:
+; SSE-NEXT:    movdqa {{.*#+}} xmm1 = [63,63,63,63]
+; SSE-NEXT:    pand %xmm0, %xmm1
+; SSE-NEXT:    pextrw $1, %xmm0, %eax
+; SSE-NEXT:    andl $31, %eax
+; SSE-NEXT:    pinsrw $1, %eax, %xmm1
+; SSE-NEXT:    pextrw $2, %xmm0, %eax
+; SSE-NEXT:    andl $7, %eax
+; SSE-NEXT:    pinsrw $2, %eax, %xmm1
 ; SSE-NEXT:    pextrw $3, %xmm0, %eax
 ; SSE-NEXT:    imull $44151, %eax, %ecx # imm = 0xAC77
 ; SSE-NEXT:    shrl $22, %ecx
 ; SSE-NEXT:    imull $95, %ecx, %ecx
 ; SSE-NEXT:    subl %ecx, %eax
-; SSE-NEXT:    pextrw $1, %xmm0, %ecx
-; SSE-NEXT:    andl $31, %ecx
-; SSE-NEXT:    movd %xmm0, %edx
-; SSE-NEXT:    andl $63, %edx
-; SSE-NEXT:    movd %edx, %xmm1
-; SSE-NEXT:    pinsrw $1, %ecx, %xmm1
-; SSE-NEXT:    pextrw $2, %xmm0, %ecx
-; SSE-NEXT:    andl $7, %ecx
-; SSE-NEXT:    pinsrw $2, %ecx, %xmm1
 ; SSE-NEXT:    pinsrw $3, %eax, %xmm1
 ; SSE-NEXT:    movdqa %xmm1, %xmm0
 ; SSE-NEXT:    retq
 ;
-; AVX-LABEL: dont_fold_urem_power_of_two:
-; AVX:       # %bb.0:
-; AVX-NEXT:    vpextrw $3, %xmm0, %eax
-; AVX-NEXT:    imull $44151, %eax, %ecx # imm = 0xAC77
-; AVX-NEXT:    shrl $22, %ecx
-; AVX-NEXT:    imull $95, %ecx, %ecx
-; AVX-NEXT:    subl %ecx, %eax
-; AVX-NEXT:    vpextrw $1, %xmm0, %ecx
-; AVX-NEXT:    andl $31, %ecx
-; AVX-NEXT:    vmovd %xmm0, %edx
-; AVX-NEXT:    andl $63, %edx
-; AVX-NEXT:    vmovd %edx, %xmm1
-; AVX-NEXT:    vpinsrw $1, %ecx, %xmm1, %xmm1
-; AVX-NEXT:    vpextrw $2, %xmm0, %ecx
-; AVX-NEXT:    andl $7, %ecx
-; AVX-NEXT:    vpinsrw $2, %ecx, %xmm1, %xmm0
-; AVX-NEXT:    vpinsrw $3, %eax, %xmm0, %xmm0
-; AVX-NEXT:    retq
+; AVX1-LABEL: dont_fold_urem_power_of_two:
+; AVX1:       # %bb.0:
+; AVX1-NEXT:    vpand {{\.?LCPI[0-9]+_[0-9]+}}(%rip), %xmm0, %xmm1
+; AVX1-NEXT:    vpextrw $1, %xmm0, %eax
+; AVX1-NEXT:    andl $31, %eax
+; AVX1-NEXT:    vpinsrw $1, %eax, %xmm1, %xmm1
+; AVX1-NEXT:    vpextrw $2, %xmm0, %eax
+; AVX1-NEXT:    andl $7, %eax
+; AVX1-NEXT:    vpinsrw $2, %eax, %xmm1, %xmm1
+; AVX1-NEXT:    vpextrw $3, %xmm0, %eax
+; AVX1-NEXT:    imull $44151, %eax, %ecx # imm = 0xAC77
+; AVX1-NEXT:    shrl $22, %ecx
+; AVX1-NEXT:    imull $95, %ecx, %ecx
+; AVX1-NEXT:    subl %ecx, %eax
+; AVX1-NEXT:    vpinsrw $3, %eax, %xmm1, %xmm0
+; AVX1-NEXT:    retq
+;
+; AVX2-LABEL: dont_fold_urem_power_of_two:
+; AVX2:       # %bb.0:
+; AVX2-NEXT:    vpbroadcastd {{.*#+}} xmm1 = [63,63,63,63]
+; AVX2-NEXT:    vpand %xmm1, %xmm0, %xmm1
+; AVX2-NEXT:    vpextrw $1, %xmm0, %eax
+; AVX2-NEXT:    andl $31, %eax
+; AVX2-NEXT:    vpinsrw $1, %eax, %xmm1, %xmm1
+; AVX2-NEXT:    vpextrw $2, %xmm0, %eax
+; AVX2-NEXT:    andl $7, %eax
+; AVX2-NEXT:    vpinsrw $2, %eax, %xmm1, %xmm1
+; AVX2-NEXT:    vpextrw $3, %xmm0, %eax
+; AVX2-NEXT:    imull $44151, %eax, %ecx # imm = 0xAC77
+; AVX2-NEXT:    shrl $22, %ecx
+; AVX2-NEXT:    imull $95, %ecx, %ecx
+; AVX2-NEXT:    subl %ecx, %eax
+; AVX2-NEXT:    vpinsrw $3, %eax, %xmm1, %xmm0
+; AVX2-NEXT:    retq
   %1 = urem <4 x i16> %x, <i16 64, i16 32, i16 8, i16 95>
   ret <4 x i16> %1
 }
 
 ; Don't fold if the divisor is one.
 define <4 x i16> @dont_fold_urem_one(<4 x i16> %x) {
 ; SSE-LABEL: dont_fold_urem_one:
 ; SSE:       # %bb.0:

llvm/test/CodeGen/X86/vec_shift5.ll

; X64-NEXT:    retq
   ret <4 x i32> %res
 }
 
 ; PR39482
 
 define <4 x i32> @extelt0_sub_pslli_v4i32(<4 x i32> %x, <4 x i32> %y){
 ; CHECK-LABEL: extelt0_sub_pslli_v4i32:
 ; CHECK:       # %bb.0:
-; CHECK-NEXT:    movd %xmm1, %eax
-; CHECK-NEXT:    movl $32, %ecx
-; CHECK-NEXT:    subl %eax, %ecx
-; CHECK-NEXT:    movd %ecx, %xmm1
+; CHECK-NEXT:    movdqa {{.*#+}} xmm2 = [32,32,32,32]
+; CHECK-NEXT:    psubd %xmm1, %xmm2
+; CHECK-NEXT:    pxor %xmm1, %xmm1
+; CHECK-NEXT:    movss {{.*#+}} xmm1 = xmm2[0],xmm1[1,2,3]
 ; CHECK-NEXT:    pslld %xmm1, %xmm0
 ; CHECK-NEXT:    ret{{[l|q]}}
   %ext = extractelement <4 x i32> %y, i64 0
   %bo = sub i32 32, %ext
   %r = tail call <4 x i32> @llvm.x86.sse2.pslli.d(<4 x i32> %x, i32 %bo)
   ret <4 x i32> %r
 }
 
 define <4 x i32> @extelt1_add_psrli_v4i32(<4 x i32> %x, <4 x i32> %y){
-; CHECK-LABEL: extelt1_add_psrli_v4i32:
-; CHECK:       # %bb.0:
-; CHECK-NEXT:    pshufd {{.*#+}} xmm1 = xmm1[1,1,1,1]
-; CHECK-NEXT:    movd %xmm1, %eax
-; CHECK-NEXT:    addl $3, %eax
-; CHECK-NEXT:    movd %eax, %xmm1
-; CHECK-NEXT:    psrld %xmm1, %xmm0
-; CHECK-NEXT:    ret{{[l|q]}}
+; X86-LABEL: extelt1_add_psrli_v4i32:
+; X86:       # %bb.0:
+; X86-NEXT:    pshufd {{.*#+}} xmm1 = xmm1[1,1,1,1]
+; X86-NEXT:    paddd {{\.?LCPI[0-9]+_[0-9]+}}, %xmm1
+; X86-NEXT:    xorps %xmm2, %xmm2
+; X86-NEXT:    movss {{.*#+}} xmm2 = xmm1[0],xmm2[1,2,3]
+; X86-NEXT:    psrld %xmm2, %xmm0
+; X86-NEXT:    retl
+;
+; X64-LABEL: extelt1_add_psrli_v4i32:
+; X64:       # %bb.0:
+; X64-NEXT:    pshufd {{.*#+}} xmm1 = xmm1[1,1,1,1]
+; X64-NEXT:    paddd {{\.?LCPI[0-9]+_[0-9]+}}(%rip), %xmm1
+; X64-NEXT:    xorps %xmm2, %xmm2
+; X64-NEXT:    movss {{.*#+}} xmm2 = xmm1[0],xmm2[1,2,3]
+; X64-NEXT:    psrld %xmm2, %xmm0
+; X64-NEXT:    retq
   %ext = extractelement <4 x i32> %y, i64 1
   %bo = add i32 %ext, 3
   %r = tail call <4 x i32> @llvm.x86.sse2.psrli.d(<4 x i32> %x, i32 %bo)
   ret <4 x i32> %r
 }
 
 define i32 @extelt1_add_psrai_v4i32_uses(<4 x i32> %x, <4 x i32> %y){
 ; CHECK-LABEL: extelt1_add_psrai_v4i32_uses: