[CallSiteSplitting] Refactor creating callsites.

[fhahn]

Summary

This change makes the call site creation more general if any of the
arguments is predicated on a condition in the call site's predecessors.

If we find a callsite, that potentially can be split, we collect the set
of conditions for the call site's predecessors (currently only 2
predecessors are allowed). To do that, we traverse each predecessor's
predecessors as long as it only has single predecessors and record the
condition, if it is relevant to the call site. For each condition, we
also check if the condition is taken or not. In case it is not taken,
we record the inverse predicate.

We use the recorded conditions to create the new call sites and split
the basic block.

This has 2 benefits: (1) it is slightly easier to see what is going on
(IMO) and (2) we can easily extend it to handle more complex control
flow.

[junbuml]

Took a quick look at this change. Look like you need to run clang-format, other than that I didn't see any issue with it. Let me take another look next week.
I believe it make this pass even better shape.

[fhahn]

Thanks for having a look! I ran clang-format now.

[junbuml]

LGTM. However, I believe we need more test cases including eq_eq_eq_taken in order to prevent any future change from breaking this pass. It will be also good to add a test case where the same argument is used multiple ICMPs.

test/Transforms/CallSiteSplitting/callsite-split-or-phi.ll

354

In order to prevent any future change from breaking this pass, I believe we need more test cases including eq_eq_eq_taken. It will be also good to add a test case where the same argument is used multiple ICMPs.

[fhahn]

test/Transforms/CallSiteSplitting/callsite-split-or-phi.ll

354

Agreed, I'll add such a test cases tomorrow.

[junbuml]

lib/Transforms/Scalar/CallSiteSplitting.cpp

139

We should not revisit the same predecessor. Otherwise, it might get into infinite loop if header is unreachable from the entry and we have a backedge from TBB to header :

define i32 @test(i32* %a, i32 %v, i32 %p) {
Entry:
  br label %End
Header:
  %tobool2 = icmp eq i32 %p, 10
  br i1 %tobool2, label %Tail, label %TBB
TBB:
  %cmp = icmp eq i32 %v, 1
  br i1 %cmp, label %Tail, label %Header
Tail:
  %r = call i32 @callee(i32* %a, i32 %v, i32 %p)
  ret i32 %r
End:
  ret i32 %v
}

[fhahn]

Added more test cases and fixed infinite loop/

[fhahn]

lib/Transforms/Scalar/CallSiteSplitting.cpp

139

Indeed, should be fixed now. I've added this example to the tests.

test/Transforms/CallSiteSplitting/callsite-split-or-phi.ll

354

I've added more test cases as suggested. I've also added tests where there are multiple constraints along a path (*constrain_same*). If there are conflicting conditions, e.g. x == 111 and x == 222, the first one seen will be used. I think that should be fine, as in that case the block should be unreachable. I've also mentioned that in the comment for recordConditions

[junbuml]

test/Transforms/CallSiteSplitting/callsite-split-or-phi.ll

346

unrechable -> unreachable

It might be good to add CHECK as this is also transformed with some comment about no infinite loop even in this unreachable case.

[fhahn]

Thank you very much, I've committed the patch with the fix to the test case.

lib/Transforms/Scalar/CallSiteSplitting.cpp

 
 using namespace llvm;
 using namespace PatternMatch;
 
 #define DEBUG_TYPE "callsite-splitting"
 
 STATISTIC(NumCallSiteSplit, "Number of call-site split");
 
-static void addNonNullAttribute(Instruction *CallI, Instruction *&NewCallI,
+static void addNonNullAttribute(Instruction *CallI, Instruction *NewCallI,
                                 Value *Op) {
-  if (!NewCallI)
-    NewCallI = CallI->clone();
   CallSite CS(NewCallI);
   unsigned ArgNo = 0;
   for (auto &I : CS.args()) {
     if (&*I == Op)
       CS.addParamAttr(ArgNo, Attribute::NonNull);
     ++ArgNo;
   }
 }
 
-static void setConstantInArgument(Instruction *CallI, Instruction *&NewCallI,
+static void setConstantInArgument(Instruction *CallI, Instruction *NewCallI,
                                   Value *Op, Constant *ConstValue) {
-  if (!NewCallI)
-    NewCallI = CallI->clone();
   CallSite CS(NewCallI);
   unsigned ArgNo = 0;
   for (auto &I : CS.args()) {
     if (&*I == Op)
       CS.setArgument(ArgNo, ConstValue);
     ++ArgNo;
   }
 }
if (isa<Constant>(*I) || CS.paramHasAttr(ArgNo, Attribute::NonNull))
       continue;
 
     if (*I == Op0)
       return true;
   }
   return false;
 }
 
-static SmallVector<BranchInst *, 2>
-findOrCondRelevantToCallArgument(CallSite CS) {
-  SmallVector<BranchInst *, 2> BranchInsts;
-  for (auto PredBB : predecessors(CS.getInstruction()->getParent())) {
-    auto *PBI = dyn_cast<BranchInst>(PredBB->getTerminator());
-    if (!PBI || !PBI->isConditional())
-      continue;
+/// If From has a conditional jump to To, add the condition to Conditions,
+/// if it is relevant to any argument at CS.
+void recordCondition(const CallSite &CS, BasicBlock *From, BasicBlock *To, SmallVectorImpl<std::pair<ICmpInst*, unsigned>> &Conditions) {
+  auto *BI = dyn_cast<BranchInst>(From->getTerminator());
+  if (!BI || !BI->isConditional())
+    return;
 
-    CmpInst::Predicate Pred;
-    Value *Cond = PBI->getCondition();
-    if (!match(Cond, m_ICmp(Pred, m_Value(), m_Constant())))
-      continue;
-    ICmpInst *Cmp = cast<ICmpInst>(Cond);
-    if (Pred == ICmpInst::ICMP_EQ || Pred == ICmpInst::ICMP_NE)
-      if (isCondRelevantToAnyCallArgument(Cmp, CS))
-        BranchInsts.push_back(PBI);
-  }
-  return BranchInsts;
+  Value *Cond = BI->getCondition();
+    return;
+
+      Conditions.push_back({Cmp, From->getTerminator()->getSuccessor(0) == To ? Pred : Cmp->getInversePredicate()});
 }
 
-static bool tryCreateCallSitesOnOrPredicatedArgument(
-    CallSite CS, Instruction *&NewCSTakenFromHeader,
-    Instruction *&NewCSTakenFromNextCond, BasicBlock *HeaderBB) {
-  auto BranchInsts = findOrCondRelevantToCallArgument(CS);
-  assert(BranchInsts.size() <= 2 &&
-         "Unexpected number of blocks in the OR predicated condition");
-  Instruction *Instr = CS.getInstruction();
-  BasicBlock *CallSiteBB = Instr->getParent();
-  TerminatorInst *HeaderTI = HeaderBB->getTerminator();
-  bool IsCSInTakenPath = CallSiteBB == HeaderTI->getSuccessor(0);
-
-  for (auto *PBI : BranchInsts) {
-    assert(isa<ICmpInst>(PBI->getCondition()) &&
-           "Unexpected condition in a conditional branch.");
-    ICmpInst *Cmp = cast<ICmpInst>(PBI->getCondition());
-    Value *Arg = Cmp->getOperand(0);
-    assert(isa<Constant>(Cmp->getOperand(1)) &&
-           "Expected op1 to be a constant.");
-    Constant *ConstVal = cast<Constant>(Cmp->getOperand(1));
-    CmpInst::Predicate Pred = Cmp->getPredicate();
-
-    if (PBI->getParent() == HeaderBB) {
-      Instruction *&CallTakenFromHeader =
-          IsCSInTakenPath ? NewCSTakenFromHeader : NewCSTakenFromNextCond;
-      Instruction *&CallUntakenFromHeader =
-          IsCSInTakenPath ? NewCSTakenFromNextCond : NewCSTakenFromHeader;
-
-      assert((Pred == ICmpInst::ICMP_EQ || Pred == ICmpInst::ICMP_NE) &&
-             "Unexpected predicate in an OR condition");
-
-      // Set the constant value for agruments in the call predicated based on
-      // the OR condition.
-      Instruction *&CallToSetConst = Pred == ICmpInst::ICMP_EQ
-                                         ? CallTakenFromHeader
-                                         : CallUntakenFromHeader;
-      setConstantInArgument(Instr, CallToSetConst, Arg, ConstVal);
-
-      // Add the NonNull attribute if compared with the null pointer.
-      if (ConstVal->getType()->isPointerTy() && ConstVal->isNullValue()) {
-        Instruction *&CallToSetAttr = Pred == ICmpInst::ICMP_EQ
-                                          ? CallUntakenFromHeader
-                                          : CallTakenFromHeader;
-        addNonNullAttribute(Instr, CallToSetAttr, Arg);
-      }
-      continue;
-    }
+/// Record ICmp conditions relevant to any argument in CS following Pred's
+/// single successors.
+void recordConditions(const CallSite &CS, BasicBlock *Pred, SmallVectorImpl<std::pair<ICmpInst*, unsigned>> &Conditions) {
+  recordCondition(CS, Pred, CS.getInstruction()->getParent(), Conditions);
+  BasicBlock *From = Pred;
+  BasicBlock *To = Pred;
+  while ((From = From->getSinglePredecessor())) {
+    recordCondition(CS, From, To, Conditions);
+    To = From;
 
-    if (Pred == ICmpInst::ICMP_EQ) {
-      if (PBI->getSuccessor(0) == Instr->getParent()) {
-        // Set the constant value for the call taken from the second block in
-        // the OR condition.
-        setConstantInArgument(Instr, NewCSTakenFromNextCond, Arg, ConstVal);
-      } else {
-        // Add the NonNull attribute if compared with the null pointer for the
-        // call taken from the second block in the OR condition.
-        if (ConstVal->getType()->isPointerTy() && ConstVal->isNullValue())
-          addNonNullAttribute(Instr, NewCSTakenFromNextCond, Arg);
-      }
-    } else {
-      if (PBI->getSuccessor(0) == Instr->getParent()) {
-        // Add the NonNull attribute if compared with the null pointer for the
-        // call taken from the second block in the OR condition.
-        if (ConstVal->getType()->isPointerTy() && ConstVal->isNullValue())
-          addNonNullAttribute(Instr, NewCSTakenFromNextCond, Arg);
-      } else if (Pred == ICmpInst::ICMP_NE) {
-        // Set the constant value for the call in the untaken path from the
-        // header block.
-        setConstantInArgument(Instr, NewCSTakenFromNextCond, Arg, ConstVal);
-      } else
-        llvm_unreachable("Unexpected condition");
-    }
-  }
-  return NewCSTakenFromHeader || NewCSTakenFromNextCond;
+Instruction *addConditions(CallSite &CS, SmallVectorImpl<std::pair<ICmpInst*, unsigned>> &Conditions) {
+  if (Conditions.empty())
+    return nullptr;
+
+  Instruction *NewCI = CS.getInstruction()->clone();
+  for (auto &Cond : Conditions) {
+    Value *Arg = Cond.first->getOperand(0);
+    Constant *ConstVal = cast<Constant>(Cond.first->getOperand(1));
+    if (Cond.second == ICmpInst::ICMP_EQ)
+      setConstantInArgument(CS.getInstruction(), NewCI, Arg, ConstVal);
+    else if (ConstVal->getType()->isPointerTy() && ConstVal->isNullValue()) {
+      assert(Cond.second == ICmpInst::ICMP_NE);
+      addNonNullAttribute(CS.getInstruction(), NewCI, Arg);
+    }
+  return NewCI;
+
+static SmallVector<BasicBlock *, 2> getTwoPredecessors(BasicBlock *BB) {
+  SmallVector<BasicBlock *, 2> Preds(predecessors((BB)));
+  assert(Preds.size() == 2 && "Expected exactly 2 predecessors!");
+  return Preds;
 }
 
 static bool canSplitCallSite(CallSite CS) {
   // FIXME: As of now we handle only CallInst. InvokeInst could be handled
   // without too much effort.
   Instruction *Instr = CS.getInstruction();
   if (!isa<CallInst>(Instr))
     return false;
static bool canSplitCallSite(CallSite CS) {
   if (Preds.size() != 2 || isa<IndirectBrInst>(Preds[0]->getTerminator()) ||
       isa<IndirectBrInst>(Preds[1]->getTerminator()))
     return false;
 
   return CallSiteBB->canSplitPredecessors();
 }
 
 /// Return true if the CS is split into its new predecessors which are directly
+
 /// hooked to each of its orignial predecessors pointed by PredBB1 and PredBB2.
 /// In OR predicated case, PredBB1 will point the header, and PredBB2 will point
 /// to the second compare block. CallInst1 and CallInst2 will be the new
 /// call-sites placed in the new predecessors split for PredBB1 and PredBB2,
 /// repectively. Therefore, CallInst1 will be the call-site placed
 /// between Header and Tail, and CallInst2 will be the call-site between TBB and
 /// Tail. For example, in the IR below with an OR condition, the call-site can
 /// be split
if (PHINode *PN = dyn_cast<PHINode>(&BI)) {
             return true;
         }
     }
     break;
   }
   return false;
 }
 
-static SmallVector<BasicBlock *, 2> getTwoPredecessors(BasicBlock *BB) {
-  SmallVector<BasicBlock *, 2> Preds(predecessors((BB)));
-  assert(Preds.size() == 2 && "Expected exactly 2 predecessors!");
-  return Preds;
-}
-
 static bool tryToSplitOnPHIPredicatedArgument(CallSite CS) {
   if (!isPredicatedOnPHI(CS))
     return false;
 
   auto Preds = getTwoPredecessors(CS.getInstruction()->getParent());
   splitCallSite(CS, Preds[0], Preds[1], nullptr, nullptr);
   return true;
 }
 // Check if one of the predecessors is a single predecessors of the other.
 // This is a requirement for control flow modeling an OR. HeaderBB points to
 // the single predecessor and OrBB points to other node. HeaderBB potentially
 // contains the first compare of the OR and OrBB the second.
 static bool isOrHeader(BasicBlock *HeaderBB, BasicBlock *OrBB) {
   return OrBB->getSinglePredecessor() == HeaderBB &&
          HeaderBB->getTerminator()->getNumSuccessors() == 2;
 }
 
 static bool tryToSplitOnOrPredicatedArgument(CallSite CS) {
   auto Preds = getTwoPredecessors(CS.getInstruction()->getParent());
-  BasicBlock *HeaderBB = nullptr;
-  BasicBlock *OrBB = nullptr;
-  if (isOrHeader(Preds[0], Preds[1])) {
-    HeaderBB = Preds[0];
-    OrBB = Preds[1];
-  } else if (isOrHeader(Preds[1], Preds[0])) {
-    HeaderBB = Preds[1];
-    OrBB = Preds[0];
-  } else
-    return false;
-
-  Instruction *CallInst1 = nullptr;
-  Instruction *CallInst2 = nullptr;
-  if (!tryCreateCallSitesOnOrPredicatedArgument(CS, CallInst1, CallInst2,
-                                                HeaderBB)) {
-    assert(!CallInst1 && !CallInst2 && "Unexpected new call-sites cloned.");
+  if (!isOrHeader(Preds[0], Preds[1]) && !isOrHeader(Preds[1], Preds[0]))
+    return false;
+
+  SmallVector<std::pair<ICmpInst*, unsigned>, 2> C1, C2;
+  recordConditions(CS, Preds[0], C1);
+  recordConditions(CS, Preds[1], C2);
+
+  Instruction *CallInst1 = addConditions(CS, C1);
+  Instruction *CallInst2 = addConditions(CS, C2);
+  if (!CallInst1 && !CallInst2)
     return false;
-  }
 
-  splitCallSite(CS, HeaderBB, OrBB, CallInst1, CallInst2);
+  splitCallSite(CS, Preds[1], Preds[0], CallInst2, CallInst1);
   return true;
 }
 
 static bool tryToSplitCallSite(CallSite CS) {
   if (!CS.arg_size() || !canSplitCallSite(CS))
     return false;
   return tryToSplitOnOrPredicatedArgument(CS) ||
          tryToSplitOnPHIPredicatedArgument(CS);

test/Transforms/CallSiteSplitting/callsite-split-or-phi.ll

Tail:
   %p = phi i32[1,%Header], [2, %TBB]
   %r = call i32 @callee(i32* %a, i32 %v, i32 %p)
   ret i32 %r
 
 End:
   ret i32 %v
 }
 
+
+;CHECK-LABEL: @test_eq_eq_eq_untaken
+;CHECK-LABEL: Tail.predBB1.split:
+;CHECK: %[[CALL1:.*]] = call i32 @callee(i32* nonnull %a, i32 %v, i32 10)
+;CHECK-LABEL: Tail.predBB2.split:
+;CHECK: %[[CALL2:.*]] = call i32 @callee(i32* nonnull %a, i32 1, i32 %p)
+;CHECK-LABEL: Tail
+;CHECK: %[[MERGED:.*]] = phi i32 [ %[[CALL1]], %Tail.predBB1.split ], [ %[[CALL2]], %Tail.predBB2.split ]
+;CHECK: ret i32 %[[MERGED]]
+define i32 @test_eq_eq_eq_untaken(i32* %a, i32 %v, i32 %p) {
+Header:
+  %tobool1 = icmp eq i32* %a, null
+  br i1 %tobool1, label %End, label %Header2
+
+Header2:
+  %tobool2 = icmp eq i32 %p, 10
+  br i1 %tobool2, label %Tail, label %TBB
+
+TBB:
+  %cmp = icmp eq i32 %v, 1
+  br i1 %cmp, label %Tail, label %End
+
+Tail:
+  %r = call i32 @callee(i32* %a, i32 %v, i32 %p)
+  ret i32 %r
+
+End:
+  ret i32 %v
+}
+
 define i32 @callee(i32* %a, i32 %v, i32 %p) {
 entry:
   %c = icmp ne i32* %a, null
   br i1 %c, label %BB1, label %BB2
 
 BB1:
   call void @dummy(i32* %a, i32 %p)
   br label %End