[AArch64][SVE][WIP] Add support for vscale constants (?)

[efriedma]

Summary

This patch allows the following IR to compile:

define i32 @vscale() {
  ret i32 ptrtoint (<vscale x 1 x i8>* getelementptr (<vscale x 1 x
i8>, <vscale x 1 x i8>* null, i64 1) to i32)
}

to the following, which produces a value equal to the runtime value of vscale:

cntd    x8
lsr     x0, x8, #1
                                // kill: def $w0 killed $w0 killed $x0
ret

Posting this patch mainly to pose the question of whether the input IR is supposed to be valid. Some people have expressed the opinion that vscale should not be an llvm::Constant.

[lattner]

This is a very clever approach, I agree it should work - nice job! That said, I don't think it follows that we should accept these constantexprs in a ShuffleVector mask.

[efriedma]

My current thinking for shuffles is that we shouldn't represent the shuffle mask of shufflevector as a Constant at all. Currently, we basically treat it as an ArrayRef<int> anyway (and in fact, that's the only representation at the SelectionDAG level). And once we're not using Constants, we can pick an appropriate representation for vscale'ed shuffles without worrying about shoving the result into a Constant. I think Sander is planning to write up a proposal along these lines.

[lattner]

My current thinking for shuffles is that we shouldn't represent the shuffle mask of shufflevector as a Constant at all. Currently, we basically treat it as an ArrayRef<int> anyway

I completely agree, this was a mistake in the representation and it would be really great to fix it! I'd recommend moving the existing functionality before proposing extending shufflevector to vscale though. One step at a time. Thanks :-)

[paulwalker-arm]

Given this establishes that adding the scalable vector type has inadvertently created a way to represent the runtime constant of vscale as a constant, what's the downside to having a clearer representation in the form of an explicit constant?

[lattner]

What's the advantage of a redundant way to represent the same thing? LLVM could have a 'sizeof' ConstantExpr for example, but never did (people use the same gep trick). It is generally good to have fewer more canonical ways to represent a thing if possible.

[paulwalker-arm]

Can we guarantee "sizeof(<vscale x 1 x i8>) == vscale" for all future targets? much like how "sizeof(<vscale x 1 x i1>) == vscale" doesn't hold true for SVE. If not then this isn't a canonical form as the exact pattern becomes target specific. Or to put another, why should vscale be linked to data layout?

[efriedma]

Can we guarantee "sizeof(<vscale x 1 x i8>) == vscale" for all future targets?

I don't think the current datalayout code actually prevents a target from specifying some sort of excessive alignment for vectors, like "v8:128:128". Granted, I can't think of any reason you'd want to do that.

[simoll]

define i32 @vscale() {
  ret i32 ptrtoint (<vscale x 1 x i8>* getelementptr (<vscale x 1 x
i8>, <vscale x 1 x i8>* null, i64 1) to i32)
}

Whether this returns vscale or something else depends on the ptrtoint mapping of the default address space.. and <vscale x 1 x i8>* null might not be i32 0 for some targets. You might get away with subtracting the integer value of null ;-) (.. i guess @arsenm is more familiar with the topic of address spaces and what may be assumed about them given his recent DevMtg talk).

Personally, I'd prefer to have an explicit and well-defined vscale constant. Having scalable vector types without vscale, to me, seems like stopping half way.

My current thinking for shuffles is that we shouldn't represent the shuffle mask of shufflevector as a Constant at all.

If shuffle masks are no longer IR entities, this means ruling out computed shuffle masks in the future, which are available on some targets (X86). Whether we want computed shuffle masks at all is a different question.

[sdesmalen]

In D71636#1807522, @simoll wrote:

define i32 @vscale() {
  ret i32 ptrtoint (<vscale x 1 x i8>* getelementptr (<vscale x 1 x
i8>, <vscale x 1 x i8>* null, i64 1) to i32)
}

Whether this returns vscale or something else depends on the ptrtoint mapping of the default address space.. and <vscale x 1 x i8>* null might not be i32 0 for some targets. You might get away with subtracting the integer value of null ;-) (.. i guess @arsenm is more familiar with the topic of address spaces and what may be assumed about them given his recent DevMtg talk).

My understanding was that for the default address space, nullptr always has the bit value 0. I don't really understand this topic enough to know the intricacies around the ptrtoint mapping of the default address space, but if this is used for sizeof, I don't see why it couldn't be used for vscale (taking into account the check for larger alignment).

The main argument for wanting vscale as a constant was to avoid copy propagation (because it is inherently a constant value) to benefit ISel and to use VScale in shuffle masks to e.g. represent zips or concats. We can drop the latter requirement if we come up with alternative ways to address shuffles for scalable vectors. With the right safe-guards with respect to datalayout, we can still use Eli's suggestion to implement the copy propagation and still benefit from the constant expression with some help from PatternMatch. I've updated my original vscale patch to reflect that: D68203.

[efriedma]

In D71636#1807522, @simoll wrote:

My current thinking for shuffles is that we shouldn't represent the shuffle mask of shufflevector as a Constant at all.

If shuffle masks are no longer IR entities, this means ruling out computed shuffle masks in the future, which are available on some targets (X86). Whether we want computed shuffle masks at all is a different question.

Yes, this is kind of a decision point: either we allow computed shufflemasks, or we get rid of the operand altogether. Or we can introduce some intrinsics and put off the decision, I guess.

I don't think there's much incentive to support computed shuffle masks. Yes, x86 has pshufb, but that doesn't generalize to other element widths/two source vectors/etc. easily. And there aren't very many practical use cases for computed shuffles in automatic vectorization.

If we allow computed shuffle masks, there's also the minor complication that we'd have to change shufflevector to produce poison, not undef, for undef indexes. But we can likely change that with some work...

[lattner]

In D71636#1813285, @efriedma wrote:

In D71636#1807522, @simoll wrote:

If shuffle masks are no longer IR entities, this means ruling out computed shuffle masks in the future, which are available on some targets (X86). Whether we want computed shuffle masks at all is a different question.

Yes, this is kind of a decision point: either we allow computed shufflemasks, or we get rid of the operand altogether. Or we can introduce some intrinsics and put off the decision, I guess.

I don't think there's much incentive to support computed shuffle masks. Yes, x86 has pshufb, but that doesn't generalize to other element widths/two source vectors/etc. easily. And there aren't very many practical use cases for computed shuffles in automatic vectorization.

If we allow computed shuffle masks, there's also the minor complication that we'd have to change shufflevector to produce poison, not undef, for undef indexes. But we can likely change that with some work...

Altivec also supports computed shuffles with vperm IIRC. The problem with this in a shufflevector context is that each of these are very different in how they model a computed shuffle. I would personally rather (keep!) shufflevector to static shuffle masks as a target independent concept. If there becomes a good abstraction for computed shuffles, we can introduce a new intrinsic or instruction at that time, which is separate from it.

[efriedma]

The main thrust of this patch was committed in D68203. But we still need the changes to SelectionDAGBuilder for general GEPs over vscale'ed types.

[efriedma]

Remaining changes are now in https://reviews.llvm.org/D73602 .

llvm/include/llvm/CodeGen/ISDOpcodes.h

enum NodeType {
     VECREDUCE_FMAX, VECREDUCE_FMIN,
     /// Integer reductions may have a result type larger than the vector element
     /// type. However, the reduction is performed using the vector element type
     /// and the value in the top bits is unspecified.
     VECREDUCE_ADD, VECREDUCE_MUL,
     VECREDUCE_AND, VECREDUCE_OR, VECREDUCE_XOR,
     VECREDUCE_SMAX, VECREDUCE_SMIN, VECREDUCE_UMAX, VECREDUCE_UMIN,
 
+    // The scale of scalable vectors.
+    VSCALE,
+
     /// BUILTIN_OP_END - This must be the last enum value in this list.
     /// The target-specific pre-isel opcode values start here.
     BUILTIN_OP_END
   };
 
   /// FIRST_TARGET_MEMORY_OPCODE - Target-specific pre-isel operations
   /// which do not reference a specific memory location should be less than
   /// this value. Those that do must not be less than this value, and can

llvm/lib/Analysis/ConstantFolding.cpp

Constant *SymbolicallyEvaluateGEP(const GEPOperator *GEP,
   if (Constant *C = CastGEPIndices(SrcElemTy, Ops, ResTy,
                                    GEP->getInRangeIndex(), DL, TLI))
     return C;
 
   Constant *Ptr = Ops[0];
   if (!Ptr->getType()->isPointerTy())
     return nullptr;
 
+  if (isa<VectorType>(SrcElemTy) && SrcElemTy->getVectorIsScalable()) {
+    // A GEP over a scalable vector produces a result not known at compile-time.
+    // FIXME: We could fold some three-operand GEPs where the first index
+    // is zero.
+    return nullptr;
+  }
+
   Type *IntPtrTy = DL.getIntPtrType(Ptr->getType());
 
   // If this is a constant expr gep that is effectively computing an
   // "offsetof", fold it into 'cast int Size to T*' instead of 'gep 0, 0, 12'
   for (unsigned i = 1, e = Ops.size(); i != e; ++i)
       if (!isa<ConstantInt>(Ops[i])) {
 
         // If this is "gep i8* Ptr, (sub 0, V)", fold this as:

llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

if (StructType *StTy = GTI.getStructTypeOrNull()) {
           Flags.setNoUnsignedWrap(true);
 
         N = DAG.getNode(ISD::ADD, dl, N.getValueType(), N,
                         DAG.getConstant(Offset, dl, N.getValueType()), Flags);
       }
     } else {
       unsigned IdxSize = DAG.getDataLayout().getIndexSizeInBits(AS);
       MVT IdxTy = MVT::getIntegerVT(IdxSize);
-      APInt ElementSize(IdxSize, DL->getTypeAllocSize(GTI.getIndexedType()));
+      TypeSize ElementSize = DL->getTypeAllocSize(GTI.getIndexedType());
 
       // 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;
+      const auto *CI = dyn_cast_or_null<ConstantInt>(C);
+      if (CI && CI->isZero())
+      if (CI && !ElementSize.isScalable()) {
         APInt Offs = ElementSize * 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);
 
         // In an inbounds GEP with an offset that is nonnegative even when
         // interpreted as signed, assume there is no unsigned overflow.
if (StructType *StTy = GTI.getStructTypeOrNull()) {
       }
 
       // If the index is smaller or larger than intptr_t, truncate or extend
       // it.
       IdxN = DAG.getSExtOrTrunc(IdxN, dl, N.getValueType());
 
       // 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 (ElementSize.getKnownMinSize() != 1) {
+        if (isPowerOf2_64(ElementSize.getKnownMinSize())) {
+          unsigned Amt = Log2_64(ElementSize.getKnownMinSize());
           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(ElementSize.getKnownMinSize(), dl,
                                           IdxN.getValueType());
           IdxN = DAG.getNode(ISD::MUL, dl,
                              N.getValueType(), IdxN, Scale);
         }
       }
+      if (ElementSize.isScalable()) {
+        SDValue VScale = DAG.getNode(ISD::VSCALE, dl, N.getValueType());
+        IdxN = DAG.getNode(ISD::MUL, dl, N.getValueType(), IdxN, VScale);
+      }
 
       N = DAG.getNode(ISD::ADD, dl,
                       N.getValueType(), N, IdxN);
     }
   }
 
   if (PtrMemTy != PtrTy && !cast<GEPOperator>(I).isInBounds())
     N = DAG.getPtrExtendInReg(N, dl, PtrMemTy);

llvm/lib/CodeGen/SelectionDAG/SelectionDAGDumper.cpp

#endif
   case ISD::VECREDUCE_OR:               return "vecreduce_or";
   case ISD::VECREDUCE_XOR:              return "vecreduce_xor";
   case ISD::VECREDUCE_SMAX:             return "vecreduce_smax";
   case ISD::VECREDUCE_SMIN:             return "vecreduce_smin";
   case ISD::VECREDUCE_UMAX:             return "vecreduce_umax";
   case ISD::VECREDUCE_UMIN:             return "vecreduce_umin";
   case ISD::VECREDUCE_FMAX:             return "vecreduce_fmax";
   case ISD::VECREDUCE_FMIN:             return "vecreduce_fmin";
+  case ISD::VSCALE:                     return "vscale";
   }
 }
 
 const char *SDNode::getIndexedModeName(ISD::MemIndexedMode AM) {
   switch (AM) {
   default:              return "";
   case ISD::PRE_INC:    return "<pre-inc>";
   case ISD::PRE_DEC:    return "<pre-dec>";

llvm/lib/Target/AArch64/AArch64ISelLowering.h

private:
   SDValue LowerFSINCOS(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerVECREDUCE(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerATOMIC_LOAD_SUB(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerATOMIC_LOAD_AND(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerDYNAMIC_STACKALLOC(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerWindowsDYNAMIC_STACKALLOC(SDValue Op, SDValue Chain,
                                          SDValue &Size,
                                          SelectionDAG &DAG) const;
+  SDValue LowerVSCALE(SDValue Op, SelectionDAG &DAG) const;
 
   SDValue BuildSDIVPow2(SDNode *N, const APInt &Divisor, SelectionDAG &DAG,
                         SmallVectorImpl<SDNode *> &Created) const override;
   SDValue getSqrtEstimate(SDValue Operand, SelectionDAG &DAG, int Enabled,
                           int &ExtraSteps, bool &UseOneConst,
                           bool Reciprocal) const override;
   SDValue getRecipEstimate(SDValue Operand, SelectionDAG &DAG, int Enabled,
                            int &ExtraSteps) const override;

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

if (Subtarget->hasSVE()) {
     // splat of 0 or undef) once vector selects supported in SVE codegen. See
     // D68877 for more details.
     for (MVT VT : MVT::integer_scalable_vector_valuetypes()) {
       if (isTypeLegal(VT))
         setOperationAction(ISD::SPLAT_VECTOR, VT, Custom);
     }
     setOperationAction(ISD::INTRINSIC_WO_CHAIN, MVT::i8, Custom);
     setOperationAction(ISD::INTRINSIC_WO_CHAIN, MVT::i16, Custom);
+    setOperationAction(ISD::VSCALE, MVT::i64, Custom);
   }
 
   PredictableSelectIsExpensive = Subtarget->predictableSelectIsExpensive();
 }
 
 void AArch64TargetLowering::addTypeForNEON(MVT VT, MVT PromotedBitwiseVT) {
   assert(VT.isVector() && "VT should be a vector type");
 
SDValue AArch64TargetLowering::LowerSTORE(SDValue Op,
 
   if (StoreNode->isTruncatingStore()) {
     return LowerTruncateVectorStore(Dl, StoreNode, VT, MemVT, DAG);
   }
 
   return SDValue();
 }
 
+SDValue AArch64TargetLowering::LowerVSCALE(SDValue Op,
+                                           SelectionDAG &DAG) const {
+  SDLoc DL(Op);
+  SDValue CNTD = DAG.getNode(
+        ISD::INTRINSIC_WO_CHAIN, DL, MVT::i64,
+        DAG.getConstant(Intrinsic::aarch64_sve_cntd, DL, MVT::i32),
+        DAG.getConstant(31, DL, MVT::i32));
+  return DAG.getNode(ISD::SRL, DL, MVT::i64, CNTD,
+                     DAG.getConstant(1, DL, MVT::i64));
+}
+
+
 SDValue AArch64TargetLowering::LowerOperation(SDValue Op,
                                               SelectionDAG &DAG) const {
   LLVM_DEBUG(dbgs() << "Custom lowering: ");
   LLVM_DEBUG(Op.dump());
 
   switch (Op.getOpcode()) {
   default:
     llvm_unreachable("unimplemented operand");
SDValue AArch64TargetLowering::LowerOperation(SDValue Op,
   case ISD::VECREDUCE_FMIN:
     return LowerVECREDUCE(Op, DAG);
   case ISD::ATOMIC_LOAD_SUB:
     return LowerATOMIC_LOAD_SUB(Op, DAG);
   case ISD::ATOMIC_LOAD_AND:
     return LowerATOMIC_LOAD_AND(Op, DAG);
   case ISD::DYNAMIC_STACKALLOC:
     return LowerDYNAMIC_STACKALLOC(Op, DAG);
+  case ISD::VSCALE:
+    return LowerVSCALE(Op, DAG);
   }
 }
 
 //===----------------------------------------------------------------------===//
 //                      Calling Convention Implementation
 //===----------------------------------------------------------------------===//
 
 /// Selects the correct CCAssignFn for a given CallingConvention value.