[SVE] Fix bug in simplification of scalable vector instructions

[ctetreau]

Summary

• Most of the simplifications in SimplifyShuffleVectorInst depend on the

concrete value of, or the length of the mask vector. For scalable
vectors, this cannot be known at compile time.

• for these tests, detect if the vector is scalable before attempting

the transformation

• The functions ShuffleVectorInst::getMaskValue and

ShuffleVectorInst::getShuffleMask access the value of the constant mask.
However, since the length of the mask is unknown at compile time, these
function do not work for scalable vectors. Add asserts to ensure that
the input mask is not scalable

[ctetreau]

Fixed permissions on test file

[efriedma]

llvm/lib/Analysis/InstructionSimplify.cpp

4079

Missing testcase?

4463

Redundant "if"?

4488

ConstantFoldShuffleVectorInstruction should be fine? Or is this temporary to work around some other issue?

4529

Can you rearrange this code so you can put the early return earlier?

[ctetreau]

llvm/lib/Analysis/InstructionSimplify.cpp

4463

The inner computation walks the indices array, so we do need to test it. I suppose I could have put the !Scalable check inside the match conditional, but I don't think that'd be any clearer. I guess it'd make the diff smaller?

4488

It could be made to be fine with the current implementation that uses a Constant*; The beginning of this function checks if (isa<UndefValue>(Mask)), and returns an undef. This currently contains a bug that I just now realized that I failed to address, but it's fixable. However, once https://reviews.llvm.org/D72467 is merged, it will be unsalvageable due to the fact that the mask is passed as an array of int, and to determine if it's an undef, you must walk the array. I wonder if it's worth fixing, since your patch seems destined to land soon.

4529

I did not rearrange the code because the current implementation forms an ordering of which simplification should be taken, and I assumed that this order was selected on purpose.

[ctetreau]

llvm/lib/Analysis/InstructionSimplify.cpp

4463

Actually, I mistakenly thought we were discussing the conditional at 4502. I believe that this is not redundant; this will walk from 0 .. MaskEltCount.Min, and will call operator[](i) on Indices, which (assuming SmallVector behaves like std::vector) will result in an out of bounds array access.

[efriedma]

llvm/lib/Analysis/InstructionSimplify.cpp

4488

With new shuffles, we'll eventually want this to work, but sure, we can leave it out for now.

4529

It didn't look like the transforms overlapped, but maybe it's more complicated than I thought at first glance?

[ctetreau]

Address review issues. Resolve additional issue discovered in the IR parser

[efriedma]

LGTM. Please request commit access (http://llvm.org/docs/DeveloperPolicy.html#new-contributors)

[rnk]

This test failed for me locally, so I reverted this change in rGa05441038a3a4. Here is the stack trace I got:
https://reviews.llvm.org/P8190

[ctetreau]

[SVE] Fix bug in simplification of scalable vector instructions

• Most of the simplifications in SimplifyShuffleVectorInst depend on the

concrete value of, or the length of the mask vector. For scalable
vectors, this cannot be known at compile time.

• for these tests, detect if the vector is scalable before attempting

the transformation
-The functions ShuffleVectorInst::getMaskValue and
ShuffleVectorInst::getShuffleMask access the value of the constant mask.
However, since the length of the mask is unknown at compile time, these
function do not work for scalable vectors. Add asserts to ensure that
the input mask is not scalable

[ctetreau]

Fixed test failure caused by new patch https://reviews.llvm.org/D73435 that was merged after this patch was approved.

[ctetreau]

Removed added extra whitespace in otherwise untouched code

[ctetreau]

llvm/lib/Analysis/ValueTracking.cpp

169 ↗

(On Diff #242157)

@RKSimon This is the new code added. It modifies your recent patch, so I figured you'd like to review it.

[RKSimon]

llvm/lib/Analysis/ValueTracking.cpp

169 ↗

(On Diff #242157)

Would it be viable to set all DemandedLHS/DemandedRHS instead so we can still handle all element cases?

[ctetreau]

llvm/lib/Analysis/ValueTracking.cpp

169 ↗

(On Diff #242157)

The problem is that for scalable vectors, we don't know how many elements there are. For a scalable vector, getVectorNumElements() returns the minimum number of elements there can be. If we have <vscale x 4 x i32>, it will return 4, which is the number of elements that would be in the vector if vscale = 1. But the value of vscale isn't known until runtime, so this doesn't really help us.

[RKSimon]

llvm/lib/Analysis/ValueTracking.cpp

169 ↗

(On Diff #242157)

OK, it's a shame we can't do anything in this case. No more comments.

[ctetreau]

llvm/lib/Analysis/ValueTracking.cpp

169 ↗

(On Diff #242157)

Scalable shufflevector is still kind of a work in progress. There's an RFC up (http://lists.llvm.org/pipermail/llvm-dev/2020-January/138762.html) to expand it, so hopefully this can be revisited after that.

[apazos]

llvm/lib/Analysis/ValueTracking.cpp

172 ↗

(On Diff #242157)

I see this line changed because of the conflict with https://reviews.llvm.org/D73435. Thanks for fixing it. LGTM.

llvm/include/llvm/IR/Instructions.h

public:
   /// Returns the pointer type returned by the GEP
   /// instruction, which may be a vector of pointers.
   static Type *getGEPReturnType(Type *ElTy, Value *Ptr,
                                 ArrayRef<Value *> IdxList) {
     Type *PtrTy = PointerType::get(checkGEPType(getIndexedType(ElTy, IdxList)),
                                    Ptr->getType()->getPointerAddressSpace());
     // Vector GEP
     if (Ptr->getType()->isVectorTy()) {
-      unsigned NumElem = Ptr->getType()->getVectorNumElements();
-      return VectorType::get(PtrTy, NumElem);
+      ElementCount EltCount = Ptr->getType()->getVectorElementCount();
+      return VectorType::get(PtrTy, EltCount);
     }
     for (Value *Index : IdxList)
       if (Index->getType()->isVectorTy()) {
-        unsigned NumElem = Index->getType()->getVectorNumElements();
-        return VectorType::get(PtrTy, NumElem);
+        ElementCount EltCount = Index->getType()->getVectorElementCount();
+        return VectorType::get(PtrTy, EltCount);
       }
     // Scalar GEP
     return PtrTy;
   }
 
   unsigned getNumIndices() const {  // Note: always non-negative
     return getNumOperands() - 1;
   }

llvm/lib/Analysis/InstructionSimplify.cpp

static Value *SimplifyGEPInst(Type *SrcTy, ArrayRef<Value *> Ops,
   // getelementptr P -> P.
   if (Ops.size() == 1)
     return Ops[0];
 
   // Compute the (pointer) type returned by the GEP instruction.
   Type *LastType = GetElementPtrInst::getIndexedType(SrcTy, Ops.slice(1));
   Type *GEPTy = PointerType::get(LastType, AS);
   if (VectorType *VT = dyn_cast<VectorType>(Ops[0]->getType()))
-    GEPTy = VectorType::get(GEPTy, VT->getNumElements());
+    GEPTy = VectorType::get(GEPTy, VT->getElementCount());
   else if (VectorType *VT = dyn_cast<VectorType>(Ops[1]->getType()))
-    GEPTy = VectorType::get(GEPTy, VT->getNumElements());
+    GEPTy = VectorType::get(GEPTy, VT->getElementCount());
 
   if (isa<UndefValue>(Ops[0]))
     return UndefValue::get(GEPTy);
 
   if (Ops.size() == 2) {
     // getelementptr P, 0 -> P.
     if (match(Ops[1], m_Zero()) && Ops[0]->getType() == GEPTy)
       return Ops[0];

 
 static Value *SimplifyShuffleVectorInst(Value *Op0, Value *Op1, Constant *Mask,
                                         Type *RetTy, const SimplifyQuery &Q,
                                         unsigned MaxRecurse) {
   if (isa<UndefValue>(Mask))
     return UndefValue::get(RetTy);
 
   Type *InVecTy = Op0->getType();
-  unsigned MaskNumElts = Mask->getType()->getVectorNumElements();
-  unsigned InVecNumElts = InVecTy->getVectorNumElements();
+  ElementCount MaskEltCount = Mask->getType()->getVectorElementCount();
+  ElementCount InVecEltCount = InVecTy->getVectorElementCount();
+
+  assert(MaskEltCount.Scalable == InVecEltCount.Scalable &&
+         "vscale mismatch between input vector and mask");
+
+  bool Scalable = MaskEltCount.Scalable;
 
   SmallVector<int, 32> Indices;
+  if (!Scalable) {
     ShuffleVectorInst::getShuffleMask(Mask, Indices);
-  assert(MaskNumElts == Indices.size() &&
+    assert(MaskEltCount.Min == Indices.size() &&
            "Size of Indices not same as number of mask elements?");
+  }
 
+  if (!Scalable) {
     // Canonicalization: If mask does not select elements from an input vector,
     // replace that input vector with undef.
     bool MaskSelects0 = false, MaskSelects1 = false;
-  for (unsigned i = 0; i != MaskNumElts; ++i) {
+    for (unsigned i = 0; i != MaskEltCount.Min; ++i) {
       if (Indices[i] == -1)
         continue;
-    if ((unsigned)Indices[i] < InVecNumElts)
+      if ((unsigned)Indices[i] < InVecEltCount.Min)
         MaskSelects0 = true;
       else
         MaskSelects1 = true;
     }
     if (!MaskSelects0)
       Op0 = UndefValue::get(InVecTy);
     if (!MaskSelects1)
       Op1 = UndefValue::get(InVecTy);
+  }
 
   auto *Op0Const = dyn_cast<Constant>(Op0);
   auto *Op1Const = dyn_cast<Constant>(Op1);
 
-  // If all operands are constant, constant fold the shuffle.
-  if (Op0Const && Op1Const)
+  // If all operands are constant, constant fold the shuffle. This
+  // transformation depends on the value of the mask which is not known at
+  // compile time for scalable vectors
+  if (!Scalable && Op0Const && Op1Const)
     return ConstantFoldShuffleVectorInstruction(Op0Const, Op1Const, Mask);
 
   // Canonicalization: if only one input vector is constant, it shall be the
-  // second one.
-  if (Op0Const && !Op1Const) {
+  // second one. This transformation depends on the value of the mask which
+  // is not known at compile time for scalable vectors
+  if (!Scalable && Op0Const && !Op1Const) {
     std::swap(Op0, Op1);
-    ShuffleVectorInst::commuteShuffleMask(Indices, InVecNumElts);
+    ShuffleVectorInst::commuteShuffleMask(Indices, InVecEltCount.Min);
   }
 
   // A splat of an inserted scalar constant becomes a vector constant:
   // shuf (inselt ?, C, IndexC), undef, <IndexC, IndexC...> --> <C, C...>
   // NOTE: We may have commuted above, so analyze the updated Indices, not the
   //       original mask constant.
+  // NOTE: This transformation depends on the value of the mask which is not
+  // known at compile time for scalable vectors
   Constant *C;
   ConstantInt *IndexC;
-  if (match(Op0, m_InsertElement(m_Value(), m_Constant(C),
+  if (!Scalable && match(Op0, m_InsertElement(m_Value(), m_Constant(C),
                                               m_ConstantInt(IndexC)))) {
     // Match a splat shuffle mask of the insert index allowing undef elements.
     int InsertIndex = IndexC->getZExtValue();
     if (all_of(Indices, [InsertIndex](int MaskElt) {
           return MaskElt == InsertIndex || MaskElt == -1;
         })) {
       assert(isa<UndefValue>(Op1) && "Expected undef operand 1 for splat");
 
       // Shuffle mask undefs become undefined constant result elements.
-      SmallVector<Constant *, 16> VecC(MaskNumElts, C);
-      for (unsigned i = 0; i != MaskNumElts; ++i)
+      SmallVector<Constant *, 16> VecC(MaskEltCount.Min, C);
+      for (unsigned i = 0; i != MaskEltCount.Min; ++i)
         if (Indices[i] == -1)
           VecC[i] = UndefValue::get(C->getType());
       return ConstantVector::get(VecC);
     }
   }
 
   // A shuffle of a splat is always the splat itself. Legal if the shuffle's
   // value type is same as the input vectors' type.
   if (auto *OpShuf = dyn_cast<ShuffleVectorInst>(Op0))
     if (isa<UndefValue>(Op1) && RetTy == InVecTy &&
         OpShuf->getMask()->getSplatValue())
       return Op0;
 
+  // All remaining transformation depend on the value of the mask, which is
+  // not known at compile time for scalable vectors.
+  if (Scalable)
+    return nullptr;
+
   // Don't fold a shuffle with undef mask elements. This may get folded in a
   // better way using demanded bits or other analysis.
   // TODO: Should we allow this?
   if (find(Indices, -1) != Indices.end())
     return nullptr;
 
   // Check if every element of this shuffle can be mapped back to the
   // corresponding element of a single root vector. If so, we don't need this
   // shuffle. This handles simple identity shuffles as well as chains of
   // shuffles that may widen/narrow and/or move elements across lanes and back.
   Value *RootVec = nullptr;
-  for (unsigned i = 0; i != MaskNumElts; ++i) {
+  for (unsigned i = 0; i != MaskEltCount.Min; ++i) {
     // Note that recursion is limited for each vector element, so if any element
     // exceeds the limit, this will fail to simplify.
     RootVec =
         foldIdentityShuffles(i, Op0, Op1, Indices[i], RootVec, MaxRecurse);
 
     // We can't replace a widening/narrowing shuffle with one of its operands.
     if (!RootVec || RootVec->getType() != RetTy)
       return nullptr;

llvm/lib/AsmParser/LLParser.cpp

int LLParser::ParseGetElementPtr(Instruction *&Inst, PerFunctionState &PFS) {
   if (Ty != BasePointerType->getElementType())
     return Error(ExplicitTypeLoc,
                  "explicit pointee type doesn't match operand's pointee type");
 
   SmallVector<Value*, 16> Indices;
   bool AteExtraComma = false;
   // GEP returns a vector of pointers if at least one of parameters is a vector.
   // All vector parameters should have the same vector width.
-  unsigned GEPWidth = BaseType->isVectorTy() ?
-    BaseType->getVectorNumElements() : 0;
+  ElementCount GEPWidth = BaseType->isVectorTy() ?
+    BaseType->getVectorElementCount() : ElementCount(0, false);
 
   while (EatIfPresent(lltok::comma)) {
     if (Lex.getKind() == lltok::MetadataVar) {
       AteExtraComma = true;
       break;
     }
     if (ParseTypeAndValue(Val, EltLoc, PFS)) return true;
     if (!Val->getType()->isIntOrIntVectorTy())
       return Error(EltLoc, "getelementptr index must be an integer");
 
     if (Val->getType()->isVectorTy()) {
-      unsigned ValNumEl = Val->getType()->getVectorNumElements();
-      if (GEPWidth && GEPWidth != ValNumEl)
+      ElementCount ValNumEl = Val->getType()->getVectorElementCount();
+      if (GEPWidth != ElementCount(0, false) && GEPWidth != ValNumEl)
         return Error(EltLoc,
           "getelementptr vector index has a wrong number of elements");
       GEPWidth = ValNumEl;
     }
     Indices.push_back(Val);
   }
 
   SmallPtrSet<Type*, 4> Visited;

llvm/lib/IR/ConstantFold.cpp

Constant *llvm::ConstantFoldInsertElementInstruction(Constant *Val,
   }
 
   return ConstantVector::get(Result);
 }
 
 Constant *llvm::ConstantFoldShuffleVectorInstruction(Constant *V1,
                                                      Constant *V2,
                                                      Constant *Mask) {
-  unsigned MaskNumElts = Mask->getType()->getVectorNumElements();
+  ElementCount MaskEltCount = Mask->getType()->getVectorElementCount();
   Type *EltTy = V1->getType()->getVectorElementType();
 
   // Undefined shuffle mask -> undefined value.
   if (isa<UndefValue>(Mask))
-    return UndefValue::get(VectorType::get(EltTy, MaskNumElts));
+    return UndefValue::get(VectorType::get(EltTy, MaskEltCount));
 
   // Don't break the bitcode reader hack.
   if (isa<ConstantExpr>(Mask)) return nullptr;
 
   // Do not iterate on scalable vector. The num of elements is unknown at
   // compile-time.
   VectorType *ValTy = cast<VectorType>(V1->getType());
   if (ValTy->isScalable())
     return nullptr;
 
+  unsigned MaskNumElts = MaskEltCount.Min;
   unsigned SrcNumElts = V1->getType()->getVectorNumElements();
 
   // Loop over the shuffle mask, evaluating each element.
   SmallVector<Constant*, 32> Result;
   for (unsigned i = 0; i != MaskNumElts; ++i) {
     int Elt = ShuffleVectorInst::getMaskValue(Mask, i);
     if (Elt == -1) {
       Result.push_back(UndefValue::get(EltTy));

llvm/lib/IR/Instructions.cpp

if (const auto *CE = dyn_cast<ConstantExpr>(Mask))
     if (CE->getOpcode() == Instruction::UserOp1)
       return true;
 
   return false;
 }
 
 int ShuffleVectorInst::getMaskValue(const Constant *Mask, unsigned i) {
   assert(i < Mask->getType()->getVectorNumElements() && "Index out of range");
+  assert(!Mask->getType()->getVectorElementCount().Scalable &&
+    "Length of scalable vectors unknown at compile time");
   if (auto *CDS = dyn_cast<ConstantDataSequential>(Mask))
     return CDS->getElementAsInteger(i);
   Constant *C = Mask->getAggregateElement(i);
   if (isa<UndefValue>(C))
     return -1;
   return cast<ConstantInt>(C)->getZExtValue();
 }
 
 void ShuffleVectorInst::getShuffleMask(const Constant *Mask,
                                        SmallVectorImpl<int> &Result) {
+  assert(!Mask->getType()->getVectorElementCount().Scalable &&
+    "Length of scalable vectors unknown at compile time");
   unsigned NumElts = Mask->getType()->getVectorNumElements();
 
   if (auto *CDS = dyn_cast<ConstantDataSequential>(Mask)) {
     for (unsigned i = 0; i != NumElts; ++i)
       Result.push_back(CDS->getElementAsInteger(i));
     return;
   }
   for (unsigned i = 0; i != NumElts; ++i) {

llvm/test/Analysis/ConstantFolding/vscale-getelementptr.ll

+; RUN: opt -early-cse -S < %s | FileCheck %s
+
+target datalayout = "e-m:e-i8:8:32-i16:16:32-i64:64-i128:128-n32:64-S128"
+target triple = "aarch64"
+
+; CHECK-LABEL: define <4 x i32*> @fixed_length_version_first() {
+; CHECK-NEXT:  ret <4 x i32*> undef
+define <4 x i32*> @fixed_length_version_first() {
+  %ptr = getelementptr i32, <4 x i32*> undef, <4 x i64> undef
+  ret <4 x i32*> %ptr
+}
+
+; CHECK-LABEL: define <4 x <4 x i32>*> @fixed_length_version_second() {
+; CHECK-NEXT:  ret <4 x <4 x i32>*> undef
+define <4 x <4 x i32>*> @fixed_length_version_second() {
+  %ptr = getelementptr <4 x i32>, <4 x i32>* undef, <4 x i64> undef
+  ret <4 x <4 x i32>*> %ptr
+}
+
+; CHECK-LABEL: define <vscale x 4 x i32*> @vscale_version_first() {
+; CHECK-NEXT:  ret <vscale x 4 x i32*> undef
+define <vscale x 4 x i32*> @vscale_version_first() {
+  %ptr = getelementptr i32, <vscale x 4 x i32*> undef, <vscale x 4 x i64> undef
+  ret <vscale x 4 x i32*> %ptr
+}
+
+; CHECK-LABEL: define <vscale x 4 x <vscale x 4 x i32>*> @vscale_version_second() {
+; CHECK-NEXT:  ret <vscale x 4 x <vscale x 4 x i32>*> undef
+define <vscale x 4 x <vscale x 4 x i32>*> @vscale_version_second() {
+  %ptr = getelementptr <vscale x 4 x i32>, <vscale x 4 x i32>* undef, <vscale x 4 x i64> undef
+  ret <vscale x 4 x <vscale x 4 x i32>*> %ptr
+}

llvm/test/Analysis/ConstantFolding/vscale-shufflevector.ll

+; RUN: opt -early-cse -S < %s | FileCheck %s
+
+target datalayout = "e-m:e-i8:8:32-i16:16:32-i64:64-i128:128-n32:64-S128"
+target triple = "aarch64"
+
+; This test checks that SimplifyInstruction does not blow up in the face of
+; a scalable shufflevector. vscale is a constant value known only at runtime.
+; Therefore, it is not possible to know the concrete value of, or the length
+; of the mask at compile time. Simplifications that depend on the value
+; of the mask cannot be performed.
+
+; Given the fact that the value of the mask is unknown at compile time for
+; scalable vectors, very few simplifications will be done. Here, we want to
+; see that the instruction can be passed to SimplifyInstruction and not crash
+; the compiler. It happens to be the case that this will be the result.
+
+; CHECK-LABEL: define <vscale x 8 x i1> @vscale_version()
+; CHECK-NEXT: %splatter = insertelement <vscale x 8 x i1> undef, i1 true, i32 0
+; CHECK-NEXT: %foo = shufflevector <vscale x 8 x i1> %splatter, <vscale x 8 x i1> undef, <vscale x 8 x i32> zeroinitializer
+; CHECK-NEXT: ret <vscale x 8 x i1> %foo
+
+define <vscale x 8 x i1> @vscale_version() {
+  %splatter = insertelement <vscale x 8 x i1> undef, i1 true, i32 0
+  %foo = shufflevector <vscale x 8 x i1> %splatter,
+                       <vscale x 8 x i1> undef,
+                       <vscale x 8 x i32> zeroinitializer
+  ret <vscale x 8 x i1> %foo
+}
+
+; The non-scalable version should be optimized as normal.
+
+; CHECK-LABEL: define <8 x i1> @fixed_length_version() {
+; CHECK-NEXT:  ret <8 x i1> <i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true, i1 true>
+define <8 x i1> @fixed_length_version() {
+  %splatter = insertelement <8 x i1> undef, i1 true, i32 0
+  %foo = shufflevector <8 x i1> %splatter,
+                       <8 x i1> undef,
+                       <8 x i32> zeroinitializer
+  ret <8 x i1> %foo
+}
+