[RISCV] Don't use zero-stride vector load for gather if not optimized

[• pcwang-thead]

Summary

We may form a zero-stride vector load when lowering gather to strided
load. As what D137699 has done, we use load+splat for this form if
there is no optimized implementation.
We restrict this to unmasked loads currently in consideration of the
complexity of hanlding all falses masks.

[reames]

This isn't correct. The strided load can be masked. For the case where all lanes are masked off, executing the scalar load is unsound and could introduce a fault.

You could allow any mask where you can prove at least one lane active, or make the scalar store conditional, but there's a bunch of complexity there. As a starting point, I suggest you restrict your transformation to when the instruction is unmasked.

[• pcwang-thead]

Restrict this to unmasked loads only.

[• pcwang-thead]

In D137931#3925757, @reames wrote:

This isn't correct. The strided load can be masked. For the case where all lanes are masked off, executing the scalar load is unsound and could introduce a fault.

You could allow any mask where you can prove at least one lane active, or make the scalar store conditional, but there's a bunch of complexity there. As a starting point, I suggest you restrict your transformation to when the instruction is unmasked.

Thanks. I didn't consider these situations before.

I tried to transform masked loads to scalar load + splat + vrgather, but it may not worth doing it as three instructions are needed. And for runtime all-zeros-mask, we may take some costs to handle it. So, I think we can only handle unmasked loads for now. :-)

[reames]

LGTM

llvm/test/CodeGen/RISCV/rvv/fixed-vector-strided-load-store-asm.ll

230

Hm, this test case is interesting.

This is a case where even with a fast broadcast load, putting the value into a scalar allows the splat to be folded into the using instruction. This trades a scalar register for a vector one, and might be generally interesting.

Maybe a case to give some further thought, definitely not blocking for this patch.

[• pcwang-thead]

llvm/test/CodeGen/RISCV/rvv/fixed-vector-strided-load-store-asm.ll

230

After thinking, I think there are some opportunities to fold splat into .vx instructions, see D138101.

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

case Intrinsic::riscv_masked_strided_load: {
     if (!IsUnmasked) {
       MVT MaskVT = getMaskTypeFor(ContainerVT);
       if (VT.isFixedLengthVector()) {
         Mask = convertToScalableVector(MaskVT, Mask, DAG, Subtarget);
         PassThru = convertToScalableVector(ContainerVT, PassThru, DAG, Subtarget);
       }
     }
 
+    auto *Load = cast<MemIntrinsicSDNode>(Op);
     SDValue VL = getDefaultVLOps(VT, ContainerVT, DL, DAG, Subtarget).second;
-
+    SDValue Ptr = Op.getOperand(3);
+    SDValue Stride = Op.getOperand(4);
+    SDValue Result, Chain;
+    if (!Subtarget.hasOptimizedZeroStrideLoad() && isNullConstant(Stride)) {
+      MVT ScalarVT = ContainerVT.getVectorElementType();
+      SDValue ScalarLoad =
+          DAG.getExtLoad(ISD::ZEXTLOAD, DL, XLenVT, Load->getChain(), Ptr,
+                         ScalarVT, Load->getMemOperand());
+      Chain = ScalarLoad.getValue(1);
+      Result = lowerScalarSplat(IsUnmasked ? SDValue() : PassThru, ScalarLoad,
+                                VL, ContainerVT, DL, DAG, Subtarget);
+    } else {
       SDValue IntID = DAG.getTargetConstant(
           IsUnmasked ? Intrinsic::riscv_vlse : Intrinsic::riscv_vlse_mask, DL,
           XLenVT);
 
-    auto *Load = cast<MemIntrinsicSDNode>(Op);
       SmallVector<SDValue, 8> Ops{Load->getChain(), IntID};
       if (IsUnmasked)
         Ops.push_back(DAG.getUNDEF(ContainerVT));
       else
         Ops.push_back(PassThru);
-    Ops.push_back(Op.getOperand(3)); // Ptr
-    Ops.push_back(Op.getOperand(4)); // Stride
+      Ops.push_back(Ptr);    // Ptr
+      Ops.push_back(Stride); // Stride
       if (!IsUnmasked)
         Ops.push_back(Mask);
       Ops.push_back(VL);
       if (!IsUnmasked) {
-      SDValue Policy = DAG.getTargetConstant(RISCVII::TAIL_AGNOSTIC, DL, XLenVT);
+        SDValue Policy =
+            DAG.getTargetConstant(RISCVII::TAIL_AGNOSTIC, DL, XLenVT);
         Ops.push_back(Policy);
       }
 
       SDVTList VTs = DAG.getVTList({ContainerVT, MVT::Other});
-    SDValue Result =
+      Result =
           DAG.getMemIntrinsicNode(ISD::INTRINSIC_W_CHAIN, DL, VTs, Ops,
                                   Load->getMemoryVT(), Load->getMemOperand());
-    SDValue Chain = Result.getValue(1);
+      Chain = Result.getValue(1);
+    }
     if (VT.isFixedLengthVector())
       Result = convertFromScalableVector(VT, Result, DAG, Subtarget);
     return DAG.getMergeValues({Result, Chain}, DL);
   }
   case Intrinsic::riscv_seg2_load:
   case Intrinsic::riscv_seg3_load:
   case Intrinsic::riscv_seg4_load:
   case Intrinsic::riscv_seg5_load:

llvm/test/CodeGen/RISCV/rvv/fixed-vector-strided-load-store-asm.ll

 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
 ; RUN: llc < %s -mtriple=riscv64 -mattr=+m,+v -riscv-v-vector-bits-min=256 | FileCheck %s --check-prefixes=CHECK,V
 ; RUN: llc < %s -mtriple=riscv64 -mattr=+m,+f,+zve32f -riscv-v-vector-bits-min=256 | FileCheck %s --check-prefixes=CHECK,ZVE32F
+; RUN: llc < %s -mtriple=riscv64 -mattr=+m,+v,+no-optimized-zero-stride-load -riscv-v-vector-bits-min=256 | FileCheck %s --check-prefixes=CHECK,NOT-OPTIMIZED
+; RUN: llc < %s -mtriple=riscv64 -mattr=+m,+f,+zve32f,+no-optimized-zero-stride-load -riscv-v-vector-bits-min=256 | FileCheck %s --check-prefixes=CHECK,NOT-OPTIMIZED
 
 %struct.foo = type { i32, i32, i32, i32 }
 
 ; void gather(signed char * __restrict  A, signed char * __restrict B) {
 ;   for (int i = 0; i != 1024; ++i)
 ;       A[i] += B[i * 5];
 ; }
 define void @gather(i8* noalias nocapture %A, i8* noalias nocapture readonly %B) {
vector.body:                                      ; preds = %vector.body, %entry
   br i1 %6, label %for.cond.cleanup, label %vector.body
 
 for.cond.cleanup:                                 ; preds = %vector.body
   ret void
 }
 
 define void @gather_zero_stride(i8* noalias nocapture %A, i8* noalias nocapture readonly %B) {
 ;
-; CHECK-LABEL: gather_zero_stride:
-; CHECK:       # %bb.0: # %entry
-; CHECK-NEXT:    li a2, 0
-; CHECK-NEXT:    li a3, 32
-; CHECK-NEXT:    li a4, 1024
-; CHECK-NEXT:  .LBB3_1: # %vector.body
-; CHECK-NEXT:    # =>This Inner Loop Header: Depth=1
-; CHECK-NEXT:    vsetvli zero, a3, e8, m1, ta, ma
-; CHECK-NEXT:    vlse8.v v8, (a1), zero
-; CHECK-NEXT:    add a5, a0, a2
-; CHECK-NEXT:    vle8.v v9, (a5)
-; CHECK-NEXT:    vadd.vv v8, v9, v8
-; CHECK-NEXT:    vse8.v v8, (a5)
-; CHECK-NEXT:    addi a2, a2, 32
-; CHECK-NEXT:    addi a1, a1, 160
-; CHECK-NEXT:    bne a2, a4, .LBB3_1
-; CHECK-NEXT:  # %bb.2: # %for.cond.cleanup
-; CHECK-NEXT:    ret
+; V-LABEL: gather_zero_stride:
+; V:       # %bb.0: # %entry
+; V-NEXT:    li a2, 0
+; V-NEXT:    li a3, 32
+; V-NEXT:    li a4, 1024
+; V-NEXT:  .LBB3_1: # %vector.body
+; V-NEXT:    # =>This Inner Loop Header: Depth=1
+; V-NEXT:    vsetvli zero, a3, e8, m1, ta, ma
+; V-NEXT:    vlse8.v v8, (a1), zero
+; V-NEXT:    add a5, a0, a2
+; V-NEXT:    vle8.v v9, (a5)
+; V-NEXT:    vadd.vv v8, v9, v8
+; V-NEXT:    vse8.v v8, (a5)
+; V-NEXT:    addi a2, a2, 32
+; V-NEXT:    addi a1, a1, 160
+; V-NEXT:    bne a2, a4, .LBB3_1
+; V-NEXT:  # %bb.2: # %for.cond.cleanup
+; V-NEXT:    ret
+;
+; ZVE32F-LABEL: gather_zero_stride:
+; ZVE32F:       # %bb.0: # %entry
+; ZVE32F-NEXT:    li a2, 0
+; ZVE32F-NEXT:    li a3, 32
+; ZVE32F-NEXT:    li a4, 1024
+; ZVE32F-NEXT:  .LBB3_1: # %vector.body
+; ZVE32F-NEXT:    # =>This Inner Loop Header: Depth=1
+; ZVE32F-NEXT:    vsetvli zero, a3, e8, m1, ta, ma
+; ZVE32F-NEXT:    vlse8.v v8, (a1), zero
+; ZVE32F-NEXT:    add a5, a0, a2
+; ZVE32F-NEXT:    vle8.v v9, (a5)
+; ZVE32F-NEXT:    vadd.vv v8, v9, v8
+; ZVE32F-NEXT:    vse8.v v8, (a5)
+; ZVE32F-NEXT:    addi a2, a2, 32
+; ZVE32F-NEXT:    addi a1, a1, 160
+; ZVE32F-NEXT:    bne a2, a4, .LBB3_1
+; ZVE32F-NEXT:  # %bb.2: # %for.cond.cleanup
+; ZVE32F-NEXT:    ret
+;
+; NOT-OPTIMIZED-LABEL: gather_zero_stride:
+; NOT-OPTIMIZED:       # %bb.0: # %entry
+; NOT-OPTIMIZED-NEXT:    li a2, 0
+; NOT-OPTIMIZED-NEXT:    li a3, 32
+; NOT-OPTIMIZED-NEXT:    li a4, 1024
+; NOT-OPTIMIZED-NEXT:  .LBB3_1: # %vector.body
+; NOT-OPTIMIZED-NEXT:    # =>This Inner Loop Header: Depth=1
+; NOT-OPTIMIZED-NEXT:    lbu a5, 0(a1)
+; NOT-OPTIMIZED-NEXT:    add a6, a0, a2
+; NOT-OPTIMIZED-NEXT:    vsetvli zero, a3, e8, m1, ta, ma
+; NOT-OPTIMIZED-NEXT:    vle8.v v8, (a6)
+; NOT-OPTIMIZED-NEXT:    vadd.vx v8, v8, a5
+; NOT-OPTIMIZED-NEXT:    vse8.v v8, (a6)
+; NOT-OPTIMIZED-NEXT:    addi a2, a2, 32
+; NOT-OPTIMIZED-NEXT:    addi a1, a1, 160
+; NOT-OPTIMIZED-NEXT:    bne a2, a4, .LBB3_1
+; NOT-OPTIMIZED-NEXT:  # %bb.2: # %for.cond.cleanup
+; NOT-OPTIMIZED-NEXT:    ret
 entry:
   br label %vector.body
 
 vector.body:                                      ; preds = %vector.body, %entry
   %index = phi i64 [ 0, %entry ], [ %index.next, %vector.body ]
   %vec.ind = phi <32 x i64> [ zeroinitializer, %entry ], [ %vec.ind.next, %vector.body ]
   %0 = mul nuw nsw <32 x i64> %vec.ind, <i64 5, i64 5, i64 5, i64 5, i64 5, i64 5, i64 5, i64 5, i64 5, i64 5, i64 5, i64 5, i64 5, i64 5, i64 5, i64 5, i64 5, i64 5, i64 5, i64 5, i64 5, i64 5, i64 5, i64 5, i64 5, i64 5, i64 5, i64 5, i64 5, i64 5, i64 5, i64 5>
   %1 = getelementptr inbounds i8, i8* %B, <32 x i64> %0