[ConstantFold][NFC] Compile time optimization for large vectors

[ThomasRaoux]

Summary

Optimize the common case of splat vector constant. For large vector going through all the elements is expensive. For splat/broadcast cases we can skip going through all the elements.

[majnemer]

llvm/include/llvm/IR/Constants.h

770

I'd move this into the initializer.

llvm/lib/Analysis/ValueTracking.cpp

179

Is it OK to set DemandedLHS bit 0 when DemandedElts[0] is false?

llvm/lib/IR/ConstantFold.cpp

69–72

Please add a comment for this.

892–897

Please add a comment here.

1388–1389

This doesn't depend on C1.

[ThomasRaoux]

Thanks David.

llvm/include/llvm/IR/Constants.h

770

Done.

llvm/lib/Analysis/ValueTracking.cpp

179

It is okay as long as DemandedElts has at least one bit set. With current code this function would never be called with an empty DemandedElts but I added an early check for this case to handle this case.

llvm/lib/IR/ConstantFold.cpp

69–72

Done.

892–897

Done.

1388–1389

Right, I pulled this part out of the if.

[majnemer]

Please upload with full context.

[ThomasRaoux]

In D76664#1939581, @majnemer wrote:

Please upload with full context.

Done, sorry about that.

[kariddi]

llvm/include/llvm/IR/Constants.h

770

This introduces a linear scan in the creation of the ConstantDataVector() class, I wonder if this might be wasteful if "isSplat()" never endsup to be called on a particular vector.

I wonder if moving this scan in "isSplat()" and caching the result might be better (as isSplat() already was an expensive method, while the construction of ConstantDataVector was not)

[kariddi]

Thanks for the patch! Some speedup in this area was dearly needed!

[majnemer]

llvm/lib/Analysis/ValueTracking.cpp

177–184

I'd float this to the top before we create more APInts.

179

Is it OK that the loop bellow only sets DemandedLHS[0] if DemandedElts[i] is true for some multiple of NumElts while this logic will set it unconditionally?

[ThomasRaoux]

llvm/lib/Analysis/ValueTracking.cpp

177–184

I believe we still need to set DemandedLHS and DemandedRHS to APInt::getNullValue(NumElts) so I cannot really move it up?

179

DemandedLHS[0] is set if for any DemandedElts[i] Shuf->getMaskValue(i) is 0 or NumElts.
In this case we know all the getMaskValue(i) are equal to 0 (since the shuffle mask is aggregateZero) so if any DemandedElts[i] is set it's associated Shuf->getMaskValue(i) will be 0 and DemandedLHS[0] will be set. Since we handle the case where all the bits of DemandedElts are 0 above we know at this point that at least one bit is set.

[ThomasRaoux]

llvm/include/llvm/IR/Constants.h

770

Here is an alternative solution were we lazily set IsSplat. I assume it is going to be rare that we create a constant and not have any pass call isSplat on it but it could make sense to try to keep creation cheap. What do you think of this solution?

[kariddi]

llvm/include/llvm/IR/Constants.h

770

Exactly what I was thinking about thanks!

I think there are potentially cases where isSplat() is not called when a Constant is created. An example is serialization and deserialization of bitcode, where the bitcode is read into an LLVM representation is constructed and then it is just printed out (or viceversa).
Or very late legalization passes before instruction selection.

[ThomasRaoux]

llvm/include/llvm/IR/Constants.h

770

Sounds good, let's leave it that way then.

[majnemer]

llvm/include/llvm/IR/Constants.h

771

You could save a byte by rolling your own optional. I think this is valuable because there may be many ConstantDataVector.

Maybe something like:

bool IsSplatSet : 1;
bool IsSplat : 1;

llvm/lib/Analysis/ValueTracking.cpp

177–184

I was just thinking of floating this bit up:

if (DemandedElts.isNullValue())
  return true;

llvm/lib/IR/ConstantFold.cpp

1389

Now that these are lazy, maybe only do C1->getSplatValue() if C2Splat is true.

llvm/lib/IR/Constants.cpp

2907

Formatting.

[ThomasRaoux]

llvm/include/llvm/IR/Constants.h

771

Done.

llvm/lib/Analysis/ValueTracking.cpp

177–184

What I meant is that DemandedLHS DemandedRHS are out parameters and do need to be set before return true. So we cannot move this above the APInt above.

llvm/lib/IR/ConstantFold.cpp

1389

Done.

llvm/lib/IR/Constants.cpp

2907

Done.

[majnemer]

LGTM

llvm/lib/IR/Instructions.cpp

1965

Instead of resize, I'd reserve + push_back. resize fills the vector with values you will overwrite anyway.

[ThomasRaoux]

Thanks David!

llvm/lib/IR/Instructions.cpp

1965

Done.

llvm/include/llvm/IR/Constants.h

 /// float/double, and whose elements are just simple data values
 /// (i.e. ConstantInt/ConstantFP). This Constant node has no operands because it
 /// stores all of the elements of the constant as densely packed data, instead
 /// of as Value*'s.
 class ConstantDataVector final : public ConstantDataSequential {
   friend class ConstantDataSequential;
 
   explicit ConstantDataVector(Type *ty, const char *Data)
-      : ConstantDataSequential(ty, ConstantDataVectorVal, Data) {}
+      : ConstantDataSequential(ty, ConstantDataVectorVal, Data),
+        IsSplatSet(false) {}
+  // Cache whether or not the constant is a splat.
+  mutable bool IsSplatSet : 1;
+  mutable bool IsSplat : 1;
+  bool isSplatData() const;
 
 public:
   ConstantDataVector(const ConstantDataVector &) = delete;
 
   /// get() constructors - Return a constant with vector type with an element
   /// count and element type matching the ArrayRef passed in.  Note that this
   /// can return a ConstantAggregateZero object.
   static Constant *get(LLVMContext &Context, ArrayRef<uint8_t> Elts);

llvm/lib/Analysis/ValueTracking.cpp

static bool getShuffleDemandedElts(const ShuffleVectorInst *Shuf,
   // The length of scalable vectors is unknown at compile time, thus we
   // cannot check their values
   if (Shuf->getMask()->getType()->getVectorElementCount().Scalable)
     return false;
 
   int NumElts = Shuf->getOperand(0)->getType()->getVectorNumElements();
   int NumMaskElts = Shuf->getMask()->getType()->getVectorNumElements();
   DemandedLHS = DemandedRHS = APInt::getNullValue(NumElts);
-
+  if (DemandedElts.isNullValue())
+    return true;
+  // Simple case of a shuffle with zeroinitializer.
+  if (isa<ConstantAggregateZero>(Shuf->getMask())) {
+    DemandedLHS.setBit(0);
+    return true;
+  }
   for (int i = 0; i != NumMaskElts; ++i) {
     if (!DemandedElts[i])
       continue;
     int M = Shuf->getMaskValue(i);
     assert(M < (NumElts * 2) && "Invalid shuffle mask constant");
 
     // For undef elements, we don't know anything about the common state of
     // the shuffle result.
     if (M == -1)

llvm/lib/IR/ConstantFold.cpp

static Constant *BitCastConstantVector(Constant *CV, VectorType *DstTy) {
   // 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 != CV->getType()->getVectorNumElements())
     return nullptr;
 
   Type *DstEltTy = DstTy->getElementType();
+  // Fast path for splatted constants.
+  if (Constant *Splat = CV->getSplatValue()) {
+    return ConstantVector::getSplat(DstTy->getVectorElementCount(),
+                                    ConstantExpr::getBitCast(Splat, DstEltTy));
+  }
 
   SmallVector<Constant*, 16> Result;
   Type *Ty = IntegerType::get(CV->getContext(), 32);
   for (unsigned i = 0; i != NumElts; ++i) {
     Constant *C =
       ConstantExpr::getExtractElement(CV, ConstantInt::get(Ty, i));
     C = ConstantExpr::getBitCast(C, DstEltTy);
     Result.push_back(C);
   }
 
   return ConstantVector::get(Result);
 }
 
Constant *llvm::ConstantFoldCastInstruction(unsigned opc, Constant *V,
   }
 
   // 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() &&
       DestTy->getVectorNumElements() == V->getType()->getVectorNumElements()) {
-    SmallVector<Constant*, 16> res;
     VectorType *DestVecTy = cast<VectorType>(DestTy);
     Type *DstEltTy = DestVecTy->getElementType();
+    // Fast path for splatted constants.
+    if (Constant *Splat = V->getSplatValue()) {
+      return ConstantVector::getSplat(
+          DestTy->getVectorElementCount(),
+          ConstantExpr::getCast(opc, Splat, DstEltTy));
+    }
+    SmallVector<Constant *, 16> res;
     Type *Ty = IntegerType::get(V->getContext(), 32);
     for (unsigned i = 0, e = V->getType()->getVectorNumElements(); 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::ConstantFoldShuffleVectorInstruction(Constant *V1,
 
   // Undefined shuffle mask -> undefined value.
   if (isa<UndefValue>(Mask))
     return UndefValue::get(VectorType::get(EltTy, MaskEltCount));
 
   // Don't break the bitcode reader hack.
   if (isa<ConstantExpr>(Mask)) return nullptr;
 
+  // If the mask is all zeros this is a splat, no need to go through all
+  // elements.
+  if (isa<ConstantAggregateZero>(Mask) && !MaskEltCount.Scalable) {
+    Type *Ty = IntegerType::get(V1->getContext(), 32);
+    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.
   VectorType *ValTy = cast<VectorType>(V1->getType());
   if (ValTy->isScalable())
     return nullptr;
 
   unsigned MaskNumElts = MaskEltCount.Min;
   unsigned SrcNumElts = V1->getType()->getVectorNumElements();
if (ConstantFP *CFP = dyn_cast<ConstantFP>(C)) {
     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;
+    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;
-    Type *Ty = IntegerType::get(VTy->getContext(), 32);
     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)) {
         return ConstantFP::get(C1->getContext(), C3V);
       }
     }
   } else if (VectorType *VTy = dyn_cast<VectorType>(C1->getType())) {
     // Do not iterate on scalable vector. The number of elements is unknown at
     // compile-time.
     if (IsScalableVector)
       return nullptr;
+    // 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->getVectorElementCount(),
+            ConstantExpr::get(Opcode, C1Splat, C2Splat));
+      }
+    }
 
     // Fold each element and create a vector constant from those constants.
     SmallVector<Constant*, 16> Result;
     Type *Ty = IntegerType::get(VTy->getContext(), 32);
     for (unsigned i = 0, e = VTy->getNumElements(); i != e; ++i) {
       Constant *ExtractIdx = ConstantInt::get(Ty, i);
       Constant *LHS = ConstantExpr::getExtractElement(C1, ExtractIdx);
       Constant *RHS = ConstantExpr::getExtractElement(C2, ExtractIdx);
case FCmpInst::FCMP_OGE:
       return ConstantInt::get(ResultTy, R==APFloat::cmpGreaterThan ||
                                         R==APFloat::cmpEqual);
     }
   } else if (C1->getType()->isVectorTy()) {
     // Do not iterate on scalable vector. The number of elements is unknown at
     // compile-time.
     if (C1->getType()->getVectorIsScalable())
       return nullptr;
+    // Fast path for splatted constants.
+    if (Constant *C1Splat = C1->getSplatValue())
+      if (Constant *C2Splat = C2->getSplatValue())
+        return ConstantVector::getSplat(
+            C1->getType()->getVectorElementCount(),
+            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 = C1->getType()->getVectorNumElements(); i != e;++i){
       Constant *C1E =

llvm/lib/IR/Constants.cpp

bool ConstantDataSequential::isCString() const {
 
   // The last value must be nul.
   if (Str.back() != 0) return false;
 
   // Other elements must be non-nul.
   return Str.drop_back().find(0) == StringRef::npos;
 }
 
-bool ConstantDataVector::isSplat() const {
+bool ConstantDataVector::isSplatData() const {
   const char *Base = getRawDataValues().data();
 
   // Compare elements 1+ to the 0'th element.
   unsigned EltSize = getElementByteSize();
   for (unsigned i = 1, e = getNumElements(); i != e; ++i)
     if (memcmp(Base, Base+i*EltSize, EltSize))
       return false;
 
   return true;
 }
 
+bool ConstantDataVector::isSplat() const {
+  if (!IsSplatSet) {
+    IsSplatSet = true;
+    IsSplat = isSplatData();
+  }
+  return IsSplat;
+}
+
 Constant *ConstantDataVector::getSplatValue() const {
   // If they're all the same, return the 0th one as a representative.
   return isSplat() ? getElementAsConstant(0) : nullptr;
 }
 
 //===----------------------------------------------------------------------===//
 //                handleOperandChange implementations
 

llvm/lib/IR/Instructions.cpp

int ShuffleVectorInst::getMaskValue(const Constant *Mask, unsigned i) {
   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 (isa<ConstantAggregateZero>(Mask)) {
+    Result.resize(NumElts, 0);
+    return;
+  }
+  Result.reserve(NumElts);
   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) {
     Constant *C = Mask->getAggregateElement(i);
-    Result.push_back(isa<UndefValue>(C) ? -1 :
                      cast<ConstantInt>(C)->getZExtValue());
   }
 }
 
 static bool isSingleSourceMaskImpl(ArrayRef<int> Mask, int NumOpElts) {
   assert(!Mask.empty() && "Shuffle mask must contain elements");
   bool UsesLHS = false;
   bool UsesRHS = false;

llvm/lib/Transforms/InstCombine/InstCombineSimplifyDemanded.cpp

Value *InstCombiner::SimplifyDemandedVectorElts(Value *V, APInt DemandedElts,
   }
   case Instruction::ShuffleVector: {
     auto *Shuffle = cast<ShuffleVectorInst>(I);
     assert(Shuffle->getOperand(0)->getType() ==
            Shuffle->getOperand(1)->getType() &&
            "Expected shuffle operands to have same type");
     unsigned OpWidth =
         Shuffle->getOperand(0)->getType()->getVectorNumElements();
+    // Handle trivial case of a splat. Only check the first element of LHS
+    // operand.
+    if (isa<ConstantAggregateZero>(Shuffle->getMask()) &&
+        DemandedElts.isAllOnesValue()) {
+      if (!isa<UndefValue>(I->getOperand(1))) {
+        I->setOperand(1, UndefValue::get(I->getOperand(1)->getType()));
+        MadeChange = true;
+      }
+      APInt LeftDemanded(OpWidth, 1);
+      APInt LHSUndefElts(OpWidth, 0);
+      simplifyAndSetOp(I, 0, LeftDemanded, LHSUndefElts);
+      if (LHSUndefElts[0])
+        UndefElts = EltMask;
+      else
+        UndefElts.clearAllBits();
+      break;
+    }
+
     APInt LeftDemanded(OpWidth, 0), RightDemanded(OpWidth, 0);
     for (unsigned i = 0; i < VWidth; i++) {
       if (DemandedElts[i]) {
         unsigned MaskVal = Shuffle->getMaskValue(i);
         if (MaskVal != -1u) {
           assert(MaskVal < OpWidth * 2 &&
                  "shufflevector mask index out of range!");
           if (MaskVal < OpWidth)