From fef619781b4d6c680b0ff8030f77edf7fc7dfa74 Mon Sep 17 00:00:00 2001
From: Peizhao Ou <peizhaoo@uci.edu>
Date: Mon, 21 May 2018 21:09:57 -0700
Subject: [PATCH] Taints load part of RMWs unconditionally

---
 include/llvm/IR/Instruction.h    | 15 +++++++++++++++
 lib/CodeGen/AtomicExpandPass.cpp |  6 ++++++
 lib/CodeGen/CodeGenPrepare.cpp   |  2 ++
 lib/IR/Instruction.cpp           |  6 ++++--
 4 files changed, 27 insertions(+), 2 deletions(-)

diff --git a/include/llvm/IR/Instruction.h b/include/llvm/IR/Instruction.h
index 03c45497fa9..c9de9e20055 100644
--- a/include/llvm/IR/Instruction.h
+++ b/include/llvm/IR/Instruction.h
@@ -41,11 +41,16 @@ class Instruction : public User,
   BasicBlock *Parent;
   DebugLoc DbgLoc;                         // 'dbg' Metadata cache.
 
+  // XXX-modified: Indicate whether this instruction should be tainted. Used for
+  // tainting the load parts of RMW.
+  bool NeedTainted;
+
   enum {
     /// HasMetadataBit - This is a bit stored in the SubClassData field which
     /// indicates whether this instruction has metadata attached to it or not.
     HasMetadataBit = 1 << 15
   };
+
 public:
   // Out of line virtual method, so the vtable, etc has a home.
   ~Instruction() override;
@@ -288,6 +293,16 @@ public:
   /// Copy I's fast-math flags
   void copyFastMathFlags(const Instruction *I);
 
+  // XXX-modified: Indicate whether this instruction should be tainted. Used for
+  // tainting the load parts of RMW.
+  bool getNeedTainted() {
+    return NeedTainted;
+  }
+
+  bool setNeedTainted() {
+    NeedTainted = true;
+  }
+
 private:
   /// hasMetadataHashEntry - Return true if we have an entry in the on-the-side
   /// metadata hash.
diff --git a/lib/CodeGen/AtomicExpandPass.cpp b/lib/CodeGen/AtomicExpandPass.cpp
index 077c52b19a7..eed71cb5ecf 100644
--- a/lib/CodeGen/AtomicExpandPass.cpp
+++ b/lib/CodeGen/AtomicExpandPass.cpp
@@ -515,6 +515,12 @@ bool AtomicExpand::expandAtomicOpToLLSC(
   // Start the main loop block now that we've taken care of the preliminaries.
   Builder.SetInsertPoint(LoopBB);
   Value *Loaded = TLI->emitLoadLinked(Builder, Addr, MemOpOrder);
+  auto* LoadedPartInst = dyn_cast<Instruction>(Loaded);
+  assert(LoadedPartInst && "Load part of RMW should be an instruction!");
+  if (MemOpOrder != Acquire && MemOpOrder != AcquireRelease &&
+      MemOpOrder != SequentiallyConsistent) {
+    LoadedPartInst->setNeedTainted();
+  }
 
   Value *NewVal = PerformOp(Builder, Loaded);
 
diff --git a/lib/CodeGen/CodeGenPrepare.cpp b/lib/CodeGen/CodeGenPrepare.cpp
index 7674e5cff92..abaa07cdf7b 100644
--- a/lib/CodeGen/CodeGenPrepare.cpp
+++ b/lib/CodeGen/CodeGenPrepare.cpp
@@ -1438,6 +1438,8 @@ bool CodeGenPrepare::runOnFunction(Function &F) {
           break;
         }
       }
+    } else if (I->getNeedTainted()) {
+      TaintRelaxedLoads(&*I, I->getNextNode());
     }
   }
   EverMadeChange |=
diff --git a/lib/IR/Instruction.cpp b/lib/IR/Instruction.cpp
index 4b33d2e66ea..226f0d62d11 100644
--- a/lib/IR/Instruction.cpp
+++ b/lib/IR/Instruction.cpp
@@ -22,7 +22,8 @@ using namespace llvm;
 
 Instruction::Instruction(Type *ty, unsigned it, Use *Ops, unsigned NumOps,
                          Instruction *InsertBefore)
-  : User(ty, Value::InstructionVal + it, Ops, NumOps), Parent(nullptr) {
+  : User(ty, Value::InstructionVal + it, Ops, NumOps), Parent(nullptr),
+    NeedTainted(false) {
 
   // If requested, insert this instruction into a basic block...
   if (InsertBefore) {
@@ -34,7 +35,8 @@ Instruction::Instruction(Type *ty, unsigned it, Use *Ops, unsigned NumOps,
 
 Instruction::Instruction(Type *ty, unsigned it, Use *Ops, unsigned NumOps,
                          BasicBlock *InsertAtEnd)
-  : User(ty, Value::InstructionVal + it, Ops, NumOps), Parent(nullptr) {
+  : User(ty, Value::InstructionVal + it, Ops, NumOps), Parent(nullptr),
+    NeedTainted(false) {
 
   // append this instruction into the basic block
   assert(InsertAtEnd && "Basic block to append to may not be NULL!");
-- 
2.34.1