[llvm] Improve performance of programUndefinedIfUndefOrPoison

[gezalore]

Summary

programUndefinedIfUndefOrPoison used to eagerly propagate the fact that
a value is poison to the users of the value. The problem is that if the
value has a lot of uses (orders of magnitude more than the scanning
limit we use in this function), then we spend the bulk of our time in
eagerly propagating the poison property, which we will mostly never use
later anyway due to the scanning limit.

I have a test case (of ~50k lines of machine generated C++), where this
results in ~60% of 35s compilation time being spent doing just this
eager propagation.

This patch changes programUndefinedIfUndefOrPoison to only propagate to
instructions actually visited, looking back to see if their operands are
poison. This should be equivalent and no functional change is intended,
but we regain virtually all of the 60% compilation time spent in this
function in my test case (i.e.: a 2.5x total compilation speedup).

[gezalore]

Hello, this is my first patch to LLVM so please educate me as you feel is appropriate.

@nikic @aqjune git tells me you might be appropriate reviewers, could you take a look please?

[nikic]

Not a fan of the mix of lazy and eager propagation here. I think the way you implemented the lazy isPoison check we may also run into a problem where the poison check will try to recurse past the initial instruction, and might explore a large tree that way, introducing a different compile-time problem.

Can we simply propagate poison as we visit instructions? I.e. when we visit an instruction, check whether it propagates poison and has a poison operand, if so mark it as poison itself?

[gezalore]

Simplify according to feedback

[gezalore]

You are right about the recursion problem, and I think your suggestion works (and the test suite and my benchmark agrees), so I updated the patch.

Thanks for the speedy review and suggestions!

[nikic]

LGTM, thanks!

llvm/lib/Analysis/ValueTracking.cpp

5800

You can use for (const Value *Op : I.operands()) here and use Op in place of U.get(), as you don't care about the specific position of the use in this case.

[nikic]

Can you please also share Name <email> to use for the commit (assuming you don't have commit access)?

[gezalore]

Thank you. I don't have commit access, I am:

Geza Lore <[e-mail]>

[gezalore]

Fix code style as recommended. Hopefully patch is good to go now.

llvm/lib/Analysis/ValueTracking.cpp

bool llvm::mustTriggerUB(const Instruction *I,
 
   for (const auto *V : NonPoisonOps)
     if (KnownPoison.count(V))
       return true;
 
   return false;
 }
 
+static bool isPoison(const Value *V, SmallSet<const Value *, 16> &KnownPoison,
+                     SmallSet<const Value *, 16> &KnownNonPoison) {
+  // Check if we already know whether V is poison.
+  if (KnownPoison.count(V))
+    return true;
+  if (KnownNonPoison.count(V))
+    return false;
+
+  // Not yet known whether V is poison, compute it lazily.
+  if (const Instruction *I = dyn_cast<const Instruction>(V)) {
+    if (propagatesPoison(cast<Operator>(I))) {
+      for (const Use &U : I->operands()) {
+        if (isPoison(U.get(), KnownPoison, KnownNonPoison)) {
+          KnownPoison.insert(V);
+          return true;
+        }
+      }
+    }
+  }
+
+  KnownNonPoison.insert(V);
+  return false;
+}
+
 static bool programUndefinedIfUndefOrPoison(const Value *V,
                                             bool PoisonOnly) {
   // We currently only look for uses of values within the same basic
   // block, as that makes it easier to guarantee that the uses will be
   // executed given that Inst is executed.
   //
   // FIXME: Expand this to consider uses beyond the same basic block. To do
   // this, look out for the distinction between post-dominance and strong
for (const auto &I : make_range(Begin, End)) {
         return true;
 
       if (!isGuaranteedToTransferExecutionToSuccessor(&I))
         break;
     }
     return false;
   }
 
-  // Set of instructions that we have proved will yield poison if Inst
-  // does.
-  SmallSet<const Value *, 16> YieldsPoison;
-  SmallSet<const BasicBlock *, 4> Visited;
-
-  YieldsPoison.insert(V);
-  auto Propagate = [&](const User *User) {
-    if (propagatesPoison(cast<Operator>(User)))
-      YieldsPoison.insert(User);
+  // Instructions we have proven they will or will not yield poison if V does
+  SmallSet<const Value *, 16> KnownPoison;
+  SmallSet<const Value *, 16> KnownNonPoison;
+
+  // Eagerly mark the first few uses of the poison Value as poison. This
+  // helps with the common case of having a handful of uses, but because we
+  // are only scanning instructions up to ScanLimit below, in cases where
+  // there are a lot of uses of this value (orders of magnitude more than
+  // ScanLimit), it can save a lot of time to compute the rest lazily in
+  // 'isPoison' instead.
+  auto PoisonAFewUsersEagerly = [&](const Value *Value) {
+    assert(KnownPoison.count(Value) && "Eagerly poisoning uses of non-poison");
+    unsigned EagerLimit = 3;
+    for (const User *U : Value->users()) {
+      if (propagatesPoison(cast<Operator>(U)))
+        KnownPoison.insert(U);
+      else
+        KnownNonPoison.insert(U);
+      if (--EagerLimit == 0)
+        return;
+    }
   };
-  for_each(V->users(), Propagate);
+
+  KnownPoison.insert(V);
+  PoisonAFewUsersEagerly(V);
+
+  SmallSet<const BasicBlock *, 4> Visited;
   Visited.insert(BB);
 
   while (true) {
     for (const auto &I : make_range(Begin, End)) {
       if (isa<DbgInfoIntrinsic>(I))
         continue;
       if (--ScanLimit == 0)
         return false;
-      if (mustTriggerUB(&I, YieldsPoison))
+
+      // Like llvm::mustTriggerUB, but with lazy evaluation via 'isPoison'
+      SmallPtrSet<const Value *, 4> NonPoisonOps;
+      getGuaranteedNonPoisonOps(&I, NonPoisonOps);
+      for (const Value *O : NonPoisonOps)
+        if (isPoison(O, KnownPoison, KnownNonPoison))
           return true;
+
       if (!isGuaranteedToTransferExecutionToSuccessor(&I))
         return false;
 
       // Mark poison that propagates from I through uses of I.
-      if (YieldsPoison.count(&I))
-        for_each(I.users(), Propagate);
+      if (isPoison(&I, KnownPoison, KnownNonPoison))
+        PoisonAFewUsersEagerly(&I);
     }
 
     BB = BB->getSingleSuccessor();
     if (!BB || !Visited.insert(BB).second)
       break;
 
     Begin = BB->getFirstNonPHI()->getIterator();
     End = BB->end();
   }
   return false;