[RISCV] Match strided loads with reversed indexing sequences

[reames]

Summary

This extends the concat_vector of loads to strided_load transform to handle reversed index pattern. The previous code expected indexing of the form (a0, a1+S, a2+S,...). However, we can also see indexing of the form (a1+S, a2+S, a3+S, .., aS). This form is a strided load starting at address aN + S*(n-1) with stride -S.

Note that this is also fixing what looks to be a bug in the memory location reasoning for forward strided case. A strided load with negative stride access eltsize bytes past base ptr, and then bytes *before* base ptr. (That is, the range should extend from before base ptr to after base ptr.)

[luke]

Makes sense to me

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

12900

Can we do for (SDValue Op : N->ops()) instead to remove the need for the extra isSimple() and isNormalLoad() checks above?

12915

Just an idea, can we use one lambda and pass in a reverse iterator of the Ptrs vector?

llvm/test/CodeGen/RISCV/rvv/fixed-vectors-strided-load-combine.ll

497

We can strike off this todo

[reames]

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

12900

We can, but we still need the BaseLd and BaseLdVT variables. Given that, I didn't think folding only part of the checks into the loop was worthwhile.

12915

Not unless you know how to do a reverse enumerate with a "first element" check. We need the index check and the index + or - one bits. I played with this and it was the best variant I found.

llvm/test/CodeGen/RISCV/rvv/fixed-vectors-strided-load-combine.ll

497

Good catch. I need to add this to my mental checklist, it's one of my regular oops.

[reames]

Address review comment

[luke]

LGTM

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

static SDValue performCONCAT_VECTORSCombine(SDNode *N, SelectionDAG &DAG,
   if (!BaseLd || !BaseLd->isSimple() || !ISD::isNormalLoad(BaseLd) ||
       !SDValue(BaseLd, 0).hasOneUse())
     return SDValue();
 
   EVT BaseLdVT = BaseLd->getValueType(0);
   SDValue BasePtr = BaseLd->getBasePtr();
 
   // Go through the loads and check that they're strided
-  SDValue CurPtr = BasePtr;
-  SDValue Stride;
+  SmallVector<SDValue> Ptrs;
+  Ptrs.push_back(BasePtr);
   Align Align = BaseLd->getAlign();
-
   for (SDValue Op : N->ops().drop_front()) {
     auto *Ld = dyn_cast<LoadSDNode>(Op);
     if (!Ld || !Ld->isSimple() || !Op.hasOneUse() ||
         Ld->getChain() != BaseLd->getChain() || !ISD::isNormalLoad(Ld) ||
         Ld->getValueType(0) != BaseLdVT)
       return SDValue();
 
-    SDValue Ptr = Ld->getBasePtr();
-    // Check that each load's pointer is (add CurPtr, Stride)
-    if (Ptr.getOpcode() != ISD::ADD || Ptr.getOperand(0) != CurPtr)
+    Ptrs.push_back(Ld->getBasePtr());
+
+    // The common alignment is the most restrictive (smallest) of all the loads
+    Align = std::min(Align, Ld->getAlign());
+  }
+
+  auto matchForwardStrided = [](ArrayRef<SDValue> Ptrs) {
+    SDValue Stride;
+    for (auto Idx : enumerate(Ptrs)) {
+      if (Idx.index() == 0)
+        continue;
+      SDValue Ptr = Idx.value();
+      // Check that each load's pointer is (add LastPtr, Stride)
+      if (Ptr.getOpcode() != ISD::ADD ||
+          Ptr.getOperand(0) != Ptrs[Idx.index()-1])
         return SDValue();
       SDValue Offset = Ptr.getOperand(1);
       if (!Stride)
         Stride = Offset;
       else if (Offset != Stride)
         return SDValue();
+    }
+    return Stride;
+  };
+  auto matchReverseStrided = [](ArrayRef<SDValue> Ptrs) {
+    SDValue Stride;
+    for (auto Idx : enumerate(Ptrs)) {
+      if (Idx.index() == Ptrs.size() - 1)
+        continue;
+      SDValue Ptr = Idx.value();
+      // Check that each load's pointer is (add NextPtr, Stride)
+      if (Ptr.getOpcode() != ISD::ADD ||
+          Ptr.getOperand(0) != Ptrs[Idx.index()+1])
+        return SDValue();
+      SDValue Offset = Ptr.getOperand(1);
+      if (!Stride)
+        Stride = Offset;
+      else if (Offset != Stride)
+        return SDValue();
+    }
+    return Stride;
+  };
 
-    // The common alignment is the most restrictive (smallest) of all the loads
-    Align = std::min(Align, Ld->getAlign());
-
-    CurPtr = Ptr;
+  bool Reversed = false;
+  SDValue Stride = matchForwardStrided(Ptrs);
+  if (!Stride) {
+    Stride = matchReverseStrided(Ptrs);
+    Reversed = true;
+    // TODO: At this point, we've successfully matched a generalized gather
+    // load.  Maybe we should emit that, and then move the specialized
+    // matchers above and below into a DAG combine?
+    if (!Stride)
+      return SDValue();
   }
 
   // A special case is if the stride is exactly the width of one of the loads,
   // in which case it's contiguous and can be combined into a regular vle
   // without changing the element size
   if (auto *ConstStride = dyn_cast<ConstantSDNode>(Stride);
-      ConstStride &&
+      ConstStride && !Reversed &&
       ConstStride->getZExtValue() == BaseLdVT.getFixedSizeInBits() / 8) {
     MachineMemOperand *MMO = DAG.getMachineFunction().getMachineMemOperand(
         BaseLd->getPointerInfo(), BaseLd->getMemOperand()->getFlags(),
         VT.getStoreSize(), Align);
     // Can't do the combine if the load isn't naturally aligned with the element
     // type
     if (!TLI.allowsMemoryAccessForAlignment(*DAG.getContext(),
                                             DAG.getDataLayout(), VT, *MMO))
if (!TLI.isLegalStridedLoadStore(WideVecVT, Align))
     return SDValue();
 
   MVT ContainerVT = TLI.getContainerForFixedLengthVector(WideVecVT);
   SDValue VL =
       getDefaultVLOps(WideVecVT, ContainerVT, DL, DAG, Subtarget).second;
   SDVTList VTs = DAG.getVTList({ContainerVT, MVT::Other});
   SDValue IntID =
       DAG.getTargetConstant(Intrinsic::riscv_vlse, DL, Subtarget.getXLenVT());
+  if (Reversed)
+    Stride = DAG.getNegative(Stride, DL, Stride->getValueType(0));
   SDValue Ops[] = {BaseLd->getChain(),
                    IntID,
                    DAG.getUNDEF(ContainerVT),
                    BasePtr,
                    Stride,
                    VL};
 
   uint64_t MemSize;
-  if (auto *ConstStride = dyn_cast<ConstantSDNode>(Stride))
+  if (auto *ConstStride = dyn_cast<ConstantSDNode>(Stride);
+      ConstStride && !Reversed && ConstStride->getSExtValue() >= 0)
     // total size = (elsize * n) + (stride - elsize) * (n-1)
     //            = elsize + stride * (n-1)
     MemSize = WideScalarVT.getSizeInBits() +
               ConstStride->getSExtValue() * (N->getNumOperands() - 1);
   else
     // If Stride isn't constant, then we can't know how much it will load
     MemSize = MemoryLocation::UnknownSize;
 

llvm/test/CodeGen/RISCV/rvv/fixed-vectors-strided-load-combine.ll

; CHECK-NEXT:    ret
   %d = load <2 x float>, ptr %d.gep
   %e.0 = shufflevector <2 x float> %a, <2 x float> %b, <4 x i32> <i32 0, i32 1, i32 2, i32 3>
   %e.1 = shufflevector <2 x float> %c, <2 x float> %d, <4 x i32> <i32 0, i32 1, i32 2, i32 3>
   %e.2 = shufflevector <4 x float> %e.0, <4 x float> %e.1, <8 x i32> <i32 0, i32 1, i32 2, i32 3, i32 4, i32 5, i32 6, i32 7>
   store <8 x float> %e.2, ptr %z
   ret void
 }
 
-; TODO: This is a strided load with a negative stride
+; This is a strided load with a negative stride
 define void @reverse_strided_constant_pos_4xv2f32(ptr %x, ptr %z, i64 %s) {
 ; CHECK-LABEL: reverse_strided_constant_pos_4xv2f32:
 ; CHECK:       # %bb.0:
-; CHECK-NEXT:    addi a2, a0, 64
-; CHECK-NEXT:    addi a3, a0, 128
-; CHECK-NEXT:    addi a4, a0, 192
-; CHECK-NEXT:    vsetivli zero, 2, e32, mf2, ta, ma
-; CHECK-NEXT:    vle32.v v8, (a4)
-; CHECK-NEXT:    vle32.v v10, (a3)
-; CHECK-NEXT:    vle32.v v12, (a2)
-; CHECK-NEXT:    vle32.v v14, (a0)
-; CHECK-NEXT:    vsetivli zero, 4, e32, m2, tu, ma
-; CHECK-NEXT:    vslideup.vi v8, v10, 2
-; CHECK-NEXT:    vsetivli zero, 6, e32, m2, tu, ma
-; CHECK-NEXT:    vslideup.vi v8, v12, 4
-; CHECK-NEXT:    vsetivli zero, 8, e32, m2, ta, ma
-; CHECK-NEXT:    vslideup.vi v8, v14, 6
-; CHECK-NEXT:    vse32.v v8, (a1)
+; CHECK-NEXT:    addi a0, a0, 192
+; CHECK-NEXT:    li a2, -64
+; CHECK-NEXT:    vsetivli zero, 4, e64, m2, ta, ma
+; CHECK-NEXT:    vlse64.v v8, (a0), a2
+; CHECK-NEXT:    vse64.v v8, (a1)
 ; CHECK-NEXT:    ret
   %x.1 = getelementptr i8, ptr %x, i64 64
   %x.2 = getelementptr i8, ptr %x.1, i64 64
   %x.3 = getelementptr i8, ptr %x.2, i64 64
   %a = load <2 x float>, ptr %x.3
   %b = load <2 x float>, ptr %x.2
   %c = load <2 x float>, ptr %x.1
   %d = load <2 x float>, ptr %x
   %e.0 = shufflevector <2 x float> %a, <2 x float> %b, <4 x i32> <i32 0, i32 1, i32 2, i32 3>
   %e.1 = shufflevector <2 x float> %c, <2 x float> %d, <4 x i32> <i32 0, i32 1, i32 2, i32 3>
   %e.2 = shufflevector <4 x float> %e.0, <4 x float> %e.1, <8 x i32> <i32 0, i32 1, i32 2, i32 3, i32 4, i32 5, i32 6, i32 7>
   store <8 x float> %e.2, ptr %z
   ret void
 }
 
 define void @reverse_strided_constant_neg_4xv2f32(ptr %x, ptr %z, i64 %s) {
 ; CHECK-LABEL: reverse_strided_constant_neg_4xv2f32:
 ; CHECK:       # %bb.0:
-; CHECK-NEXT:    addi a2, a0, -64
-; CHECK-NEXT:    addi a3, a0, -128
-; CHECK-NEXT:    addi a4, a0, -192
-; CHECK-NEXT:    vsetivli zero, 2, e32, mf2, ta, ma
-; CHECK-NEXT:    vle32.v v8, (a4)
-; CHECK-NEXT:    vle32.v v10, (a3)
-; CHECK-NEXT:    vle32.v v12, (a2)
-; CHECK-NEXT:    vle32.v v14, (a0)
-; CHECK-NEXT:    vsetivli zero, 4, e32, m2, tu, ma
-; CHECK-NEXT:    vslideup.vi v8, v10, 2
-; CHECK-NEXT:    vsetivli zero, 6, e32, m2, tu, ma
-; CHECK-NEXT:    vslideup.vi v8, v12, 4
-; CHECK-NEXT:    vsetivli zero, 8, e32, m2, ta, ma
-; CHECK-NEXT:    vslideup.vi v8, v14, 6
-; CHECK-NEXT:    vse32.v v8, (a1)
+; CHECK-NEXT:    addi a0, a0, -192
+; CHECK-NEXT:    li a2, 64
+; CHECK-NEXT:    vsetivli zero, 4, e64, m2, ta, ma
+; CHECK-NEXT:    vlse64.v v8, (a0), a2
+; CHECK-NEXT:    vse64.v v8, (a1)
 ; CHECK-NEXT:    ret
   %x.1 = getelementptr i8, ptr %x, i64 -64
   %x.2 = getelementptr i8, ptr %x.1, i64 -64
   %x.3 = getelementptr i8, ptr %x.2, i64 -64
   %a = load <2 x float>, ptr %x.3
   %b = load <2 x float>, ptr %x.2
   %c = load <2 x float>, ptr %x.1
   %d = load <2 x float>, ptr %x
   %e.0 = shufflevector <2 x float> %a, <2 x float> %b, <4 x i32> <i32 0, i32 1, i32 2, i32 3>
   %e.1 = shufflevector <2 x float> %c, <2 x float> %d, <4 x i32> <i32 0, i32 1, i32 2, i32 3>
   %e.2 = shufflevector <4 x float> %e.0, <4 x float> %e.1, <8 x i32> <i32 0, i32 1, i32 2, i32 3, i32 4, i32 5, i32 6, i32 7>
   store <8 x float> %e.2, ptr %z
   ret void
 }
 
-; TODO: This is a strided load with a negative stride
+; This is a strided load with a negative stride
 define void @reverse_strided_runtime_4xv2f32(ptr %x, ptr %z, i64 %s) {
 ; CHECK-LABEL: reverse_strided_runtime_4xv2f32:
 ; CHECK:       # %bb.0:
-; CHECK-NEXT:    add a3, a0, a2
-; CHECK-NEXT:    add a4, a3, a2
-; CHECK-NEXT:    add a2, a4, a2
-; CHECK-NEXT:    vsetivli zero, 2, e32, mf2, ta, ma
-; CHECK-NEXT:    vle32.v v8, (a2)
-; CHECK-NEXT:    vle32.v v10, (a4)
-; CHECK-NEXT:    vle32.v v12, (a3)
-; CHECK-NEXT:    vle32.v v14, (a0)
-; CHECK-NEXT:    vsetivli zero, 4, e32, m2, tu, ma
-; CHECK-NEXT:    vslideup.vi v8, v10, 2
-; CHECK-NEXT:    vsetivli zero, 6, e32, m2, tu, ma
-; CHECK-NEXT:    vslideup.vi v8, v12, 4
-; CHECK-NEXT:    vsetivli zero, 8, e32, m2, ta, ma
-; CHECK-NEXT:    vslideup.vi v8, v14, 6
-; CHECK-NEXT:    vse32.v v8, (a1)
+; CHECK-NEXT:    add a0, a0, a2
+; CHECK-NEXT:    add a3, a2, a2
+; CHECK-NEXT:    add a0, a0, a3
+; CHECK-NEXT:    neg a2, a2
+; CHECK-NEXT:    vsetivli zero, 4, e64, m2, ta, ma
+; CHECK-NEXT:    vlse64.v v8, (a0), a2
+; CHECK-NEXT:    vse64.v v8, (a1)
 ; CHECK-NEXT:    ret
   %x.1 = getelementptr i8, ptr %x, i64 %s
   %x.2 = getelementptr i8, ptr %x.1, i64 %s
   %x.3 = getelementptr i8, ptr %x.2, i64 %s
   %a = load <2 x float>, ptr %x.3
   %b = load <2 x float>, ptr %x.2
   %c = load <2 x float>, ptr %x.1
   %d = load <2 x float>, ptr %x
   %e.0 = shufflevector <2 x float> %a, <2 x float> %b, <4 x i32> <i32 0, i32 1, i32 2, i32 3>
   %e.1 = shufflevector <2 x float> %c, <2 x float> %d, <4 x i32> <i32 0, i32 1, i32 2, i32 3>
   %e.2 = shufflevector <4 x float> %e.0, <4 x float> %e.1, <8 x i32> <i32 0, i32 1, i32 2, i32 3, i32 4, i32 5, i32 6, i32 7>
   store <8 x float> %e.2, ptr %z
   ret void
 }