[VectorUtils] Attempt to truncate icmps

[jmolloy]

Summary

Hi Hal,

This is a toned-down version of r268921, which was reverted due to PR27690.

The intent is the same, which is to allow loops containing icmps to be vectorized at smaller bitwidths. DemandedBits doesn't return any useful information for icmp operands. For example:

%1 = zext i8 %0 to i32
%2 = xor i32 %1, 255
%3 = icmp ult i32 %2, 24

This can obviously become:

%2 = xor i8 %0, 255
%3 = icmp ult i8 %2, 24

But demandedbits will not and can not tell us this, because while the upper 24 bits can be *truncated* and still return the same result, they cannot be arbitrarily changed (undef), which is the criterion DemandedBits uses.

This patch adds a fairly simple version of DemandedBits called DemandedBitsWithAssumptions (bikeshedding accepted). This is a thin wrapper around DemandedBits that allows a user to supply "facts" about the bits demanded in the dataflow graph that perhaps DemandedBits would be unable to deduce itself.

We do need to recompute DemandedBits afresh when doing this; there is possibly room in the future for more caching if this turns out to have an impact.

[jmolloy]

Ping!

[jmolloy]

Ping!

include/llvm/Analysis/DemandedBits.h

 class DominatorTree;
 class AssumptionCache;
 
 class DemandedBits {
 public:
   DemandedBits(Function &F, AssumptionCache &AC, DominatorTree &DT) :
     F(F), AC(AC), DT(DT), Analyzed(false) {}
 
+  virtual ~DemandedBits() {}
+
   /// Return the bits demanded from instruction I.
   APInt getDemandedBits(Instruction *I);
 
   /// Return true if, during analysis, I could not be reached.
   bool isInstructionDead(Instruction *I);
   
   void print(raw_ostream &OS);
 
-private:
+protected:
   Function &F;
   AssumptionCache ∾
   DominatorTree &DT;
 
   void performAnalysis();
-  void determineLiveOperandBits(const Instruction *UserI,
+  virtual void determineLiveOperandBits(const Instruction *UserI,
     const Instruction *I, unsigned OperandNo,
     const APInt &AOut, APInt &AB,
     APInt &KnownZero, APInt &KnownOne,
     APInt &KnownZero2, APInt &KnownOne2);
 
   bool Analyzed;
 
   // The set of visited instructions (non-integer-typed only).
   SmallPtrSet<Instruction*, 32> Visited;
   DenseMap<Instruction *, APInt> AliveBits;
 };
 
+/// A version of DemandedBits that can be fed facts about edges in
+/// the use-def graph that it may not be able to compute itself.
+class DemandedBitsWithAssumptions : public DemandedBits {
+public:
+  /// A fact about an edge in the use-def graph. Read as "<User> demands
+  /// <DemandedBits> from <Operand>".
+  struct EdgeAssumption {
+    Value *User, *Operand;
+    APInt DemandedBits;
+  };
+
+  /// Create a new DemandedBits implementation based upon \c DB, but with
+  /// facts from \c Assumptions.
+  DemandedBitsWithAssumptions(DemandedBits &DB,
+                              ArrayRef<EdgeAssumption> Assumptions);
+
+private:
+  SmallVector<EdgeAssumption, 4> Assumptions;
+
+  virtual void determineLiveOperandBits(const Instruction *UserI,
+    const Instruction *I, unsigned OperandNo,
+    const APInt &AOut, APInt &AB,
+    APInt &KnownZero, APInt &KnownOne,
+    APInt &KnownZero2, APInt &KnownOne2);
+};
+
 class DemandedBitsWrapperPass : public FunctionPass {
 private:
   mutable Optional<DemandedBits> DB;
 public:
   static char ID; // Pass identification, replacement for typeid
   DemandedBitsWrapperPass();
 
   bool runOnFunction(Function &F) override;

lib/Analysis/DemandedBits.cpp

 void DemandedBits::print(raw_ostream &OS) {
   performAnalysis();
   for (auto &KV : AliveBits) {
     OS << "DemandedBits: 0x" << utohexstr(KV.second.getLimitedValue()) << " for "
        << *KV.first << "\n";
   }
 }
 
+DemandedBitsWithAssumptions::DemandedBitsWithAssumptions(
+    DemandedBits &DB, ArrayRef<EdgeAssumption> Assumptions_)
+    : DemandedBits(DB) {
+  // Always recompute the DB graph.
+  Analyzed = false;
+
+  for (auto &A : Assumptions_)
+    Assumptions.push_back(A);
+}
+
+void DemandedBitsWithAssumptions::determineLiveOperandBits(
+    const Instruction *UserI, const Instruction *I, unsigned OperandNo,
+    const APInt &AOut, APInt &AB, APInt &KnownZero, APInt &KnownOne,
+    APInt &KnownZero2, APInt &KnownOne2) {
+  for (auto &A : Assumptions) {
+    if (A.User == UserI && A.Operand == I) {
+      AB = A.DemandedBits;
+      return;
+    }
+  }
+
+  return DemandedBits::determineLiveOperandBits(UserI, I, OperandNo, AOut, AB,
+                                                KnownZero, KnownOne, KnownZero2,
+                                                KnownOne2);
+}
+
 FunctionPass *llvm::createDemandedBitsWrapperPass() {
   return new DemandedBitsWrapperPass();
 }
 
 char DemandedBitsAnalysis::PassID;
 
 DemandedBits DemandedBitsAnalysis::run(Function &F,
                                              AnalysisManager<Function> &AM) {

lib/Analysis/VectorUtils.cpp

llvm::computeMinimumValueSizes(ArrayRef<BasicBlock *> Blocks, DemandedBits &DB,
   // of connected values must have the same minimum bitwidth.
   EquivalenceClasses<Value *> ECs;
   SmallVector<Value *, 16> Worklist;
   SmallPtrSet<Value *, 4> Roots;
   SmallPtrSet<Value *, 16> Visited;
   DenseMap<Value *, uint64_t> DBits;
   SmallPtrSet<Instruction *, 4> InstructionSet;
   MapVector<Instruction *, uint64_t> MinBWs;
+  SmallVector<DemandedBitsWithAssumptions::EdgeAssumption, 4> Assumptions;
+
+  assert(!Blocks.empty() && "Need at least one block!");
+  const DataLayout &DL = Blocks[0]->getParent()->getParent()->getDataLayout();
   
   // Determine the roots. We work bottom-up, from truncs or icmps.
   bool SeenExtFromIllegalType = false;
   for (auto *BB : Blocks)
     for (auto &I : *BB) {
       InstructionSet.insert(&I);
 
       if (TTI && (isa<ZExtInst>(&I) || isa<SExtInst>(&I)) &&
           !TTI->isTypeLegal(I.getOperand(0)->getType()))
         SeenExtFromIllegalType = true;
 
       // Only deal with non-vector integers up to 64-bits wide.
       if ((isa<TruncInst>(&I) || isa<ICmpInst>(&I)) &&
           !I.getType()->isVectorTy() &&
           I.getOperand(0)->getType()->getScalarSizeInBits() <= 64) {
         // Don't make work for ourselves. If we know the loaded type is legal,
         // don't add it to the worklist.
         if (TTI && isa<TruncInst>(&I) && TTI->isTypeLegal(I.getType()))
           continue;
 
         Worklist.push_back(&I);
         Roots.insert(&I);
+
+        // DemandedBits can never truncate ICmps on its own; its model
+        // is that any bit not demanded becomes undefined. Some ICmps
+        // have unneeded high bits, but those bits can only be
+        // removed, not arbitrarily modified.
+        //
+        // ICmps where both operands have more than one leading sign
+        // bit (for signed comparisons) or at least one leading zero
+        // (for unsigned comparisons) can be executed on truncated
+        // types.
+        //
+        // We identify these here and pass the information to DemandedBits
+        // using the helper class DemandedBitsWithAssumptions.
+        if (auto *ICI = dyn_cast<ICmpInst>(&I)) {
+          if (isa<IntegerType>(ICI->getOperand(0)->getType())) {
+            unsigned NSB0 = ComputeNumSignBits(ICI->getOperand(0), DL);
+            unsigned NSB1 = ComputeNumSignBits(ICI->getOperand(1), DL);
+            unsigned NSB = std::min(NSB0, NSB1);
+
+            bool CanTruncate = true;
+            if (ICI->isUnsigned()) {
+              // We know how many sign bits there are. But do we know that the
+              // sign bit is zero? If it's not, then an unsigned comparison
+              // cannot be truncated.
+              bool KnownZero0 = false, KnownZero1 = false, KnownOne = false;
+              ComputeSignBit(ICI->getOperand(0), KnownZero0, KnownOne, DL);
+              ComputeSignBit(ICI->getOperand(1), KnownZero1, KnownOne, DL);
+
+              CanTruncate = KnownZero0 && KnownZero1;
+            } else {
+              // Make sure at least one sign bit remains after the truncation!
+              --NSB;
+            }
+
+            if (CanTruncate) {
+              auto BW = ICI->getOperand(0)->getType()->getIntegerBitWidth();
+              APInt B = APInt::getLowBitsSet(BW, BW - NSB);
+
+              Assumptions.push_back({ICI, ICI->getOperand(0), B});
+              Assumptions.push_back({ICI, ICI->getOperand(1), B});
+            }
           }
         }
+
+      }
+    }
+
   // Early exit.
   if (Worklist.empty() || (TTI && !SeenExtFromIllegalType))
     return MinBWs;
 
+  DemandedBitsWithAssumptions DBA(DB, Assumptions);
+  
   // Now proceed breadth-first, unioning values together.
   while (!Worklist.empty()) {
     Value *Val = Worklist.pop_back_val();
     Value *Leader = ECs.getOrInsertLeaderValue(Val);
 
     if (Visited.count(Val))
       continue;
     Visited.insert(Val);
 
     // Non-instructions terminate a chain successfully.
     if (!isa<Instruction>(Val))
       continue;
     Instruction *I = cast<Instruction>(Val);
 
     // If we encounter a type that is larger than 64 bits, we can't represent
     // it so bail out.
-    if (DB.getDemandedBits(I).getBitWidth() > 64)
+    if (DBA.getDemandedBits(I).getBitWidth() > 64)
       return MapVector<Instruction *, uint64_t>();
 
-    uint64_t V = DB.getDemandedBits(I).getZExtValue();
+    uint64_t V = DBA.getDemandedBits(I).getZExtValue();
     DBits[Leader] |= V;
     DBits[I] = V;
 
     // Casts, loads and instructions outside of our range terminate a chain
     // successfully.
     if (isa<SExtInst>(I) || isa<ZExtInst>(I) || isa<LoadInst>(I) ||
         !InstructionSet.count(I))
       continue;

test/Transforms/LoopVectorize/AArch64/loop-vectorization-factors.ll

for.body:                                         ; preds = %entry, %for.body
   %arrayidx3 = getelementptr inbounds i8, i8* %q, i64 %indvars.iv
   store i8 %conv1, i8* %arrayidx3
   %indvars.iv.next = add nuw nsw i64 %indvars.iv, 1
   %lftr.wideiv = trunc i64 %indvars.iv.next to i32
   %exitcond = icmp eq i32 %lftr.wideiv, %len
   br i1 %exitcond, label %for.cond.cleanup, label %for.body
 }
 
+; CHECK-LABEL: @add_g
+; CHECK: load <16 x i8>
+; CHECK: xor <16 x i8>
+; CHECK: icmp ult <16 x i8>
+; CHECK: select <16 x i1> {{.*}}, <16 x i8>
+; CHECK: store <16 x i8>
+define void @add_g(i8* noalias nocapture readonly %p, i8* noalias nocapture readonly %q, i8* noalias nocapture
+%r, i8 %arg1, i32 %len) #0 {
+  %1 = icmp sgt i32 %len, 0
+  br i1 %1, label %.lr.ph, label %._crit_edge
+
+.lr.ph:                                           ; preds = %0
+  %2 = sext i8 %arg1 to i64
+  br label %3
+
+._crit_edge:                                      ; preds = %3, %0
+  ret void
+
+; <label>:3                                       ; preds = %3, %.lr.ph
+  %indvars.iv = phi i64 [ 0, %.lr.ph ], [ %indvars.iv.next, %3 ]
+  %x4 = getelementptr inbounds i8, i8* %p, i64 %indvars.iv
+  %x5 = load i8, i8* %x4
+  %x7 = getelementptr inbounds i8, i8* %q, i64 %indvars.iv
+  %x8 = load i8, i8* %x7
+  %x9 = zext i8 %x5 to i32
+  %x10 = xor i32 %x9, 255
+  %x11 = icmp ult i32 %x10, 24
+  %x12 = select i1 %x11, i32 %x10, i32 24
+  %x13 = trunc i32 %x12 to i8
+  store i8 %x13, i8* %x4
+  %indvars.iv.next = add nuw nsw i64 %indvars.iv, 1
+  %lftr.wideiv = trunc i64 %indvars.iv.next to i32
+  %exitcond = icmp eq i32 %lftr.wideiv, %len
+  br i1 %exitcond, label %._crit_edge, label %3
+}
+
 attributes #0 = { nounwind }