[SVE] Add support for lowering GEPs involving scalable vectors.

[efriedma]

Summary

This includes both GEPs where the indexed type is a scalable vector, and GEPs where the result type is a scalable vector.

I had to make a couple related fixes to allow all the testcases to compile.

[efriedma]

Just noticed CodeGen/X86/large-gep-scale.ll is failing; need to tweak the ElementSize handling a bit.

[efriedma]

Fix regression test failure. Add a few comments to try to explain the convoluted logic for the various different bitwidths.

[fpetrogalli]

Hi @eli.friedman ,

thank you for working on this, I am working on adding the addressing modes for SVE, this is required to be able to test what I am doing in the back-end.

I have the following questions/requests.

1. Could you please describe in a comment in the tests the math behind operations like the following? %d = getelementptr <vscale x 2 x i64>, <vscale x 2 x i64>* %base, <2 x i64> <i64 1, i64 1>. This could be determined by looking at the codegen, but having it explained in a target independent representation would make things easier to undertsand.
2. Could you please tests expressions like the following? getelementptr <vscale x 2 x i64>, <vscale x 2 x i64> * %base, i64 4, getelementptr <vscale x 2 x i64>, <vscale x 2 x i64> * %base, [i32|i64] %offset and getelementptr i8, <vscale x 2 x i8*> %bases, i32 1?
3. You tests the cases for <vscale x 2 x i64>. Could you please test more configurations of <vscale x N x Ty>, at least covering all cases in which sizeof( N x Ty) == 16?

All these requests will probably require you to add more patterns in AArch64SVEInstrInfo.td, other than the one you have already added. If that's the case, I think it will be worth splitting the patch in two patches, one for the CodeGen of the patterns and one for the codegen of the GEP (of course, unless it is not possible to reproduce such patterns without the GEP on scalable types...)

Thanks!

Francesco

[fpetrogalli]

Hello @efriedma ,

I have started adding SVE reg+imm addressing mode on top of your patch (see https://reviews.llvm.org/D74254).

To make it work, I had to add handling of getelementpr when adding a constant int. The hunk I have added could be useful to let you implement the extra cases I have asked you to add in this patch. Feel free to grab it as it is, or modify it if you find a better way to handle such cases.

Meanwhile, I will keep working on the addressing modes.

Thank you!

Francesco

[efriedma]

Could you please describe in a comment in the tests the math behind operations like the following?

I'm not sure what you mean by "the math". A scalar GEP produces a scalar result, a vector GEP produces a vector result. A scalar operand to a vector GEP is splatted. I'm not sure what useful explanation I could add; I could add a comment "This is a vector GEP", but I'm not sure that's helpful; the result type is already explicitly written in the "ret" instruction.

Could you please tests expressions like the following?

Sure.

You tests the cases for <vscale x 2 x i64>. Could you please test more configurations of <vscale x N x Ty>, at least covering all cases in which sizeof( N x Ty) == 16?

For vector GEPs, all the math always has to happen in N x i64 (since pointers are 64 bits wide). So if N>2, we have to run type legalization, and that probably explodes in unrelated ways, so I don't really want to do that. I can add more tests for scalar GEPs over vector pointee types; that shouldn't require more patterns, I think.

To make it work, I had to add handling of getelementpr when adding a constant int

What's the difference between the DAG produced by your patch, vs. my patch alone?

[efriedma]

Rebased. Added more tests. Optimized emitted DAG for GEPs over scalable vectors.

[sdesmalen]

llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

3917

I'm not sure I understand this comment; why would ElementSize not fit into IdxTy?

[fpetrogalli]

In D73602#1864968, @efriedma wrote:

Could you please describe in a comment in the tests the math behind operations like the following?

I'm not sure what you mean by "the math". A scalar GEP produces a scalar result, a vector GEP produces a vector result. A scalar operand to a vector GEP is splatted. I'm not sure what useful explanation I could add; I could add a comment "This is a vector GEP", but I'm not sure that's helpful; the result type is already explicitly written in the "ret" instruction.

Yeah, I was confused by what the GEP was doing, I thought you where somehow adding the values of the fixed length vector to the first lanes of the scalable one. My bad. No need to add any comment.

Could you please tests expressions like the following?

Sure.

Thanks!

You tests the cases for <vscale x 2 x i64>. Could you please test more configurations of <vscale x N x Ty>, at least covering all cases in which sizeof( N x Ty) == 16?

For vector GEPs, all the math always has to happen in N x i64 (since pointers are 64 bits wide). So if N>2, we have to run type legalization, and that probably explodes in unrelated ways, so I don't really want to do that. I can add more tests for scalar GEPs over vector pointee types; that shouldn't require more patterns, I think.

Yes, I was not asking to modify the types of the RHS argument, I was asking to modify the type of the base pointer, so that you tests things like getelementptr <vscale x 2 x i8>, <vscale x 2 x i8>*, i64 %whatever

To make it work, I had to add handling of getelementpr when adding a constant int

What's the difference between the DAG produced by your patch, vs. my patch alone?

I have explained it in the comment of the first example.

llvm/test/CodeGen/AArch64/sve-gep.ll

7–8

These two instructions could be replaced by a single one: addvl x0, x0, #4. Do you think you can change your patch to generate this sequence, or should this be done in the back-end?

Also, the code I wrote in the patch at https://reviews.llvm.org/D74254 generates the following node for the base address out of this example (well, my patches uses size in bits, but I recon it should be bytes):

(ISD::ADD, %base (ISD::VSCALE 16))

It essentially encodes the multiplicative constant in the VSCALE node directly (provided that the argument of the MUL is an immediate constant). That makes it very easy to detect in the back-end because it doesn't require to check for MUL/SUB and SHL. Do you think you can emit such DAG in your patch?

[efriedma]

llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

3917

If you're on a 32-bit target, and you have a GEP over a type larger than 4GB. (Not sure this is something that *should* be supported, but we have at least one testcase in the LLVM tree.)

llvm/test/CodeGen/AArch64/sve-gep.ll

7–8

We probably need to add patterns for "add(x, vscale) -> addvl" in TableGen; that's not really related to this patch, though.

I guess I could modify the current patch where it does if (CI && !ElementScalable) { to generate an appropriate vscale node directly. But we probably want a mul(vscale(n1), n2) -> vscale(n1*n2) dagcombine anyway.

[fpetrogalli]

llvm/test/CodeGen/AArch64/sve-gep.ll

7–8

We probably need to add patterns for "add(x, vscale) -> addvl" in TableGen; that's not really related to this patch, though.

OK

I guess I could modify the current patch where it does if (CI && !ElementScalable) { to generate an appropriate vscale node directly. But we probably want a mul(vscale(n1), n2) -> vscale(n1*n2) dagcombine anyway.

Ah, DAG Combine. Right. Yep, we can do that there.

42

Maybe try two different constants in the <2 x i64> vector? And also a version with one of the elements of the vector coming from an input variable of the function?

53

Same here. Worth using two different constants and also a version in which the elements of the index vector are not constant but runtime values.

[andwar]

llvm/test/CodeGen/AArch64/sve-gep.ll

78

This test case seems identical to scalable_of_fixed_1.

[sdesmalen]

LGTM!

llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

3917

Ah okay, sounds a bit obscure indeed, but that makes sense.

llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

void SelectionDAGBuilder::visitGetElementPtr(const User &I) {
   SDValue N = getValue(Op0);
   SDLoc dl = getCurSDLoc();
   auto &TLI = DAG.getTargetLoweringInfo();
   MVT PtrTy = TLI.getPointerTy(DAG.getDataLayout(), AS);
   MVT PtrMemTy = TLI.getPointerMemTy(DAG.getDataLayout(), AS);
 
   // Normalize Vector GEP - all scalar operands should be converted to the
   // splat vector.
-  unsigned VectorWidth = I.getType()->isVectorTy() ?
-    I.getType()->getVectorNumElements() : 0;
+  bool IsVectorGEP = I.getType()->isVectorTy();
+  ElementCount VectorElementCount = IsVectorGEP ?
+    I.getType()->getVectorElementCount() : ElementCount(0, false);
 
-  if (VectorWidth && !N.getValueType().isVector()) {
+  if (IsVectorGEP && !N.getValueType().isVector()) {
     LLVMContext &Context = *DAG.getContext();
-    EVT VT = EVT::getVectorVT(Context, N.getValueType(), VectorWidth);
+    EVT VT = EVT::getVectorVT(Context, N.getValueType(), VectorElementCount);
+    if (VectorElementCount.Scalable)
+      N = DAG.getSplatVector(VT, dl, N);
+    else
       N = DAG.getSplatBuildVector(VT, dl, N);
   }
 
   for (gep_type_iterator GTI = gep_type_begin(&I), E = gep_type_end(&I);
        GTI != E; ++GTI) {
     const Value *Idx = GTI.getOperand();
     if (StructType *StTy = GTI.getStructTypeOrNull()) {
       unsigned Field = cast<Constant>(Idx)->getUniqueInteger().getZExtValue();
       if (Field) {
         // N = N + Offset
         uint64_t Offset = DL->getStructLayout(StTy)->getElementOffset(Field);
 
         // In an inbounds GEP with an offset that is nonnegative even when
         // interpreted as signed, assume there is no unsigned overflow.
         SDNodeFlags Flags;
         if (int64_t(Offset) >= 0 && cast<GEPOperator>(I).isInBounds())
           Flags.setNoUnsignedWrap(true);
 
         N = DAG.getNode(ISD::ADD, dl, N.getValueType(), N,
                         DAG.getConstant(Offset, dl, N.getValueType()), Flags);
       }
     } else {
+      // IdxSize is the width of the arithmetic according to IR semantics.
+      // In SelectionDAG, we may prefer to do arithmetic in a wider bitwidth
+      // (and fix up the result later).
       unsigned IdxSize = DAG.getDataLayout().getIndexSizeInBits(AS);
       MVT IdxTy = MVT::getIntegerVT(IdxSize);
-      APInt ElementSize(IdxSize, DL->getTypeAllocSize(GTI.getIndexedType()));
+      TypeSize ElementSize = DL->getTypeAllocSize(GTI.getIndexedType());
+      // We intentionally mask away the high bits here; ElementSize may not
+      // fit in IdxTy.
+      APInt ElementMul(IdxSize, ElementSize.getKnownMinSize());
+      bool ElementScalable = ElementSize.isScalable();
 
       // If this is a scalar constant or a splat vector of constants,
       // handle it quickly.
       const auto *C = dyn_cast<Constant>(Idx);
       if (C && isa<VectorType>(C->getType()))
         C = C->getSplatValue();
 
-      if (const auto *CI = dyn_cast_or_null<ConstantInt>(C)) {
-        if (CI->isZero())
-          continue;
-        APInt Offs = ElementSize * CI->getValue().sextOrTrunc(IdxSize);
+      const auto *CI = dyn_cast_or_null<ConstantInt>(C);
+      if (CI && CI->isZero())
+      if (CI && !ElementScalable) {
+        APInt Offs = ElementMul * CI->getValue().sextOrTrunc(IdxSize);
         LLVMContext &Context = *DAG.getContext();
-        SDValue OffsVal = VectorWidth ?
-          DAG.getConstant(Offs, dl, EVT::getVectorVT(Context, IdxTy, VectorWidth)) :
-          DAG.getConstant(Offs, dl, IdxTy);
+        SDValue OffsVal;
+        if (IsVectorGEP)
+          OffsVal = DAG.getConstant(
+              Offs, dl, EVT::getVectorVT(Context, IdxTy, VectorElementCount));
+        else
+          OffsVal = DAG.getConstant(Offs, dl, IdxTy);
 
         // In an inbounds GEP with an offset that is nonnegative even when
         // interpreted as signed, assume there is no unsigned overflow.
         SDNodeFlags Flags;
         if (Offs.isNonNegative() && cast<GEPOperator>(I).isInBounds())
           Flags.setNoUnsignedWrap(true);
 
         OffsVal = DAG.getSExtOrTrunc(OffsVal, dl, N.getValueType());
 
         N = DAG.getNode(ISD::ADD, dl, N.getValueType(), N, OffsVal, Flags);
         continue;
       }
 
-      // N = N + Idx * ElementSize;
+      // N = N + Idx * ElementMul;
       SDValue IdxN = getValue(Idx);
 
-      if (!IdxN.getValueType().isVector() && VectorWidth) {
-        EVT VT = EVT::getVectorVT(*Context, IdxN.getValueType(), VectorWidth);
+      if (!IdxN.getValueType().isVector() && IsVectorGEP) {
+        EVT VT = EVT::getVectorVT(*Context, IdxN.getValueType(),
+                                  VectorElementCount);
+        if (VectorElementCount.Scalable)
+          IdxN = DAG.getSplatVector(VT, dl, IdxN);
+        else
           IdxN = DAG.getSplatBuildVector(VT, dl, IdxN);
       }
 
       // If the index is smaller or larger than intptr_t, truncate or extend
       // it.
       IdxN = DAG.getSExtOrTrunc(IdxN, dl, N.getValueType());
 
+      if (ElementScalable) {
+        EVT VScaleTy = N.getValueType().getScalarType();
+        SDValue VScale = DAG.getNode(
+            ISD::VSCALE, dl, VScaleTy,
+            DAG.getConstant(ElementMul.getZExtValue(), dl, VScaleTy));
+        if (IsVectorGEP)
+          VScale = DAG.getSplatVector(N.getValueType(), dl, VScale);
+        IdxN = DAG.getNode(ISD::MUL, dl, N.getValueType(), IdxN, VScale);
+      } else {
         // If this is a multiply by a power of two, turn it into a shl
         // immediately.  This is a very common case.
-      if (ElementSize != 1) {
-        if (ElementSize.isPowerOf2()) {
-          unsigned Amt = ElementSize.logBase2();
+        if (ElementMul != 1) {
+          if (ElementMul.isPowerOf2()) {
+            unsigned Amt = ElementMul.logBase2();
             IdxN = DAG.getNode(ISD::SHL, dl,
                                N.getValueType(), IdxN,
                                DAG.getConstant(Amt, dl, IdxN.getValueType()));
           } else {
-          SDValue Scale = DAG.getConstant(ElementSize.getZExtValue(), dl,
+            SDValue Scale = DAG.getConstant(ElementMul.getZExtValue(), dl,
                                             IdxN.getValueType());
             IdxN = DAG.getNode(ISD::MUL, dl,
                                N.getValueType(), IdxN, Scale);
           }
         }
+      }
 
       N = DAG.getNode(ISD::ADD, dl,
                       N.getValueType(), N, IdxN);
     }
   }
 
   if (PtrMemTy != PtrTy && !cast<GEPOperator>(I).isInBounds())
     N = DAG.getPtrExtendInReg(N, dl, PtrMemTy);

llvm/test/CodeGen/AArch64/sve-gep.ll

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc -mtriple=aarch64-linux-gnu -mattr=+sve < %s | FileCheck %s
+
+define <vscale x 2 x i64>* @scalar_of_scalable_1(<vscale x 2 x i64>* %base) {
+; CHECK-LABEL: scalar_of_scalable_1:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    rdvl x8, #1
+; CHECK-NEXT:    add x0, x0, x8, lsl #2
+; CHECK-NEXT:    ret
+  %d = getelementptr <vscale x 2 x i64>, <vscale x 2 x i64>* %base, i64 4
+  ret <vscale x 2 x i64>* %d
+}
+
+define <vscale x 2 x i64>* @scalar_of_scalable_2(<vscale x 2 x i64>* %base, i64 %offset) {
+; CHECK-LABEL: scalar_of_scalable_2:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    rdvl x8, #1
+; CHECK-NEXT:    madd x0, x1, x8, x0
+; CHECK-NEXT:    ret
+  %d = getelementptr <vscale x 2 x i64>, <vscale x 2 x i64>* %base, i64 %offset
+  ret <vscale x 2 x i64>* %d
+}
+
+define <vscale x 2 x i32>* @scalar_of_scalable_3(<vscale x 2 x i32>* %base, i64 %offset) {
+; CHECK-LABEL: scalar_of_scalable_3:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    cnth x8
+; CHECK-NEXT:    madd x0, x1, x8, x0
+; CHECK-NEXT:    ret
+  %d = getelementptr <vscale x 2 x i32>, <vscale x 2 x i32>* %base, i64 %offset
+  ret <vscale x 2 x i32>* %d
+}
+
+define <2 x <vscale x 2 x i64>*> @fixed_of_scalable_1(<vscale x 2 x i64>* %base) {
+; CHECK-LABEL: fixed_of_scalable_1:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    rdvl x8, #1
+; CHECK-NEXT:    dup v0.2d, x8
+; CHECK-NEXT:    dup v1.2d, x0
+; CHECK-NEXT:    add v0.2d, v1.2d, v0.2d
+; CHECK-NEXT:    ret
+  %d = getelementptr <vscale x 2 x i64>, <vscale x 2 x i64>* %base, <2 x i64> <i64 1, i64 1>
+  ret <2 x <vscale x 2 x i64>*> %d
+}
+
+define <2 x <vscale x 2 x i64>*> @fixed_of_scalable_2(<2 x <vscale x 2 x i64>*> %base) {
+; CHECK-LABEL: fixed_of_scalable_2:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    rdvl x8, #1
+; CHECK-NEXT:    dup v1.2d, x8
+; CHECK-NEXT:    add v0.2d, v0.2d, v1.2d
+; CHECK-NEXT:    ret
+  %d = getelementptr <vscale x 2 x i64>, <2 x <vscale x 2 x i64>*> %base, <2 x i64> <i64 1, i64 1>
+  ret <2 x <vscale x 2 x i64>*> %d
+}
+
+define <vscale x 2 x i8*> @scalable_of_fixed_1(i8* %base) {
+; CHECK-LABEL: scalable_of_fixed_1:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    mov z0.d, x0
+; CHECK-NEXT:    add z0.d, z0.d, #1 // =0x1
+; CHECK-NEXT:    ret
+  %idx = shufflevector <vscale x 2 x i64> insertelement (<vscale x 2 x i64> undef, i64 1, i32 0), <vscale x 2 x i64> zeroinitializer, <vscale x 2 x i32> zeroinitializer
+  %d = getelementptr i8, i8* %base, <vscale x 2 x i64> %idx
+  ret <vscale x 2 x i8*> %d
+}
+
+define <vscale x 2 x i8*> @scalable_of_fixed_2(<vscale x 2 x i8*> %base) {
+; CHECK-LABEL: scalable_of_fixed_2:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    add z0.d, z0.d, #1 // =0x1
+; CHECK-NEXT:    ret
+  %idx = shufflevector <vscale x 2 x i64> insertelement (<vscale x 2 x i64> undef, i64 1, i32 0), <vscale x 2 x i64> zeroinitializer, <vscale x 2 x i32> zeroinitializer
+  %d = getelementptr i8, <vscale x 2 x i8*> %base, <vscale x 2 x i64> %idx
+  ret <vscale x 2 x i8*> %d
+}
+
+define <vscale x 2 x i8*> @scalable_of_fixed_3(i8* %base) {
+; CHECK-LABEL: scalable_of_fixed_3:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    mov z0.d, x0
+; CHECK-NEXT:    add z0.d, z0.d, #1 // =0x1
+; CHECK-NEXT:    ret
+  %idx = shufflevector <vscale x 2 x i64> insertelement (<vscale x 2 x i64> undef, i64 1, i32 0), <vscale x 2 x i64> zeroinitializer, <vscale x 2 x i32> zeroinitializer
+  %d = getelementptr i8, i8* %base, <vscale x 2 x i64> %idx
+  ret <vscale x 2 x i8*> %d
+}
+
+define <vscale x 2 x i8*> @scalable_of_fixed_4(i8* %base, <vscale x 2 x i32> %idx) {
+; CHECK-LABEL: scalable_of_fixed_4:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    ptrue p0.d
+; CHECK-NEXT:    sxtw z0.d, p0/m, z0.d
+; CHECK-NEXT:    mov z1.d, x0
+; CHECK-NEXT:    add z0.d, z1.d, z0.d
+; CHECK-NEXT:    ret
+  %d = getelementptr i8, i8* %base, <vscale x 2 x i32> %idx
+  ret <vscale x 2 x i8*> %d
+}
+
+define <vscale x 2 x <vscale x 2 x i64>*> @scalable_of_scalable_1(<vscale x 2 x i64>* %base) {
+; CHECK-LABEL: scalable_of_scalable_1:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    rdvl x8, #1
+; CHECK-NEXT:    mov z1.d, x8
+; CHECK-NEXT:    mov z0.d, x0
+; CHECK-NEXT:    mul z1.d, z1.d, #1
+; CHECK-NEXT:    add z0.d, z0.d, z1.d
+; CHECK-NEXT:    ret
+  %idx = shufflevector <vscale x 2 x i64> insertelement (<vscale x 2 x i64> undef, i64 1, i32 0), <vscale x 2 x i64> zeroinitializer, <vscale x 2 x i32> zeroinitializer
+  %d = getelementptr <vscale x 2 x i64>, <vscale x 2 x i64>* %base, <vscale x 2 x i64> %idx
+  ret <vscale x 2 x <vscale x 2 x i64>*> %d
+}
+
+define <vscale x 2 x <vscale x 2 x i64>*> @scalable_of_scalable_2(<vscale x 2 x <vscale x 2 x i64>*> %base) {
+; CHECK-LABEL: scalable_of_scalable_2:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    rdvl x8, #1
+; CHECK-NEXT:    mov z1.d, x8
+; CHECK-NEXT:    mul z1.d, z1.d, #1
+; CHECK-NEXT:    add z0.d, z0.d, z1.d
+; CHECK-NEXT:    ret
+  %idx = shufflevector <vscale x 2 x i64> insertelement (<vscale x 2 x i64> undef, i64 1, i32 0), <vscale x 2 x i64> zeroinitializer, <vscale x 2 x i32> zeroinitializer
+  %d = getelementptr <vscale x 2 x i64>, <vscale x 2 x <vscale x 2 x i64>*> %base, <vscale x 2 x i64> %idx
+  ret <vscale x 2 x <vscale x 2 x i64>*> %d
+}
+
+define <vscale x 2 x <vscale x 2 x i64>*> @scalable_of_scalable_3(<vscale x 2 x <vscale x 2 x i64>*> %base, <vscale x 2 x i32> %idx) {
+; CHECK-LABEL: scalable_of_scalable_3:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    ptrue p0.d
+; CHECK-NEXT:    rdvl x8, #1
+; CHECK-NEXT:    sxtw z1.d, p0/m, z1.d
+; CHECK-NEXT:    mov z2.d, x8
+; CHECK-NEXT:    mul z1.d, p0/m, z1.d, z2.d
+; CHECK-NEXT:    add z0.d, z0.d, z1.d
+; CHECK-NEXT:    ret
+  %d = getelementptr <vscale x 2 x i64>, <vscale x 2 x <vscale x 2 x i64>*> %base, <vscale x 2 x i32> %idx
+  ret <vscale x 2 x <vscale x 2 x i64>*> %d
+}