[X86][AVX] Attempt to fold a scaled index into a gather/scatter scale immediate (PR13310)

[RKSimon]

Summary

If the index operand for a gather/scatter intrinsic is being scaled (self-addition or a shl-by-immediate) then we may be able to fold that scaling into the intrinsic scale immediate value instead.

Fixes PR13310.

[RKSimon]

Fix clang-format warnings

[pengfei]

llvm/test/CodeGen/X86/masked_gather_scatter.ll

774–775

Why do we need an extra xor now?

[RKSimon]

I'm going to see whether this would be better moving to matchVectorAddressRecursively now that D111595 has landed.

llvm/test/CodeGen/X86/masked_gather_scatter.ll

774–775

D112505 - for dependency breaking I guess :)

[pengfei]

llvm/test/CodeGen/X86/masked_gather_scatter.ll

774–775

Interesting :)

[craig.topper]

llvm/test/CodeGen/X86/masked_gather_scatter.ll

774–775

I guess this is because we were calling getAVX2Gather which adds dependency breaking?

[RKSimon]

rebase

[RKSimon]

rebase

[pengfei]

LGTM.

[craig.topper]

llvm/lib/Target/X86/X86ISelLowering.cpp

50239

Does the Index element size need to be pointer width for this to be legal? The implicit sign extend is applied before the scale.

[RKSimon]

llvm/lib/Target/X86/X86ISelLowering.cpp

50239

Nice catch! I'll add a check/test.

[craig.topper]

llvm/lib/Target/X86/X86ISelLowering.cpp

50252

Is there a common helper function we can split off of getAVX2GatherNode and getGatherNode to use here. Most of the code in getAVX2GatherNode isn't needed for this path and the Mask handling in getAVX2GatherNode happens to work for AVX512 gathers, but wasn't intentional.

The helper could probably be shared with the getMemIntrinsicNode call in LowerMGATHER as well.

Similar for getScatterNode and the getMemIntrinsicNode call in LowerMSCATTER.

llvm/lib/Target/X86/X86ISelLowering.cpp

   APInt DemandedMask(APInt::getAllOnesValue(NumBits));
   if (TLI.SimplifyDemandedBits(SDValue(N, 0), DemandedMask, DCI))
     return SDValue(N, 0);
 
   return SDValue();
 }
 
 static SDValue combineX86GatherScatter(SDNode *N, SelectionDAG &DAG,
-                                       TargetLowering::DAGCombinerInfo &DCI) {
+                                       TargetLowering::DAGCombinerInfo &DCI,
+                                       const X86Subtarget &Subtarget) {
+  auto *MemOp = cast<X86MaskedGatherScatterSDNode>(N);
+  SDValue Index = MemOp->getIndex();
+  SDValue Scale = MemOp->getScale();
+  SDValue Mask = MemOp->getMask();
+
+  // Attempt to fold an index scale into the scale value directly.
+  // TODO: Could we move this into a gather/scatter equivalent of
+  //       X86DAGToDAGISel::matchAddressRecursively?
+  if ((Index.getOpcode() == X86ISD::VSHLI ||
+       (Index.getOpcode() == ISD::ADD &&
+        Index.getOperand(0) == Index.getOperand(1))) &&
+      isa<ConstantSDNode>(Scale)) {
+    unsigned ShiftAmt =
+        Index.getOpcode() == ISD::ADD ? 1 : Index.getConstantOperandVal(1);
+    uint64_t ScaleAmt = cast<ConstantSDNode>(Scale)->getZExtValue();
+    uint64_t NewScaleAmt = ScaleAmt * (1ULL << ShiftAmt);
+    if (isPowerOf2_64(NewScaleAmt) && NewScaleAmt <= 8) {
+      SDValue NewIndex = Index.getOperand(0);
+      SDValue NewScale = DAG.getTargetConstant(NewScaleAmt, SDLoc(N),
+                                               Scale.getValueType());
+      if (N->getOpcode() == X86ISD::MGATHER)
+        return getAVX2GatherNode(N->getOpcode(), SDValue(N, 0), DAG,
+                                 MemOp->getOperand(1), Mask,
+                                 MemOp->getBasePtr(), NewIndex, NewScale,
+                                 MemOp->getChain(), Subtarget);
+      if (N->getOpcode() == X86ISD::MSCATTER)
+        return getScatterNode(N->getOpcode(), SDValue(N, 0), DAG,
+                              MemOp->getOperand(1), Mask, MemOp->getBasePtr(),
+                              NewIndex, NewScale, MemOp->getChain(), Subtarget);
+    }
+  }
+
   // With vector masks we only demand the upper bit of the mask.
-  SDValue Mask = cast<X86MaskedGatherScatterSDNode>(N)->getMask();
   if (Mask.getScalarValueSizeInBits() != 1) {
     const TargetLowering &TLI = DAG.getTargetLoweringInfo();
     APInt DemandedMask(APInt::getSignMask(Mask.getScalarValueSizeInBits()));
     if (TLI.SimplifyDemandedBits(Mask, DemandedMask, DCI)) {
       if (N->getOpcode() != ISD::DELETED_NODE)
         DCI.AddToWorklist(N);
       return SDValue(N, 0);
     }

   bool IsSub = N->getOpcode() == ISD::SUB;
   SDValue X = N->getOperand(0);
   SDValue Y = N->getOperand(1);
 
   // If this is an add, canonicalize a zext operand to the RHS.
   // TODO: Incomplete? What if both sides are zexts?
   if (!IsSub && X.getOpcode() == ISD::ZERO_EXTEND &&
       Y.getOpcode() != ISD::ZERO_EXTEND)
     std::swap(X, Y);
 
   // Look through a one-use zext.
   bool PeekedThroughZext = false;
   if (Y.getOpcode() == ISD::ZERO_EXTEND && Y.hasOneUse()) {
     Y = Y.getOperand(0);
     PeekedThroughZext = true;
   }
 
   // If this is an add, canonicalize a setcc operand to the RHS.
   // TODO: Incomplete? What if both sides are setcc?
   // TODO: Should we allow peeking through a zext of the other operand?
   if (!IsSub && !PeekedThroughZext && X.getOpcode() == X86ISD::SETCC &&
       Y.getOpcode() != X86ISD::SETCC)
     std::swap(X, Y);
 
   if (Y.getOpcode() != X86ISD::SETCC || !Y.hasOneUse())
     return SDValue();
 
   SDLoc DL(N);
   EVT VT = N->getValueType(0);
   X86::CondCode CC = (X86::CondCode)Y.getConstantOperandVal(0);

   case ISD::FMA:
   case ISD::STRICT_FMA:     return combineFMA(N, DAG, DCI, Subtarget);
   case X86ISD::FMADDSUB_RND:
   case X86ISD::FMSUBADD_RND:
   case X86ISD::FMADDSUB:
   case X86ISD::FMSUBADD:    return combineFMADDSUB(N, DAG, DCI);
   case X86ISD::MOVMSK:      return combineMOVMSK(N, DAG, DCI, Subtarget);
   case X86ISD::MGATHER:
-  case X86ISD::MSCATTER:    return combineX86GatherScatter(N, DAG, DCI);
+  case X86ISD::MSCATTER:    return combineX86GatherScatter(N, DAG, DCI, Subtarget);
   case ISD::MGATHER:
   case ISD::MSCATTER:       return combineGatherScatter(N, DAG, DCI);
   case X86ISD::PCMPEQ:
   case X86ISD::PCMPGT:      return combineVectorCompare(N, DAG, Subtarget);
   case X86ISD::PMULDQ:
   case X86ISD::PMULUDQ:     return combinePMULDQ(N, DAG, DCI, Subtarget);
   case X86ISD::KSHIFTL:
   case X86ISD::KSHIFTR:     return combineKSHIFT(N, DAG, DCI);

llvm/test/CodeGen/X86/masked_gather_scatter.ll

 ; The base pointer is not splat, can't find unform base
 define <16 x float> @test14(float* %base, i32 %ind, <16 x float*> %vec) {
 ; KNL_64-LABEL: test14:
 ; KNL_64:       # %bb.0:
 ; KNL_64-NEXT:    vmovq %xmm0, %rax
 ; KNL_64-NEXT:    vmovd %esi, %xmm0
 ; KNL_64-NEXT:    vpbroadcastd %xmm0, %ymm0
 ; KNL_64-NEXT:    vpmovsxdq %ymm0, %zmm0
-; KNL_64-NEXT:    vpsllq $2, %zmm0, %zmm0
 ; KNL_64-NEXT:    kxnorw %k0, %k0, %k1
-; KNL_64-NEXT:    vgatherqps (%rax,%zmm0), %ymm1 {%k1}
+; KNL_64-NEXT:    vxorps %xmm1, %xmm1, %xmm1
+; KNL_64-NEXT:    vgatherqps (%rax,%zmm0,4), %ymm1 {%k1}
 ; KNL_64-NEXT:    vinsertf64x4 $1, %ymm1, %zmm1, %zmm0
 ; KNL_64-NEXT:    retq
 ;
 ; KNL_32-LABEL: test14:
 ; KNL_32:       # %bb.0:
 ; KNL_32-NEXT:    vmovd %xmm0, %eax
-; KNL_32-NEXT:    vpslld $2, {{[0-9]+}}(%esp){1to16}, %zmm1
+; KNL_32-NEXT:    vbroadcastss {{[0-9]+}}(%esp), %zmm1
 ; KNL_32-NEXT:    kxnorw %k0, %k0, %k1
-; KNL_32-NEXT:    vgatherdps (%eax,%zmm1), %zmm0 {%k1}
+; KNL_32-NEXT:    vpxor %xmm0, %xmm0, %xmm0
+; KNL_32-NEXT:    vgatherdps (%eax,%zmm1,4), %zmm0 {%k1}
 ; KNL_32-NEXT:    retl
 ;
 ; SKX-LABEL: test14:
 ; SKX:       # %bb.0:
 ; SKX-NEXT:    vmovq %xmm0, %rax
 ; SKX-NEXT:    vpbroadcastd %esi, %ymm0
 ; SKX-NEXT:    vpmovsxdq %ymm0, %zmm0
-; SKX-NEXT:    vpsllq $2, %zmm0, %zmm0
 ; SKX-NEXT:    kxnorw %k0, %k0, %k1
-; SKX-NEXT:    vgatherqps (%rax,%zmm0), %ymm1 {%k1}
+; SKX-NEXT:    vxorps %xmm1, %xmm1, %xmm1
+; SKX-NEXT:    vgatherqps (%rax,%zmm0,4), %ymm1 {%k1}
 ; SKX-NEXT:    vinsertf64x4 $1, %ymm1, %zmm1, %zmm0
 ; SKX-NEXT:    retq
 ;
 ; SKX_32-LABEL: test14:
 ; SKX_32:       # %bb.0:
 ; SKX_32-NEXT:    vmovd %xmm0, %eax
-; SKX_32-NEXT:    vpslld $2, {{[0-9]+}}(%esp){1to16}, %zmm1
+; SKX_32-NEXT:    vbroadcastss {{[0-9]+}}(%esp), %zmm1
 ; SKX_32-NEXT:    kxnorw %k0, %k0, %k1
-; SKX_32-NEXT:    vgatherdps (%eax,%zmm1), %zmm0 {%k1}
+; SKX_32-NEXT:    vpxor %xmm0, %xmm0, %xmm0
+; SKX_32-NEXT:    vgatherdps (%eax,%zmm1,4), %zmm0 {%k1}
 ; SKX_32-NEXT:    retl
 
   %broadcast.splatinsert = insertelement <16 x float*> %vec, float* %base, i32 1
   %broadcast.splat = shufflevector <16 x float*> %broadcast.splatinsert, <16 x float*> undef, <16 x i32> zeroinitializer
 
   %gep.random = getelementptr float, <16 x float*> %broadcast.splat, i32 %ind
 
   %res = call <16 x float> @llvm.masked.gather.v16f32.v16p0f32(<16 x float*> %gep.random, i32 4, <16 x i1> <i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true>, <16 x float> undef)
; SKX_32-NEXT:    retl
   %1 = insertelement <4 x i32*> undef, i32* %ptr, i32 0
   %2 = shufflevector <4 x i32*> %1, <4 x i32*> undef, <4 x i32> zeroinitializer
   call void @llvm.masked.scatter.v4i32.v4p0i32(<4 x i32> %val, <4 x i32*> %2, i32 4, <4 x i1> %mask)
   ret void
 }
 
 ;
 ; PR13310
-; FIXME: Failure to fold scaled-index into gather/scatter scale operand.
+; Attempt to fold scaled-index into gather/scatter scale operand.
 ;
 
 define <8 x float> @scaleidx_x86gather(float* %base, <8 x i32> %index, <8 x i32> %imask) nounwind {
 ; KNL_64-LABEL: scaleidx_x86gather:
 ; KNL_64:       # %bb.0:
-; KNL_64-NEXT:    vpslld $2, %ymm0, %ymm2
-; KNL_64-NEXT:    vpxor %xmm0, %xmm0, %xmm0
-; KNL_64-NEXT:    vgatherdps %ymm1, (%rdi,%ymm2), %ymm0
+; KNL_64-NEXT:    vxorps %xmm2, %xmm2, %xmm2
+; KNL_64-NEXT:    vgatherdps %ymm1, (%rdi,%ymm0,4), %ymm2
+; KNL_64-NEXT:    vmovaps %ymm2, %ymm0
 ; KNL_64-NEXT:    retq
 ;
 ; KNL_32-LABEL: scaleidx_x86gather:
 ; KNL_32:       # %bb.0:
 ; KNL_32-NEXT:    movl {{[0-9]+}}(%esp), %eax
-; KNL_32-NEXT:    vpslld $2, %ymm0, %ymm2
-; KNL_32-NEXT:    vpxor %xmm0, %xmm0, %xmm0
-; KNL_32-NEXT:    vgatherdps %ymm1, (%eax,%ymm2), %ymm0
+; KNL_32-NEXT:    vxorps %xmm2, %xmm2, %xmm2
+; KNL_32-NEXT:    vgatherdps %ymm1, (%eax,%ymm0,4), %ymm2
+; KNL_32-NEXT:    vmovaps %ymm2, %ymm0
 ; KNL_32-NEXT:    retl
 ;
 ; SKX-LABEL: scaleidx_x86gather:
 ; SKX:       # %bb.0:
-; SKX-NEXT:    vpslld $2, %ymm0, %ymm2
-; SKX-NEXT:    vpxor %xmm0, %xmm0, %xmm0
-; SKX-NEXT:    vgatherdps %ymm1, (%rdi,%ymm2), %ymm0
+; SKX-NEXT:    vxorps %xmm2, %xmm2, %xmm2
+; SKX-NEXT:    vgatherdps %ymm1, (%rdi,%ymm0,4), %ymm2
+; SKX-NEXT:    vmovaps %ymm2, %ymm0
 ; SKX-NEXT:    retq
 ;
 ; SKX_32-LABEL: scaleidx_x86gather:
 ; SKX_32:       # %bb.0:
 ; SKX_32-NEXT:    movl {{[0-9]+}}(%esp), %eax
-; SKX_32-NEXT:    vpslld $2, %ymm0, %ymm2
-; SKX_32-NEXT:    vpxor %xmm0, %xmm0, %xmm0
-; SKX_32-NEXT:    vgatherdps %ymm1, (%eax,%ymm2), %ymm0
+; SKX_32-NEXT:    vxorps %xmm2, %xmm2, %xmm2
+; SKX_32-NEXT:    vgatherdps %ymm1, (%eax,%ymm0,4), %ymm2
+; SKX_32-NEXT:    vmovaps %ymm2, %ymm0
 ; SKX_32-NEXT:    retl
   %ptr = bitcast float* %base to i8*
   %mask = bitcast <8 x i32> %imask to <8 x float>
   %scaledindex = mul <8 x i32> %index, <i32 4, i32 4, i32 4, i32 4, i32 4, i32 4, i32 4, i32 4>
   %v = call <8 x float> @llvm.x86.avx2.gather.d.ps.256(<8 x float> undef, i8* %ptr, <8 x i32> %scaledindex, <8 x float> %mask, i8 1) nounwind
   ret <8 x float> %v
 }
 
; SKX_32-NEXT:    retl
   ret <8 x float> %v
 }
 declare <8 x float> @llvm.x86.avx2.gather.d.ps.256(<8 x float>, i8*, <8 x i32>, <8 x float>, i8) nounwind readonly
 
 define void @scaleidx_x86scatter(<16 x float> %value, float* %base, <16 x i32> %index, i16 %imask) nounwind {
 ; KNL_64-LABEL: scaleidx_x86scatter:
 ; KNL_64:       # %bb.0:
 ; KNL_64-NEXT:    kmovw %esi, %k1
-; KNL_64-NEXT:    vpaddd %zmm1, %zmm1, %zmm1
-; KNL_64-NEXT:    vscatterdps %zmm0, (%rdi,%zmm1,2) {%k1}
+; KNL_64-NEXT:    vscatterdps %zmm0, (%rdi,%zmm1,4) {%k1}
 ; KNL_64-NEXT:    vzeroupper
 ; KNL_64-NEXT:    retq
 ;
 ; KNL_32-LABEL: scaleidx_x86scatter:
 ; KNL_32:       # %bb.0:
 ; KNL_32-NEXT:    movl {{[0-9]+}}(%esp), %eax
 ; KNL_32-NEXT:    kmovw {{[0-9]+}}(%esp), %k1
-; KNL_32-NEXT:    vpaddd %zmm1, %zmm1, %zmm1
-; KNL_32-NEXT:    vscatterdps %zmm0, (%eax,%zmm1,2) {%k1}
+; KNL_32-NEXT:    vscatterdps %zmm0, (%eax,%zmm1,4) {%k1}
 ; KNL_32-NEXT:    vzeroupper
 ; KNL_32-NEXT:    retl
 ;
 ; SKX-LABEL: scaleidx_x86scatter:
 ; SKX:       # %bb.0:
 ; SKX-NEXT:    kmovw %esi, %k1
-; SKX-NEXT:    vpaddd %zmm1, %zmm1, %zmm1
-; SKX-NEXT:    vscatterdps %zmm0, (%rdi,%zmm1,2) {%k1}
+; SKX-NEXT:    vscatterdps %zmm0, (%rdi,%zmm1,4) {%k1}
 ; SKX-NEXT:    vzeroupper
 ; SKX-NEXT:    retq
 ;
 ; SKX_32-LABEL: scaleidx_x86scatter:
 ; SKX_32:       # %bb.0:
 ; SKX_32-NEXT:    movl {{[0-9]+}}(%esp), %eax
 ; SKX_32-NEXT:    kmovw {{[0-9]+}}(%esp), %k1
-; SKX_32-NEXT:    vpaddd %zmm1, %zmm1, %zmm1
-; SKX_32-NEXT:    vscatterdps %zmm0, (%eax,%zmm1,2) {%k1}
+; SKX_32-NEXT:    vscatterdps %zmm0, (%eax,%zmm1,4) {%k1}
 ; SKX_32-NEXT:    vzeroupper
 ; SKX_32-NEXT:    retl
   %ptr = bitcast float* %base to i8*
   %mask = bitcast i16 %imask to <16 x i1>
   %scaledindex = shl <16 x i32> %index, <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>
   call void @llvm.x86.avx512.mask.scatter.dps.512(i8* %ptr, <16 x i1> %mask, <16 x i32> %scaledindex, <16 x float> %value, i32 2)
   ret void
 }

 ; KNL_32-NEXT:    kshiftlw $8, %k0, %k0
 ; KNL_32-NEXT:    kshiftrw $8, %k0, %k1
 ; KNL_32-NEXT:    vscatterdps %zmm0, (%eax,%zmm1,4) {%k1}
 ; KNL_32-NEXT:    vzeroupper
 ; KNL_32-NEXT:    retl
 ;
 ; SKX-LABEL: scaleidx_scatter:
 ; SKX:       # %bb.0:
-; SKX-NEXT:    vpaddd %ymm1, %ymm1, %ymm1
 ; SKX-NEXT:    kmovw %esi, %k1
-; SKX-NEXT:    vscatterdps %ymm0, (%rdi,%ymm1,4) {%k1}
+; SKX-NEXT:    vscatterdps %ymm0, (%rdi,%ymm1,8) {%k1}
 ; SKX-NEXT:    vzeroupper
 ; SKX-NEXT:    retq
 ;
 ; SKX_32-LABEL: scaleidx_scatter:
 ; SKX_32:       # %bb.0:
 ; SKX_32-NEXT:    movl {{[0-9]+}}(%esp), %eax
-; SKX_32-NEXT:    vpaddd %ymm1, %ymm1, %ymm1
 ; SKX_32-NEXT:    kmovb {{[0-9]+}}(%esp), %k1
-; SKX_32-NEXT:    vscatterdps %ymm0, (%eax,%ymm1,4) {%k1}
+; SKX_32-NEXT:    vscatterdps %ymm0, (%eax,%ymm1,8) {%k1}
 ; SKX_32-NEXT:    vzeroupper
 ; SKX_32-NEXT:    retl
   %scaledindex = mul <8 x i32> %index, <i32 2, i32 2, i32 2, i32 2, i32 2, i32 2, i32 2, i32 2>
   %ptrs = getelementptr float, float* %base, <8 x i32> %scaledindex
   %mask = bitcast i8 %imask to <8 x i1>
   call void @llvm.masked.scatter.v8f32.v8p0f32(<8 x float> %value, <8 x float*> %ptrs, i32 1, <8 x i1> %mask)
   ret void
 }