merge
[c11llvm.git] / CDSPass.cpp
index 9158baef37e6d41cc8e0bb82003c08995f95a095..0cebdf01ac647d65fcaf9d437ef4e6e01ce4358b 100644 (file)
@@ -1,7 +1,8 @@
-//===-- CdsPass.cpp - xxx -------------------------------===//
-//
-//                     The LLVM Compiler Infrastructure
+//===-- CDSPass.cpp - xxx -------------------------------===//
 //
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
@@ -25,7 +26,6 @@
 #include "llvm/Analysis/ValueTracking.h"
 #include "llvm/Analysis/CaptureTracking.h"
 #include "llvm/IR/BasicBlock.h"
-#include "llvm/IR/CFG.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/Transforms/Utils/Local.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
 #include "llvm/Transforms/IPO/PassManagerBuilder.h"
-#include <list>
 #include <vector>
-// #include "llvm/Support/MathExtras.h"
 
-#define DEBUG_TYPE "CDS"
 using namespace llvm;
 
-#define FUNCARRAYSIZE 4
+#define DEBUG_TYPE "CDS"
+#include <llvm/IR/DebugLoc.h>
+
+Value *getPosition( Instruction * I, IRBuilder <> IRB)
+{
+       const DebugLoc & debug_location = I->getDebugLoc ();
+       std::string position_string;
+       {
+               llvm::raw_string_ostream position_stream (position_string);
+               debug_location . print (position_stream);
+       }
+
+       return IRB . CreateGlobalStringPtr (position_string);
+}
 
 STATISTIC(NumInstrumentedReads, "Number of instrumented reads");
 STATISTIC(NumInstrumentedWrites, "Number of instrumented writes");
+STATISTIC(NumAccessesWithBadSize, "Number of accesses with bad size");
 // STATISTIC(NumInstrumentedVtableWrites, "Number of vtable ptr writes");
 // STATISTIC(NumInstrumentedVtableReads, "Number of vtable ptr reads");
 
@@ -63,10 +74,6 @@ STATISTIC(NumOmittedReadsFromConstantGlobals,
 STATISTIC(NumOmittedReadsFromVtable, "Number of vtable reads");
 STATISTIC(NumOmittedNonCaptured, "Number of accesses ignored due to capturing");
 
-Type * Int8Ty;
-Type * Int16Ty;
-Type * Int32Ty;
-Type * Int64Ty;
 Type * OrdTy;
 
 Type * Int8PtrTy;
@@ -76,13 +83,16 @@ Type * Int64PtrTy;
 
 Type * VoidTy;
 
-Constant * CdsLoad[FUNCARRAYSIZE];
-Constant * CdsStore[FUNCARRAYSIZE];
-Constant * CdsAtomicLoad[FUNCARRAYSIZE];
-Constant * CdsAtomicStore[FUNCARRAYSIZE];
-Constant * CdsAtomicRMW[AtomicRMWInst::LAST_BINOP + 1][FUNCARRAYSIZE];
-Constant * CdsAtomicCAS[FUNCARRAYSIZE];
-Constant * CdsAtomicThreadFence;
+static const size_t kNumberOfAccessSizes = 4;
+Constant * CDSLoad[kNumberOfAccessSizes];
+Constant * CDSStore[kNumberOfAccessSizes];
+Constant * CDSAtomicInit[kNumberOfAccessSizes];
+Constant * CDSAtomicLoad[kNumberOfAccessSizes];
+Constant * CDSAtomicStore[kNumberOfAccessSizes];
+Constant * CDSAtomicRMW[AtomicRMWInst::LAST_BINOP + 1][kNumberOfAccessSizes];
+Constant * CDSAtomicCAS_V1[kNumberOfAccessSizes];
+Constant * CDSAtomicCAS_V2[kNumberOfAccessSizes];
+Constant * CDSAtomicThreadFence;
 
 int getAtomicOrderIndex(AtomicOrdering order){
   switch (order) {
@@ -104,126 +114,355 @@ int getAtomicOrderIndex(AtomicOrdering order){
   }
 }
 
-int getTypeSize(Type* type) {
-  if (type==Int32PtrTy) {
-    return sizeof(int)*8;
-  } else if (type==Int8PtrTy) {
-    return sizeof(char)*8;
-  } else if (type==Int16PtrTy) {
-    return sizeof(short)*8;
-  } else if (type==Int64PtrTy) {
-    return sizeof(long long int)*8;
-  } else {
-    return sizeof(void*)*8;
-  }
-
-  return -1;
-}
-
-static int sizetoindex(int size) {
-  switch(size) {
-    case 8:     return 0;
-    case 16:    return 1;
-    case 32:    return 2;
-    case 64:    return 3;
-  }
-  return -1;
-}
-
 namespace {
-  struct CdsPass : public FunctionPass {
+  struct CDSPass : public FunctionPass {
     static char ID;
-    CdsPass() : FunctionPass(ID) {}
+    CDSPass() : FunctionPass(ID) {}
     bool runOnFunction(Function &F) override; 
 
   private:
     void initializeCallbacks(Module &M);
     bool instrumentLoadOrStore(Instruction *I, const DataLayout &DL);
-    bool instrumentAtomic(Instruction *I);
+    bool instrumentAtomic(Instruction *I, const DataLayout &DL);
+    bool instrumentAtomicCall(CallInst *CI, const DataLayout &DL);
     void chooseInstructionsToInstrument(SmallVectorImpl<Instruction *> &Local,
                                       SmallVectorImpl<Instruction *> &All,
                                       const DataLayout &DL);
     bool addrPointsToConstantData(Value *Addr);
+    int getMemoryAccessFuncIndex(Value *Addr, const DataLayout &DL);
   };
 }
 
-void CdsPass::initializeCallbacks(Module &M) {
-  LLVMContext &Ctx = M.getContext();
+static bool isVtableAccess(Instruction *I) {
+  if (MDNode *Tag = I->getMetadata(LLVMContext::MD_tbaa))
+    return Tag->isTBAAVtableAccess();
+  return false;
+}
+
+void CDSPass::initializeCallbacks(Module &M) {
+       LLVMContext &Ctx = M.getContext();
 
-  Int8Ty  = Type::getInt8Ty(Ctx);
-  Int16Ty = Type::getInt16Ty(Ctx);
-  Int32Ty = Type::getInt32Ty(Ctx);
-  Int64Ty = Type::getInt64Ty(Ctx);
-  OrdTy = Type::getInt32Ty(Ctx);
+       Type * Int1Ty = Type::getInt1Ty(Ctx);
+       OrdTy = Type::getInt32Ty(Ctx);
 
-  Int8PtrTy  = Type::getInt8PtrTy(Ctx);
-  Int16PtrTy = Type::getInt16PtrTy(Ctx);
-  Int32PtrTy = Type::getInt32PtrTy(Ctx);
-  Int64PtrTy = Type::getInt64PtrTy(Ctx);
+       Int8PtrTy  = Type::getInt8PtrTy(Ctx);
+       Int16PtrTy = Type::getInt16PtrTy(Ctx);
+       Int32PtrTy = Type::getInt32PtrTy(Ctx);
+       Int64PtrTy = Type::getInt64PtrTy(Ctx);
 
-  VoidTy = Type::getVoidTy(Ctx);
+       VoidTy = Type::getVoidTy(Ctx);
   
+       // Get the function to call from our untime library.
+       for (unsigned i = 0; i < kNumberOfAccessSizes; i++) {
+               const unsigned ByteSize = 1U << i;
+               const unsigned BitSize = ByteSize * 8;
+
+               std::string ByteSizeStr = utostr(ByteSize);
+               std::string BitSizeStr = utostr(BitSize);
+
+               Type *Ty = Type::getIntNTy(Ctx, BitSize);
+               Type *PtrTy = Ty->getPointerTo();
+
+               // uint8_t cds_atomic_load8 (void * obj, int atomic_index)
+               // void cds_atomic_store8 (void * obj, int atomic_index, uint8_t val)
+               SmallString<32> LoadName("cds_load" + BitSizeStr);
+               SmallString<32> StoreName("cds_store" + BitSizeStr);
+               SmallString<32> AtomicInitName("cds_atomic_init" + BitSizeStr);
+               SmallString<32> AtomicLoadName("cds_atomic_load" + BitSizeStr);
+               SmallString<32> AtomicStoreName("cds_atomic_store" + BitSizeStr);
+
+               CDSLoad[i]  = M.getOrInsertFunction(LoadName, VoidTy, PtrTy);
+               CDSStore[i] = M.getOrInsertFunction(StoreName, VoidTy, PtrTy);
+               CDSAtomicInit[i] = M.getOrInsertFunction(AtomicInitName, 
+                                                               VoidTy, PtrTy, Ty, Int8PtrTy);
+               CDSAtomicLoad[i]  = M.getOrInsertFunction(AtomicLoadName, 
+                                                               Ty, PtrTy, OrdTy, Int8PtrTy);
+               CDSAtomicStore[i] = M.getOrInsertFunction(AtomicStoreName, 
+                                                               VoidTy, PtrTy, Ty, OrdTy, Int8PtrTy);
+
+               for (int op = AtomicRMWInst::FIRST_BINOP; 
+                       op <= AtomicRMWInst::LAST_BINOP; ++op) {
+                       CDSAtomicRMW[op][i] = nullptr;
+                       std::string NamePart;
+
+                       if (op == AtomicRMWInst::Xchg)
+                               NamePart = "_exchange";
+                       else if (op == AtomicRMWInst::Add) 
+                               NamePart = "_fetch_add";
+                       else if (op == AtomicRMWInst::Sub)
+                               NamePart = "_fetch_sub";
+                       else if (op == AtomicRMWInst::And)
+                               NamePart = "_fetch_and";
+                       else if (op == AtomicRMWInst::Or)
+                               NamePart = "_fetch_or";
+                       else if (op == AtomicRMWInst::Xor)
+                               NamePart = "_fetch_xor";
+                       else
+                               continue;
+
+                       SmallString<32> AtomicRMWName("cds_atomic" + NamePart + BitSizeStr);
+                       CDSAtomicRMW[op][i] = M.getOrInsertFunction(AtomicRMWName, 
+                                                                               Ty, PtrTy, Ty, OrdTy, Int8PtrTy);
+               }
+
+               // only supportes strong version
+               SmallString<32> AtomicCASName_V1("cds_atomic_compare_exchange" + BitSizeStr + "_v1");
+               SmallString<32> AtomicCASName_V2("cds_atomic_compare_exchange" + BitSizeStr + "_v2");
+               CDSAtomicCAS_V1[i] = M.getOrInsertFunction(AtomicCASName_V1, 
+                                                               Ty, PtrTy, Ty, Ty, OrdTy, OrdTy, Int8PtrTy);
+               CDSAtomicCAS_V2[i] = M.getOrInsertFunction(AtomicCASName_V2, 
+                                                               Int1Ty, PtrTy, PtrTy, Ty, OrdTy, OrdTy, Int8PtrTy);
+       }
+
+       CDSAtomicThreadFence = M.getOrInsertFunction("cds_atomic_thread_fence", 
+                                                                                                       VoidTy, OrdTy, Int8PtrTy);
+}
 
-  // Get the function to call from our untime library.
-  for (unsigned i = 0; i < FUNCARRAYSIZE; i++) {
-    const unsigned ByteSize = 1U << i;
-    const unsigned BitSize = ByteSize * 8;
-//    errs() << BitSize << "\n";
-    std::string ByteSizeStr = utostr(ByteSize);
-    std::string BitSizeStr = utostr(BitSize);
+void printArgs(CallInst *);
 
-    Type *Ty = Type::getIntNTy(Ctx, BitSize);
-    Type *PtrTy = Ty->getPointerTo();
+bool isAtomicCall(Instruction *I) {
+       if ( auto *CI = dyn_cast<CallInst>(I) ) {
+               Function *fun = CI->getCalledFunction();
+               if (fun == NULL)
+                       return false;
 
-    // uint8_t cds_atomic_load8 (void * obj, int atomic_index)
-    // void cds_atomic_store8 (void * obj, int atomic_index, uint8_t val)
-    SmallString<32> LoadName("cds_load" + BitSizeStr);
-    SmallString<32> StoreName("cds_store" + BitSizeStr);
-    SmallString<32> AtomicLoadName("cds_atomic_load" + BitSizeStr);
-    SmallString<32> AtomicStoreName("cds_atomic_store" + BitSizeStr);
-
-//    CdsLoad[i]  = M.getOrInsertFunction(LoadName, Ty, PtrTy);
-//    CdsStore[i] = M.getOrInsertFunction(StoreName, VoidTy, PtrTy, Ty);
-    CdsLoad[i]  = M.getOrInsertFunction(LoadName, VoidTy, PtrTy);
-    CdsStore[i] = M.getOrInsertFunction(StoreName, VoidTy, PtrTy);
-    CdsAtomicLoad[i]  = M.getOrInsertFunction(AtomicLoadName, Ty, PtrTy, OrdTy);
-    CdsAtomicStore[i] = M.getOrInsertFunction(AtomicStoreName, VoidTy, PtrTy, OrdTy, Ty);
-
-    for (int op = AtomicRMWInst::FIRST_BINOP; op <= AtomicRMWInst::LAST_BINOP; ++op) {
-      CdsAtomicRMW[op][i] = nullptr;
-      std::string NamePart;
-
-      if (op == AtomicRMWInst::Xchg)
-        NamePart = "_exchange";
-      else if (op == AtomicRMWInst::Add) 
-        NamePart = "_fetch_add";
-      else if (op == AtomicRMWInst::Sub)
-        NamePart = "_fetch_sub";
-      else if (op == AtomicRMWInst::And)
-        NamePart = "_fetch_and";
-      else if (op == AtomicRMWInst::Or)
-        NamePart = "_fetch_or";
-      else if (op == AtomicRMWInst::Xor)
-        NamePart = "_fetch_xor";
-      else
-        continue;
+               StringRef funName = fun->getName();
+                // todo: come up with better rules for function name checking
+                if ( funName.contains("atomic_") ) {
+                        return true;
+                } else if (funName.contains("atomic") ) {
+                       return true;
+               }
+       }
 
-      SmallString<32> AtomicRMWName("cds_atomic" + NamePart + BitSizeStr);
-      CdsAtomicRMW[op][i] = M.getOrInsertFunction(AtomicRMWName, Ty, PtrTy, OrdTy, Ty);
-    }
+       return false;
+}
 
-    // only supportes strong version
-    SmallString<32> AtomicCASName("cds_atomic_compare_exchange" + BitSizeStr);    
-    CdsAtomicCAS[i]   = M.getOrInsertFunction(AtomicCASName, Ty, PtrTy, Ty, Ty, OrdTy, OrdTy);
-  }
+void printArgs (CallInst *CI) {
+       Function *fun = CI->getCalledFunction();
+       StringRef funName = fun->getName();
+
+       User::op_iterator begin = CI->arg_begin();
+       User::op_iterator end = CI->arg_end();
+
+       if ( funName.contains("atomic_") ) {
+               std::vector<Value *> parameters;
+
+               for (User::op_iterator it = begin; it != end; ++it) {
+                       Value *param = *it;
+                       parameters.push_back(param);
+                       errs() << *param << " type: " << *param->getType()  << "\n";
+               }
+       }
 
-  CdsAtomicThreadFence = M.getOrInsertFunction("cds_atomic_thread_fence", VoidTy, OrdTy);
 }
 
-static bool isVtableAccess(Instruction *I) {
-  if (MDNode *Tag = I->getMetadata(LLVMContext::MD_tbaa))
-    return Tag->isTBAAVtableAccess();
-  return false;
+bool CDSPass::instrumentAtomicCall(CallInst *CI, const DataLayout &DL) {
+       IRBuilder<> IRB(CI);
+       Function *fun = CI->getCalledFunction();
+       StringRef funName = fun->getName();
+       std::vector<Value *> parameters;
+
+       User::op_iterator begin = CI->arg_begin();
+       User::op_iterator end = CI->arg_end();
+       for (User::op_iterator it = begin; it != end; ++it) {
+               Value *param = *it;
+               parameters.push_back(param);
+       }
+
+       // obtain source line number of the CallInst
+       Value *position = getPosition(CI, IRB);
+
+       // the pointer to the address is always the first argument
+       Value *OrigPtr = parameters[0];
+       int Idx = getMemoryAccessFuncIndex(OrigPtr, DL);
+       if (Idx < 0)
+               return false;
+
+       const unsigned ByteSize = 1U << Idx;
+       const unsigned BitSize = ByteSize * 8;
+       Type *Ty = Type::getIntNTy(IRB.getContext(), BitSize);
+       Type *PtrTy = Ty->getPointerTo();
+
+       // atomic_init; args = {obj, order}
+       if (funName.contains("atomic_init")) {
+               Value *ptr = IRB.CreatePointerCast(OrigPtr, PtrTy);
+               Value *val = IRB.CreateBitOrPointerCast(parameters[1], Ty);
+               Value *args[] = {ptr, val, position};
+
+               Instruction* funcInst=CallInst::Create(CDSAtomicInit[Idx], args);
+               ReplaceInstWithInst(CI, funcInst);
+
+               return true;
+       }
+
+       // atomic_load; args = {obj, order}
+       if (funName.contains("atomic_load")) {
+               bool isExplicit = funName.contains("atomic_load_explicit");
+
+               Value *ptr = IRB.CreatePointerCast(OrigPtr, PtrTy);
+               Value *order;
+               if (isExplicit)
+                       order = IRB.CreateBitOrPointerCast(parameters[1], OrdTy);
+               else 
+                       order = ConstantInt::get(OrdTy, 
+                                                       (int) AtomicOrderingCABI::seq_cst);
+               Value *args[] = {ptr, order, position};
+               
+               Instruction* funcInst=CallInst::Create(CDSAtomicLoad[Idx], args);
+               ReplaceInstWithInst(CI, funcInst);
+
+               return true;
+       } else if (funName.contains("atomic") && 
+                                funName.contains("load")) {
+                // does this version of call always have an atomic order as an argument?
+                Value *ptr = IRB.CreatePointerCast(OrigPtr, PtrTy);
+                Value *order = IRB.CreateBitOrPointerCast(parameters[1], OrdTy);
+                Value *args[] = {ptr, order, position};
+
+                //Instruction* funcInst=CallInst::Create(CDSAtomicLoad[Idx], args);
+                CallInst *funcInst = IRB.CreateCall(CDSAtomicLoad[Idx], args);
+                Value *RetVal = IRB.CreateIntToPtr(funcInst, CI->getType());
+
+                CI->replaceAllUsesWith(RetVal);
+                CI->eraseFromParent();
+                
+                return true;
+        }
+
+       // atomic_store; args = {obj, val, order}
+       if (funName.contains("atomic_store")) {
+               bool isExplicit = funName.contains("atomic_store_explicit");
+               Value *OrigVal = parameters[1];
+
+               Value *ptr = IRB.CreatePointerCast(OrigPtr, PtrTy);
+               Value *val = IRB.CreatePointerCast(OrigVal, Ty);
+               Value *order;
+               if (isExplicit)
+                       order = IRB.CreateBitOrPointerCast(parameters[2], OrdTy);
+               else 
+                       order = ConstantInt::get(OrdTy, 
+                                                       (int) AtomicOrderingCABI::seq_cst);
+               Value *args[] = {ptr, val, order, position};
+               
+               Instruction* funcInst=CallInst::Create(CDSAtomicStore[Idx], args);
+               ReplaceInstWithInst(CI, funcInst);
+
+               return true;
+       } else if (funName.contains("atomic") && 
+                                funName.contains("EEEE5store")) {
+                // does this version of call always have an atomic order as an argument?
+                Value *OrigVal = parameters[1];
+
+                Value *ptr = IRB.CreatePointerCast(OrigPtr, PtrTy);
+                Value *val = IRB.CreatePointerCast(OrigVal, Ty);
+                Value *order = IRB.CreateBitOrPointerCast(parameters[1], OrdTy);
+                Value *args[] = {ptr, val, order, position};
+
+                Instruction* funcInst = CallInst::Create(CDSAtomicStore[Idx], args);
+                ReplaceInstWithInst(CI, funcInst);
+
+                return true;
+        }
+
+
+       // atomic_fetch_*; args = {obj, val, order}
+       if (funName.contains("atomic_fetch_") || 
+                       funName.contains("atomic_exchange") ) {
+               bool isExplicit = funName.contains("_explicit");
+               Value *OrigVal = parameters[1];
+
+               int op;
+               if ( funName.contains("_fetch_add") )
+                       op = AtomicRMWInst::Add;
+               else if ( funName.contains("_fetch_sub") )
+                       op = AtomicRMWInst::Sub;
+               else if ( funName.contains("_fetch_and") )
+                       op = AtomicRMWInst::And;
+               else if ( funName.contains("_fetch_or") )
+                       op = AtomicRMWInst::Or;
+               else if ( funName.contains("_fetch_xor") )
+                       op = AtomicRMWInst::Xor;
+               else if ( funName.contains("atomic_exchange") )
+                       op = AtomicRMWInst::Xchg;
+               else {
+                       errs() << "Unknown atomic read modify write operation\n";
+                       return false;
+               }
+
+               Value *ptr = IRB.CreatePointerCast(OrigPtr, PtrTy);
+               Value *val = IRB.CreatePointerCast(OrigVal, Ty);
+               Value *order;
+               if (isExplicit)
+                       order = IRB.CreateBitOrPointerCast(parameters[2], OrdTy);
+               else 
+                       order = ConstantInt::get(OrdTy, 
+                                                       (int) AtomicOrderingCABI::seq_cst);
+               Value *args[] = {ptr, val, order, position};
+               
+               Instruction* funcInst=CallInst::Create(CDSAtomicRMW[op][Idx], args);
+               ReplaceInstWithInst(CI, funcInst);
+
+               return true;
+       } else if (funName.contains("fetch")) {
+                errs() << "atomic exchange captured. Not implemented yet. ";
+                errs() << "See source file :";
+                getPositionPrint(CI, IRB);
+        } else if (funName.contains("exchange") &&
+                                !funName.contains("compare_exchange") ) {
+                errs() << "atomic exchange captured. Not implemented yet. ";
+                errs() << "See source file :";
+                getPositionPrint(CI, IRB);
+        }
+
+       /* atomic_compare_exchange_*; 
+          args = {obj, expected, new value, order1, order2}
+       */
+       if ( funName.contains("atomic_compare_exchange_") ) {
+               bool isExplicit = funName.contains("_explicit");
+
+               Value *Addr = IRB.CreatePointerCast(OrigPtr, PtrTy);
+               Value *CmpOperand = IRB.CreatePointerCast(parameters[1], PtrTy);
+               Value *NewOperand = IRB.CreateBitOrPointerCast(parameters[2], Ty);
+
+               Value *order_succ, *order_fail;
+               if (isExplicit) {
+                       order_succ = IRB.CreateBitOrPointerCast(parameters[3], OrdTy);
+                       order_fail = IRB.CreateBitOrPointerCast(parameters[4], OrdTy);
+               } else  {
+                       order_succ = ConstantInt::get(OrdTy, 
+                                                       (int) AtomicOrderingCABI::seq_cst);
+                       order_fail = ConstantInt::get(OrdTy, 
+                                                       (int) AtomicOrderingCABI::seq_cst);
+               }
+
+               Value *args[] = {Addr, CmpOperand, NewOperand, 
+                                                       order_succ, order_fail, position};
+               
+               Instruction* funcInst=CallInst::Create(CDSAtomicCAS_V2[Idx], args);
+               ReplaceInstWithInst(CI, funcInst);
+
+               return true;
+       } else if ( funName.contains("compare_exchange_strong") || 
+                                funName.contains("compare_exchange_weak") ) {
+                Value *Addr = IRB.CreatePointerCast(OrigPtr, PtrTy);
+                Value *CmpOperand = IRB.CreatePointerCast(parameters[1], PtrTy);
+                Value *NewOperand = IRB.CreateBitOrPointerCast(parameters[2], Ty);
+
+                Value *order_succ, *order_fail;
+                order_succ = IRB.CreateBitOrPointerCast(parameters[3], OrdTy);
+                order_fail = IRB.CreateBitOrPointerCast(parameters[4], OrdTy);
+
+                Value *args[] = {Addr, CmpOperand, NewOperand, 
+                                                        order_succ, order_fail, position};
+                Instruction* funcInst=CallInst::Create(CDSAtomicCAS_V2[Idx], args);
+                ReplaceInstWithInst(CI, funcInst);
+
+                return true;
+        }
+
+
+       return false;
 }
 
 static bool shouldInstrumentReadWriteFromAddress(const Module *M, Value *Addr) {
@@ -257,7 +496,7 @@ static bool shouldInstrumentReadWriteFromAddress(const Module *M, Value *Addr) {
   return true;
 }
 
-bool CdsPass::addrPointsToConstantData(Value *Addr) {
+bool CDSPass::addrPointsToConstantData(Value *Addr) {
   // If this is a GEP, just analyze its pointer operand.
   if (GetElementPtrInst *GEP = dyn_cast<GetElementPtrInst>(Addr))
     Addr = GEP->getPointerOperand();
@@ -278,50 +517,59 @@ bool CdsPass::addrPointsToConstantData(Value *Addr) {
   return false;
 }
 
-bool CdsPass::runOnFunction(Function &F) {
-  if (F.getName() == "main") 
+bool CDSPass::runOnFunction(Function &F) {
+  if (F.getName() == "main") {
     F.setName("user_main");
+    errs() << "main replaced by user_main\n";
+  }
 
-  initializeCallbacks( *F.getParent() );
+  if (true) {
+    initializeCallbacks( *F.getParent() );
 
-  SmallVector<Instruction*, 8> AllLoadsAndStores;
-  SmallVector<Instruction*, 8> LocalLoadsAndStores;
-  SmallVector<Instruction*, 8> AtomicAccesses;
+    SmallVector<Instruction*, 8> AllLoadsAndStores;
+    SmallVector<Instruction*, 8> LocalLoadsAndStores;
+    SmallVector<Instruction*, 8> AtomicAccesses;
 
-  std::vector<Instruction *> worklist;
+    std::vector<Instruction *> worklist;
 
-  bool Res = false;
-  const DataLayout &DL = F.getParent()->getDataLayout();
-  
-  errs() << "Before\n";
-  F.dump();
-
-  for (auto &B : F) {
-    for (auto &I : B) {
-      if ( (&I)->isAtomic() ) {
-        AtomicAccesses.push_back(&I);
-      } else if (isa<LoadInst>(I) || isa<StoreInst>(I)) {
-        LocalLoadsAndStores.push_back(&I);
+    bool Res = false;
+    const DataLayout &DL = F.getParent()->getDataLayout();
+
+    errs() << "--- " << F.getName() << "---\n";
+
+    for (auto &B : F) {
+      for (auto &I : B) {
+        if ( (&I)->isAtomic() || isAtomicCall(&I) ) {
+          AtomicAccesses.push_back(&I);
+        } else if (isa<LoadInst>(I) || isa<StoreInst>(I)) {
+          LocalLoadsAndStores.push_back(&I);
+        } else if (isa<CallInst>(I) || isa<InvokeInst>(I)) {
+          // not implemented yet
+        }
       }
+
+      chooseInstructionsToInstrument(LocalLoadsAndStores, AllLoadsAndStores, DL);
     }
-    chooseInstructionsToInstrument(LocalLoadsAndStores, AllLoadsAndStores, DL);
-  }
 
-  for (auto Inst : AllLoadsAndStores) {
-    Res |= instrumentLoadOrStore(Inst, DL);
-  }
+    for (auto Inst : AllLoadsAndStores) {
+//      Res |= instrumentLoadOrStore(Inst, DL);
+//      errs() << "load and store are replaced\n";
+    }
 
-  for (auto Inst : AtomicAccesses) {
-    Res |= instrumentAtomic(Inst);
-  } 
+    for (auto Inst : AtomicAccesses) {
+      Res |= instrumentAtomic(Inst, DL);
+    }
+
+    if (F.getName() == "user_main") {
+      // F.dump();
+    }
+
+  }
 
-  errs() << "After\n";
-  F.dump();
-  
   return false;
 }
 
-void CdsPass::chooseInstructionsToInstrument(
+void CDSPass::chooseInstructionsToInstrument(
     SmallVectorImpl<Instruction *> &Local, SmallVectorImpl<Instruction *> &All,
     const DataLayout &DL) {
   SmallPtrSet<Value*, 8> WriteTargets;
@@ -364,7 +612,7 @@ void CdsPass::chooseInstructionsToInstrument(
 }
 
 
-bool CdsPass::instrumentLoadOrStore(Instruction *I,
+bool CDSPass::instrumentLoadOrStore(Instruction *I,
                                             const DataLayout &DL) {
   IRBuilder<> IRB(I);
   bool IsWrite = isa<StoreInst>(*I);
@@ -378,10 +626,10 @@ bool CdsPass::instrumentLoadOrStore(Instruction *I,
   if (Addr->isSwiftError())
     return false;
 
-  int size = getTypeSize(Addr->getType());
-  int index = sizetoindex(size);
+  int Idx = getMemoryAccessFuncIndex(Addr, DL);
 
-//  not supported by Cds yet
+
+//  not supported by CDS yet
 /*  if (IsWrite && isVtableAccess(I)) {
     LLVM_DEBUG(dbgs() << "  VPTR : " << *I << "\n");
     Value *StoredValue = cast<StoreInst>(I)->getValueOperand();
@@ -410,69 +658,66 @@ bool CdsPass::instrumentLoadOrStore(Instruction *I,
 */
 
   Value *OnAccessFunc = nullptr;
-  OnAccessFunc = IsWrite ? CdsStore[index] : CdsLoad[index];
-
+  OnAccessFunc = IsWrite ? CDSStore[Idx] : CDSLoad[Idx];
+  
+  Type *ArgType = IRB.CreatePointerCast(Addr, Addr->getType())->getType();
+
+  if ( ArgType != Int8PtrTy && ArgType != Int16PtrTy && 
+               ArgType != Int32PtrTy && ArgType != Int64PtrTy ) {
+        //errs() << "A load or store of type ";
+        //errs() << *ArgType;
+        //errs() << " is passed in\n";
+        return false;  // if other types of load or stores are passed in
+  }
   IRB.CreateCall(OnAccessFunc, IRB.CreatePointerCast(Addr, Addr->getType()));
   if (IsWrite) NumInstrumentedWrites++;
   else         NumInstrumentedReads++;
   return true;
 }
 
-
-bool CdsPass::instrumentAtomic(Instruction * I) {
+bool CDSPass::instrumentAtomic(Instruction * I, const DataLayout &DL) {
   IRBuilder<> IRB(I);
   // LLVMContext &Ctx = IRB.getContext();
 
-  if (StoreInst *SI = dyn_cast<StoreInst>(I)) {
-    int atomic_order_index = getAtomicOrderIndex(SI->getOrdering());
-
-    Value *val = SI->getValueOperand();
-    Value *ptr = SI->getPointerOperand();
-    Value *order = ConstantInt::get(OrdTy, atomic_order_index);
-    Value *args[] = {ptr, order, val};
+  if (auto *CI = dyn_cast<CallInst>(I)) {
+    return instrumentAtomicCall(CI, DL);
+  }
 
-    int size=getTypeSize(ptr->getType());
-    int index=sizetoindex(size);
+  Value *position = getPosition(I, IRB);
 
-    Instruction* funcInst=CallInst::Create(CdsAtomicStore[index], args,"");
-    ReplaceInstWithInst(SI, funcInst);
-    errs() << "Store replaced\n";
-  } else if (LoadInst *LI = dyn_cast<LoadInst>(I)) {
+  if (LoadInst *LI = dyn_cast<LoadInst>(I)) {
+    Value *Addr = LI->getPointerOperand();
+    int Idx=getMemoryAccessFuncIndex(Addr, DL);
     int atomic_order_index = getAtomicOrderIndex(LI->getOrdering());
-
-    Value *ptr = LI->getPointerOperand();
     Value *order = ConstantInt::get(OrdTy, atomic_order_index);
-    Value *args[] = {ptr, order};
-
-    int size=getTypeSize(ptr->getType());
-    int index=sizetoindex(size);
-
-    Instruction* funcInst=CallInst::Create(CdsAtomicLoad[index], args, "");
+    Value *args[] = {Addr, order, position};
+    Instruction* funcInst=CallInst::Create(CDSAtomicLoad[Idx], args);
     ReplaceInstWithInst(LI, funcInst);
-    errs() << "Load Replaced\n";
+  } else if (StoreInst *SI = dyn_cast<StoreInst>(I)) {
+    Value *Addr = SI->getPointerOperand();
+    int Idx=getMemoryAccessFuncIndex(Addr, DL);
+    int atomic_order_index = getAtomicOrderIndex(SI->getOrdering());
+    Value *val = SI->getValueOperand();
+    Value *order = ConstantInt::get(OrdTy, atomic_order_index);
+    Value *args[] = {Addr, val, order, position};
+    Instruction* funcInst=CallInst::Create(CDSAtomicStore[Idx], args);
+    ReplaceInstWithInst(SI, funcInst);
   } else if (AtomicRMWInst *RMWI = dyn_cast<AtomicRMWInst>(I)) {
+    Value *Addr = RMWI->getPointerOperand();
+    int Idx=getMemoryAccessFuncIndex(Addr, DL);
     int atomic_order_index = getAtomicOrderIndex(RMWI->getOrdering());
-
     Value *val = RMWI->getValOperand();
-    Value *ptr = RMWI->getPointerOperand();
     Value *order = ConstantInt::get(OrdTy, atomic_order_index);
-    Value *args[] = {ptr, order, val};
-
-    int size = getTypeSize(ptr->getType());
-    int index = sizetoindex(size);
-
-    Instruction* funcInst = CallInst::Create(CdsAtomicRMW[RMWI->getOperation()][index], args, "");
+    Value *args[] = {Addr, val, order, position};
+    Instruction* funcInst = CallInst::Create(CDSAtomicRMW[RMWI->getOperation()][Idx], args);
     ReplaceInstWithInst(RMWI, funcInst);
-    errs() << RMWI->getOperationName(RMWI->getOperation());
-    errs() << " replaced\n";
   } else if (AtomicCmpXchgInst *CASI = dyn_cast<AtomicCmpXchgInst>(I)) {
     IRBuilder<> IRB(CASI);
 
     Value *Addr = CASI->getPointerOperand();
+    int Idx=getMemoryAccessFuncIndex(Addr, DL);
 
-    int size = getTypeSize(Addr->getType());
-    int index = sizetoindex(size);
-    const unsigned ByteSize = 1U << index;
+    const unsigned ByteSize = 1U << Idx;
     const unsigned BitSize = ByteSize * 8;
     Type *Ty = Type::getIntNTy(IRB.getContext(), BitSize);
     Type *PtrTy = Ty->getPointerTo();
@@ -487,9 +732,9 @@ bool CdsPass::instrumentAtomic(Instruction * I) {
 
     Value *Args[] = {IRB.CreatePointerCast(Addr, PtrTy),
                      CmpOperand, NewOperand,
-                     order_succ, order_fail};
+                     order_succ, order_fail, position};
 
-    CallInst *funcInst = IRB.CreateCall(CdsAtomicCAS[index], Args);
+    CallInst *funcInst = IRB.CreateCall(CDSAtomicCAS_V1[Idx], Args);
     Value *Success = IRB.CreateICmpEQ(funcInst, CmpOperand);
 
     Value *OldVal = funcInst;
@@ -508,25 +753,40 @@ bool CdsPass::instrumentAtomic(Instruction * I) {
   } else if (FenceInst *FI = dyn_cast<FenceInst>(I)) {
     int atomic_order_index = getAtomicOrderIndex(FI->getOrdering());
     Value *order = ConstantInt::get(OrdTy, atomic_order_index);
-    Value *Args[] = {order};
+    Value *Args[] = {order, position};
 
-    CallInst *funcInst = CallInst::Create(CdsAtomicThreadFence, Args);
+    CallInst *funcInst = CallInst::Create(CDSAtomicThreadFence, Args);
     ReplaceInstWithInst(FI, funcInst);
-    errs() << "Thread Fences replaced\n";
+//    errs() << "Thread Fences replaced\n";
   }
   return true;
 }
 
+int CDSPass::getMemoryAccessFuncIndex(Value *Addr,
+                                              const DataLayout &DL) {
+  Type *OrigPtrTy = Addr->getType();
+  Type *OrigTy = cast<PointerType>(OrigPtrTy)->getElementType();
+  assert(OrigTy->isSized());
+  uint32_t TypeSize = DL.getTypeStoreSizeInBits(OrigTy);
+  if (TypeSize != 8  && TypeSize != 16 &&
+      TypeSize != 32 && TypeSize != 64 && TypeSize != 128) {
+    NumAccessesWithBadSize++;
+    // Ignore all unusual sizes.
+    return -1;
+  }
+  size_t Idx = countTrailingZeros(TypeSize / 8);
+  assert(Idx < kNumberOfAccessSizes);
+  return Idx;
+}
 
 
-char CdsPass::ID = 0;
+char CDSPass::ID = 0;
 
 // Automatically enable the pass.
-// http://adriansampson.net/blog/clangpass.html
-static void registerCdsPass(const PassManagerBuilder &,
+static void registerCDSPass(const PassManagerBuilder &,
                          legacy::PassManagerBase &PM) {
-  PM.add(new CdsPass());
+  PM.add(new CDSPass());
 }
 static RegisterStandardPasses 
-       RegisterMyPass(PassManagerBuilder::EP_EarlyAsPossible,
-registerCdsPass);
\ No newline at end of file
+       RegisterMyPass(PassManagerBuilder::EP_OptimizerLast,
+registerCDSPass);