1 files changed, 23 insertions, 12 deletions

diff --git a/llvm/lib/Transforms/IPO/IPConstantPropagation.cpp b/llvm/lib/Transforms/IPO/IPConstantPropagation.cpp
index 7d55ebecbf9..fecbc809f8b 100644
--- a/llvm/lib/Transforms/IPO/IPConstantPropagation.cpp
+++ b/llvm/lib/Transforms/IPO/IPConstantPropagation.cpp
@@ -62,32 +62,43 @@ static bool PropagateConstantsIntoArguments(Function &F) {
     // Ignore blockaddress uses.
     if (isa<BlockAddress>(UR)) continue;
 
-    // Used by a non-instruction, or not the callee of a function, do not
-    // transform.
-    if (!isa<CallInst>(UR) && !isa<InvokeInst>(UR))
-      return false;
-
-    CallSite CS(cast<Instruction>(UR));
-    if (!CS.isCallee(&U))
+    // If no abstract call site was created we did not understand the use, bail.
+    AbstractCallSite ACS(&U);
+    if (!ACS)
       return false;
 
     // Check out all of the potentially constant arguments.  Note that we don't
     // inspect varargs here.
-    CallSite::arg_iterator AI = CS.arg_begin();
     Function::arg_iterator Arg = F.arg_begin();
-    for (unsigned i = 0, e = ArgumentConstants.size(); i != e;
-         ++i, ++AI, ++Arg) {
+    for (unsigned i = 0, e = ArgumentConstants.size(); i != e; ++i, ++Arg) {
 
       // If this argument is known non-constant, ignore it.
       if (ArgumentConstants[i].second)
         continue;
 
-      Constant *C = dyn_cast<Constant>(*AI);
+      Value *V = ACS.getCallArgOperand(i);
+      Constant *C = dyn_cast_or_null<Constant>(V);
+
+      // We can only propagate thread independent values through callbacks.
+      // This is different to direct/indirect call sites because for them we
+      // know the thread executing the caller and callee is the same. For
+      // callbacks this is not guaranteed, thus a thread dependent value could
+      // be different for the caller and callee, making it invalid to propagate.
+      if (C && ACS.isCallbackCall() && C->isThreadDependent()) {
+        // Argument became non-constant. If all arguments are non-constant now,
+        // give up on this function.
+        if (++NumNonconstant == ArgumentConstants.size())
+          return false;
+
+        ArgumentConstants[i].second = true;
+        continue;
+      }
+
       if (C && ArgumentConstants[i].first == nullptr) {
         ArgumentConstants[i].first = C;   // First constant seen.
       } else if (C && ArgumentConstants[i].first == C) {
         // Still the constant value we think it is.
-      } else if (*AI == &*Arg) {
+      } else if (V == &*Arg) {
         // Ignore recursive calls passing argument down.
       } else {
         // Argument became non-constant.  If all arguments are non-constant now,