[SVE] Eliminate bad VectorType::getNumElements() calls from ConstantFold

[ctetreau]

Summary

Assume all usages of this function are explicitly fixed-width operations
and cast to FixedVectorType

[ctetreau]

rebase

[sdesmalen]

llvm/lib/IR/ConstantFold.cpp

574

The FIXME is probably unnecessary as it should never be possible to bitcast a fixed-width vector to a scalable vector.

813

How about extractelement <vscale x 16 x i8> zeroinitializer, i64 0 ?

(This probably cannot be a Constant anymore when the index > 16)

1014–1015

This condition/branch has become superfluous given the use of dyn_cast<FixedVectorType> below, and the default of return nullptr.

1381

Although not as superfluous as the case above, I still think it can be removed.

2302

Is it worth just fixing this for scalable vectors instead of adding the FIXME and then doing the wrong thing? It is possible to construct a test-case that breaks this assert.

I think you can use VectorType::get(..., VT->getElementCount()) here (but it would need a test-case)

[ctetreau]

rebase

[ctetreau]

address code review issues

[ctetreau]

llvm/lib/IR/ConstantFold.cpp

813

I don't quite remember what (if anything) we decided to do about this case last thursday. I believe I was on the "it might be out of bounds, but we can't know so we shouldn't do anything" side of the argument, but I don't remember which side won.

Regardless, it would probably be more appropriate to just skip this particular transformation rather than letting it explode. I think it makes sense to just do that for now and revisit when we have all the correctness issues sorted out. I suppose this represents a behavior change, and needs a test...

1014–1015

good catch

1381

This could potentially result in a behavior change, which I am trying to avoid as much as possible in this patch. If you're sure it doesn't, I can change it. Otherwise, I can add a FIXME.

2302

If I had the authority to make the call, I'd say we should just replace all instances of VectorType::get(..., SomeVec->getNumElements() /*, false */) with VectorType::get(..., SomeVec->getElementCount()) without tests, as the original form is almost certainly a bug.

Unfortunately, others disagree with this. While it pains me to write this code, my current goal is to just get rid of the bad usages of VectorType::getNumElements in as frictionless a manner as possible, which means taking the coward's way out so as not to get bogged down writing a billion tests.

[ctetreau]

Add support and test for for constant folding extractelement (<vscale x n x ty> zeroinitializer, i64 i) for i = 0 .. n - 1

[ctetreau]

llvm/lib/IR/ConstantFold.cpp

813

I went ahead and explicitly added support for this case when the index is less than the minNumElements.

[sdesmalen]

LGTM.

llvm/lib/IR/ConstantFold.cpp

813

I went ahead and explicitly added support for this case when the index is less than the minNumElements

Thanks, I think that's the right fix.

This change probably sufficient, because AIUI there is no issue with this form if index > 16. The LangRef indeed specifies the result is 'Poison' if this turns out to be out-of-bounds, so that will lead to undefined behaviour in its propagated uses, there just isn't necessarily a way to know that it is poison at compiletime. The ISD nodes have comments saying that this case should be guarded so that it doesn't cause a runcrash at the point of doing the extract (e.g.if the extract is implemented by doing a vector store + scalar load). This behaviour is not really any different from if it wasn't constant.

1381

maybe better to err on the safe side here. Fine to leave it as is.

2302

I see your point, but there's the danger that without tests we don't guard the code-path now working for scalable vectors . If you make the change to getElementCount, we'll probably never revisit this to make a test for it. If we add tests while going through the code-base, we'll incrementally guard other code-paths as well, because it's difficult to test these things entirely in isolation.

I tried constructing a test for this case, but this leads to an unrelated issue parsing the IR, error: constant expression type mismatch.

define <vscale x 2 x i64*> @foo() {
  ret <vscale x 2 x i64*> getelementptr ([1 x i64], [1 x i64]* null, i64 0, <vscale x 2 x i64> zeroinitializer)
}

whereas removing vscale x causes the expression to be folded.

I agree that this (together with parser error) is better fixed in a separate patch.

llvm/lib/IR/ConstantFold.cpp

static Constant *BitCastConstantVector(Constant *CV, VectorType *DstTy) {
   // Do not iterate on scalable vector. The num of elements is unknown at
   // compile-time.
   if (isa<ScalableVectorType>(DstTy))
     return nullptr;
 
   // If this cast changes element count then we can't handle it here:
   // doing so requires endianness information.  This should be handled by
   // Analysis/ConstantFolding.cpp
-  unsigned NumElts = DstTy->getNumElements();
-  if (NumElts != cast<VectorType>(CV->getType())->getNumElements())
+  unsigned NumElts = cast<FixedVectorType>(DstTy)->getNumElements();
+  if (NumElts != cast<FixedVectorType>(CV->getType())->getNumElements())
     return nullptr;
 
   Type *DstEltTy = DstTy->getElementType();
   // Fast path for splatted constants.
   if (Constant *Splat = CV->getSplatValue()) {
     return ConstantVector::getSplat(DstTy->getElementCount(),
                                     ConstantExpr::getBitCast(Splat, DstEltTy));
   }
if (CE->isCast()) {
         // This is casting one pointer type to another, always BitCast
         return ConstantExpr::getPointerCast(CE->getOperand(0), DestTy);
     }
   }
 
   // If the cast operand is a constant vector, perform the cast by
   // operating on each element. In the cast of bitcasts, the element
   // count may be mismatched; don't attempt to handle that here.
   if ((isa<ConstantVector>(V) || isa<ConstantDataVector>(V)) &&
       DestTy->isVectorTy() &&
-      cast<VectorType>(DestTy)->getNumElements() ==
-          cast<VectorType>(V->getType())->getNumElements()) {
+      cast<FixedVectorType>(DestTy)->getNumElements() ==
+          cast<FixedVectorType>(V->getType())->getNumElements()) {
     VectorType *DestVecTy = cast<VectorType>(DestTy);
     Type *DstEltTy = DestVecTy->getElementType();
     // Fast path for splatted constants.
     if (Constant *Splat = V->getSplatValue()) {
       return ConstantVector::getSplat(
           cast<VectorType>(DestTy)->getElementCount(),
           ConstantExpr::getCast(opc, Splat, DstEltTy));
     }
     SmallVector<Constant *, 16> res;
     Type *Ty = IntegerType::get(V->getContext(), 32);
-    for (unsigned i = 0, e = cast<VectorType>(V->getType())->getNumElements();
+    for (unsigned i = 0,
+                  e = cast<FixedVectorType>(V->getType())->getNumElements();
          i != e; ++i) {
       Constant *C =
         ConstantExpr::getExtractElement(V, ConstantInt::get(Ty, i));
       res.push_back(ConstantExpr::getCast(opc, C, DstEltTy));
     }
     return ConstantVector::get(res);
   }
 
Constant *llvm::ConstantFoldExtractElementInstruction(Constant *Val,
   // extractelt C, undef -> undef
   if (isa<UndefValue>(Val) || isa<UndefValue>(Idx))
     return UndefValue::get(ValVTy->getElementType());
 
   auto *CIdx = dyn_cast<ConstantInt>(Idx);
   if (!CIdx)
     return nullptr;
 
+  if (auto *ValFVTy = dyn_cast<FixedVectorType>(Val->getType())) {
     // ee({w,x,y,z}, wrong_value) -> undef
-  if (CIdx->uge(ValVTy->getNumElements()))
-    return UndefValue::get(ValVTy->getElementType());
+    if (CIdx->uge(ValFVTy->getNumElements()))
+      return UndefValue::get(ValFVTy->getElementType());
+  }
 
   // ee (gep (ptr, idx0, ...), idx) -> gep (ee (ptr, idx), ee (idx0, idx), ...)
   if (auto *CE = dyn_cast<ConstantExpr>(Val)) {
     if (CE->getOpcode() == Instruction::GetElementPtr) {
       SmallVector<Constant *, 8> Ops;
       Ops.reserve(CE->getNumOperands());
       for (unsigned i = 0, e = CE->getNumOperands(); i != e; ++i) {
         Constant *Op = CE->getOperand(i);
         if (Op->getType()->isVectorTy()) {
           Constant *ScalarOp = ConstantExpr::getExtractElement(Op, Idx);
           if (!ScalarOp)
-            return  nullptr;
+            return nullptr;
           Ops.push_back(ScalarOp);
         } else
           Ops.push_back(Op);
       }
       return CE->getWithOperands(Ops, ValVTy->getElementType(), false,
                                  Ops[0]->getType()->getPointerElementType());
     }
   }
 
+  // CAZ of type ScalableVectorType and n < CAZ->getMinNumElements() =>
+  //   extractelt CAZ, n -> 0
+  if (auto *ValSVTy = dyn_cast<ScalableVectorType>(Val->getType())) {
+    if (!CIdx->uge(ValSVTy->getMinNumElements())) {
+      if (auto *CAZ = dyn_cast<ConstantAggregateZero>(Val))
+        return CAZ->getElementValue(CIdx->getZExtValue());
+    }
+    return nullptr;
+  }
+
   return Val->getAggregateElement(CIdx);
 }
 
 Constant *llvm::ConstantFoldInsertElementInstruction(Constant *Val,
                                                      Constant *Elt,
                                                      Constant *Idx) {
   if (isa<UndefValue>(Idx))
     return UndefValue::get(Val->getType());
 
   ConstantInt *CIdx = dyn_cast<ConstantInt>(Idx);
   if (!CIdx) return nullptr;
 
   // Do not iterate on scalable vector. The num of elements is unknown at
   // compile-time.
-  VectorType *ValTy = cast<VectorType>(Val->getType());
-  if (isa<ScalableVectorType>(ValTy))
+  if (isa<ScalableVectorType>(Val->getType()))
     return nullptr;
 
-  unsigned NumElts = cast<VectorType>(Val->getType())->getNumElements();
+  auto *ValTy = cast<FixedVectorType>(Val->getType());
+
+  unsigned NumElts = ValTy->getNumElements();
   if (CIdx->uge(NumElts))
     return UndefValue::get(Val->getType());
 
   SmallVector<Constant*, 16> Result;
   Result.reserve(NumElts);
   auto *Ty = Type::getInt32Ty(Val->getContext());
   uint64_t IdxVal = CIdx->getZExtValue();
   for (unsigned i = 0; i != NumElts; ++i) {
Constant *Elt =
         ConstantExpr::getExtractElement(V1, ConstantInt::get(Ty, 0));
     return ConstantVector::getSplat(MaskEltCount, Elt);
   }
   // Do not iterate on scalable vector. The num of elements is unknown at
   // compile-time.
   if (isa<ScalableVectorType>(V1VTy))
     return nullptr;
 
-  unsigned SrcNumElts = V1VTy->getNumElements();
+  unsigned SrcNumElts = V1VTy->getElementCount().Min;
 
   // Loop over the shuffle mask, evaluating each element.
   SmallVector<Constant*, 32> Result;
   for (unsigned i = 0; i != MaskNumElts; ++i) {
     int Elt = Mask[i];
     if (Elt == -1) {
       Result.push_back(UndefValue::get(EltTy));
       continue;
Constant *llvm::ConstantFoldUnaryInstruction(unsigned Opcode, Constant *C) {
 
   if (ConstantFP *CFP = dyn_cast<ConstantFP>(C)) {
     const APFloat &CV = CFP->getValueAPF();
     switch (Opcode) {
     default:
       break;
     case Instruction::FNeg:
       return ConstantFP::get(C->getContext(), neg(CV));
     }
-  } else if (VectorType *VTy = dyn_cast<VectorType>(C->getType())) {
-    // Do not iterate on scalable vector. The number of elements is unknown at
-    // compile-time.
-    if (IsScalableVector)
-      return nullptr;
+  } else if (auto *VTy = dyn_cast<FixedVectorType>(C->getType())) {
+
     Type *Ty = IntegerType::get(VTy->getContext(), 32);
     // Fast path for splatted constants.
     if (Constant *Splat = C->getSplatValue()) {
       Constant *Elt = ConstantExpr::get(Opcode, Splat);
       return ConstantVector::getSplat(VTy->getElementCount(), Elt);
     }
 
     // Fold each element and create a vector constant from those constants.
-    SmallVector<Constant*, 16> Result;
+    SmallVector<Constant *, 16> Result;
     for (unsigned i = 0, e = VTy->getNumElements(); i != e; ++i) {
       Constant *ExtractIdx = ConstantInt::get(Ty, i);
       Constant *Elt = ConstantExpr::getExtractElement(C, ExtractIdx);
 
       Result.push_back(ConstantExpr::get(Opcode, Elt));
     }
 
     return ConstantVector::get(Result);
if (ConstantFP *CFP2 = dyn_cast<ConstantFP>(C2)) {
       case Instruction::FDiv:
         (void)C3V.divide(C2V, APFloat::rmNearestTiesToEven);
         return ConstantFP::get(C1->getContext(), C3V);
       case Instruction::FRem:
         (void)C3V.mod(C2V);
         return ConstantFP::get(C1->getContext(), C3V);
       }
     }
-  } else if (VectorType *VTy = dyn_cast<VectorType>(C1->getType())) {
+  } else if (IsScalableVector) {
     // Do not iterate on scalable vector. The number of elements is unknown at
     // compile-time.
-    if (IsScalableVector)
-      return nullptr;
+    // FIXME: this branch can potentially be removed
+  } else if (auto *VTy = dyn_cast<FixedVectorType>(C1->getType())) {
     // Fast path for splatted constants.
     if (Constant *C2Splat = C2->getSplatValue()) {
       if (Instruction::isIntDivRem(Opcode) && C2Splat->isNullValue())
         return UndefValue::get(VTy);
       if (Constant *C1Splat = C1->getSplatValue()) {
         return ConstantVector::getSplat(
             VTy->getElementCount(),
             ConstantExpr::get(Opcode, C1Splat, C2Splat));
if (Constant *C1Splat = C1->getSplatValue())
             C1VTy->getElementCount(),
             ConstantExpr::getCompare(pred, C1Splat, C2Splat));
 
     // If we can constant fold the comparison of each element, constant fold
     // the whole vector comparison.
     SmallVector<Constant*, 4> ResElts;
     Type *Ty = IntegerType::get(C1->getContext(), 32);
     // Compare the elements, producing an i1 result or constant expr.
-    for (unsigned i = 0, e = C1VTy->getNumElements(); i != e; ++i) {
+    for (unsigned i = 0, e = C1VTy->getElementCount().Min; i != e; ++i) {
       Constant *C1E =
         ConstantExpr::getExtractElement(C1, ConstantInt::get(Ty, i));
       Constant *C2E =
         ConstantExpr::getExtractElement(C2, ConstantInt::get(Ty, i));
 
       ResElts.push_back(ConstantExpr::getCompare(pred, C1E, C2E));
     }
 
for (unsigned i = 0, e = Idxs.size(); i != e; ++i)
       }
     if (isNull) {
       PointerType *PtrTy = cast<PointerType>(C->getType()->getScalarType());
       Type *Ty = GetElementPtrInst::getIndexedType(PointeeTy, Idxs);
 
       assert(Ty && "Invalid indices for GEP!");
       Type *OrigGEPTy = PointerType::get(Ty, PtrTy->getAddressSpace());
       Type *GEPTy = PointerType::get(Ty, PtrTy->getAddressSpace());
-      if (VectorType *VT = dyn_cast<VectorType>(C->getType()))
-        GEPTy = FixedVectorType::get(OrigGEPTy, VT->getNumElements());
-
+      if (VectorType *VT = dyn_cast<VectorType>(C->getType())) {
+        // FIXME: handle scalable vectors (use getElementCount())
+        GEPTy = FixedVectorType::get(
+            OrigGEPTy, cast<FixedVectorType>(VT)->getNumElements());
+      }
       // The GEP returns a vector of pointers when one of more of
       // its arguments is a vector.
       for (unsigned i = 0, e = Idxs.size(); i != e; ++i) {
         if (auto *VT = dyn_cast<VectorType>(Idxs[i]->getType())) {
-          GEPTy = FixedVectorType::get(OrigGEPTy, VT->getNumElements());
+          // FIXME: handle scalable vectors
+          GEPTy = FixedVectorType::get(
+              OrigGEPTy, cast<FixedVectorType>(VT)->getNumElements());
           break;
         }
       }
 
       return Constant::getNullValue(GEPTy);
     }
   }
 
for (unsigned i = 1, e = Idxs.size(); i != e;
     auto *PrevIdx =
         NewIdxs[i - 1] ? NewIdxs[i - 1] : cast<Constant>(Idxs[i - 1]);
     bool IsCurrIdxVector = CurrIdx->getType()->isVectorTy();
     bool IsPrevIdxVector = PrevIdx->getType()->isVectorTy();
     bool UseVector = IsCurrIdxVector || IsPrevIdxVector;
 
     if (!IsCurrIdxVector && IsPrevIdxVector)
       CurrIdx = ConstantDataVector::getSplat(
-          cast<VectorType>(PrevIdx->getType())->getNumElements(), CurrIdx);
+          cast<FixedVectorType>(PrevIdx->getType())->getNumElements(), CurrIdx);
 
     if (!IsPrevIdxVector && IsCurrIdxVector)
       PrevIdx = ConstantDataVector::getSplat(
-          cast<VectorType>(CurrIdx->getType())->getNumElements(), PrevIdx);
+          cast<FixedVectorType>(CurrIdx->getType())->getNumElements(), PrevIdx);
 
     Constant *Factor =
         ConstantInt::get(CurrIdx->getType()->getScalarType(), NumElements);
     if (UseVector)
       Factor = ConstantDataVector::getSplat(
           IsPrevIdxVector
-              ? cast<VectorType>(PrevIdx->getType())->getNumElements()
-              : cast<VectorType>(CurrIdx->getType())->getNumElements(),
+              ? cast<FixedVectorType>(PrevIdx->getType())->getNumElements()
+              : cast<FixedVectorType>(CurrIdx->getType())->getNumElements(),
           Factor);
 
     NewIdxs[i] = ConstantExpr::getSRem(CurrIdx, Factor);
 
     Constant *Div = ConstantExpr::getSDiv(CurrIdx, Factor);
 
     unsigned CommonExtendedWidth =
         std::max(PrevIdx->getType()->getScalarSizeInBits(),
                  Div->getType()->getScalarSizeInBits());
     CommonExtendedWidth = std::max(CommonExtendedWidth, 64U);
 
     // Before adding, extend both operands to i64 to avoid
     // overflow trouble.
     Type *ExtendedTy = Type::getIntNTy(Div->getContext(), CommonExtendedWidth);
     if (UseVector)
       ExtendedTy = FixedVectorType::get(
           ExtendedTy,
           IsPrevIdxVector
-              ? cast<VectorType>(PrevIdx->getType())->getNumElements()
-              : cast<VectorType>(CurrIdx->getType())->getNumElements());
+              ? cast<FixedVectorType>(PrevIdx->getType())->getNumElements()
+              : cast<FixedVectorType>(CurrIdx->getType())->getNumElements());
 
     if (!PrevIdx->getType()->isIntOrIntVectorTy(CommonExtendedWidth))
       PrevIdx = ConstantExpr::getSExt(PrevIdx, ExtendedTy);
 
     if (!Div->getType()->isIntOrIntVectorTy(CommonExtendedWidth))
       Div = ConstantExpr::getSExt(Div, ExtendedTy);
 
     NewIdxs[i - 1] = ConstantExpr::getAdd(PrevIdx, Div);

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

+; RUN: opt -instcombine -S < %s | FileCheck %s
+
+; CHECK-LABEL: definitely_in_bounds
+; CHECK: ret i8 0
+define i8 @definitely_in_bounds() {
+  ret i8 extractelement (<vscale x 16 x i8> zeroinitializer, i64 15)
+}
+
+; CHECK-LABEL: maybe_in_bounds
+; CHECK: ret i8 extractelement (<vscale x 16 x i8> zeroinitializer, i64 16)
+define i8 @maybe_in_bounds() {
+  ret i8 extractelement (<vscale x 16 x i8> zeroinitializer, i64 16)
+}