[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

64–67

Please add a comment for this.

891–896

Please add a comment here.

1385–1386

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

64–67

Done.

891–896

Done.

1385–1386

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–178

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–178

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–178

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–178

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

   friend class ConstantDataSequential;
 
   explicit ConstantDataVector(Type *ty, const char *Data)
-      : ConstantDataSequential(ty, ConstantDataVectorVal, Data) {}
+      : ConstantDataSequential(ty, ConstantDataVectorVal, Data),
+        IsSplat(isSplatData(ty, Data)) {}
+  const bool IsSplat;
+  bool isSplatData(Type *ty, const char *Data) const;
 
 public:
   ConstantDataVector(const ConstantDataVector &) = delete;

llvm/lib/Analysis/ValueTracking.cpp

   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;

llvm/lib/IR/ConstantFold.cpp

     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);

   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 =

   // 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());

     // 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);

     // compile-time.
     if (IsScalableVector)
       return nullptr;
+    // Fast path for splatted constants.
+    Constant *C1Splat = C1->getSplatValue();
+    Constant *C2Splat = C2->getSplatValue();
+    if (C2Splat && Instruction::isIntDivRem(Opcode) && C2Splat->isNullValue())
+      return UndefValue::get(VTy);
+    if (C1Splat && C2Splat) {
+      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;

     // compile-time.
     if (C1->getType()->getVectorIsScalable())
       return nullptr;
+    // Fast path for splatted constants.
+    Constant *C1Splat = C1->getSplatValue();
+    Constant *C2Splat = C2->getSplatValue();
+    if (C1Splat && C2Splat) {
+      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.

llvm/lib/IR/Constants.cpp

   return Str.drop_back().find(0) == StringRef::npos;
 }
 
-bool ConstantDataVector::isSplat() const {
-  const char *Base = getRawDataValues().data();
+bool ConstantDataVector::isSplatData(Type *Ty, const char *Data) const {
+  const char *Base = Data;
 
   // Compare elements 1+ to the 0'th element.
-  unsigned EltSize = getElementByteSize();
-  for (unsigned i = 1, e = getNumElements(); i != e; ++i)
+  unsigned EltSize = Ty->getVectorElementType()->getPrimitiveSizeInBits() / 8;
+  for (unsigned i = 1, e = Ty->getVectorNumElements(); i != e; ++i)
     if (memcmp(Base, Base+i*EltSize, EltSize))
       return false;
 
   return true;
 }
 
+bool ConstantDataVector::isSplat() const { return IsSplat; }
+
 Constant *ConstantDataVector::getSplatValue() const {
   // If they're all the same, return the 0th one as a representative.
   return isSplat() ? getElementAsConstant(0) : nullptr;

llvm/lib/IR/Instructions.cpp

   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.resize(NumElts);
   if (auto *CDS = dyn_cast<ConstantDataSequential>(Mask)) {
     for (unsigned i = 0; i != NumElts; ++i)
-      Result.push_back(CDS->getElementAsInteger(i));
+      Result[i] = 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());
+    Result[i] = isa<UndefValue>(C) ? -1 : cast<ConstantInt>(C)->getZExtValue();
   }
 }
 

llvm/lib/Transforms/InstCombine/InstCombineSimplifyDemanded.cpp

            "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]) {