[CallSiteSplitting] Remove some indirection (NFC).

[fhahn]

Summary

With this patch I tried to reduce the complexity of the code sightly, by
removing some indirection. Please let me know what you think.

[junbuml]

Thanks Florian for cleaning up this pass. Overall look good, but this doesn't seem a NFC and can miss an opportunity for a constant phi in non-OR structure. Please see my inline comments.

lib/Transforms/Scalar/CallSiteSplitting.cpp

374

In case HeaderBB has more than two successors (e.g, if a switch is it's terminator), this change could allow splitting. If you want to keep this change NFC, the number of successors of HeaderBB may need to be checked.

381

We also split a call-site if it use a PHI with all constant incoming values. In this case, Preds are not necessarily an OR structure. For example, in the diamond pattern below, we should split the call site.

define i32 @test(i32* %a,  i32 %v) {
entry:
  br i1 undef, label %TBB0, label %TBB1
TBB0:
  br i1 undef, label %Tail, label %End
TBB1:
  br i1 undef, label %Tail, label %End
Tail:
  %p = phi i32[1,%TBB0], [2, %TBB1]
  %r = call i32 @callee(i32* %a, i32 %v, i32 %p)
  ret i32 %r
End:
  ret i32 %v
}

[fhahn]

Added tryToSplitOnOrPredicatedArgument and tryToSplitOnPHIPredicatedArgument to make clear which checks are required for each. This changed the order of the predecessors in 2 test cases, as now we do not determine the HeaderBB for PHI predicated arguments.

I have also added the test case Jun suggested to make sure this change does not regress things.

[fhahn]

lib/Transforms/Scalar/CallSiteSplitting.cpp

374

I think for the OR case the checks for ICMP would catch this case later (as switch for i1 only can have 2 successors I think) , but I could add an explicit check.

381

Done! I've separated the code paths splitting on PHI and on Or, as they have different legality checks.

[junbuml]

lib/Transforms/Scalar/CallSiteSplitting.cpp

374

Currently, the below test case is not handled, but with this change, we can split the call-site. Although it's not the OR structure, it doesn't seem to break the original intention of this pass. Looks like we can even extend this pass to support switch better. However, if we want to support this in this change, I think we should run tests.

define i32 @test(i32* %a, i32 %i, i32 %v) {
Header:
  switch i32 %i, label %End [
    i32 0, label %BB0
    i32 1, label %BB1
    i32 2, label %BB2
    i32 3, label %Tail
    i32 4, label %Tail
  ]

BB0:
  br label %End

BB1:
  br label %End

BB2:
  %cmp = icmp eq i32 %v, 1
  br i1 %cmp, label %Tail, label %End

Tail:
  %r = call i32 @callee(i32* %a, i32 %v)
  ret i32 %r

End:
  ret i32 %v
}

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

232 ↗

(On Diff #123012)

Maybe test_cfg_no_or --> test_cfg_no_or_phi ?

[davidxl]

In D40037#925314, @junbuml wrote:

Thanks Florian for cleaning up this pass. Overall look good, but this doesn't seem a NFC and can miss an opportunity for a constant phi in non-OR structure. Please see my inline comments.

Just a side comment -- this looks more and more like more general code cloning optimization to enable constant/range propagations -- may want to consider rename the pass and make it more general.

[fhahn]

In D40037#926426, @davidxl wrote:

Just a side comment -- this looks more and more like more general code cloning optimization to enable constant/range propagations -- may want to consider rename the pass and make it more general.

Yep, one of the reasons I tried reduce some indirection is to make it easier to extend the pass :)

[fhahn]

lib/Transforms/Scalar/CallSiteSplitting.cpp

374

I think it will be best to keep this change as a NFC and go from there, what do you think?

I would plan to update the patch to only do the transform if the header has 2 successors in the OR case and post a follow-up patch to enable the additional transform after that.

[junbuml]

lib/Transforms/Scalar/CallSiteSplitting.cpp

374

Sounds good to me.

[fhahn]

Add check to ensure we only split call sites in tryToSplitOnOrPredicatedArgument if the header block has exactly 2 successors. I plan to post a patch relaxing that soonish.

[junbuml]

LGTM. Please run clang-format before committing.

lib/Transforms/Scalar/CallSiteSplitting.cpp

117

Looks like more than 80?

lib/Transforms/Scalar/CallSiteSplitting.cpp

static void setConstantInArgument(Instruction *CallI, Instruction *&NewCallI,
   unsigned ArgNo = 0;
   for (auto &I : CS.args()) {
     if (&*I == Op)
       CS.setArgument(ArgNo, ConstValue);
     ++ArgNo;
   }
 }
 
-static bool createCallSitesOnOrPredicatedArgument(
+static bool isCondRelevantToAnyCallArgument(ICmpInst *Cmp, CallSite CS) {
+  assert(isa<Constant>(Cmp->getOperand(1)) && "Expected a constant operand.");
+  Value *Op0 = Cmp->getOperand(0);
+  unsigned ArgNo = 0;
+  for (CallSite::arg_iterator I = CS.arg_begin(), E = CS.arg_end(); I != E;
+       ++I, ++ArgNo) {
+    // Don't consider constant or arguments that are already known non-null.
+    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;
+
+    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;
+}
+
+static bool tryCreateCallSitesOnOrPredicatedArgument(
     CallSite CS, Instruction *&NewCSTakenFromHeader,
-    Instruction *&NewCSTakenFromNextCond,
-    SmallVectorImpl<BranchInst *> &BranchInsts, BasicBlock *HeaderBB) {
+    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 (unsigned I = 0, E = BranchInsts.size(); I != E; ++I) {
-    BranchInst *PBI = BranchInsts[I];
+  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();
static bool canSplitCallSite(CallSite CS) {
   // Allow splitting a call-site only when there is no instruction before the
   // call-site in the basic block. Based on this constraint, we only clone the
   // call instruction, and we do not move a call-site across any other
   // instruction.
   BasicBlock *CallSiteBB = Instr->getParent();
   if (Instr != CallSiteBB->getFirstNonPHI())
     return false;
 
-  pred_iterator PII = pred_begin(CallSiteBB);
-  pred_iterator PIE = pred_end(CallSiteBB);
-  unsigned NumPreds = std::distance(PII, PIE);
-
-  // Allow only one extra call-site. No more than two from one call-site.
-  if (NumPreds != 2)
-    return false;
-
   // Cannot split an edge from an IndirectBrInst.
-  BasicBlock *Preds[2] = {*PII++, *PII};
+  SmallVector<BasicBlock*, 2> Preds(predecessors(CallSiteBB));
   if (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.
-/// Note that PredBB1 and PredBB2 are decided in findPredicatedArgument(),
-/// especially for the OR predicated case where PredBB1 will point the header,
-/// and PredBB2 will point the 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
+/// 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
 ///
 /// from :
 ///
 ///   Header:
 ///     %c = icmp eq i32* %a, null
static void splitCallSite(CallSite CS, BasicBlock *PredBB1, BasicBlock *PredBB2,
   DEBUG(dbgs() << "    " << *CallInst1 << " in " << SplitBlock1->getName()
                << "\n");
   DEBUG(dbgs() << "    " << *CallInst2 << " in " << SplitBlock2->getName()
                << "\n");
   Instr->eraseFromParent();
   NumCallSiteSplit++;
 }
 
-static bool isCondRelevantToAnyCallArgument(ICmpInst *Cmp, CallSite CS) {
-  assert(isa<Constant>(Cmp->getOperand(1)) && "Expected a constant operand.");
-  Value *Op0 = Cmp->getOperand(0);
-  unsigned ArgNo = 0;
-  for (CallSite::arg_iterator I = CS.arg_begin(), E = CS.arg_end(); I != E;
-       ++I, ++ArgNo) {
-    // Don't consider constant or arguments that are already known non-null.
-    if (isa<Constant>(*I) || CS.paramHasAttr(ArgNo, Attribute::NonNull))
-      continue;
-
-    if (*I == Op0)
-      return true;
-  }
-  return false;
-}
-
-static void findOrCondRelevantToCallArgument(
-    CallSite CS, BasicBlock *PredBB, BasicBlock *OtherPredBB,
-    SmallVectorImpl<BranchInst *> &BranchInsts, BasicBlock *&HeaderBB) {
-  auto *PBI = dyn_cast<BranchInst>(PredBB->getTerminator());
-  if (!PBI || !PBI->isConditional())
-    return;
-
-  if (PBI->getSuccessor(0) == OtherPredBB ||
-      PBI->getSuccessor(1) == OtherPredBB)
-    if (PredBB == OtherPredBB->getSinglePredecessor()) {
-      assert(!HeaderBB && "Expect to find only a single header block");
-      HeaderBB = PredBB;
-    }
-
-  CmpInst::Predicate Pred;
-  Value *Cond = PBI->getCondition();
-  if (!match(Cond, m_ICmp(Pred, m_Value(), m_Constant())))
-    return;
-  ICmpInst *Cmp = cast<ICmpInst>(Cond);
-  if (Pred == ICmpInst::ICMP_EQ || Pred == ICmpInst::ICMP_NE)
-    if (isCondRelevantToAnyCallArgument(Cmp, CS))
-      BranchInsts.push_back(PBI);
-}
-
 // Return true if the call-site has an argument which is a PHI with only
 // constant incoming values.
 static bool isPredicatedOnPHI(CallSite CS) {
   Instruction *Instr = CS.getInstruction();
   BasicBlock *Parent = Instr->getParent();
   if (Instr != Parent->getFirstNonPHI())
     return false;
 
if (PHINode *PN = dyn_cast<PHINode>(&BI)) {
             return true;
         }
     }
     break;
   }
   return false;
 }
 
-// Return true if an agument in CS is predicated on an 'or' condition.
-// Create new call-site with arguments constrained based on the OR condition.
-static bool findPredicatedOnOrCondition(CallSite CS, BasicBlock *PredBB1,
-                                        BasicBlock *PredBB2,
-                                        Instruction *&NewCallTakenFromHeader,
-                                        Instruction *&NewCallTakenFromNextCond,
-                                        BasicBlock *&HeaderBB) {
-  SmallVector<BranchInst *, 4> BranchInsts;
-  findOrCondRelevantToCallArgument(CS, PredBB1, PredBB2, BranchInsts, HeaderBB);
-  findOrCondRelevantToCallArgument(CS, PredBB2, PredBB1, BranchInsts, HeaderBB);
-  if (BranchInsts.empty() || !HeaderBB)
-    return false;
-
-  // If an OR condition is detected, try to create call sites with constrained
-  // arguments (e.g., NonNull attribute or constant value).
-  return createCallSitesOnOrPredicatedArgument(CS, NewCallTakenFromHeader,
-                                               NewCallTakenFromNextCond,
-                                               BranchInsts, HeaderBB);
-}
-
-static bool findPredicatedArgument(CallSite CS, Instruction *&CallInst1,
-                                   Instruction *&CallInst2,
-                                   BasicBlock *&PredBB1, BasicBlock *&PredBB2) {
-  BasicBlock *CallSiteBB = CS.getInstruction()->getParent();
-  pred_iterator PII = pred_begin(CallSiteBB);
-  pred_iterator PIE = pred_end(CallSiteBB);
-  assert(std::distance(PII, PIE) == 2 && "Expect only two predecessors.");
-  (void)PIE;
-  BasicBlock *Preds[2] = {*PII++, *PII};
-  BasicBlock *&HeaderBB = PredBB1;
-  if (!findPredicatedOnOrCondition(CS, Preds[0], Preds[1], CallInst1, CallInst2,
-                                   HeaderBB) &&
-      !isPredicatedOnPHI(CS))
+static bool tryToSplitCallSite(CallSite CS) {
+  if (!CS.arg_size())
     return false;
 
-  if (!PredBB1)
-    PredBB1 = Preds[0];
-
-  PredBB2 = PredBB1 == Preds[0] ? Preds[1] : Preds[0];
-  return true;
-}
+  BasicBlock *HeaderBB = nullptr;
+  BasicBlock *OrBB = nullptr;
+  BasicBlock *CallSiteBB = CS.getInstruction()->getParent();
+  SmallVector<BasicBlock*, 2> Preds(predecessors(CallSiteBB));
+  if (Preds.size() != 2)
+    return false;
 
-static bool tryToSplitCallSite(CallSite CS) {
-  if (!CS.arg_size())
+  // 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.
+  if (Preds[1]->getSinglePredecessor() == Preds[0]) {
+    HeaderBB = Preds[0];
+    OrBB = Preds[1];
+  } else if (Preds[0]->getSinglePredecessor() == Preds[1]) {
+    HeaderBB = Preds[1];
+    OrBB = Preds[0];
+  } else
     return false;
 
-  BasicBlock *PredBB1 = nullptr;
-  BasicBlock *PredBB2 = nullptr;
   Instruction *CallInst1 = nullptr;
   Instruction *CallInst2 = nullptr;
-  if (!canSplitCallSite(CS) ||
-      !findPredicatedArgument(CS, CallInst1, CallInst2, PredBB1, PredBB2)) {
+  if (!canSplitCallSite(CS))
+    return false;
+
+  if (!tryCreateCallSitesOnOrPredicatedArgument(CS, CallInst1, CallInst2,
+                                                 HeaderBB) &&
+      !isPredicatedOnPHI(CS)) {
     assert(!CallInst1 && !CallInst2 && "Unexpected new call-sites cloned.");
     return false;
   }
-  splitCallSite(CS, PredBB1, PredBB2, CallInst1, CallInst2);
+  splitCallSite(CS, HeaderBB, OrBB, CallInst1, CallInst2);
   return true;
 }
 
 static bool doCallSiteSplitting(Function &F, TargetLibraryInfo &TLI) {
   bool Changed = false;
   for (Function::iterator BI = F.begin(), BE = F.end(); BI != BE;) {
     BasicBlock &BB = *BI++;
     for (BasicBlock::iterator II = BB.begin(), IE = BB.end(); II != IE;) {