[DAGCombiner][x86] add transform/hook to load a scalar directly for use in a vector binop

[spatel]

Summary

This is the reversal for the proposed IR canonicalization in D50992 (insert+vector op --> scalar op+insert).

I've enabled x86 in the minimal way because this looks like a close call for most recent Intel. They have fast (1uop / 1 cycle latency) transfer from GPR to *MM according to Agner / llvm-mca. And the current AVX2 code is likely too broadcast-happy as seen in the test diffs.

So this doesn't exercise the recent broadcast improvements from D51125 / D51186 yet.

Enabling more opcodes/types (the test file covers all 18 IR-equivalent binops) looks tricky. For example, we should have gotten the 'and' with i32 test, but v4i32 is promoted to v2i64. Pre-SSE4, we don't have pmulld, so we can't uniformly allow isLegalOrCustom().

[spatel]

Ping.

[lebedev.ri]

This looks reasonable to me, some nits.

include/llvm/CodeGen/TargetLowering.h

535

You probably want to add some explanatory comment here.

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

15147

Do we need to check that BO.getOpcode() is actually valid for VecVT?
Or are we ok with scalarization?

test/CodeGen/X86/load-scalar-as-vector.ll

23

Why is there a vpbroadcastd here, as compared to AVX1 version?

[spatel]

include/llvm/CodeGen/TargetLowering.h

535

Oops - yes, I forgot that.

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

15147

Scalarization would almost certainly defeat the point of this transform, but I think it's better to have all of those checks in one place - in the target hook.

test/CodeGen/X86/load-scalar-as-vector.ll

23

For AVX2 targets, we assume that a load+broadcast costs no more than a full vector load, but it saves data space, so that's the default. I suspect that we need to revisit that per-CPU and see if that's actually true.

[spatel]

test/CodeGen/X86/load-scalar-as-vector.ll

23

Oops - in this case, we loaded a scalar though, so there must be something going wrong with our undef knowledge.

[spatel]

Patch updated:
Added documentation comment for the new target hook.

[RKSimon]

My main concern is the growth in the constant pool due to vector constants.

[spatel]

In D51553#1234619, @RKSimon wrote:

My main concern is the growth in the constant pool due to vector constants.

I agree that that is a trade-off, but I don't think we can make an informed decision about it at this level. We can fine-tune the hook in this patch if we see regressions, but we're still just guessing at the eventual global outcome.
So without that kind of knowledge, we have to look at the micro-motivation.

Here's llvm-mca output using scalar math on btver2 (Jaguar):

[0,0]     DeeeeeER  .    ..   movl	(%rdi), %eax
[0,1]     D=====eER .    ..   addl	$42, %eax
[0,2]     .D=====eeeeeeeeER   movd	%eax, %xmm0  <--- move to vector register has 8-cycle latency?

vs. vector op:

[0,0]     DeeeeeER  .   movd	(%rdi), %xmm0
[0,1]     D==eeeeeeER   paddd	32(%rip), %xmm0

Seems like the right choice to do the transform here?

[spatel]

Ping.

[andreadb]

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

15103–15105

I wonder if this should be disabled for optsize/minsize.
On x86, this is going to increase .text, and it also potentially introduces a new constant in the pool.

mov (%rdi), %eax
add $0x2a, %eax
vmovd %eax, %xmm0
retq

is 10 bytes.

vmovd (%rdi), %xmm0
vpaddd 0x0(%rip), %xmm0, %xmm0  ## 4 bytes PC relative relocation
retq

is 13 bytes.

Plus, 16B in the constant pool (if we didn't see that constant before)...

[spatel]

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

15103–15105

Yes, that seems like a good constraint for x86. I'm not sure that other targets would want the same limitation though, so I think we should make that another option in the target hook itself rather than here in generic code.

[spatel]

Patch updated:
Disallow the transform on x86 if we're optimizing for size. Test added at rL342755.

[spatel]

Ping.

[RKSimon]

test/CodeGen/X86/load-scalar-as-vector.ll

5

Please can you rebase as I added avx512 tests at rL342773 and possibly add some tests with 256/512 vectors as well?

[spatel]

test/CodeGen/X86/load-scalar-as-vector.ll

5

Sure - there's no difference for AVX512 with the current set of tests on this patch, but there should be after rL343491.

[spatel]

Patch updated:
No code changes, but I added 256/512-bit versions of the 128-bit tests that are affected by this patch.

The new results show that we are only transforming in the cases where we have legal ops for the type. Ie, this transform does not fire after we split the big vectors into the smaller legal vector type supported by the target. That's because we don't have an insert_vector_elt at that stage; it's a scalar_to_vector post-legalization.

So we can view the wider cases as conservatively ok here, but we should handle the pattern with scalar_to_vector as a follow-up. Or we might say that we don't want this transform to ever fire for bigger vector types? It's harder to justify this transform with the longer vectors if those instructions are executed as multi-uop sequences and/or have effects that we can't model (throttling the clock rate). OTOH, the transform is limited to insertion into element 0 currently, so the fact that we're using wide vector ops in these tests isn't necessary in the first place.

[efriedma]

For an insert into lane zero, it's never profitable to generate ymm/zmm instructions; we can just generate an xmm instruction and implicitly extend the result.

[spatel]

In D51553#1251450, @efriedma wrote:

For an insert into lane zero, it's never profitable to generate ymm/zmm instructions; we can just generate an xmm instruction and implicitly extend the result.

Right - that's what I was suggesting at the end of my last reply. So do we want to put this patch on hold to improve the undef/splat behaviors seen here first, or proceed here and file bugs for those problems?

[lebedev.ri]

What's the status here?

[spatel]

Patch rebased - at first glance, we haven't made much progress on the issues raised months ago despite dozens of narrowing/shrinking/demanded bits/demanded elts patches.

[lebedev.ri]

(just removing from review queue)

[RKSimon]

@spatel reverse ping

[spatel]

Rebased to apply cleanly to current master, but the test diffs remain unchanged. We are still creating vector ops that are too wide.

This seems easier to implement in vector-combine now that we have it. We will need to make sure that we are not fighting the recently added scalarization transforms.

[lebedev.ri]

This review seems to be stuck/dead, consider abandoning if no longer relevant.

include/llvm/CodeGen/TargetLowering.h

public:
   /// Return true if we should transform.
   /// Return false if there is no preference.
   virtual bool shouldTransformSignedTruncationCheck(EVT XVT,
                                                     unsigned KeptBits) const {
     // By default, let's assume that no one prefers shifts.
     return false;
   }
 
+  /// Given an insert-element vector (InsElt) and a scalar operation (Op),
+  /// return true if it would be profitable to convert the scalar operation into
+  /// a vector operation. This would normally be true if:
+  /// 1. The vector operation does not cost much more than a scalar version.
+  /// 2. The target can avoid a costly transfer from scalar to vector register
+  ///    by loading a scalar operand directly into a vector register.
+  virtual bool shouldLoadScalarIntoVectorOp(SDValue InsElt, SDValue Op) const {
+    return false;
+  }
+
   /// Return true if the target wants to use the optimization that
   /// turns ext(promotableInst1(...(promotableInstN(load)))) into
   /// promotedInst1(...(promotedInstN(ext(load)))).
   bool enableExtLdPromotion() const { return EnableExtLdPromotion; }
 
   /// Return true if the target can combine store(extractelement VectorTy,
   /// Idx).
   /// \p Cost[out] gives the cost of that transformation when this is true.

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

SDValue DAGCombiner::combineInsertEltToShuffle(SDNode *N, unsigned InsIndex) {
   SDValue DestVecBC = DAG.getBitcast(ShufVT, DestVec);
   SDValue Shuf = DAG.getVectorShuffle(ShufVT, DL, DestVecBC, PaddedSubV, Mask);
   AddToWorklist(PaddedSubV.getNode());
   AddToWorklist(DestVecBC.getNode());
   AddToWorklist(Shuf.getNode());
   return DAG.getBitcast(VT, Shuf);
 }
 
+/// Try to convert a load of a scalar + scalar binop + insert element into a
+/// load + insert element + vector binop. If we can load the scalar directly
+/// into a vector register, this eliminates a potentially expensive transfer
+/// from scalar register to vector register.
+static SDValue combineLoadBinopInsElt(SDNode *Ins, SelectionDAG &DAG) {
+  // TODO: This can be loosened to allow insertion into any constant vector.
+  SDValue UndefVec = Ins->getOperand(0);
+  SDValue BO = Ins->getOperand(1);
+  if (!UndefVec.isUndef() || !ISD::isBinaryOp(BO.getNode()) || !BO.hasOneUse())
+    return SDValue();
+
+  const TargetLowering &TLI = DAG.getTargetLoweringInfo();
+  if (!TLI.shouldLoadScalarIntoVectorOp(SDValue(Ins, 0), BO))
+    return SDValue();
+
+  // We are matching a binop that has a loaded operand and a constant operand.
+  // This is complicated because the operands can be in either order (and we
+  // must capture that fact), and the constant can be either integer or FP.
+  EVT VecVT = Ins->getOperand(0).getValueType();
+  SDValue Ld, C;
+  SDLoc DL(Ins);
+  auto matchLoadAndConstant = [&](SDValue Op0, SDValue Op1) {
+    if (Op0.getOpcode() != ISD::LOAD)
+      return false;
+    // Splat a scalar constant operand for use in a vector op. The caller can
+    // adjust (blend) this constant with the original insertion vector constant.
+    if (auto *CInt = dyn_cast<ConstantSDNode>(Op1)) {
+      // Account for scalar operand size differences. For example, scalar shift
+      // amount may have a different type than the other operand/result.
+      unsigned Width = VecVT.getScalarSizeInBits();
+      C = DAG.getConstant(CInt->getAPIntValue().zextOrTrunc(Width), DL, VecVT);
+    } else if (auto *CFP = dyn_cast<ConstantFPSDNode>(Op1)) {
+      C = DAG.getConstantFP(CFP->getValueAPF(), DL, VecVT);
+    } else {
+      return false;
+    }
+    Ld = Op0;
+    return true;
+  };
+
+  SDValue BO0 = BO.getOperand(0), BO1 = BO.getOperand(1);
+  bool Op0IsLoad = matchLoadAndConstant(BO0, BO1);
+  if (!Op0IsLoad && !matchLoadAndConstant(BO1, BO0))
+    return SDValue();
+
+  // ins undef, (bo (load X), C), index --> bo (ins undef, (load X), index), C'
+  // ins undef, (bo C, (load X)), index --> bo C', (ins undef, (load X), index)
+  SDValue NewInsert = DAG.getNode(ISD::INSERT_VECTOR_ELT, DL, VecVT,
+                                  UndefVec, Op0IsLoad ? BO0 : BO1,
+                                  Ins->getOperand(2));
+
+  // TODO: We created a splat constant above because we did not check the
+  //       insert index. If the insert index is a constant and/or we were not
+  //       originally inserting into an undef constant, we should compute the
+  //       other constant elements as needed.
+  return Op0IsLoad ? DAG.getNode(BO.getOpcode(), DL, VecVT, NewInsert, C)
+                   : DAG.getNode(BO.getOpcode(), DL, VecVT, C, NewInsert);
+}
+
 SDValue DAGCombiner::visitINSERT_VECTOR_ELT(SDNode *N) {
   SDValue InVec = N->getOperand(0);
   SDValue InVal = N->getOperand(1);
   SDValue EltNo = N->getOperand(2);
   SDLoc DL(N);
 
   // If the inserted element is an UNDEF, just use the input vector.
   if (InVal.isUndef())
if (!IndexC) {
     // inselt undef, InVal, EltNo --> build_vector < InVal, InVal, ... >
     if (InVec.isUndef() && TLI.shouldSplatInsEltVarIndex(VT)) {
       SmallVector<SDValue, 8> Ops(VT.getVectorNumElements(), InVal);
       return DAG.getBuildVector(VT, DL, Ops);
     }
     return SDValue();
   }
 
+  if (SDValue BO = combineLoadBinopInsElt(N, DAG))
+    return BO;
+
   // We must know which element is being inserted for folds below here.
   unsigned Elt = IndexC->getZExtValue();
   if (SDValue Shuf = combineInsertEltToShuffle(N, Elt))
     return Shuf;
 
   // Canonicalize insert_vector_elt dag nodes.
   // Example:
   // (insert_vector_elt (insert_vector_elt A, Idx0), Idx1)

lib/Target/X86/X86ISelLowering.h

shouldTransformSignedTruncationCheck(EVT XVT,
       // We are ok with KeptBitsVT being byte/word/dword, what MOVS supports.
       // XVT will be larger than KeptBitsVT.
       MVT KeptBitsVT = MVT::getIntegerVT(KeptBits);
       return VTIsOk(XVT) && VTIsOk(KeptBitsVT);
     }
 
     bool shouldSplatInsEltVarIndex(EVT VT) const override;
 
+    bool shouldLoadScalarIntoVectorOp(SDValue Ins, SDValue Op) const override;
+
     bool convertSetCCLogicToBitwiseLogic(EVT VT) const override {
       return VT.isScalarInteger();
     }
 
     /// Vector-sized comparisons are fast using PCMPEQ + PMOVMSK or PTEST.
     MVT hasFastEqualityCompare(unsigned NumBits) const override;
 
     /// Allow multiple load pairs per block for smaller and faster code.

lib/Target/X86/X86ISelLowering.cpp

 }
 
 bool X86TargetLowering::shouldSplatInsEltVarIndex(EVT VT) const {
   // Any legal vector type can be splatted more efficiently than
   // loading/spilling from memory.
   return isTypeLegal(VT);
 }
 
+/// If we can load an integer scalar into a vector register with minimal
+/// shuffling and the vector operation is supported, then avoiding a transfer
+/// from GPR to vector is probably a win.
+bool X86TargetLowering::shouldLoadScalarIntoVectorOp(SDValue InsElt,
+                                                     SDValue BinOp) const {
+  // Without SSE2, we only have movss and no integer vector ops; don't bother.
+  // FP uses the same registers, so the transform would never help for an FP op.
+  EVT VecVT = InsElt.getValueType();
+  if (!Subtarget.hasSSE2() || VecVT.isFloatingPoint())
+    return false;
+
+  // Loading into the 0-index lane is possible with SSE2 using movd/movq.
+  // TODO: AVX1 and AVX2 can splat (broadcast) various scalar types.
+  EVT ScalarVT = BinOp.getValueType();
+  SDValue InsIdx = InsElt.getOperand(2);
+  if ((ScalarVT != MVT::i32 && ScalarVT != MVT::i64) || !isNullConstant(InsIdx))
+    return false;
+
+  // Filter out illegal vector types, ISA holes, and unsupported vector ops.
+  // TODO: This eliminates custom/promoted ops that are probably ok.
+  auto Opcode = BinOp.getOpcode();
+  if (!isTypeDesirableForOp(Opcode, VecVT) || !isOperationLegal(Opcode, VecVT))
+    return false;
+
+  return true;
+}
+
 MVT X86TargetLowering::hasFastEqualityCompare(unsigned NumBits) const {
   MVT VT = MVT::getIntegerVT(NumBits);
   if (isTypeLegal(VT))
     return VT;
 
   // PMOVMSKB can handle this.
   if (NumBits == 128 && isTypeLegal(MVT::v16i8))
     return MVT::v16i8;

test/CodeGen/X86/load-scalar-as-vector.ll

 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
 ; RUN: llc < %s -mtriple=x86_64-- -mattr=+sse2       | FileCheck %s --check-prefixes=SSE,SSE2
 ; RUN: llc < %s -mtriple=x86_64-- -mattr=+sse4.2     | FileCheck %s --check-prefixes=SSE,SSE4
 ; RUN: llc < %s -mtriple=x86_64-- -mattr=+avx        | FileCheck %s --check-prefixes=AVX,AVX1
 ; RUN: llc < %s -mtriple=x86_64-- -mattr=+avx2       | FileCheck %s --check-prefixes=AVX,AVX2
 
 define <4 x i32> @add_op1_constant(i32* %p) nounwind {
 ; SSE-LABEL: add_op1_constant:
 ; SSE:       # %bb.0:
-; SSE-NEXT:    movl (%rdi), %eax
-; SSE-NEXT:    addl $42, %eax
-; SSE-NEXT:    movd %eax, %xmm0
+; SSE-NEXT:    movd {{.*#+}} xmm0 = mem[0],zero,zero,zero
+; SSE-NEXT:    paddd {{.*}}(%rip), %xmm0
 ; SSE-NEXT:    retq
 ;
-; AVX-LABEL: add_op1_constant:
-; AVX:       # %bb.0:
-; AVX-NEXT:    movl (%rdi), %eax
-; AVX-NEXT:    addl $42, %eax
-; AVX-NEXT:    vmovd %eax, %xmm0
-; AVX-NEXT:    retq
+; AVX1-LABEL: add_op1_constant:
+; AVX1:       # %bb.0:
+; AVX1-NEXT:    vmovd {{.*#+}} xmm0 = mem[0],zero,zero,zero
+; AVX1-NEXT:    vpaddd {{.*}}(%rip), %xmm0, %xmm0
+; AVX1-NEXT:    retq
+;
+; AVX2-LABEL: add_op1_constant:
+; AVX2:       # %bb.0:
+; AVX2-NEXT:    vmovd {{.*#+}} xmm0 = mem[0],zero,zero,zero
+; AVX2-NEXT:    vpbroadcastd {{.*#+}} xmm1 = [42,42,42,42]
+; AVX2-NEXT:    vpaddd %xmm1, %xmm0, %xmm0
+; AVX2-NEXT:    retq
   %x = load i32, i32* %p
   %b = add i32 %x, 42
   %r = insertelement <4 x i32> undef, i32 %b, i32 0
   ret <4 x i32> %r
 }
 
 define <8 x i16> @add_op0_constant(i16* %p) nounwind {
 ; SSE-LABEL: add_op0_constant:
; AVX-NEXT:    retq
   %b = add i16 42, %x
   %r = insertelement <8 x i16> undef, i16 %b, i32 0
   ret <8 x i16> %r
 }
 
 define <2 x i64> @sub_op0_constant(i64* %p) nounwind {
 ; SSE-LABEL: sub_op0_constant:
 ; SSE:       # %bb.0:
-; SSE-NEXT:    movl $42, %eax
-; SSE-NEXT:    subq (%rdi), %rax
-; SSE-NEXT:    movq %rax, %xmm0
+; SSE-NEXT:    movq {{.*#+}} xmm1 = mem[0],zero
+; SSE-NEXT:    movdqa {{.*#+}} xmm0 = [42,42]
+; SSE-NEXT:    psubq %xmm1, %xmm0
 ; SSE-NEXT:    retq
 ;
 ; AVX-LABEL: sub_op0_constant:
 ; AVX:       # %bb.0:
-; AVX-NEXT:    movl $42, %eax
-; AVX-NEXT:    subq (%rdi), %rax
-; AVX-NEXT:    vmovq %rax, %xmm0
+; AVX-NEXT:    vmovq {{.*#+}} xmm0 = mem[0],zero
+; AVX-NEXT:    vmovdqa {{.*#+}} xmm1 = [42,42]
+; AVX-NEXT:    vpsubq %xmm0, %xmm1, %xmm0
 ; AVX-NEXT:    retq
   %x = load i64, i64* %p
   %b = sub i64 42, %x
   %r = insertelement <2 x i64> undef, i64 %b, i32 0
   ret <2 x i64> %r
 }
 
 define <16 x i8> @sub_op1_constant(i8* %p) nounwind {

 ; AVX-NEXT:    retq
   %x = load i8, i8* %p
   %b = sub i8 %x, 42
   %r = insertelement <16 x i8> undef, i8 %b, i32 0
   ret <16 x i8> %r
 }
 
 define <4 x i32> @mul_op1_constant(i32* %p) nounwind {
-; SSE-LABEL: mul_op1_constant:
-; SSE:       # %bb.0:
-; SSE-NEXT:    imull $42, (%rdi), %eax
-; SSE-NEXT:    movd %eax, %xmm0
-; SSE-NEXT:    retq
-;
-; AVX-LABEL: mul_op1_constant:
-; AVX:       # %bb.0:
-; AVX-NEXT:    imull $42, (%rdi), %eax
-; AVX-NEXT:    vmovd %eax, %xmm0
-; AVX-NEXT:    retq
+; SSE2-LABEL: mul_op1_constant:
+; SSE2:       # %bb.0:
+; SSE2-NEXT:    imull $42, (%rdi), %eax
+; SSE2-NEXT:    movd %eax, %xmm0
+; SSE2-NEXT:    retq
+;
+; SSE4-LABEL: mul_op1_constant:
+; SSE4:       # %bb.0:
+; SSE4-NEXT:    movd {{.*#+}} xmm0 = mem[0],zero,zero,zero
+; SSE4-NEXT:    pmulld {{.*}}(%rip), %xmm0
+; SSE4-NEXT:    retq
+;
+; AVX1-LABEL: mul_op1_constant:
+; AVX1:       # %bb.0:
+; AVX1-NEXT:    vmovd {{.*#+}} xmm0 = mem[0],zero,zero,zero
+; AVX1-NEXT:    vpmulld {{.*}}(%rip), %xmm0, %xmm0
+; AVX1-NEXT:    retq
+;
+; AVX2-LABEL: mul_op1_constant:
+; AVX2:       # %bb.0:
+; AVX2-NEXT:    vmovd {{.*#+}} xmm0 = mem[0],zero,zero,zero
+; AVX2-NEXT:    vpbroadcastd {{.*#+}} xmm1 = [42,42,42,42]
+; AVX2-NEXT:    vpmulld %xmm1, %xmm0, %xmm0
+; AVX2-NEXT:    retq
   %x = load i32, i32* %p
   %b = mul i32 %x, 42
   %r = insertelement <4 x i32> undef, i32 %b, i32 0
   ret <4 x i32> %r
 }
 
 define <8 x i16> @mul_op0_constant(i16* %p) nounwind {
 ; SSE-LABEL: mul_op0_constant:
; AVX-NEXT:    retq
   %b = and i32 %x, 42
   %r = insertelement <4 x i32> undef, i32 %b, i32 0
   ret <4 x i32> %r
 }
 
 define <2 x i64> @or_op1_constant(i64* %p) nounwind {
 ; SSE-LABEL: or_op1_constant:
 ; SSE:       # %bb.0:
-; SSE-NEXT:    movq (%rdi), %rax
-; SSE-NEXT:    orq $42, %rax
-; SSE-NEXT:    movq %rax, %xmm0
+; SSE-NEXT:    movsd {{.*#+}} xmm0 = mem[0],zero
+; SSE-NEXT:    orps {{.*}}(%rip), %xmm0
 ; SSE-NEXT:    retq
 ;
 ; AVX-LABEL: or_op1_constant:
 ; AVX:       # %bb.0:
-; AVX-NEXT:    movq (%rdi), %rax
-; AVX-NEXT:    orq $42, %rax
-; AVX-NEXT:    vmovq %rax, %xmm0
+; AVX-NEXT:    vmovsd {{.*#+}} xmm0 = mem[0],zero
+; AVX-NEXT:    vorps {{.*}}(%rip), %xmm0, %xmm0
 ; AVX-NEXT:    retq
   %x = load i64, i64* %p
   %b = or i64 %x, 42
   %r = insertelement <2 x i64> undef, i64 %b, i32 0
   ret <2 x i64> %r
 }
 
 define <8 x i16> @xor_op1_constant(i16* %p) nounwind {