From 9c6cd879346afe6bc0e29e117393158fd89f5feb Mon Sep 17 00:00:00 2001
From: Peizhao Ou <peizhaoo@uci.edu>
Date: Wed, 7 Mar 2018 00:21:40 -0800
Subject: [PATCH] Adds bogus conditional branch after all relaxed loads

---
 lib/CodeGen/CodeGenPrepare.cpp | 68 +++++++++++++++++++++++++++-------
 1 file changed, 54 insertions(+), 14 deletions(-)

diff --git a/lib/CodeGen/CodeGenPrepare.cpp b/lib/CodeGen/CodeGenPrepare.cpp
index ac8fbbf9c76..5f22c14f90f 100644
--- a/lib/CodeGen/CodeGenPrepare.cpp
+++ b/lib/CodeGen/CodeGenPrepare.cpp
@@ -785,15 +785,16 @@ void AddFakeConditionalBranch(Instruction* SplitInst, Value* Condition) {
 }
 
 // Returns true if the code is changed, and false otherwise.
-void TaintRelaxedLoads(Instruction* UsageInst) {
+void TaintRelaxedLoads(Instruction* UsageInst, Instruction* InsertPoint) {
   // For better performance, we can add a "AND X 0" instruction before the
   // condition.
   auto* BB = UsageInst->getParent();
-  auto* InsertPoint = UsageInst->getNextNode();
-  while (dyn_cast<PHINode>(InsertPoint)) {
-    InsertPoint = InsertPoint->getNextNode();
+  if (InsertPoint == nullptr) {
+    InsertPoint = UsageInst->getNextNode();
+    while (dyn_cast<PHINode>(InsertPoint)) {
+      InsertPoint = InsertPoint->getNextNode();
+    }
   }
-  IRBuilder<true, NoFolder> Builder(InsertPoint);
   // First thing is to cast 'UsageInst' to an integer type if necessary.
   Value* AndTarget = nullptr;
   if (IntegerType::classof(UsageInst->getType())) {
@@ -802,9 +803,44 @@ void TaintRelaxedLoads(Instruction* UsageInst) {
     Type* TargetIntegerType = IntegerType::get(
         UsageInst->getContext(),
         BB->getModule()->getDataLayout().getPointerSizeInBits());
+    IRBuilder<true, NoFolder> Builder(UsageInst->getNextNode());
     AndTarget = createCast(Builder, UsageInst, TargetIntegerType);
   }
 
+  // Check whether InsertPoint is a added fake conditional branch.
+  BranchInst* BI = nullptr;
+  if ((BI = dyn_cast<BranchInst>(InsertPoint)) && BI->isConditional()) {
+    auto* Cond = dyn_cast<Instruction>(BI->getOperand(0));
+    if (Cond && Cond->getOpcode() == Instruction::ICmp) {
+      auto* CmpInst = dyn_cast<ICmpInst>(Cond);
+      auto* Op0 = dyn_cast<Instruction>(Cond->getOperand(0));
+      auto* Op1 = dyn_cast<ConstantInt>(Cond->getOperand(1));
+      // %tmp = And X, 0
+      // %cmp = ICMP_NE %tmp, 0
+      // Br %cmp
+      // =>
+      // %tmp1 = And X, NewTaintedVal
+      // %tmp2 = And %tmp1, 0
+      // %cmp = ICMP_NE %tmp2, 0
+      // Br %cmp
+      if (CmpInst && CmpInst->getPredicate() == CmpInst::ICMP_NE && Op0 &&
+          Op0->getOpcode() == Instruction::And && Op1 && Op1->isZero()) {
+        auto* Op01 = dyn_cast<ConstantInt>(Op0->getOperand(1));
+        if (Op01 && Op01->isZero()) {
+          // Now we have a previously added fake cond branch.
+          auto* Op00 = Op0->getOperand(0);
+          IRBuilder<true, NoFolder> Builder(CmpInst);
+          AndTarget = Builder.CreateAnd(Op00, AndTarget);
+          auto* AndZero = dyn_cast<Instruction>(Builder.CreateAnd(
+              AndTarget, Constant::getNullValue(AndTarget->getType())));
+          CmpInst->setOperand(0, AndZero);
+          return;
+        }
+      }
+    }
+  }
+
+  IRBuilder<true, NoFolder> Builder(InsertPoint);
   auto* AndZero = dyn_cast<Instruction>(
       Builder.CreateAnd(AndTarget, Constant::getNullValue(AndTarget->getType())));
   auto* FakeCondition = dyn_cast<Instruction>(Builder.CreateICmp(
@@ -895,9 +931,11 @@ Instruction* findMostRecentDependenceUsage(LoadInst* LI, Instruction* LaterInst,
 
 // XXX-comment: Returns whether the code has been changed.
 bool AddFakeConditionalBranchAfterMonotonicLoads(
-    const SmallVector<LoadInst*, 1>& MonotonicLoadInsts, DominatorTree* DT) {
+    SmallSet<LoadInst*, 1>& MonotonicLoadInsts, DominatorTree* DT) {
   bool Changed = false;
-  for (auto* LI : MonotonicLoadInsts) {
+  while (!MonotonicLoadInsts.empty()) {
+    auto* LI = *MonotonicLoadInsts.begin();
+    MonotonicLoadInsts.erase(LI);
     SmallVector<BasicBlock*, 2> ChainedBB;
     auto* FirstInst = findFirstStoreCondBranchInst(LI, &ChainedBB);
     if (FirstInst != nullptr) {
@@ -922,7 +960,8 @@ bool AddFakeConditionalBranchAfterMonotonicLoads(
     if (FirstInst && (SI = dyn_cast<StoreInst>(FirstInst))) {
       // For immediately coming stores, taint the address of the store.
       if (SI->getParent() == LI->getParent() || DT->dominates(LI, SI)) {
-        Changed |= taintStoreAddress(SI, LI);
+        TaintRelaxedLoads(LI, SI);
+        Changed = true;
       } else {
         auto* Inst =
             findMostRecentDependenceUsage(LI, FirstInst, &ChainedBB, DT);
@@ -930,25 +969,26 @@ bool AddFakeConditionalBranchAfterMonotonicLoads(
           LI->setOrdering(Acquire);
           Changed = true;
         } else {
-          Changed |= taintStoreAddress(SI, Inst);
+          TaintRelaxedLoads(Inst, SI);
+          Changed = true;
         }
       }
     } else {
       // No upcoming branch
       if (!FirstInst) {
-        TaintRelaxedLoads(LI);
+        TaintRelaxedLoads(LI, nullptr);
         Changed = true;
       } else {
         // For immediately coming branch, directly add a fake branch.
         if (FirstInst->getParent() == LI->getParent() ||
             DT->dominates(LI, FirstInst)) {
-          TaintRelaxedLoads(LI);
+          TaintRelaxedLoads(LI, FirstInst);
           Changed = true;
         } else {
           auto* Inst =
               findMostRecentDependenceUsage(LI, FirstInst, &ChainedBB, DT);
           if (Inst) {
-            TaintRelaxedLoads(Inst);
+            TaintRelaxedLoads(Inst, FirstInst);
           } else {
             LI->setOrdering(Acquire);
           }
@@ -1332,14 +1372,14 @@ bool CodeGenPrepare::runOnFunction(Function &F) {
   // XXX-comment: Delay dealing with relaxed loads in this function to avoid
   // further changes done by other passes (e.g., SimplifyCFG).
   // Collect all the relaxed loads.
-  SmallVector<LoadInst*, 1> MonotonicLoadInsts;
+  SmallSet<LoadInst*, 1> MonotonicLoadInsts;
   for (inst_iterator I = inst_begin(F), E = inst_end(F); I != E; ++I) {
     if (I->isAtomic()) {
       switch (I->getOpcode()) {
         case Instruction::Load: {
           auto* LI = dyn_cast<LoadInst>(&*I);
           if (LI->getOrdering() == Monotonic) {
-            MonotonicLoadInsts.push_back(LI);
+            MonotonicLoadInsts.insert(LI);
           }
           break;
         }
-- 
2.34.1