[llvm] [RISCV] Match strided vector bases in RISCVGatherScatterLowering (PR #93972)

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-backend-risc-v

Author: Luke Lau (lukel97)

<details>
<summary>Changes</summary>

Currently we only match GEPs with a scalar base pointer, but a common pattern that's emitted from the loop vectorizer is a strided vector base plus some sort of scalar offset:

    %base = getelementptr i64, ptr %p, <vscale x 1 x i64> %step
    %gep = getelementptr i64, <vscale x 1 x ptr> %base, i64 %offset

This is common for accesses into a struct e.g. f[i].b below:

    struct F { int a; char b; };

    void foo(struct F *f) {
      for (int i = 0; i < 1024; i += 2) {
        f[i].a++;
        f[i].b++;
      }
    }

This patch handles this case in RISCVGatherScatterLowering by recursing on the base pointer if it's a vector.

With this we can convert roughly 80% of the indexed loads and stores emitted to strided loads and stores on SPEC CPU 2017, -O3 -march=rva22u64_v


---
Full diff: https://github.com/llvm/llvm-project/pull/93972.diff


2 Files Affected:

- (modified) llvm/lib/Target/RISCV/RISCVGatherScatterLowering.cpp (+16) 
- (modified) llvm/test/CodeGen/RISCV/rvv/strided-load-store.ll (+2-4) 


``````````diff
diff --git a/llvm/lib/Target/RISCV/RISCVGatherScatterLowering.cpp b/llvm/lib/Target/RISCV/RISCVGatherScatterLowering.cpp
index f0bd25f167d80..f7cca854d2767 100644
--- a/llvm/lib/Target/RISCV/RISCVGatherScatterLowering.cpp
+++ b/llvm/lib/Target/RISCV/RISCVGatherScatterLowering.cpp
@@ -349,6 +349,22 @@ RISCVGatherScatterLowering::determineBaseAndStride(Instruction *Ptr,
 
   SmallVector<Value *, 2> Ops(GEP->operands());
 
+  // If the base pointer is a vector, check if it's strided.
+  if (GEP->getPointerOperand()->getType()->isVectorTy()) {
+    auto [BaseBase, Stride] = determineBaseAndStride(
+        cast<Instruction>(GEP->getPointerOperand()), Builder);
+    // If GEP's offset is scalar then we can add it to the base pointer's base.
+    auto IsScalar = [](Value *Idx) { return !Idx->getType()->isVectorTy(); };
+    if (BaseBase && all_of(GEP->indices(), IsScalar)) {
+      Builder.SetInsertPoint(GEP);
+      SmallVector<Value *> Indices(GEP->indices());
+      Value *OffsetBase =
+          Builder.CreateGEP(GEP->getSourceElementType(), BaseBase, Indices, "",
+                            GEP->isInBounds());
+      return {OffsetBase, Stride};
+    }
+  }
+
   // Base pointer needs to be a scalar.
   Value *ScalarBase = Ops[0];
   if (ScalarBase->getType()->isVectorTy()) {
diff --git a/llvm/test/CodeGen/RISCV/rvv/strided-load-store.ll b/llvm/test/CodeGen/RISCV/rvv/strided-load-store.ll
index 4feecbbdef94f..53b20161cdaea 100644
--- a/llvm/test/CodeGen/RISCV/rvv/strided-load-store.ll
+++ b/llvm/test/CodeGen/RISCV/rvv/strided-load-store.ll
@@ -301,10 +301,8 @@ define void @constant_stride(<vscale x 1 x i64> %x, ptr %p, i64 %stride) {
 
 define <vscale x 1 x i64> @vector_base_scalar_offset(ptr %p, i64 %offset) {
 ; CHECK-LABEL: @vector_base_scalar_offset(
-; CHECK-NEXT:    [[STEP:%.*]] = call <vscale x 1 x i64> @llvm.experimental.stepvector.nxv1i64()
-; CHECK-NEXT:    [[PTRS1:%.*]] = getelementptr i64, ptr [[P:%.*]], <vscale x 1 x i64> [[STEP]]
-; CHECK-NEXT:    [[PTRS2:%.*]] = getelementptr i64, <vscale x 1 x ptr> [[PTRS1]], i64 [[OFFSET:%.*]]
-; CHECK-NEXT:    [[X:%.*]] = call <vscale x 1 x i64> @llvm.masked.gather.nxv1i64.nxv1p0(<vscale x 1 x ptr> [[PTRS2]], i32 8, <vscale x 1 x i1> shufflevector (<vscale x 1 x i1> insertelement (<vscale x 1 x i1> poison, i1 true, i64 0), <vscale x 1 x i1> poison, <vscale x 1 x i32> zeroinitializer), <vscale x 1 x i64> poison)
+; CHECK-NEXT:    [[TMP1:%.*]] = getelementptr i64, ptr [[P:%.*]], i64 [[OFFSET:%.*]]
+; CHECK-NEXT:    [[X:%.*]] = call <vscale x 1 x i64> @llvm.riscv.masked.strided.load.nxv1i64.p0.i64(<vscale x 1 x i64> poison, ptr [[TMP1]], i64 8, <vscale x 1 x i1> shufflevector (<vscale x 1 x i1> insertelement (<vscale x 1 x i1> poison, i1 true, i64 0), <vscale x 1 x i1> poison, <vscale x 1 x i32> zeroinitializer))
 ; CHECK-NEXT:    ret <vscale x 1 x i64> [[X]]
 ;
   %step = call <vscale x 1 x i64> @llvm.experimental.stepvector.nxv1i64()

``````````

</details>


https://github.com/llvm/llvm-project/pull/93972