[SLP] Fix sign-extends for type-shrinking

[mssimpso]

Summary

This patch ensures the correct minimum bit width during type-shrinking. Previously when type-shrinking, we always sign-extended values back to their original width. However, if we are going to sign-extend, and the sign bit is unknown, we have to increase the minimum bit width by one bit so the sign-extend will fill the upper bits correctly. If the sign bit is known to be zero, we can perform a zero-extend instead. This should fix PR31243.

Reference: https://llvm.org/bugs/show_bug.cgi?id=31243

[mkuper]

test/Transforms/SLPVectorizer/X86/minimum-sizes.ll

13

Can you make the test more explicit by checking the extends we actually want to appear?

[mssimpso]

test/Transforms/SLPVectorizer/X86/minimum-sizes.ll

13

Sure. I'll upload a new version.

[mssimpso]

Addressed Michael's comments.

• Made expected test output more explicit.

[mkuper]

test/Transforms/SLPVectorizer/X86/minimum-sizes.ll

13

Shouldn't we also have a test where we increase the bit-width because we don't know the sign bit?
I'm a bit surprised this didn't affect any of our existing tests.

[mssimpso]

test/Transforms/SLPVectorizer/X86/minimum-sizes.ll

13

That's a good idea - I'll add another test for the sext case.

[mssimpso]

Addressed Michael's comments.

• Reduced the existing test case. The test now uses -slp-threshold to force vectorization.
• Added another test for the sext case.
• Added FIXME's in both the source and test explaining why the sext case is currently suboptimal, but correct.

[mssimpso]

Renamed test variables.

[mkuper]

LGTM with a couple of nits.

lib/Transforms/Vectorize/SLPVectorizer.cpp

1936

Could you rename this? ExTy makes it look like a type (yes, it's the "type of the extensions", but...)

3533

Can you use all_of?

[mssimpso]

lib/Transforms/Vectorize/SLPVectorizer.cpp

1936

Sure, I'll just go with Extend.

3533

Yes, good point.

[Gerolf]

Michael, could you also lend this your expert eye? Do you agree with the extra memory (pair vs vector) and compile-time (look at the loop tree) investment for this fix? Shortening the types may not be worth it at least on some architectures.

Thanks
Gerolf

lib/Transforms/Vectorize/SLPVectorizer.cpp

1939–1940

I would expect the TTI function to make use of the extra information of the sign bit, but can't see that code. Or is this intentional?

3521

This looks convoluted. You define a lambda and then use the inverted value. I think what you want i hasSignBitSet[] which returns is KnownOne. Then you use this as IsSigned = hasSignBitSet(). I think in your code you mixup concepts and intended use.

[mssimpso]

In D27466#618754, @Gerolf wrote:

Michael, could you also lend this your expert eye? Do you agree with the extra memory (pair vs vector) and compile-time (look at the loop tree) investment for this fix? Shortening the types may not be worth it at least on some architectures.

Gerolf,

In terms of compile-time, the fix adds a single pass over the *roots* of non-store-rooted trees. The number of roots is equal to the vectorization factor, so this will be minimal. In terms of memory usage, the fix adds a bool for each instruction in the vectorizable trees that we can type-shrink. This also seems minimal to me.

lib/Transforms/Vectorize/SLPVectorizer.cpp

1939–1940

I'm not sure how the TTI function could use sign bit information here. We know that if we have a zext, then the sign bit is zero. If we have a sext, the sign bit can be either one or zero. The TTI function returns the expect cost of the pair of instructions:

%0 = extractelement ...
%1 = zext/sext %0 ...

We're giving the TTI function the kind of extend we are performing. Is there are case you have in mind that I'm not thinking of?

3521

I'm not sure this is what we want. We're checking if the sign bit is known to be zero, not known to be one. We set IsSigned only if we can't prove the sign bit to be zero (even though in reality it may be). Perhaps the "IsSigned" name is a bit confusing? I'll rename it when rewriting the code below to use "any_of" to address Michael's comment.

To clarify what's happening, we're checking the roots to determine if we know the sign bit is zero. If so, we zext. If we can't prove the sign bit to be zero, we add one to the max bit width and sext.

[mssimpso]

Addressed comments from Michael and Gerolf. Thanks!

• Rewrote sign bit computation to use "all_of" (Michael), and renamed variable to make the logic less confusing (Gerolf).
• Renamed ExTy to Extend (Michael).

[mkuper]

Still LGTM.

Gerolf / Michael Z, are you ok with continuing the discussion post-commit, if you still have objections?

Gerolf, correct me if I'm wrong, but it sounds like your main issue was not with this patch, but rather with the whole approach to type shortening.
I believe we should keep it as is, but even if we do end up deciding to remove it, I think we should fix the functional bug we have (PR31243) first.

[mzolotukhin]

I agree with Michael K here - first we need to fix the stability issue, then we can discuss if there is anything we want to do another way.

Michael Z

[mssimpso]

Thanks everyone. I'll go ahead and commit the bug fix for now, and then we can continue the discussion.

lib/Transforms/Vectorize/SLPVectorizer.cpp

#endif
   DemandedBits *DB;
   const DataLayout *DL;
   unsigned MaxVecRegSize; // This is set by TTI or overridden by cl::opt.
   unsigned MinVecRegSize; // Set by cl::opt (default: 128).
   /// Instruction builder to construct the vectorized tree.
   IRBuilder<> Builder;
 
   /// A map of scalar integer values to the smallest bit width with which they
-  /// can legally be represented.
-  MapVector<Value *, uint64_t> MinBWs;
+  /// can legally be represented. The values map to (width, signed) pairs,
+  /// where "width" indicates the minimum bit width and "signed" is True if the
+  /// value must be signed-extended, rather than zero-extended, back to its
+  /// original width.
+  MapVector<Value *, std::pair<uint64_t, bool>> MinBWs;
 };
 
 } // end namespace llvm
 } // end namespace slpvectorizer
 
 void BoUpSLP::buildTree(ArrayRef<Value *> Roots,
                         ArrayRef<Value *> UserIgnoreLst) {
   deleteTree();
int BoUpSLP::getEntryCost(TreeEntry *E) {
   Type *ScalarTy = VL[0]->getType();
   if (StoreInst *SI = dyn_cast<StoreInst>(VL[0]))
     ScalarTy = SI->getValueOperand()->getType();
   VectorType *VecTy = VectorType::get(ScalarTy, VL.size());
 
   // If we have computed a smaller type for the expression, update VecTy so
   // that the costs will be accurate.
   if (MinBWs.count(VL[0]))
-    VecTy = VectorType::get(IntegerType::get(F->getContext(), MinBWs[VL[0]]),
-                            VL.size());
+    VecTy = VectorType::get(
+        IntegerType::get(F->getContext(), MinBWs[VL[0]].first), VL.size());
 
   if (E->NeedToGather) {
     if (allConstant(VL))
       return 0;
     if (isSplat(VL)) {
       return TTI->getShuffleCost(TargetTransformInfo::SK_Broadcast, VecTy, 0);
     }
     return getGatherCost(E->Scalars);
if (EphValues.count(EU.User))
       continue;
 
     // If we plan to rewrite the tree in a smaller type, we will need to sign
     // extend the extracted value back to the original type. Here, we account
     // for the extract and the added cost of the sign extend if needed.
     auto *VecTy = VectorType::get(EU.Scalar->getType(), BundleWidth);
     auto *ScalarRoot = VectorizableTree[0].Scalars[0];
     if (MinBWs.count(ScalarRoot)) {
-      auto *MinTy = IntegerType::get(F->getContext(), MinBWs[ScalarRoot]);
+      auto *MinTy = IntegerType::get(F->getContext(), MinBWs[ScalarRoot].first);
+      auto Extend =
+          MinBWs[ScalarRoot].second ? Instruction::SExt : Instruction::ZExt;
       VecTy = VectorType::get(MinTy, BundleWidth);
-      ExtractCost += TTI->getExtractWithExtendCost(
-          Instruction::SExt, EU.Scalar->getType(), VecTy, EU.Lane);
+      ExtractCost += TTI->getExtractWithExtendCost(Extend, EU.Scalar->getType(),
+                                                   VecTy, EU.Lane);
     } else {
       ExtractCost +=
           TTI->getVectorInstrCost(Instruction::ExtractElement, VecTy, EU.Lane);
     }
   }
 
   int SpillCost = getSpillCost();
   Cost += SpillCost + ExtractCost;
Value *BoUpSLP::vectorizeTree() {
   // If the vectorized tree can be rewritten in a smaller type, we truncate the
   // vectorized root. InstCombine will then rewrite the entire expression. We
   // sign extend the extracted values below.
   auto *ScalarRoot = VectorizableTree[0].Scalars[0];
   if (MinBWs.count(ScalarRoot)) {
     if (auto *I = dyn_cast<Instruction>(VectorRoot))
       Builder.SetInsertPoint(&*++BasicBlock::iterator(I));
     auto BundleWidth = VectorizableTree[0].Scalars.size();
-    auto *MinTy = IntegerType::get(F->getContext(), MinBWs[ScalarRoot]);
+    auto *MinTy = IntegerType::get(F->getContext(), MinBWs[ScalarRoot].first);
     auto *VecTy = VectorType::get(MinTy, BundleWidth);
     auto *Trunc = Builder.CreateTrunc(VectorRoot, VecTy);
     VectorizableTree[0].VectorizedValue = Trunc;
   }
 
   DEBUG(dbgs() << "SLP: Extracting " << ExternalUses.size() << " values .\n");
 
+  // If necessary, sign-extend or zero-extend ScalarRoot to the larger type
+  // specified by ScalarType.
+  auto extend = [&](Value *ScalarRoot, Value *Ex, Type *ScalarType) {
+    if (!MinBWs.count(ScalarRoot))
+      return Ex;
+    if (MinBWs[ScalarRoot].second)
+      return Builder.CreateSExt(Ex, ScalarType);
+    return Builder.CreateZExt(Ex, ScalarType);
+  };
+
   // Extract all of the elements with the external uses.
   for (const auto &ExternalUse : ExternalUses) {
     Value *Scalar = ExternalUse.Scalar;
     llvm::User *User = ExternalUse.User;
 
     // Skip users that we already RAUW. This happens when one instruction
     // has multiple uses of the same value.
     if (!is_contained(Scalar->users(), User))
if (auto *VecI = dyn_cast<Instruction>(Vec)) {
                 PH->getIncomingBlock(i)->getTerminator();
             if (isa<CatchSwitchInst>(IncomingTerminator)) {
               Builder.SetInsertPoint(VecI->getParent(),
                                      std::next(VecI->getIterator()));
             } else {
               Builder.SetInsertPoint(PH->getIncomingBlock(i)->getTerminator());
             }
             Value *Ex = Builder.CreateExtractElement(Vec, Lane);
-            if (MinBWs.count(ScalarRoot))
-              Ex = Builder.CreateSExt(Ex, Scalar->getType());
+            Ex = extend(ScalarRoot, Ex, Scalar->getType());
             CSEBlocks.insert(PH->getIncomingBlock(i));
             PH->setOperand(i, Ex);
           }
         }
       } else {
         Builder.SetInsertPoint(cast<Instruction>(User));
         Value *Ex = Builder.CreateExtractElement(Vec, Lane);
-        if (MinBWs.count(ScalarRoot))
-          Ex = Builder.CreateSExt(Ex, Scalar->getType());
+        Ex = extend(ScalarRoot, Ex, Scalar->getType());
         CSEBlocks.insert(cast<Instruction>(User)->getParent());
         User->replaceUsesOfWith(Scalar, Ex);
      }
     } else {
       Builder.SetInsertPoint(&F->getEntryBlock().front());
       Value *Ex = Builder.CreateExtractElement(Vec, Lane);
-      if (MinBWs.count(ScalarRoot))
-        Ex = Builder.CreateSExt(Ex, Scalar->getType());
+      Ex = extend(ScalarRoot, Ex, Scalar->getType());
       CSEBlocks.insert(&F->getEntryBlock());
       User->replaceUsesOfWith(Scalar, Ex);
     }
 
     DEBUG(dbgs() << "SLP: Replaced:" << *User << ".\n");
   }
 
   // For each vectorized value:
void BoUpSLP::computeMinimumValueSizes() {
   // We first check if all the bits of the roots are demanded. If they're not,
   // we can truncate the roots to this narrower type.
   for (auto *Root : TreeRoot) {
     auto Mask = DB->getDemandedBits(cast<Instruction>(Root));
     MaxBitWidth = std::max<unsigned>(
         Mask.getBitWidth() - Mask.countLeadingZeros(), MaxBitWidth);
   }
 
+  // True if the roots can be zero-extended back to their original type, rather
+  // than sign-extended. We know that if the leading bits are not demanded, we
+  // can safely zero-extend. So we initialize IsKnownPositive to True.
+  bool IsKnownPositive = true;
+
   // If all the bits of the roots are demanded, we can try a little harder to
   // compute a narrower type. This can happen, for example, if the roots are
   // getelementptr indices. InstCombine promotes these indices to the pointer
   // width. Thus, all their bits are technically demanded even though the
   // address computation might be vectorized in a smaller type.
   //
   // We start by looking at each entry that can be demoted. We compute the
   // maximum bit width required to store the scalar by using ValueTracking to
   // compute the number of high-order bits we can truncate.
   if (MaxBitWidth == DL->getTypeSizeInBits(TreeRoot[0]->getType())) {
     MaxBitWidth = 8u;
+
+    // Determine if the sign bit of all the roots is known to be zero. If not,
+    // IsKnownPositive is set to False.
+    IsKnownPositive = all_of(TreeRoot, [&](Value *R) {
+      bool KnownZero = false;
+      bool KnownOne = false;
+      ComputeSignBit(R, KnownZero, KnownOne, *DL);
+      return KnownZero;
+    });
+
+    // Determine the maximum number of bits required to store the scalar
+    // values.
     for (auto *Scalar : ToDemote) {
       auto NumSignBits = ComputeNumSignBits(Scalar, *DL, 0, AC, 0, DT);
       auto NumTypeBits = DL->getTypeSizeInBits(Scalar->getType());
       MaxBitWidth = std::max<unsigned>(NumTypeBits - NumSignBits, MaxBitWidth);
     }
+
+    // If we can't prove that the sign bit is zero, we must add one to the
+    // maximum bit width to account for the unknown sign bit. This preserves
+    // the existing sign bit so we can safely sign-extend the root back to the
+    // original type. Otherwise, if we know the sign bit is zero, we will
+    // zero-extend the root instead.
+    //
+    // FIXME: This is somewhat suboptimal, as there will be cases where adding
+    //        one to the maximum bit width will yield a larger-than-necessary
+    //        type. In general, we need to add an extra bit only if we can't
+    //        prove that the upper bit of the original type is equal to the
+    //        upper bit of the proposed smaller type. If these two bits are the
+    //        same (either zero or one) we know that sign-extending from the
+    //        smaller type will result in the same value. Here, since we can't
+    //        yet prove this, we are just making the proposed smaller type
+    //        larger to ensure correctness.
+    if (!IsKnownPositive)
+      ++MaxBitWidth;
   }
 
   // Round MaxBitWidth up to the next power-of-two.
   if (!isPowerOf2_64(MaxBitWidth))
     MaxBitWidth = NextPowerOf2(MaxBitWidth);
 
   // If the maximum bit width we compute is less than the with of the roots'
   // type, we can proceed with the narrowing. Otherwise, do nothing.
   if (MaxBitWidth >= TreeRootIT->getBitWidth())
     return;
 
   // If we can truncate the root, we must collect additional values that might
   // be demoted as a result. That is, those seeded by truncations we will
   // modify.
   while (!Roots.empty())
     collectValuesToDemote(Roots.pop_back_val(), Expr, ToDemote, Roots);
 
   // Finally, map the values we can demote to the maximum bit with we computed.
   for (auto *Scalar : ToDemote)
-    MinBWs[Scalar] = MaxBitWidth;
+    MinBWs[Scalar] = std::make_pair(MaxBitWidth, !IsKnownPositive);
 }
 
 namespace {
 /// The SLPVectorizer Pass.
 struct SLPVectorizer : public FunctionPass {
   SLPVectorizerPass Impl;
 
   /// Pass identification, replacement for typeid

test/Transforms/SLPVectorizer/X86/minimum-sizes.ll

+; RUN: opt -S -slp-threshold=-6 -slp-vectorizer -instcombine < %s | FileCheck %s
+
+target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux-gnu"
+
+; These tests ensure that we do not regress due to PR31243. Note that we set
+; the SLP threshold to force vectorization even when not profitable.
+
+; CHECK-LABEL: @PR31243_zext
+;
+; When computing minimum sizes, if we can prove the sign bit is zero, we can
+; zero-extend the roots back to their original sizes.
+;
+; CHECK: %[[OR:.+]] = or <2 x i8> {{.*}}, <i8 1, i8 1>
+; CHECK: %[[E0:.+]] = extractelement <2 x i8> %[[OR]], i32 0
+; CHECK: %[[Z0:.+]] = zext i8 %[[E0]] to i64
+; CHECK: getelementptr inbounds i8, i8* %ptr, i64 %[[Z0]]
+; CHECK: %[[E1:.+]] = extractelement <2 x i8> %[[OR]], i32 1
+; CHECK: %[[Z1:.+]] = zext i8 %[[E1]] to i64
+; CHECK: getelementptr inbounds i8, i8* %ptr, i64 %[[Z1]]
+;
+define i8 @PR31243_zext(i8 %v0, i8 %v1, i8 %v2, i8 %v3, i8* %ptr) {
+entry:
+  %tmp0 = zext i8 %v0 to i32
+  %tmp1 = zext i8 %v1 to i32
+  %tmp2 = or i32 %tmp0, 1
+  %tmp3 = or i32 %tmp1, 1
+  %tmp4 = getelementptr inbounds i8, i8* %ptr, i32 %tmp2
+  %tmp5 = getelementptr inbounds i8, i8* %ptr, i32 %tmp3
+  %tmp6 = load i8, i8* %tmp4
+  %tmp7 = load i8, i8* %tmp5
+  %tmp8 = add i8 %tmp6, %tmp7
+  ret i8 %tmp8
+}
+
+; CHECK-LABEL: @PR31243_sext
+;
+; When computing minimum sizes, if we cannot prove the sign bit is zero, we
+; have to include one extra bit for signedness since we will sign-extend the
+; roots.
+;
+; FIXME: This test is suboptimal since the compuation can be performed in i8.
+;        In general, we need to add an extra bit to the maximum bit width only
+;        if we can't prove that the upper bit of the original type is equal to
+;        the upper bit of the proposed smaller type. If these two bits are the
+;        same (either zero or one) we know that sign-extending from the smaller
+;        type will result in the same value. Since we don't yet perform this
+;        optimization, we make the proposed smaller type (i8) larger (i16) to
+;        ensure correctness.
+;
+; CHECK: %[[S0:.+]] = sext <2 x i8> {{.*}} to <2 x i16>
+; CHECK: %[[OR:.+]] = or <2 x i16> %[[S0]], <i16 1, i16 1>
+; CHECK: %[[E0:.+]] = extractelement <2 x i16> %[[OR]], i32 0
+; CHECK: %[[S1:.+]] = sext i16 %[[E0]] to i64
+; CHECK: getelementptr inbounds i8, i8* %ptr, i64 %[[S1]]
+; CHECK: %[[E1:.+]] = extractelement <2 x i16> %[[OR]], i32 1
+; CHECK: %[[S2:.+]] = sext i16 %[[E1]] to i64
+; CHECK: getelementptr inbounds i8, i8* %ptr, i64 %[[S2]]
+;
+define i8 @PR31243_sext(i8 %v0, i8 %v1, i8 %v2, i8 %v3, i8* %ptr) {
+entry:
+  %tmp0 = sext i8 %v0 to i32
+  %tmp1 = sext i8 %v1 to i32
+  %tmp2 = or i32 %tmp0, 1
+  %tmp3 = or i32 %tmp1, 1
+  %tmp4 = getelementptr inbounds i8, i8* %ptr, i32 %tmp2
+  %tmp5 = getelementptr inbounds i8, i8* %ptr, i32 %tmp3
+  %tmp6 = load i8, i8* %tmp4
+  %tmp7 = load i8, i8* %tmp5
+  %tmp8 = add i8 %tmp6, %tmp7
+  ret i8 %tmp8
+}