[AArch64][SVE] Add more folds to make use of gather/scatter with 32-bit indices

[CarolineConcatto]

Summary

In AArch64ISelLowering.cpp this patch implements this fold:

1. GEP (%ptr, SHL ((stepvector(A) + splat(%offset))) << splat(B)))

into GEP (%ptr + (%offset << B), step_vector (A << B))

The above transform simplifies the index operand so that it can be expressed
as i32 elements.
This allows using only one gather/scatter assembly instruction instead of two.

Patch by Paul Walker (@paulwalker-arm).

Depends on D117900

[david-arm]

Just as a suggestion, perhaps it's worth changing the title of this patch so that it's different to D117900? It can be a bit confusing otherwise trying to review them! You could change it to something like:

[AArch64][SVE] Fold gather/scatter with 32bits when possible (Part 2)

or

[AArch64][SVE] Add more folds to make use of gather/scatter with 32-bit indices

?

[CarolineConcatto]

• change commit message tittle
• rebase depends on D118459 too

[sdesmalen]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

16182

Does the following work:

if (auto *Shift = dyn_cast_or_null<ConstantSDNode>(DAG.getSplatValue(Index.getOperand(1))) {
  ..
}

?

If so, that removes the need to do auto *C = dyn_cast<ConstantSDNode>(Shift) and the return if C == nullptr.

16190

nit: the comment as-is seems redundant.

I think it's worth saying though that Stride is not scaled explicitly by 'Scale', because this is done implicitly by the gather/scatter addressing mode.

16192–16193

nit: Change to:

// BasePtr = BasePtr + ((Offset * Scale) << Shift)
Offset = DAG.getNode(ISD::MUL, DL, MVT::i64, Offset, N->getScale());

[CarolineConcatto]

• Address nits

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

16182

I can, but I will use Shift later on:
Offset = DAG.getNode(ISD::SHL, DL, MVT::i64, Offset, Shift);
My skills may not be the best. But I believe that I cannot use Shift as a ConstantSDNode when setting Offset.
I believe I need to make Shift be a splat of vscale type again.
I don't think this would be better than now.

[CarolineConcatto]

• Use dyn_cast_or_null<ConstantSDNode> to create Shift
• get Shift value with SDValue(Shift, 0)

as suggested by Sander

[sdesmalen]

Just a few final comments, otherwise looks fine to me.

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

16181

I'd suggest creating some variables here, so that further on you don't need to write something like:

Index.getOperand(0).getOperand(0).getOperand(0)

Instead, if you create some variables:

SDValue Add = Index.getOperand(0);
SDValue ShiftOp = Index.getOperand(1);
SDValue StepOp = Add.getOperand(0);
SDValue OffsetOp = Add.getOperand(1);

Then the code below becomes a bit more readable:

if (auto *Shift = dyn_cast_or_null<ConstantSDNode>(DAG.getSplatValue(ShiftOp))) {
  if (SDValue Offset = DAG.getSplatValue(OffsetOp)) {
    int64_t Step = cast<ConstantSDNode>(StepOp.getOperand(0))->getSExtValue();
    ...
  }
}

16190

nit: s/scaled/scale/

[CarolineConcatto]

• Create SDValues Add, ShiftOp ,StepOp ,OffsetOp as requested

[sdesmalen]

Thanks @CarolineConcatto, LGTM

[CarolineConcatto]

• rebase on top of main

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

if (Index.getOpcode() == ISD::ADD &&
     SDValue StepVector = Index.getOperand(0);
     if (auto Offset = DAG.getSplatValue(Index.getOperand(1))) {
       Stride = cast<ConstantSDNode>(StepVector.getOperand(0))->getSExtValue();
       Offset = DAG.getNode(ISD::MUL, DL, MVT::i64, Offset, N->getScale());
       BasePtr = DAG.getNode(ISD::ADD, DL, MVT::i64, BasePtr, Offset);
     }
   }
 
+  // Index = shl((step(const) + splat(offset))), splat(shift))
+  if (Index.getOpcode() == ISD::SHL &&
+      Index.getOperand(0).getOpcode() == ISD::ADD &&
+      Index.getOperand(0).getOperand(0).getOpcode() == ISD::STEP_VECTOR) {
+    if (auto *Shift = dyn_cast_or_null<ConstantSDNode>(
+            DAG.getSplatValue(Index.getOperand(1))))
+      if (auto Offset = DAG.getSplatValue(Index.getOperand(0).getOperand(1))) {
+        int64_t Step = cast<ConstantSDNode>(
+                           Index.getOperand(0).getOperand(0).getOperand(0))
+                           ->getSExtValue();
+        // Stride does not scaled explicitly by 'Scale', because it happens in
+        // the gather/scatter addressing mode.
+        Stride = Step << Shift->getSExtValue();
+        // BasePtr = BasePtr + ((Offset * Scale) << Shift)
+        Offset = DAG.getNode(ISD::MUL, DL, MVT::i64, Offset, N->getScale());
+        Offset = DAG.getNode(ISD::SHL, DL, MVT::i64, Offset, SDValue(Shift, 0));
+        BasePtr = DAG.getNode(ISD::ADD, DL, MVT::i64, BasePtr, Offset);
+      }
+  }
+
   // Earlier returns because there is no fold change.
   if (Stride == 0)
     return false;
 
   // Ensure all element can be represented as a signed i32 value.
   if (Stride < std::numeric_limits<int32_t>::min() ||
       Stride > std::numeric_limits<int32_t>::max())
     return false;

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

; CHECK-NEXT:    ret
   %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 the resulting load is "vscale x 4" wide, despite the offset giving the
+; impression the gather must be split due to it's <vscale x 4 x i64> offset.
+; gather_f32(base, index(offset, 8 * sizeof(float))
+define <vscale x 4 x i8> @gather_i8_index_offset_8([8 x i8]* %base, i64 %offset, <vscale x 4 x i1> %pg) #0 {
+; CHECK-LABEL: gather_i8_index_offset_8:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    add x8, x0, x1, lsl #3
+; CHECK-NEXT:    index z0.s, #0, #8
+; CHECK-NEXT:    ld1sb { z0.s }, p0/z, [x8, z0.s, sxtw]
+; 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
+  %step = call <vscale x 4 x i64> @llvm.experimental.stepvector.nxv4i64()
+  %t2 = add <vscale x 4 x i64> %t1, %step
+  %t3 = getelementptr [8 x i8], [8 x i8]* %base, <vscale x 4 x i64> %t2
+  %t4 = bitcast <vscale x 4 x [8 x i8]*> %t3 to <vscale x 4 x i8*>
+  %load = call <vscale x 4 x i8> @llvm.masked.gather.nxv4i8(<vscale x 4 x i8*> %t4, i32 4, <vscale x 4 x i1> %pg, <vscale x 4 x i8> undef)
+  ret <vscale x 4 x i8> %load
+}
+
+; Ensure the resulting load is "vscale x 4" wide, despite the offset giving the
+; impression the gather must be split due to it's <vscale x 4 x i64> offset.
+; gather_f32(base, index(offset, 8 * sizeof(float))
+define <vscale x 4 x float> @gather_f32_index_offset_8([8 x float]* %base, i64 %offset, <vscale x 4 x i1> %pg) #0 {
+; CHECK-LABEL: gather_f32_index_offset_8:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    mov w8, #32
+; CHECK-NEXT:    add x9, x0, x1, lsl #5
+; CHECK-NEXT:    index z0.s, #0, w8
+; CHECK-NEXT:    ld1w { z0.s }, p0/z, [x9, z0.s, sxtw]
+; 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
+  %step = call <vscale x 4 x i64> @llvm.experimental.stepvector.nxv4i64()
+  %t2 = add <vscale x 4 x i64> %t1, %step
+  %t3 = getelementptr [8 x float], [8 x float]* %base, <vscale x 4 x i64> %t2
+  %t4 = bitcast <vscale x 4 x [8 x float]*> %t3 to <vscale x 4 x float*>
+  %load = call <vscale x 4 x float> @llvm.masked.gather.nxv4f32(<vscale x 4 x float*> %t4, i32 4, <vscale x 4 x i1> %pg, <vscale x 4 x float> undef)
+  ret <vscale x 4 x float> %load
+}
+
+; Ensure the resulting store is "vscale x 4" wide, despite the offset giving the
+; impression the scatter must be split due to it's <vscale x 4 x i64> offset.
+; scatter_f16(base, index(offset, 8 * sizeof(i8))
+define void @scatter_i8_index_offset_8([8 x i8]* %base, i64 %offset, <vscale x 4 x i1> %pg, <vscale x 4 x i8> %data) #0 {
+; CHECK-LABEL: scatter_i8_index_offset_8:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    add x8, x0, x1, lsl #3
+; CHECK-NEXT:    index z1.s, #0, #8
+; CHECK-NEXT:    st1b { z0.s }, p0, [x8, z1.s, sxtw]
+; 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
+  %step = call <vscale x 4 x i64> @llvm.experimental.stepvector.nxv4i64()
+  %t2 = add <vscale x 4 x i64> %t1, %step
+  %t3 = getelementptr [8 x i8], [8 x i8]* %base, <vscale x 4 x i64> %t2
+  %t4 = bitcast <vscale x 4 x [8 x i8]*> %t3 to <vscale x 4 x i8*>
+  call void @llvm.masked.scatter.nxv4i8(<vscale x 4 x i8> %data, <vscale x 4 x i8*> %t4, i32 2, <vscale x 4 x i1> %pg)
+  ret void
+}
+
+; Ensure the resulting store is "vscale x 4" wide, despite the offset giving the
+; impression the scatter must be split due to it's <vscale x 4 x i64> offset.
+; scatter_f16(base, index(offset, 8 * sizeof(half))
+define void @scatter_f16_index_offset_8([8 x half]* %base, i64 %offset, <vscale x 4 x i1> %pg, <vscale x 4 x half> %data) #0 {
+; CHECK-LABEL: scatter_f16_index_offset_8:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    mov w8, #16
+; CHECK-NEXT:    add x9, x0, x1, lsl #4
+; CHECK-NEXT:    index z1.s, #0, w8
+; CHECK-NEXT:    st1h { z0.s }, p0, [x9, z1.s, sxtw]
+; 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
+  %step = call <vscale x 4 x i64> @llvm.experimental.stepvector.nxv4i64()
+  %t2 = add <vscale x 4 x i64> %t1, %step
+  %t3 = getelementptr [8 x half], [8 x half]* %base, <vscale x 4 x i64> %t2
+  %t4 = bitcast <vscale x 4 x [8 x half]*> %t3 to <vscale x 4 x half*>
+  call void @llvm.masked.scatter.nxv4f16(<vscale x 4 x half> %data, <vscale x 4 x half*> %t4, i32 2, <vscale x 4 x i1> %pg)
+  ret void
+}
+
 
 attributes #0 = { "target-features"="+sve" vscale_range(1, 16) }
 
+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 <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 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()