[AArch64][SVE]Make better use of gather/scatter when inside a loop body

[CarolineConcatto]

Summary

In the loopvectorizer it hoists out from the loop body some variables
to optimize the loop body. When the loop has a gather/scatter
it hoists out parts of the index, as the step_vector.

However, with gather/scatter when parts of the index (step_vector and index)
are hoisted out from the loop, the compiler cannot check if the index could be
folded in 32 bits (in findMoreOptimalIndexType).

This patch pulls stepvector and the index to be inside the loop body, so we
can check if the compiler can fold the gather/scatter.

[sdesmalen]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

12240

You may want to limit this recursion a bit further to some threshold, to avoid it becoming too expensive.

Maybe you can rewrite this to a loop, like this:

unsigned Threshold = 3;  // I'd expect 3 or 4 to be sufficient for most cases.
for (const User *U : I->users()) {
  for (unsigned J=0; J<Threshold; ++J) {
    if (isa<GetElementPtrInst>(U))
      return true;
    if (isa<PHINode>(U) || !U->hasOneUse())
      break;
    U = cast<Instruction>(U)->getUser();
  }
}
return false;

12301

This doesn't need to be limited to the second operand. You can just combine it with the case for splat.

12308–12310

I'm a bit confused about what this is doing, it seems like this is adding:

I->getOperand(i)->getOperandUse(0);
I->getOperandUse(i);

is that correct?

I notice it also hasn't checked the other operand yet before it returns true.

What I think you want to do is:

• First check if I is used by a GEP in some way.
• Then iterate through each of the operands, adding it to Ops if it's either a splat/stepvector.
• If any splats/stepvectors have been found, return.

[CarolineConcatto]

• use a threshold for isInstUsedByGEP to reduce future costs.
• add test for the threshold feature.
• pull step and splat only when being used by a ge.p

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

12301

Step vector will always be the second operand, because the work you did in
https://reviews.llvm.org/D118459

12308–12310

If it is a splat suffle it needs to get the shutter and the insert.
I->getOperand(i)->getOperandUse(0); -> insert element in position 0
I->getOperandUse(i); -> is the suffle of the insert
If we only push the shuffle, the optimisation does not see it as a splat further in the line.

I believe now it checks all operands that need to be checked.

[david-arm]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

12305

I think you might be able to simplify this a little with this:

auto OpI = I->getOperand(i);
if (isSplatShuffle(OpI))
  Ops.push_back(&cast<Instruction>(OpI)->getOperandUse(0));
Ops.push_back(&I->getOperandUse(i));

[david-arm]

Hi @CarolineConcatto, I just have a few more minor comments!

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

12240

nit: Just a minor comment, but perhaps it's simpler to explain this as something like:

// Use a threshold to limit the amount of recursion.

12245

nit: Perhaps these can be combined into a single statement, i.e.

if (isa<PHINode>(UI)) || !UI->hasOneUse())
  return false;

to reduce lines of code?

12249

I think it's more common to use

if (isa<GetElementPtrInst>(UI))

here.

[sdesmalen]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

12243

Because you iterate all uses here, it will look at all uses for the first entry to this function, and later on only check a single use (because of the if (!UI->hasOneUse()) return false;). I don't think that's your intention.

12301

In D118459 it canonicalizes binop(splat(x), step_vector) -> binop(step_vector, splat(x)), so that we can recognize binop(step_vector, splat(x)) directly.

At this point you'll want to push down the stepvector intrinsic regardless of which operand it is used in, because if it happens to be used in operand 0, we want to recognize that case as well. If it's in operand 1 and you push it down, SelectionDAG will swap the operands.

12301–12305

This seems to be doing things the wrong way around. It first checks the operands 0 and 1 individually, and then goes into a loop to go over all the operands to check them again.

I'm thinking more of something like this:

for (unsigned J = 0; J < I->getNumOperands(); ++J) {
  Use &U = I->getOperandUse(J);
  if (isSplatShuffle(U.get())) {
    Ops.push_back(&cast<Instruction>(U.get())->getOperandUse(0));
    Ops.push_back(&U);
  } else if (auto *II = dyn_cast<IntrinsicInst>(U.get()))
    if (II->getIntrinsicID() == Intrinsic::experimental_stepvector)
      Ops.push_back(&U);
}

12312

If it did not find anything to sink based on the above criteria, the code below may still be applicable. So maybe you can query the Ops.size() before/after this loop, and only return true if it is different (because you've added new operands).

[CarolineConcatto]

• Earlier return in isInstUsedByGEP when I has more than 1 use
• Update for loop in shouldSinkOperands to check all operands

[CarolineConcatto]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

12305

I changed a bit the code, as asked by Sander.
So now I cannot do as you suggested David

[CarolineConcatto]

Removed unnecessary for loop from isInstUsedByGEP

[tschuett]

In D120894#3375098, @CarolineConcatto wrote:

Removed unnecessary for loop from isInstUsedByGEP

It sounds more like Optional<bool>, if the threshold was exceeded. i.e. is false sound?

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

 }
 
 static bool isSplatShuffle(Value *V) {
   if (auto *Shuf = dyn_cast<ShuffleVectorInst>(V))
     return is_splat(Shuf->getShuffleMask());
   return false;
 }
 
+static bool isInstUsedByGEP(Instruction *I, unsigned Threshold) {
+  // Use a threshold to limit the amount of recursion.
+  if (!I->hasOneUse() || isa<PHINode>(I) || Threshold == 0)
+    return false;
+  Use &U = *I->use_begin();
+  return isa<GetElementPtrInst>(I) ||
+         isInstUsedByGEP(cast<Instruction>(U.getUser()), Threshold - 1);
+}
+
 /// Check if sinking \p I's operands to I's basic block is profitable, because
 /// the operands can be folded into a target instruction, e.g.
 /// shufflevectors extracts and/or sext/zext can be folded into (u,s)subl(2).
 bool AArch64TargetLowering::shouldSinkOperands(
     Instruction *I, SmallVectorImpl<Use *> &Ops) const {
   if (!I->getType()->isVectorTy())
     return false;
 
   if (IntrinsicInst *II = dyn_cast<IntrinsicInst>(I)) {
     switch (II->getIntrinsicID()) {
if (IntrinsicInst *II = dyn_cast<IntrinsicInst>(I)) {
     default:
       return false;
     }
   }
 
   switch (I->getOpcode()) {
   case Instruction::Sub:
   case Instruction::Add: {
+    unsigned OpsSize = Ops.size();
+    if (isInstUsedByGEP(I, /*Threshold*/ 4)) {
+      for (unsigned J = 0; J < I->getNumOperands(); ++J) {
+        Use &U = I->getOperandUse(J);
+        if (isSplatShuffle(U.get())) {
+          Ops.push_back(&cast<Instruction>(U.get())->getOperandUse(0));
+          Ops.push_back(&U);
+        } else if (auto *II = dyn_cast<IntrinsicInst>(U.get()))
+          if (II->getIntrinsicID() == Intrinsic::experimental_stepvector)
+            Ops.push_back(&U);
+      }
+      if (Ops.size() > OpsSize)
+        return true;
+    }
     if (!areExtractExts(I->getOperand(0), I->getOperand(1)))
       return false;
 
     // If the exts' operands extract either the lower or upper elements, we
     // can sink them too.
     auto Ext1 = cast<Instruction>(I->getOperand(0));
     auto Ext2 = cast<Instruction>(I->getOperand(1));
     if (areExtractShuffleVectors(Ext1->getOperand(0), Ext2->getOperand(0))) {

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

; CHECK-NEXT:    ret
   %splat.const8.ins = insertelement <vscale x 4 x i64> undef, i64 8, i32 0
   %splat.const8 = shufflevector <vscale x 4 x i64> %splat.const8.ins, <vscale x 4 x i64> undef, <vscale x 4 x i32> zeroinitializer
   %mul = mul <vscale x 4 x i64> %add2, %splat.const8
   %gep = getelementptr [8 x half], [8 x half]* %base, <vscale x 4 x i64> %mul
   %gep.bc = bitcast <vscale x 4 x [8 x half]*> %gep to <vscale x 4 x half*>
   call void @llvm.masked.scatter.nxv4f16(<vscale x 4 x half> %data, <vscale x 4 x half*> %gep.bc, i32 2, <vscale x 4 x i1> %pg)
   ret void
 }
+
+;; Check fold when step is outside the loop
+;; The tests below are a simplification of the llvm-ir generated by:
+;; void gather_stride16(float *dst, float *src1, float *src2, int n) {
+;;   for (int i = 0; i < n; i++) {
+;;    for (int j = 0; j < n; j++) {
+;;      dst[i + j] += src1[i + (j * (int)16)] * src2[i + j];
+
+define void @gather_nxv4f32(float* %src, float* %dst, <vscale x 4 x i1> %pg) #0 {
+; CHECK-LABEL: gather_nxv4f32:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    mov w9, #16
+; CHECK-NEXT:    mov x8, xzr
+; CHECK-NEXT:    ptrue p1.s
+; CHECK-NEXT:    index z0.s, #0, w9
+; CHECK-NEXT:    cnth x9
+; CHECK-NEXT:  .LBB13_1: // %vector.ph
+; CHECK-NEXT:    // =>This Loop Header: Depth=1
+; CHECK-NEXT:    // Child Loop BB13_2 Depth 2
+; CHECK-NEXT:    mov x10, xzr
+; CHECK-NEXT:  .LBB13_2: // %vector.body
+; CHECK-NEXT:    // Parent Loop BB13_1 Depth=1
+; CHECK-NEXT:    // => This Inner Loop Header: Depth=2
+; CHECK-NEXT:    add x11, x0, x8, lsl #2
+; CHECK-NEXT:    add x11, x11, x10, lsl #6
+; CHECK-NEXT:    ld1w { z1.s }, p0/z, [x11, z0.s, sxtw #2]
+; CHECK-NEXT:    st1w { z1.s }, p1, [x1, x10, lsl #2]
+; CHECK-NEXT:    add x10, x10, x9
+; CHECK-NEXT:    cmp x10, #10
+; CHECK-NEXT:    b.eq .LBB13_2
+; CHECK-NEXT:  // %bb.3: // %for.cond1.cleanup
+; CHECK-NEXT:    // in Loop: Header=BB13_1 Depth=1
+; CHECK-NEXT:    add x8, x8, #1
+; CHECK-NEXT:    cmp x8, #10
+; CHECK-NEXT:    b.ne .LBB13_1
+; CHECK-NEXT:  // %bb.4: // %for.cond.cleanup
+; CHECK-NEXT:    ret
+entry:
+  br label %vector.ph
+
+vector.ph:
+  %index2 = phi i64 [ 0, %entry ], [ %index2.next, %for.cond1.cleanup ]
+  %splatinsertIdx2 = insertelement <vscale x 4 x i64> poison, i64 %index2, i64 0
+  %splatIdx2 = shufflevector <vscale x 4 x i64> %splatinsertIdx2, <vscale x 4 x i64> poison, <vscale x 4 x i32> zeroinitializer
+  %step = call <vscale x 4 x i64> @llvm.experimental.stepvector.nxv4i64()
+  br label %vector.body
+
+vector.body:                                      ; preds = %vector.body, %vector.ph
+  %idx = phi i64 [ 0, %vector.ph ], [ %idx.next, %vector.body ]
+  %splatinsertIdx = insertelement <vscale x 4 x i64> poison, i64 %idx, i64 0
+  %splatIdx = shufflevector <vscale x 4 x i64> %splatinsertIdx, <vscale x 4 x i64> poison, <vscale x 4 x i32> zeroinitializer
+   %splatinsert4 = insertelement <vscale x 4 x i64> poison, i64 4, i64 0
+  %splat4 = shufflevector <vscale x 4 x i64> %splatinsert4, <vscale x 4 x i64> poison, <vscale x 4 x i32> zeroinitializer
+  %t0.0 = add <vscale x 4 x i64> %splatIdx, %step
+  %t0.1 = shl nsw <vscale x 4 x i64> %t0.0, %splat4
+  %t0.2 = add nuw nsw <vscale x 4 x i64> %t0.1, %splatIdx2
+  %t0.3 = getelementptr inbounds float, float* %src, <vscale x 4 x i64> %t0.2
+  %wide.masked.gather = call <vscale x 4 x float> @llvm.masked.gather.nxv4f32(<vscale x 4 x float*> %t0.3, i32 4, <vscale x 4 x i1> %pg, <vscale x 4 x float> undef)
+  %t8 = getelementptr inbounds float, float* %dst, i64 %idx
+  %t9 = bitcast float* %t8 to <vscale x 4 x float>*
+  store <vscale x 4 x float> %wide.masked.gather, <vscale x 4 x float>* %t9, align 4
+  %t14 = call i64 @llvm.vscale.i64()
+  %t15 = shl nuw nsw i64 %t14, 3
+  %idx.next = add nuw i64 %idx, %t15
+  %t16 = icmp eq i64 %idx.next, 10
+  br i1 %t16, label %vector.body, label %for.cond1.cleanup, !llvm.loop !0
+
+for.cond1.cleanup:
+  %index2.next = add nuw nsw i64 %index2, 1
+  %exitcond = icmp eq i64 %index2.next, 10
+  br i1 %exitcond, label %for.cond.cleanup, label  %vector.ph, !llvm.loop !0
+
+for.cond.cleanup:                                 ; preds = %middle.block, %for.body
+  ret void
+}
+
+define void @scatter_nxv4i8(i8* %base, i64 %offset, <vscale x 4 x i1> %pg, <vscale x 4 x i8> %data, i64 %n) #0 {
+; CHECK-LABEL: scatter_nxv4i8:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    mov w9, #16
+; CHECK-NEXT:    mov x8, xzr
+; CHECK-NEXT:    index z1.s, #0, w9
+; CHECK-NEXT:    cnth x9
+; CHECK-NEXT:  .LBB14_1: // %vector.ph
+; CHECK-NEXT:    // =>This Loop Header: Depth=1
+; CHECK-NEXT:    // Child Loop BB14_2 Depth 2
+; CHECK-NEXT:    mov x10, xzr
+; CHECK-NEXT:  .LBB14_2: // %vector.body
+; CHECK-NEXT:    // Parent Loop BB14_1 Depth=1
+; CHECK-NEXT:    // => This Inner Loop Header: Depth=2
+; CHECK-NEXT:    add x11, x0, x8
+; CHECK-NEXT:    add x11, x11, x10, lsl #4
+; CHECK-NEXT:    add x10, x10, x9
+; CHECK-NEXT:    cmp x10, #10
+; CHECK-NEXT:    st1b { z0.s }, p0, [x11, z1.s, sxtw]
+; CHECK-NEXT:    b.eq .LBB14_2
+; CHECK-NEXT:  // %bb.3: // %for.cond1.cleanup
+; CHECK-NEXT:    // in Loop: Header=BB14_1 Depth=1
+; CHECK-NEXT:    add x8, x8, #1
+; CHECK-NEXT:    cmp x8, #10
+; CHECK-NEXT:    b.ne .LBB14_1
+; CHECK-NEXT:  // %bb.4: // %for.cond.cleanup
+; CHECK-NEXT:    ret
+entry:
+ br label %vector.ph
+
+vector.ph:
+  %index2 = phi i64 [ 0, %entry ], [ %index2.next, %for.cond1.cleanup ]
+  %splatinsertIdx2 = insertelement <vscale x 4 x i64> poison, i64 %index2, i64 0
+  %splatIdx2 = shufflevector <vscale x 4 x i64> %splatinsertIdx2, <vscale x 4 x i64> poison, <vscale x 4 x i32> zeroinitializer
+  %step = call <vscale x 4 x i64> @llvm.experimental.stepvector.nxv4i64()
+  br label %vector.body
+
+vector.body:                                      ; preds = %vector.body, %vector.ph
+  %index = phi i64 [ 0, %vector.ph ], [ %index.next, %vector.body ]
+  %.splatinsert = insertelement <vscale x 4 x i64> poison, i64 %index, i64 0
+  %splatIdx = shufflevector <vscale x 4 x i64> %.splatinsert, <vscale x 4 x i64> poison, <vscale x 4 x i32> zeroinitializer
+  %splatinsert4 = insertelement <vscale x 4 x i64> poison, i64 4, i64 0
+  %splat4 = shufflevector <vscale x 4 x i64> %splatinsert4, <vscale x 4 x i64> poison, <vscale x 4 x i32> zeroinitializer
+  %t0.0 = add <vscale x 4 x i64> %splatIdx, %step
+  %t0.1 = shl <vscale x 4 x i64> %t0.0, %splat4 ; With mul it does not fold
+  %t0.2 = add nuw nsw <vscale x 4 x i64> %t0.1, %splatIdx2
+  %t0.3 = getelementptr i8, i8* %base, <vscale x 4 x i64> %t0.2
+  call void @llvm.masked.scatter.nxv4i8(<vscale x 4 x i8> %data, <vscale x 4 x i8*> %t0.3, i32 2, <vscale x 4 x i1> %pg)
+  %t14 = call i64 @llvm.vscale.i64()
+  %t15 = shl nuw nsw i64 %t14, 3
+  %index.next = add nuw i64 %index, %t15
+  %t16 = icmp eq i64 %index.next, 10
+  br i1 %t16, label %vector.body, label %for.cond1.cleanup, !llvm.loop !0
+
+for.cond1.cleanup:
+   %index2.next = add nuw nsw i64 %index2, 1
+   %exitcond = icmp eq i64 %index2.next, 10
+   br i1 %exitcond, label %for.cond.cleanup, label %vector.ph, !llvm.loop !0
+
+for.cond.cleanup:                                 ; preds = %middle.block, %for.body
+  ret void
+}
+
+define void @mul_scatter_nxv4i8(i8* %base, i64 %offset, <vscale x 4 x i1> %pg, <vscale x 4 x i8> %data1, <vscale x 4 x i8> %data2, i64 %n) #0 {
+; CHECK-LABEL: mul_scatter_nxv4i8:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    mov w9, #16
+; CHECK-NEXT:    mov w10, #256
+; CHECK-NEXT:    mov x8, xzr
+; CHECK-NEXT:    index z2.s, #0, w9
+; CHECK-NEXT:    index z3.s, #0, w10
+; CHECK-NEXT:    cnth x9
+; CHECK-NEXT:  .LBB15_1: // %vector.ph
+; CHECK-NEXT:    // =>This Loop Header: Depth=1
+; CHECK-NEXT:    // Child Loop BB15_2 Depth 2
+; CHECK-NEXT:    mov x10, xzr
+; CHECK-NEXT:  .LBB15_2: // %vector.body
+; CHECK-NEXT:    // Parent Loop BB15_1 Depth=1
+; CHECK-NEXT:    // => This Inner Loop Header: Depth=2
+; CHECK-NEXT:    add x11, x0, x8
+; CHECK-NEXT:    add x12, x11, x10, lsl #4
+; CHECK-NEXT:    add x11, x11, x10, lsl #8
+; CHECK-NEXT:    add x10, x10, x9
+; CHECK-NEXT:    cmp x10, #10
+; CHECK-NEXT:    st1b { z0.s }, p0, [x12, z2.s, sxtw]
+; CHECK-NEXT:    st1b { z1.s }, p0, [x11, z3.s, sxtw]
+; CHECK-NEXT:    b.eq .LBB15_2
+; CHECK-NEXT:  // %bb.3: // %for.cond1.cleanup
+; CHECK-NEXT:    // in Loop: Header=BB15_1 Depth=1
+; CHECK-NEXT:    add x8, x8, #1
+; CHECK-NEXT:    cmp x8, #10
+; CHECK-NEXT:    b.ne .LBB15_1
+; CHECK-NEXT:  // %bb.4: // %for.cond.cleanup
+; CHECK-NEXT:    ret
+entry:
+ br label %vector.ph
+
+vector.ph:
+  %index2 = phi i64 [ 0, %entry ], [ %index2.next, %for.cond1.cleanup ]
+  %splatinsertIdx2 = insertelement <vscale x 4 x i64> poison, i64 %index2, i64 0
+  %splatIdx2 = shufflevector <vscale x 4 x i64> %splatinsertIdx2, <vscale x 4 x i64> poison, <vscale x 4 x i32> zeroinitializer
+  %step = call <vscale x 4 x i64> @llvm.experimental.stepvector.nxv4i64()
+  br label %vector.body
+
+vector.body:                                      ; preds = %vector.body, %vector.ph
+  %index = phi i64 [ 0, %vector.ph ], [ %index.next, %vector.body ]
+  %.splatinsert = insertelement <vscale x 4 x i64> poison, i64 %index, i64 0
+  %splatIdx = shufflevector <vscale x 4 x i64> %.splatinsert, <vscale x 4 x i64> poison, <vscale x 4 x i32> zeroinitializer
+  %splatinsert4 = insertelement <vscale x 4 x i64> poison, i64 4, i64 0
+  %splat4 = shufflevector <vscale x 4 x i64> %splatinsert4, <vscale x 4 x i64> poison, <vscale x 4 x i32> zeroinitializer
+  %splatinsert8 = insertelement <vscale x 4 x i64> poison, i64 8, i64 0
+  %splat8 = shufflevector <vscale x 4 x i64> %splatinsert8, <vscale x 4 x i64> poison, <vscale x 4 x i32> zeroinitializer
+  %t0.0 = add <vscale x 4 x i64> %splatIdx, %step
+  %t0.1 = shl <vscale x 4 x i64> %t0.0, %splat4
+  %t0.2 = add nuw nsw <vscale x 4 x i64> %t0.1, %splatIdx2
+  %t0.3 = getelementptr i8, i8* %base, <vscale x 4 x i64> %t0.2
+  call void @llvm.masked.scatter.nxv4i8(<vscale x 4 x i8> %data1, <vscale x 4 x i8*> %t0.3, i32 2, <vscale x 4 x i1> %pg)
+  ;second round use of step
+  %t1.0 = add <vscale x 4 x i64> %splatIdx, %step
+  %t1.1 = shl <vscale x 4 x i64> %t1.0, %splat8
+  %t1.2 = add nuw nsw <vscale x 4 x i64> %t1.1, %splatIdx2
+  %t1.3 = getelementptr i8, i8* %base, <vscale x 4 x i64> %t1.2
+  call void @llvm.masked.scatter.nxv4i8(<vscale x 4 x i8> %data2, <vscale x 4 x i8*> %t1.3, i32 2, <vscale x 4 x i1> %pg)
+  %t14 = call i64 @llvm.vscale.i64()
+  %t15 = shl nuw nsw i64 %t14, 3
+  %index.next = add nuw i64 %index, %t15
+  %t16 = icmp eq i64 %index.next, 10
+  br i1 %t16, label %vector.body, label %for.cond1.cleanup, !llvm.loop !0
+
+for.cond1.cleanup:
+  %index2.next = add nuw nsw i64 %index2, 1
+  %exitcond = icmp eq i64 %index2.next, 10
+  br i1 %exitcond, label %for.cond.cleanup, label  %vector.ph, !llvm.loop !0
+
+for.cond.cleanup:                                 ; preds = %middle.block, %for.body
+  ret void
+}
+
+; stepvector is hidden further behind GEP two adds and a shift.
+define void @loop_scatter_f16_index_add_add_mul([8 x half]* %base, i64 %offset, i64 %offset2, <vscale x 4 x i1> %pg, <vscale x 4 x half> %data) #0 {
+; CHECK-LABEL: loop_scatter_f16_index_add_add_mul:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    mov w9, #128
+; CHECK-NEXT:    mov x8, xzr
+; CHECK-NEXT:    index z1.s, #0, w9
+; CHECK-NEXT:    cnth x9
+; CHECK-NEXT:  .LBB16_1: // %vector.ph
+; CHECK-NEXT:    // =>This Loop Header: Depth=1
+; CHECK-NEXT:    // Child Loop BB16_2 Depth 2
+; CHECK-NEXT:    mov x10, xzr
+; CHECK-NEXT:  .LBB16_2: // %vector.body
+; CHECK-NEXT:    // Parent Loop BB16_1 Depth=1
+; CHECK-NEXT:    // => This Inner Loop Header: Depth=2
+; CHECK-NEXT:    add x11, x0, x8, lsl #7
+; CHECK-NEXT:    add x11, x11, x10, lsl #7
+; CHECK-NEXT:    add x10, x10, x9
+; CHECK-NEXT:    cmp x10, #10
+; CHECK-NEXT:    st1h { z0.s }, p0, [x11, z1.s, sxtw]
+; CHECK-NEXT:    b.eq .LBB16_2
+; CHECK-NEXT:  // %bb.3: // %for.cond1.cleanup
+; CHECK-NEXT:    // in Loop: Header=BB16_1 Depth=1
+; CHECK-NEXT:    add x8, x8, #1
+; CHECK-NEXT:    cmp x8, #10
+; CHECK-NEXT:    b.ne .LBB16_1
+; CHECK-NEXT:  // %bb.4: // %for.cond.cleanup
+; CHECK-NEXT:    ret
+entry:
+ br label %vector.ph
+
+vector.ph:
+  %index2 = phi i64 [ 0, %entry ], [ %index2.next, %for.cond1.cleanup ]
+  %splatinsertIdx2 = insertelement <vscale x 4 x i64> poison, i64 %index2, i64 0
+  %splatIdx2 = shufflevector <vscale x 4 x i64> %splatinsertIdx2, <vscale x 4 x i64> poison, <vscale x 4 x i32> zeroinitializer
+  %step = call <vscale x 4 x i64> @llvm.experimental.stepvector.nxv4i64()
+  br label %vector.body
+
+vector.body:
+  %index = phi i64 [ 0, %vector.ph ], [ %index.next, %vector.body ]
+  %.splatinsert = insertelement <vscale x 4 x i64> poison, i64 %index, i64 0
+  %splatIdx = shufflevector <vscale x 4 x i64> %.splatinsert, <vscale x 4 x i64> poison, <vscale x 4 x i32> zeroinitializer
+  %add1 = add <vscale x 4 x i64> %splatIdx, %step
+  %add2 = add <vscale x 4 x i64> %add1, %splatIdx2
+  %splat.const8.ins = insertelement <vscale x 4 x i64> undef, i64 8, i32 0
+  %splat.const8 = shufflevector <vscale x 4 x i64> %splat.const8.ins, <vscale x 4 x i64> undef, <vscale x 4 x i32> zeroinitializer
+  %mul = mul <vscale x 4 x i64> %add2, %splat.const8
+  %gep = getelementptr [8 x half], [8 x half]* %base, <vscale x 4 x i64> %mul
+  %gep.bc = bitcast <vscale x 4 x [8 x half]*> %gep to <vscale x 4 x half*>
+  call void @llvm.masked.scatter.nxv4f16(<vscale x 4 x half> %data, <vscale x 4 x half*> %gep.bc, i32 2, <vscale x 4 x i1> %pg)
+  %t14 = call i64 @llvm.vscale.i64()
+  %t15 = shl nuw nsw i64 %t14, 3
+  %index.next = add nuw i64 %index, %t15
+  %t16 = icmp eq i64 %index.next, 10
+  br i1 %t16, label %vector.body, label %for.cond1.cleanup, !llvm.loop !0
+
+for.cond1.cleanup:
+  %index2.next = add nuw nsw i64 %index2, 1
+  %exitcond = icmp eq i64 %index2.next, 10
+  br i1 %exitcond, label %for.cond.cleanup, label  %vector.ph, !llvm.loop !0
+
+for.cond.cleanup:                                 ; preds = %middle.block, %for.body
+  ret void
+
+}
+
+
+;; Negative test
+
+; stepvector is hidden further behind GEP four adds and a shift.
+; Threshold is 3
+define void @neg_loop_scatter_f16_index_3add_mul([8 x half]* %base, i64 %offset, i64 %offset2, <vscale x 4 x i1> %pg, <vscale x 4 x half> %data) #0 {
+; CHECK-LABEL: neg_loop_scatter_f16_index_3add_mul:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    mov x8, xzr
+; CHECK-NEXT:    index z1.d, #0, #1
+; CHECK-NEXT:    mov z2.d, z1.d
+; CHECK-NEXT:    punpklo p1.h, p0.b
+; CHECK-NEXT:    uunpklo z3.d, z0.s
+; CHECK-NEXT:    punpkhi p0.h, p0.b
+; CHECK-NEXT:    incd z2.d
+; CHECK-NEXT:    uunpkhi z0.d, z0.s
+; CHECK-NEXT:    cnth x9
+; CHECK-NEXT:  .LBB17_1: // %vector.ph
+; CHECK-NEXT:    // =>This Loop Header: Depth=1
+; CHECK-NEXT:    // Child Loop BB17_2 Depth 2
+; CHECK-NEXT:    mov x10, xzr
+; CHECK-NEXT:  .LBB17_2: // %vector.body
+; CHECK-NEXT:    // Parent Loop BB17_1 Depth=1
+; CHECK-NEXT:    // => This Inner Loop Header: Depth=2
+; CHECK-NEXT:    mov z4.d, x10
+; CHECK-NEXT:    mov z5.d, x8
+; CHECK-NEXT:    add z6.d, z4.d, z2.d
+; CHECK-NEXT:    add z4.d, z4.d, z1.d
+; CHECK-NEXT:    add x10, x10, x9
+; CHECK-NEXT:    add z6.d, z6.d, z5.d
+; CHECK-NEXT:    add z4.d, z4.d, z5.d
+; CHECK-NEXT:    lsl z5.d, z6.d, #7
+; CHECK-NEXT:    lsl z4.d, z4.d, #7
+; CHECK-NEXT:    add z5.d, z5.d, #512 // =0x200
+; CHECK-NEXT:    add z4.d, z4.d, #512 // =0x200
+; CHECK-NEXT:    cmp x10, #10
+; CHECK-NEXT:    st1h { z3.d }, p1, [x0, z4.d]
+; CHECK-NEXT:    st1h { z0.d }, p0, [x0, z5.d]
+; CHECK-NEXT:    b.eq .LBB17_2
+; CHECK-NEXT:  // %bb.3: // %for.cond1.cleanup
+; CHECK-NEXT:    // in Loop: Header=BB17_1 Depth=1
+; CHECK-NEXT:    add x8, x8, #1
+; CHECK-NEXT:    cmp x8, #10
+; CHECK-NEXT:    b.ne .LBB17_1
+; CHECK-NEXT:  // %bb.4: // %for.cond.cleanup
+; CHECK-NEXT:    ret
+entry:
+ br label %vector.ph
+
+vector.ph:
+  %index2 = phi i64 [ 0, %entry ], [ %index2.next, %for.cond1.cleanup ]
+  %splatinsertIdx2 = insertelement <vscale x 4 x i64> poison, i64 %index2, i64 0
+  %splatIdx2 = shufflevector <vscale x 4 x i64> %splatinsertIdx2, <vscale x 4 x i64> poison, <vscale x 4 x i32> zeroinitializer
+  %step = call <vscale x 4 x i64> @llvm.experimental.stepvector.nxv4i64()
+  br label %vector.body
+
+vector.body:
+  %index = phi i64 [ 0, %vector.ph ], [ %index.next, %vector.body ]
+  %.splatinsert = insertelement <vscale x 4 x i64> poison, i64 %index, i64 0
+  %splatIdx = shufflevector <vscale x 4 x i64> %.splatinsert, <vscale x 4 x i64> poison, <vscale x 4 x i32> zeroinitializer
+  %add1 = add <vscale x 4 x i64> %splatIdx, %step
+  %add2 = add <vscale x 4 x i64> %add1, %splatIdx2
+  %splat.const4.ins = insertelement <vscale x 4 x i64> undef, i64 4 , i32 0
+  %splat.const4 = shufflevector <vscale x 4 x i64> %splat.const4.ins, <vscale x 4 x i64> undef, <vscale x 4 x i32> zeroinitializer
+  %add3 = add <vscale x 4 x i64> %add2, %splat.const4
+  %splat.const8.ins = insertelement <vscale x 4 x i64> undef, i64 8, i32 0
+  %splat.const8 = shufflevector <vscale x 4 x i64> %splat.const8.ins, <vscale x 4 x i64> undef, <vscale x 4 x i32> zeroinitializer
+  %mul = mul <vscale x 4 x i64> %add3, %splat.const8
+  %gep = getelementptr [8 x half], [8 x half]* %base, <vscale x 4 x i64> %mul
+  %gep.bc = bitcast <vscale x 4 x [8 x half]*> %gep to <vscale x 4 x half*>
+  call void @llvm.masked.scatter.nxv4f16(<vscale x 4 x half> %data, <vscale x 4 x half*> %gep.bc, i32 2, <vscale x 4 x i1> %pg)
+  %t14 = call i64 @llvm.vscale.i64()
+  %t15 = shl nuw nsw i64 %t14, 3
+  %index.next = add nuw i64 %index, %t15
+  %t16 = icmp eq i64 %index.next, 10
+  br i1 %t16, label %vector.body, label %for.cond1.cleanup, !llvm.loop !0
+
+for.cond1.cleanup:
+  %index2.next = add nuw nsw i64 %index2, 1
+  %exitcond = icmp eq i64 %index2.next, 10
+  br i1 %exitcond, label %for.cond.cleanup, label  %vector.ph, !llvm.loop !0
+
+for.cond.cleanup:                                 ; preds = %middle.block, %for.body
+  ret void
+
+}
+
+
+;; The fold does not apply for: splat + (step+splat)
+;; only for : (step+splat) + splat
+;; because we cannot guarantee position of splat in
+;;  canonicalizeCommutativeBinop
+define void @neg_scatter_nxv4i8(i8* %base, i64 %offset, <vscale x 4 x i1> %pg, <vscale x 4 x i8> %data, i64 %n) #0 {
+; CHECK-LABEL: neg_scatter_nxv4i8:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    mov x8, xzr
+; CHECK-NEXT:    index z1.d, #0, #1
+; CHECK-NEXT:    mov z2.d, z1.d
+; CHECK-NEXT:    punpklo p1.h, p0.b
+; CHECK-NEXT:    uunpklo z3.d, z0.s
+; CHECK-NEXT:    punpkhi p0.h, p0.b
+; CHECK-NEXT:    incd z2.d
+; CHECK-NEXT:    uunpkhi z0.d, z0.s
+; CHECK-NEXT:    cnth x9
+; CHECK-NEXT:  .LBB18_1: // %vector.ph
+; CHECK-NEXT:    // =>This Loop Header: Depth=1
+; CHECK-NEXT:    // Child Loop BB18_2 Depth 2
+; CHECK-NEXT:    mov x10, xzr
+; CHECK-NEXT:  .LBB18_2: // %vector.body
+; CHECK-NEXT:    // Parent Loop BB18_1 Depth=1
+; CHECK-NEXT:    // => This Inner Loop Header: Depth=2
+; CHECK-NEXT:    mov z4.d, x10
+; CHECK-NEXT:    add x10, x10, x9
+; CHECK-NEXT:    add z6.d, z2.d, z4.d
+; CHECK-NEXT:    add z4.d, z1.d, z4.d
+; CHECK-NEXT:    mov z5.d, x8
+; CHECK-NEXT:    lsl z6.d, z6.d, #4
+; CHECK-NEXT:    lsl z4.d, z4.d, #4
+; CHECK-NEXT:    add z6.d, z5.d, z6.d
+; CHECK-NEXT:    add z4.d, z5.d, z4.d
+; CHECK-NEXT:    cmp x10, #10
+; CHECK-NEXT:    st1b { z3.d }, p1, [x0, z4.d]
+; CHECK-NEXT:    st1b { z0.d }, p0, [x0, z6.d]
+; CHECK-NEXT:    b.eq .LBB18_2
+; CHECK-NEXT:  // %bb.3: // %for.cond1.cleanup
+; CHECK-NEXT:    // in Loop: Header=BB18_1 Depth=1
+; CHECK-NEXT:    add x8, x8, #1
+; CHECK-NEXT:    cmp x8, #10
+; CHECK-NEXT:    b.ne .LBB18_1
+; CHECK-NEXT:  // %bb.4: // %for.cond.cleanup
+; CHECK-NEXT:    ret
+entry:
+ br label %vector.ph
+
+vector.ph:
+  %index2 = phi i64 [ 0, %entry ], [ %index2.next, %for.cond1.cleanup ]
+  %splatinsertIdx2 = insertelement <vscale x 4 x i64> poison, i64 %index2, i64 0
+  %splatIdx2 = shufflevector <vscale x 4 x i64> %splatinsertIdx2, <vscale x 4 x i64> poison, <vscale x 4 x i32> zeroinitializer
+  %step = call <vscale x 4 x i64> @llvm.experimental.stepvector.nxv4i64()
+  br label %vector.body
+
+vector.body:                                      ; preds = %vector.body, %vector.ph
+  %index = phi i64 [ 0, %vector.ph ], [ %index.next, %vector.body ]
+  %.splatinsert = insertelement <vscale x 4 x i64> poison, i64 %index, i64 0
+  %splatIdx = shufflevector <vscale x 4 x i64> %.splatinsert, <vscale x 4 x i64> poison, <vscale x 4 x i32> zeroinitializer
+  %splatinsert4 = insertelement <vscale x 4 x i64> poison, i64 4, i64 0
+  %splat4 = shufflevector <vscale x 4 x i64> %splatinsert4, <vscale x 4 x i64> poison, <vscale x 4 x i32> zeroinitializer
+  %t0.0 = add <vscale x 4 x i64> %splatIdx, %step
+  %t0.1 = shl <vscale x 4 x i64> %t0.0, %splat4
+  %t0.2 = add nuw nsw <vscale x 4 x i64> %splatIdx2 , %t0.1
+  %t0.3 = getelementptr i8, i8* %base, <vscale x 4 x i64> %t0.2
+  call void @llvm.masked.scatter.nxv4i8(<vscale x 4 x i8> %data, <vscale x 4 x i8*> %t0.3, i32 2, <vscale x 4 x i1> %pg)
+  %t14 = call i64 @llvm.vscale.i64()
+  %t15 = shl nuw nsw i64 %t14, 3
+  %index.next = add nuw i64 %index, %t15
+  %t16 = icmp eq i64 %index.next, 10
+  br i1 %t16, label %vector.body, label %for.cond1.cleanup, !llvm.loop !0
+
+for.cond1.cleanup:
+   %index2.next = add nuw nsw i64 %index2, 1
+   %exitcond = icmp eq i64 %index2.next, 10
+   br i1 %exitcond, label %for.cond.cleanup, label %vector.ph, !llvm.loop !0
+
+for.cond.cleanup:                                 ; preds = %middle.block, %for.body
+  ret void
+}
+
+; Index is removed.
+; "vscale x 4" wide scatter when the offset is too big.
+define void @neg_scatter_nxv4i8_noIndex(i8* %base, i64 %offset, <vscale x 4 x i1> %pg, <vscale x 4 x i8> %data, i64 %n) #0 {
+; CHECK-LABEL: neg_scatter_nxv4i8_noIndex:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    mov x8, xzr
+; CHECK-NEXT:    cnth x9
+; CHECK-NEXT:    punpklo p1.h, p0.b
+; CHECK-NEXT:    uunpklo z1.d, z0.s
+; CHECK-NEXT:    punpkhi p0.h, p0.b
+; CHECK-NEXT:    uunpkhi z0.d, z0.s
+; CHECK-NEXT:  .LBB19_1: // %vector.ph
+; CHECK-NEXT:    // =>This Loop Header: Depth=1
+; CHECK-NEXT:    // Child Loop BB19_2 Depth 2
+; CHECK-NEXT:    mov x10, xzr
+; CHECK-NEXT:  .LBB19_2: // %vector.body
+; CHECK-NEXT:    // Parent Loop BB19_1 Depth=1
+; CHECK-NEXT:    // => This Inner Loop Header: Depth=2
+; CHECK-NEXT:    add x10, x10, x9
+; CHECK-NEXT:    st1b { z1.d }, p1, [x0, z0.d]
+; CHECK-NEXT:    cmp x10, #10
+; CHECK-NEXT:    st1b { z0.d }, p0, [x0, z0.d]
+; CHECK-NEXT:    b.eq .LBB19_2
+; CHECK-NEXT:  // %bb.3: // %for.cond1.cleanup
+; CHECK-NEXT:    // in Loop: Header=BB19_1 Depth=1
+; CHECK-NEXT:    add x8, x8, #1
+; CHECK-NEXT:    cmp x8, #10
+; CHECK-NEXT:    b.ne .LBB19_1
+; CHECK-NEXT:  // %bb.4: // %for.cond.cleanup
+; CHECK-NEXT:    ret
+entry:
+ br label %vector.ph
+
+vector.ph:
+  %index2 = phi i64 [ 0, %entry ], [ %index2.next, %for.cond1.cleanup ]
+  %splatinsertIdx2 = insertelement <vscale x 4 x i64> poison, i64 %index2, i64 0
+  %splatIdx2 = shufflevector <vscale x 4 x i64> %splatinsertIdx2, <vscale x 4 x i64> poison, <vscale x 4 x i32> zeroinitializer
+  %step = call <vscale x 4 x i64> @llvm.experimental.stepvector.nxv4i64()
+  br label %vector.body
+
+vector.body:                                      ; preds = %vector.body, %vector.ph
+  %index = phi i64 [ 0, %vector.ph ], [ %index.next, %vector.body ]
+  %.splatinsert = insertelement <vscale x 4 x i64> poison, i64 %index, i64 0
+  %splatIdx = shufflevector <vscale x 4 x i64> %.splatinsert, <vscale x 4 x i64> poison, <vscale x 4 x i32> zeroinitializer
+  %splatinsert4 = insertelement <vscale x 4 x i64> poison, i64 33554432, i64 0
+  %splat4 = shufflevector <vscale x 4 x i64> %splatinsert4, <vscale x 4 x i64> poison, <vscale x 4 x i32> zeroinitializer
+  %t0.0 = add <vscale x 4 x i64> %splatIdx, %step
+  %t0.1 = shl <vscale x 4 x i64> %t0.0, %splat4
+  %t0.2 = add nuw nsw <vscale x 4 x i64> %t0.1, %splatIdx2
+  %t0.3 = getelementptr i8, i8* %base, <vscale x 4 x i64> %t0.2
+  call void @llvm.masked.scatter.nxv4i8(<vscale x 4 x i8> %data, <vscale x 4 x i8*> %t0.3, i32 2, <vscale x 4 x i1> %pg)
+  %t14 = call i64 @llvm.vscale.i64()
+  %t15 = shl nuw nsw i64 %t14, 3
+  %index.next = add nuw i64 %index, %t15
+  %t16 = icmp eq i64 %index.next, 10
+  br i1 %t16, label %vector.body, label %for.cond1.cleanup, !llvm.loop !0
+
+for.cond1.cleanup:
+  %index2.next = add nuw nsw i64 %index2, 1
+  %exitcond = icmp eq i64 %index2.next, 10
+  br i1 %exitcond, label %for.cond.cleanup, label  %vector.ph, !llvm.loop !0
+
+for.cond.cleanup:                                 ; preds = %middle.block, %for.body
+  ret void
+}
+
+
 attributes #0 = { "target-features"="+sve" vscale_range(1, 16) }
+!0 = distinct !{!0, !1}
+!1 = !{!"llvm.loop.isvectorized", i32 1}
+!2 = distinct !{!2, !3, !1}
+!3 = !{!"llvm.loop.unroll.runtime.disable"}
 
+declare i64 @llvm.vscale.i64()
 declare <vscale x 4 x float> @llvm.masked.gather.nxv4f32(<vscale x 4 x float*>, i32, <vscale x 4 x i1>, <vscale x 4 x float>)
 
 declare <vscale x 4 x i8> @llvm.masked.gather.nxv4i8(<vscale x 4 x i8*>, i32, <vscale x 4 x i1>, <vscale x 4 x i8>)
 declare void @llvm.masked.scatter.nxv4i8(<vscale x 4 x i8>, <vscale x 4 x i8*>, i32, <vscale x 4 x i1>)
 declare void @llvm.masked.scatter.nxv4i16(<vscale x 4 x i16>, <vscale x 4 x i16*>, i32, <vscale x 4 x i1>)
 declare void @llvm.masked.scatter.nxv4f16(<vscale x 4 x half>, <vscale x 4 x half*>, i32, <vscale x 4 x i1>)
 declare <vscale x 4 x i64> @llvm.experimental.stepvector.nxv4i64()
+