Finished support of IPA EA & TSan, fixed bugs

* Switched TSan tests to global EA * Fixed bug in isExternalEscapedObject * Added support of variadic functions * Fixed InstrOmitted logic
focs-lab · MorthimerMcMare · Oct 14, 2024 · Oct 15, 2024 · Oct 18, 2024 · Oct 18, 2024
commit d0c40d7a23b7d5a7c7ba11b6c8512876ea1e517c
diff --git a/compiler-rt/test/tsan/lit.cfg.py b/compiler-rt/test/tsan/lit.cfg.py
@@ -48,7 +48,8 @@ def get_required_attr(config, attr_name):
     extra_cflags = []
 
 # Enable new Escape Analysis instead of default Capture Tracking
-escape_analysis_flag = ["-mllvm -tsan-use-escape-analysis"]
+#escape_analysis_flag = ["-mllvm -tsan-use-escape-analysis"]
+escape_analysis_flag = ["-mllvm -tsan-use-escape-analysis-global"]
 
 tsan_incdir = config.test_source_root + "/../"
 # Setup default compiler flags used with -fsanitize=thread option.

diff --git a/llvm/lib/Analysis/EscapeAnalysis.cpp b/llvm/lib/Analysis/EscapeAnalysis.cpp
@@ -163,9 +163,6 @@ bool EscapeAnalysisInfo::isExternalEscapedObject(const Value *V) const {
   if (const auto *CI = dyn_cast<CallInst>(V))
     return CI->getFunctionType()->getReturnType()->isPointerTy();
 
-  if (ArgsEscapes.has_value())
-    return isa<GlobalVariable>(V);
-
   return ((isa<Argument>(V) && V->getType()->isPointerTy()) ||
           isa<GlobalVariable>(V));
 }
@@ -378,6 +375,11 @@ EscapeAnalysisInfo::getEscapeKindForOpnd(const Use &U) const {
         assert(FuncIt != ArgsEscapes->get().end() &&
                "ArgsEscapes must contain information about called function");
 
+        // If it's variadic function, all arguments after fixed are escaped
+        if (Callee->isVarArg() && (Call->getDataOperandNo(&U) >=
+                                   Callee->getFunctionType()->getNumParams()))
+          return {EscapeKind::MAY_ESCAPE, std::nullopt};
+
         const auto ArgEscIt = FuncIt->second.find(Call->getDataOperandNo(&U));
         assert(ArgEscIt != FuncIt->second.end() &&
                "ArgEscapes must contain information about all arguments");
@@ -644,11 +646,14 @@ SmallVector<Value *, 8>
 EscapeAnalysisInfo::getUnderlyingMayEscObjects(const Value *V,
                                                const unsigned MaxLookup) {
   SmallVector<Value *, 8> UnderlObjs;
+  LLVM_DEBUG(dbgs() << "getUnderlyingMayEscObjects for " << *V << "\n");
   getUnderlyingObjectsForCodeGenWithoutPHIInvCheck(V, UnderlObjs, MaxLookup);
   LLVM_DEBUG(if (!UnderlObjs.empty()) {
     dbgs() << "\tgetUnderlyingMayEscObjects:";
     for (const auto *Obj : UnderlObjs)
-      dbgs() << "\t\t" << *Obj << "\n"; });
+      dbgs() << "\t\t" << *Obj << "\n"; }
+      else
+      dbgs() << "\tgetUnderlyingMayEscObjects -- empty\n"; );
   return UnderlObjs;
 }
 
@@ -675,7 +680,7 @@ void EscapeAnalysisInfo::printEscapingForBB(const BasicBlock *BB,
 
   OS << "Escaping objects for BB " << BB->getName() << ":\n";
   for (const auto *V : It->second.getEscapedObjs()) {
-    if (isExternalEscapedObject(V))
+    if (isExternalEscapedObject(V) || isa<Argument>(V))
       // I'm not sure, we should not print objects escaping by definition
       // (such as global variables or pointer arguments),
       // but let's omit them for now

diff --git a/llvm/lib/Transforms/Instrumentation/ThreadSanitizer.cpp b/llvm/lib/Transforms/Instrumentation/ThreadSanitizer.cpp
@@ -487,7 +487,7 @@ void ThreadSanitizer::chooseInstructionsToInstrument(
   DenseMap<Value *, size_t> WriteTargets; // Map of addresses to index in All
   // Iterate from the end.
   for (Instruction *I : reverse(Local)) {
-    LLVM_DEBUG(dbgs() << "\nchoose I: " << *I << "\n");
+    LLVM_DEBUG(dbgs() << "\nchooseInstructionsToInstrument I: " << *I << "\n");
 
     const bool IsWrite = isa<StoreInst>(*I);
     Value *Addr = IsWrite ? cast<StoreInst>(I)->getPointerOperand()
@@ -521,27 +521,43 @@ void ThreadSanitizer::chooseInstructionsToInstrument(
     }
 
     if (EAI.has_value()) {
-      const auto EscObjs = EscapeAnalysisInfo::getUnderlyingMayEscObjects(Addr);
-      for (const auto *Obj : EscObjs) {
+      bool InstrOmitted = false;
+      for (const auto *Obj :
+           EscapeAnalysisInfo::getUnderlyingMayEscObjects(Addr)) {
+        dbgs() << "check Obj " << *Obj << "\n";
         const bool IsEscaped = EAI.value().isEscapedForBB(I->getParent(), Obj);
         DEBUG_WITH_TYPE("tsan-ea", compareCaptureAndEA(I, Addr, IsEscaped));
-        if (!EAI.value().isEscapedForBB(I->getParent(), Obj)) {
-          LLVM_DEBUG(dbgs() << "Omit\n");
-          NumOmittedNonEscaped++;
-          continue;
+        if (IsEscaped) {
+          InstrOmitted = false;
+          break;
         }
+        InstrOmitted = true;
+      }
+      if (InstrOmitted) {
+        LLVM_DEBUG(dbgs() << "Instruction omitted\n");
+        NumOmittedNonEscaped++;
+        continue;
       }
     } else if (EAIGlobal.has_value()) {
-      const auto EscObjs = EscapeAnalysisInfo::getUnderlyingMayEscObjects(Addr);
-      for (const auto *Obj : EscObjs) {
-        bool IsEscaped = EAIGlobal.value()->isEscapedForBBInFunc(
+      bool InstrOmitted = false;
+      for (const auto *Obj :
+           EscapeAnalysisInfo::getUnderlyingMayEscObjects(Addr)) {
+        LLVM_DEBUG(dbgs() << "EscapeAnalysisInfo " << *Obj << " -- ");
+        const bool IsEscaped = EAIGlobal.value()->isEscapedForBBInFunc(
           I->getFunction(), I->getParent(), Obj);
         DEBUG_WITH_TYPE("tsan-ea", compareCaptureAndEA(I, Addr, IsEscaped));
-        if (!IsEscaped) {
-          LLVM_DEBUG(dbgs() << "Omit\n");
-          NumOmittedNonEscaped++;
-          continue;
+        if (IsEscaped) {
+          LLVM_DEBUG(dbgs() << "escaped\n");
+          InstrOmitted = false;
+          break;
         }
+        LLVM_DEBUG(dbgs() << "not escaped\n");
+        InstrOmitted = true;
+      }
+      if (InstrOmitted) {
+        LLVM_DEBUG(dbgs() << "Instruction omitted\n");
+        NumOmittedNonEscaped++;
+        continue;
       }
     } else {
       if (isa<AllocaInst>(getUnderlyingObject(Addr))) {
@@ -587,6 +603,12 @@ bool ThreadSanitizer::sanitizeFunction(
     Function &F, const TargetLibraryInfo &TLI,
     const std::optional<EscapeAnalysisInfo> &EAI,
     std::optional<EscapeAnalysisGlobalInfo*> EAIGlobal) {
+
+  LLVM_DEBUG(dbgs() <<
+    "\n%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%\n"
+    "%%%%%%%%%%%%%%%%%%%% Func " << F.getName() << "\t%%%%%%%%%%%%%%%%%%%%%%\n"
+    "%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%\n");
+
   // This is required to prevent instrumenting call to __tsan_init from within
   // the module constructor.
   if (F.getName() == kTsanModuleCtorName)