[DAG] Move one-use add of splat to base of scatter/gather

[reames]

Summary

This extends the uniform base transform used with scatter/gather to support one-use vector adds-of-splats with a non-zero base. This has the effect of essentially reassociating an add from vector to scalar domain.

The motivation is to improve the lowering of scatter/gather operations fed by complex geps.

[craig.topper]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

10671–10672

I think we typically pass SDLoc by reference

[reames]

Address review comment

[reames]

Corrected rebase

[craig.topper]

LGTM

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

assert(TopHalf && BottomHalf &&
          "same value. This should have been addressed before this function.");
   return DAG.getNode(
       ISD::CONCAT_VECTORS, DL, VT,
       BottomHalf->isZero() ? RHS->getOperand(0) : LHS->getOperand(0),
       TopHalf->isZero() ? RHS->getOperand(1) : LHS->getOperand(1));
 }
 
 bool refineUniformBase(SDValue &BasePtr, SDValue &Index, bool IndexIsScaled,
-                       SelectionDAG &DAG) {
-  if (!isNullConstant(BasePtr) || Index.getOpcode() != ISD::ADD)
+                       SelectionDAG &DAG, const SDLoc &DL) {
+  if (Index.getOpcode() != ISD::ADD)
     return false;
 
   // Only perform the transformation when existing operands can be reused.
   if (IndexIsScaled)
     return false;
 
+  if (!isNullConstant(BasePtr) && !Index.hasOneUse())
+    return false;
+
+  EVT VT = BasePtr.getValueType();
   if (SDValue SplatVal = DAG.getSplatValue(Index.getOperand(0));
-      SplatVal && SplatVal.getValueType() == BasePtr.getValueType()) {
+      SplatVal && SplatVal.getValueType() == VT) {
+    if (isNullConstant(BasePtr))
       BasePtr = SplatVal;
+    else
+      BasePtr = DAG.getNode(ISD::ADD, DL, VT, BasePtr, SplatVal);
     Index = Index.getOperand(1);
     return true;
   }
   if (SDValue SplatVal = DAG.getSplatValue(Index.getOperand(1));
-      SplatVal && SplatVal.getValueType() == BasePtr.getValueType()) {
+      SplatVal && SplatVal.getValueType() == VT) {
+    if (isNullConstant(BasePtr))
       BasePtr = SplatVal;
+    else
+      BasePtr = DAG.getNode(ISD::ADD, DL, VT, BasePtr, SplatVal);
     Index = Index.getOperand(0);
     return true;
   }
   return false;
 }
 
 // Fold sext/zext of index into index type.
 bool refineIndexType(SDValue &Index, ISD::MemIndexType &IndexType, EVT DataVT,
SDValue DAGCombiner::visitVPSCATTER(SDNode *N) {
   SDValue VL = MSC->getVectorLength();
   ISD::MemIndexType IndexType = MSC->getIndexType();
   SDLoc DL(N);
 
   // Zap scatters with a zero mask.
   if (ISD::isConstantSplatVectorAllZeros(Mask.getNode()))
     return Chain;
 
-  if (refineUniformBase(BasePtr, Index, MSC->isIndexScaled(), DAG)) {
+  if (refineUniformBase(BasePtr, Index, MSC->isIndexScaled(), DAG, DL)) {
     SDValue Ops[] = {Chain, StoreVal, BasePtr, Index, Scale, Mask, VL};
     return DAG.getScatterVP(DAG.getVTList(MVT::Other), MSC->getMemoryVT(),
                             DL, Ops, MSC->getMemOperand(), IndexType);
   }
 
   if (refineIndexType(Index, IndexType, StoreVal.getValueType(), DAG)) {
     SDValue Ops[] = {Chain, StoreVal, BasePtr, Index, Scale, Mask, VL};
     return DAG.getScatterVP(DAG.getVTList(MVT::Other), MSC->getMemoryVT(),
SDValue DAGCombiner::visitMSCATTER(SDNode *N) {
   SDValue BasePtr = MSC->getBasePtr();
   ISD::MemIndexType IndexType = MSC->getIndexType();
   SDLoc DL(N);
 
   // Zap scatters with a zero mask.
   if (ISD::isConstantSplatVectorAllZeros(Mask.getNode()))
     return Chain;
 
-  if (refineUniformBase(BasePtr, Index, MSC->isIndexScaled(), DAG)) {
+  if (refineUniformBase(BasePtr, Index, MSC->isIndexScaled(), DAG, DL)) {
     SDValue Ops[] = {Chain, StoreVal, Mask, BasePtr, Index, Scale};
     return DAG.getMaskedScatter(DAG.getVTList(MVT::Other), MSC->getMemoryVT(),
                                 DL, Ops, MSC->getMemOperand(), IndexType,
                                 MSC->isTruncatingStore());
   }
 
   if (refineIndexType(Index, IndexType, StoreVal.getValueType(), DAG)) {
     SDValue Ops[] = {Chain, StoreVal, Mask, BasePtr, Index, Scale};
SDValue DAGCombiner::visitVPGATHER(SDNode *N) {
   SDValue Chain = MGT->getChain();
   SDValue Index = MGT->getIndex();
   SDValue Scale = MGT->getScale();
   SDValue BasePtr = MGT->getBasePtr();
   SDValue VL = MGT->getVectorLength();
   ISD::MemIndexType IndexType = MGT->getIndexType();
   SDLoc DL(N);
 
-  if (refineUniformBase(BasePtr, Index, MGT->isIndexScaled(), DAG)) {
+  if (refineUniformBase(BasePtr, Index, MGT->isIndexScaled(), DAG, DL)) {
     SDValue Ops[] = {Chain, BasePtr, Index, Scale, Mask, VL};
     return DAG.getGatherVP(
         DAG.getVTList(N->getValueType(0), MVT::Other), MGT->getMemoryVT(), DL,
         Ops, MGT->getMemOperand(), IndexType);
   }
 
   if (refineIndexType(Index, IndexType, N->getValueType(0), DAG)) {
     SDValue Ops[] = {Chain, BasePtr, Index, Scale, Mask, VL};
SDValue DAGCombiner::visitMGATHER(SDNode *N) {
   SDValue BasePtr = MGT->getBasePtr();
   ISD::MemIndexType IndexType = MGT->getIndexType();
   SDLoc DL(N);
 
   // Zap gathers with a zero mask.
   if (ISD::isConstantSplatVectorAllZeros(Mask.getNode()))
     return CombineTo(N, PassThru, MGT->getChain());
 
-  if (refineUniformBase(BasePtr, Index, MGT->isIndexScaled(), DAG)) {
+  if (refineUniformBase(BasePtr, Index, MGT->isIndexScaled(), DAG, DL)) {
     SDValue Ops[] = {Chain, PassThru, Mask, BasePtr, Index, Scale};
     return DAG.getMaskedGather(
         DAG.getVTList(N->getValueType(0), MVT::Other), MGT->getMemoryVT(), DL,
         Ops, MGT->getMemOperand(), IndexType, MGT->getExtensionType());
   }
 
   if (refineIndexType(Index, IndexType, N->getValueType(0), DAG)) {
     SDValue Ops[] = {Chain, PassThru, Mask, BasePtr, Index, Scale};

llvm/test/CodeGen/AArch64/sve-gather-scatter-addr-opts.ll

 
 ; Ensure we don't use a "vscale x 4" wide scatter when the offset is too big.
 define void @scatter_i8_index_offset_maximum_plus_one(i8* %base, i64 %offset, <vscale x 4 x i1> %pg, <vscale x 4 x i8> %data) #0 {
 ; CHECK-LABEL: scatter_i8_index_offset_maximum_plus_one:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    rdvl x8, #1
 ; CHECK-NEXT:    mov w9, #67108864
 ; CHECK-NEXT:    lsr x8, x8, #4
-; CHECK-NEXT:    add x10, x0, x1
+; CHECK-NEXT:    add x11, x0, x1
+; CHECK-NEXT:    mov w10, #33554432
 ; CHECK-NEXT:    punpklo p1.h, p0.b
-; CHECK-NEXT:    uunpklo z3.d, z0.s
-; CHECK-NEXT:    mul x8, x8, x9
-; CHECK-NEXT:    mov w9, #33554432
+; CHECK-NEXT:    madd x8, x8, x9, x11
+; CHECK-NEXT:    uunpklo z2.d, z0.s
 ; CHECK-NEXT:    punpkhi p0.h, p0.b
 ; CHECK-NEXT:    uunpkhi z0.d, z0.s
-; CHECK-NEXT:    index z1.d, #0, x9
-; CHECK-NEXT:    mov z2.d, x8
-; CHECK-NEXT:    st1b { z3.d }, p1, [x10, z1.d]
-; CHECK-NEXT:    add z2.d, z1.d, z2.d
-; CHECK-NEXT:    st1b { z0.d }, p0, [x10, z2.d]
+; CHECK-NEXT:    index z1.d, #0, x10
+; CHECK-NEXT:    st1b { z2.d }, p1, [x11, z1.d]
+; CHECK-NEXT:    st1b { z0.d }, p0, [x8, z1.d]
 ; CHECK-NEXT:    ret
   %t0 = insertelement <vscale x 4 x i64> undef, i64 %offset, i32 0
   %t1 = shufflevector <vscale x 4 x i64> %t0, <vscale x 4 x i64> undef, <vscale x 4 x i32> zeroinitializer
   %t2 = insertelement <vscale x 4 x i64> undef, i64 33554432, i32 0
   %t3 = shufflevector <vscale x 4 x i64> %t2, <vscale x 4 x i64> undef, <vscale x 4 x i32> zeroinitializer
   %step = call <vscale x 4 x i64> @llvm.experimental.stepvector.nxv4i64()
   %t4 = mul <vscale x 4 x i64> %t3, %step
   %t5 = add <vscale x 4 x i64> %t1, %t4
   %t6 = getelementptr i8, i8* %base, <vscale x 4 x i64> %t5
   call void @llvm.masked.scatter.nxv4i8(<vscale x 4 x i8> %data, <vscale x 4 x i8*> %t6, i32 2, <vscale x 4 x i1> %pg)
   ret void
 }
 
 ; Ensure we don't use a "vscale x 4" wide scatter when the offset is too small.
 define void @scatter_i8_index_offset_minimum_minus_one(i8* %base, i64 %offset, <vscale x 4 x i1> %pg, <vscale x 4 x i8> %data) #0 {
 ; CHECK-LABEL: scatter_i8_index_offset_minimum_minus_one:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    rdvl x8, #1
 ; CHECK-NEXT:    mov x9, #-2
 ; CHECK-NEXT:    lsr x8, x8, #4
 ; CHECK-NEXT:    movk x9, #64511, lsl #16
-; CHECK-NEXT:    add x10, x0, x1
+; CHECK-NEXT:    add x11, x0, x1
+; CHECK-NEXT:    mov x10, #-33554433
+; CHECK-NEXT:    madd x8, x8, x9, x11
 ; CHECK-NEXT:    punpklo p1.h, p0.b
-; CHECK-NEXT:    mul x8, x8, x9
-; CHECK-NEXT:    mov x9, #-33554433
-; CHECK-NEXT:    uunpklo z3.d, z0.s
+; CHECK-NEXT:    uunpklo z2.d, z0.s
 ; CHECK-NEXT:    punpkhi p0.h, p0.b
+; CHECK-NEXT:    index z1.d, #0, x10
 ; CHECK-NEXT:    uunpkhi z0.d, z0.s
-; CHECK-NEXT:    index z1.d, #0, x9
-; CHECK-NEXT:    mov z2.d, x8
-; CHECK-NEXT:    st1b { z3.d }, p1, [x10, z1.d]
-; CHECK-NEXT:    add z2.d, z1.d, z2.d
-; CHECK-NEXT:    st1b { z0.d }, p0, [x10, z2.d]
+; CHECK-NEXT:    st1b { z2.d }, p1, [x11, z1.d]
+; CHECK-NEXT:    st1b { z0.d }, p0, [x8, z1.d]
 ; CHECK-NEXT:    ret
   %t0 = insertelement <vscale x 4 x i64> undef, i64 %offset, i32 0
   %t1 = shufflevector <vscale x 4 x i64> %t0, <vscale x 4 x i64> undef, <vscale x 4 x i32> zeroinitializer
   %t2 = insertelement <vscale x 4 x i64> undef, i64 -33554433, i32 0
   %t3 = shufflevector <vscale x 4 x i64> %t2, <vscale x 4 x i64> undef, <vscale x 4 x i32> zeroinitializer
   %step = call <vscale x 4 x i64> @llvm.experimental.stepvector.nxv4i64()
   %t4 = mul <vscale x 4 x i64> %t3, %step
   %t5 = add <vscale x 4 x i64> %t1, %t4
   %t6 = getelementptr i8, i8* %base, <vscale x 4 x i64> %t5
   call void @llvm.masked.scatter.nxv4i8(<vscale x 4 x i8> %data, <vscale x 4 x i8*> %t6, i32 2, <vscale x 4 x i1> %pg)
   ret void
 }
 
 ; Ensure we don't use a "vscale x 4" wide scatter when the stride is too big .
 define void @scatter_i8_index_stride_too_big(i8* %base, i64 %offset, <vscale x 4 x i1> %pg, <vscale x 4 x i8> %data) #0 {
 ; CHECK-LABEL: scatter_i8_index_stride_too_big:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    rdvl x8, #1
 ; CHECK-NEXT:    mov x9, #-9223372036854775808
 ; CHECK-NEXT:    lsr x8, x8, #4
-; CHECK-NEXT:    add x10, x0, x1
+; CHECK-NEXT:    add x11, x0, x1
+; CHECK-NEXT:    mov x10, #4611686018427387904
 ; CHECK-NEXT:    punpklo p1.h, p0.b
-; CHECK-NEXT:    uunpklo z3.d, z0.s
-; CHECK-NEXT:    mul x8, x8, x9
-; CHECK-NEXT:    mov x9, #4611686018427387904
+; CHECK-NEXT:    madd x8, x8, x9, x11
+; CHECK-NEXT:    uunpklo z2.d, z0.s
 ; CHECK-NEXT:    punpkhi p0.h, p0.b
 ; CHECK-NEXT:    uunpkhi z0.d, z0.s
-; CHECK-NEXT:    index z1.d, #0, x9
-; CHECK-NEXT:    mov z2.d, x8
-; CHECK-NEXT:    st1b { z3.d }, p1, [x10, z1.d]
-; CHECK-NEXT:    add z2.d, z1.d, z2.d
-; CHECK-NEXT:    st1b { z0.d }, p0, [x10, z2.d]
+; CHECK-NEXT:    index z1.d, #0, x10
+; CHECK-NEXT:    st1b { z2.d }, p1, [x11, z1.d]
+; CHECK-NEXT:    st1b { z0.d }, p0, [x8, z1.d]
 ; CHECK-NEXT:    ret
   %t0 = insertelement <vscale x 4 x i64> undef, i64 %offset, i32 0
   %t1 = shufflevector <vscale x 4 x i64> %t0, <vscale x 4 x i64> undef, <vscale x 4 x i32> zeroinitializer
   %t2 = insertelement <vscale x 4 x i64> undef, i64 4611686018427387904, i32 0
   %t3 = shufflevector <vscale x 4 x i64> %t2, <vscale x 4 x i64> undef, <vscale x 4 x i32> zeroinitializer
   %step = call <vscale x 4 x i64> @llvm.experimental.stepvector.nxv4i64()
   %t4 = mul <vscale x 4 x i64> %t3, %step
   %t5 = add <vscale x 4 x i64> %t1, %t4

llvm/test/CodeGen/RISCV/rvv/mscatter-combine.ll

 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
 ; RUN: llc -mtriple=riscv32 -mattr=+m,+d,+zfh,+experimental-zvfh,+v -target-abi=ilp32d \
 ; RUN:     -verify-machineinstrs < %s | FileCheck %s --check-prefixes=RV32
 ; RUN: llc -mtriple=riscv64 -mattr=+m,+d,+zfh,+experimental-zvfh,+v -target-abi=lp64d \
 ; RUN:     -verify-machineinstrs < %s | FileCheck %s --check-prefixes=RV64
 
 %struct = type { i64, i64, ptr, i32, i32, i32, [4 x i32] }
 
 define void @complex_gep(ptr %p, <vscale x 2 x i64> %vec.ind, <vscale x 2 x i1> %m) {
 ; RV32-LABEL: complex_gep:
 ; RV32:       # %bb.0:
 ; RV32-NEXT:    vsetvli a1, zero, e32, m1, ta, mu
-; RV32-NEXT:    vmv.v.x v10, a0
-; RV32-NEXT:    vnsrl.wi v11, v8, 0
-; RV32-NEXT:    li a0, 48
-; RV32-NEXT:    vmadd.vx v11, a0, v10
-; RV32-NEXT:    vmv.v.i v8, 0
-; RV32-NEXT:    li a0, 28
-; RV32-NEXT:    vsoxei32.v v8, (a0), v11, v0.t
+; RV32-NEXT:    vnsrl.wi v10, v8, 0
+; RV32-NEXT:    li a1, 48
+; RV32-NEXT:    vmul.vx v8, v10, a1
+; RV32-NEXT:    addi a0, a0, 28
+; RV32-NEXT:    vmv.v.i v9, 0
+; RV32-NEXT:    vsoxei32.v v9, (a0), v8, v0.t
 ; RV32-NEXT:    ret
 ;
 ; RV64-LABEL: complex_gep:
 ; RV64:       # %bb.0:
-; RV64-NEXT:    vsetvli a1, zero, e64, m2, ta, mu
-; RV64-NEXT:    vmv.v.x v10, a0
-; RV64-NEXT:    li a0, 56
-; RV64-NEXT:    vmacc.vx v10, a0, v8
+; RV64-NEXT:    li a1, 56
+; RV64-NEXT:    vsetvli a2, zero, e64, m2, ta, mu
+; RV64-NEXT:    vmul.vx v8, v8, a1
+; RV64-NEXT:    addi a0, a0, 32
 ; RV64-NEXT:    vsetvli zero, zero, e32, m1, ta, mu
-; RV64-NEXT:    vmv.v.i v8, 0
-; RV64-NEXT:    li a0, 32
-; RV64-NEXT:    vsoxei64.v v8, (a0), v10, v0.t
+; RV64-NEXT:    vmv.v.i v10, 0
+; RV64-NEXT:    vsoxei64.v v10, (a0), v8, v0.t
 ; RV64-NEXT:    ret
   %gep = getelementptr inbounds %struct, ptr %p, <vscale x 2 x i64> %vec.ind, i32 5
   call void @llvm.masked.scatter.nxv2i32.nxv2p0(<vscale x 2 x i32> zeroinitializer, <vscale x 2 x ptr> %gep, i32 8, <vscale x 2 x i1> %m)
   ret void
 }
 
 declare void @llvm.masked.scatter.nxv2i32.nxv2p0(<vscale x 2 x i32>, <vscale x 2 x ptr>, i32, <vscale x 2 x i1>)