[dfsan] Introduce further optimization to reduce the number of union queries.

Specifically, do not compute a union if it is statically known that one
shadow set subsumes the other.

llvm-svn: 213100

2 files changed, 51 insertions, 4 deletions

diff --git a/llvm/lib/Transforms/Instrumentation/DataFlowSanitizer.cpp b/llvm/lib/Transforms/Instrumentation/DataFlowSanitizer.cpp
index 6419adc85a2..35057cdd47e 100644
--- a/llvm/lib/Transforms/Instrumentation/DataFlowSanitizer.cpp
+++ b/llvm/lib/Transforms/Instrumentation/DataFlowSanitizer.cpp
@@ -63,7 +63,10 @@
 #include "llvm/Support/SpecialCaseList.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
 #include "llvm/Transforms/Utils/Local.h"
+#include <algorithm>
 #include <iterator>
+#include <set>
+#include <utility>
 
 using namespace llvm;
 
@@ -284,6 +287,7 @@ struct DFSanFunction {
   };
   DenseMap<std::pair<Value *, Value *>, CachedCombinedShadow>
       CachedCombinedShadows;
+  DenseMap<Value *, std::set<Value *>> ShadowElements;
 
   DFSanFunction(DataFlowSanitizer &DFS, Function *F, bool IsNativeABI)
       : DFS(DFS), F(F), IA(DFS.getInstrumentedABI()),
@@ -892,6 +896,24 @@ Value *DFSanFunction::combineShadows(Value *V1, Value *V2, Instruction *Pos) {
   if (V1 == V2)
     return V1;
 
+  auto V1Elems = ShadowElements.find(V1);
+  auto V2Elems = ShadowElements.find(V2);
+  if (V1Elems != ShadowElements.end() && V2Elems != ShadowElements.end()) {
+    if (std::includes(V1Elems->second.begin(), V1Elems->second.end(),
+                      V2Elems->second.begin(), V2Elems->second.end())) {
+      return V1;
+    } else if (std::includes(V2Elems->second.begin(), V2Elems->second.end(),
+                             V1Elems->second.begin(), V1Elems->second.end())) {
+      return V2;
+    }
+  } else if (V1Elems != ShadowElements.end()) {
+    if (V1Elems->second.count(V2))
+      return V1;
+  } else if (V2Elems != ShadowElements.end()) {
+    if (V2Elems->second.count(V1))
+      return V2;
+  }
+
   auto Key = std::make_pair(V1, V2);
   if (V1 > V2)
     std::swap(Key.first, Key.second);
@@ -917,6 +939,20 @@ Value *DFSanFunction::combineShadows(Value *V1, Value *V2, Instruction *Pos) {
 
   CCS.Block = Tail;
   CCS.Shadow = Phi;
+
+  std::set<Value *> UnionElems;
+  if (V1Elems != ShadowElements.end()) {
+    UnionElems = V1Elems->second;
+  } else {
+    UnionElems.insert(V1);
+  }
+  if (V2Elems != ShadowElements.end()) {
+    UnionElems.insert(V2Elems->second.begin(), V2Elems->second.end());
+  } else {
+    UnionElems.insert(V2);
+  }
+  ShadowElements[Phi] = std::move(UnionElems);
+
   return Phi;
 }
 
diff --git a/llvm/test/Instrumentation/DataFlowSanitizer/union.ll b/llvm/test/Instrumentation/DataFlowSanitizer/union.ll
index 30162e75a25..2b31081776b 100644
--- a/llvm/test/Instrumentation/DataFlowSanitizer/union.ll
+++ b/llvm/test/Instrumentation/DataFlowSanitizer/union.ll
@@ -8,10 +8,10 @@ target datalayout = "e-p:64:64:64-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:64:64-f3
 
 ; CHECK-LABEL: @"dfs$f"
 define void @f(i32 %x, i32 %y) {
-  ; CHECK: __dfsan_union
+  ; CHECK: call{{.*}}__dfsan_union
   %xay = add i32 %x, %y
   store i32 %xay, i32* @a
-  ; CHECK-NOT: __dfsan_union
+  ; CHECK-NOT: call{{.*}}__dfsan_union
   %xmy = mul i32 %x, %y
   store i32 %xmy, i32* @b
   ret void
@@ -25,13 +25,13 @@ define void @g(i1 %p, i32 %x, i32 %y) {
   br i1 %p, label %l1, label %l2
 
 l1:
-  ; CHECK: __dfsan_union
+  ; CHECK: call{{.*}}__dfsan_union
   %xay = add i32 %x, %y
   store i32 %xay, i32* @a
   br label %l3
 
 l2:
-  ; CHECK: __dfsan_union
+  ; CHECK: call{{.*}}__dfsan_union
   %xmy = mul i32 %x, %y
   store i32 %xmy, i32* @b
   br label %l3
@@ -39,3 +39,14 @@ l2:
 l3:
   ret void
 }
+
+; In this case, we know that the label for %xayax subsumes the label for %xay.
+
+; CHECK-LABEL: @"dfs$h"
+define i32 @h(i32 %x, i32 %y) {
+  ; CHECK: call{{.*}}__dfsan_union
+  %xay = add i32 %x, %y
+  ; CHECK-NOT: call{{.*}}__dfsan_union
+  %xayax = add i32 %xay, %x
+  ret i32 %xayax
+}