[InstCombine][asan] Don't speculate loads before select ptr

llvm/include/llvm/Analysis/Loads.h

@@ -69,8 +69,7 @@ bool isDereferenceableAndAlignedPointer(const Value *V, Align Alignment,
 /// quick local scan of the basic block containing ScanFrom, to determine if
 /// the address is already accessed.
 bool isSafeToLoadUnconditionally(Value *V, Align Alignment, const APInt &Size,
-                                 const DataLayout &DL,
-                                 Instruction *ScanFrom = nullptr,
+                                 const DataLayout &DL, Instruction *ScanFrom,
                                  AssumptionCache *AC = nullptr,
                                  const DominatorTree *DT = nullptr,
                                  const TargetLibraryInfo *TLI = nullptr);
@@ -100,12 +99,18 @@ bool isDereferenceableReadOnlyLoop(Loop *L, ScalarEvolution *SE,
 /// quick local scan of the basic block containing ScanFrom, to determine if
 /// the address is already accessed.
 bool isSafeToLoadUnconditionally(Value *V, Type *Ty, Align Alignment,
-                                 const DataLayout &DL,
-                                 Instruction *ScanFrom = nullptr,
+                                 const DataLayout &DL, Instruction *ScanFrom,
                                  AssumptionCache *AC = nullptr,
                                  const DominatorTree *DT = nullptr,
                                  const TargetLibraryInfo *TLI = nullptr);
 
+/// Return true if speculation of the given load must be suppressed to avoid
+/// ordering or interfering with an active sanitizer.  If not suppressed,
+/// dereferenceability and alignment must be proven separately.  Note: This
+/// is only needed for raw reasoning; if you use the interface below
+/// (isSafeToSpeculativelyExecute), this is handled internally.
+bool mustSuppressSpeculation(const LoadInst &LI);
+
 /// The default number of maximum instructions to scan in the block, used by
 /// FindAvailableLoadedValue().
 extern cl::opt<unsigned> DefMaxInstsToScan;

llvm/include/llvm/Analysis/ValueTracking.h

@@ -792,13 +792,6 @@ bool onlyUsedByLifetimeMarkers(const Value *V);
 /// droppable instructions.
 bool onlyUsedByLifetimeMarkersOrDroppableInsts(const Value *V);
 
-/// Return true if speculation of the given load must be suppressed to avoid
-/// ordering or interfering with an active sanitizer.  If not suppressed,
-/// dereferenceability and alignment must be proven separately.  Note: This
-/// is only needed for raw reasoning; if you use the interface below
-/// (isSafeToSpeculativelyExecute), this is handled internally.
-bool mustSuppressSpeculation(const LoadInst &LI);
-
 /// Return true if the instruction does not have any effects besides
 /// calculating the result and does not have undefined behavior.
 ///

llvm/lib/Analysis/Loads.cpp

@@ -345,6 +345,19 @@ bool llvm::isDereferenceableAndAlignedInLoop(LoadInst *LI, Loop *L,
                                             HeaderFirstNonPHI, AC, &DT);
 }
 
+static bool suppressSpeculativeLoadForSanitizers(const Instruction &CtxI) {
+  const Function &F = *CtxI.getFunction();
+  // Speculative load may create a race that did not exist in the source.
+  return F.hasFnAttribute(Attribute::SanitizeThread) ||
+         // Speculative load may load data from dirty regions.
+         F.hasFnAttribute(Attribute::SanitizeAddress) ||
+         F.hasFnAttribute(Attribute::SanitizeHWAddress);
+}
+
+bool llvm::mustSuppressSpeculation(const LoadInst &LI) {
+  return !LI.isUnordered() || suppressSpeculativeLoadForSanitizers(LI);
+}
+
 /// Check if executing a load of this pointer value cannot trap.
 ///
 /// If DT and ScanFrom are specified this method performs context-sensitive
@@ -365,8 +378,12 @@ bool llvm::isSafeToLoadUnconditionally(Value *V, Align Alignment, const APInt &S
   // If DT is not specified we can't make context-sensitive query
   const Instruction* CtxI = DT ? ScanFrom : nullptr;
   if (isDereferenceableAndAlignedPointer(V, Alignment, Size, DL, CtxI, AC, DT,
-                                         TLI))
-    return true;
+                                         TLI)) {
+    // With sanitizers `Dereferenceable` is not always enough for unconditional
+    // load.
+    if (!ScanFrom || !suppressSpeculativeLoadForSanitizers(*ScanFrom))
+      return true;
+  }
 
   if (!ScanFrom)
     return false;

llvm/lib/Analysis/ValueTracking.cpp

@@ -6798,17 +6798,6 @@ bool llvm::onlyUsedByLifetimeMarkersOrDroppableInsts(const Value *V) {
       V, /* AllowLifetime */ true, /* AllowDroppable */ true);
 }
 
-bool llvm::mustSuppressSpeculation(const LoadInst &LI) {
-  if (!LI.isUnordered())
-    return true;
-  const Function &F = *LI.getFunction();
-  // Speculative load may create a race that did not exist in the source.
-  return F.hasFnAttribute(Attribute::SanitizeThread) ||
-    // Speculative load may load data from dirty regions.
-    F.hasFnAttribute(Attribute::SanitizeAddress) ||
-    F.hasFnAttribute(Attribute::SanitizeHWAddress);
-}
-
 bool llvm::isSafeToSpeculativelyExecute(const Instruction *Inst,
                                         const Instruction *CtxI,
                                         AssumptionCache *AC,

llvm/test/Transforms/InstCombine/load.ll

@@ -56,6 +56,18 @@ define i32 @test5(i1 %C) {
   ret i32 %Z
 }
 
+; FIXME: Constants should be allowed for this optimization.
+define i32 @test5_asan(i1 %C) sanitize_address {
+; CHECK-LABEL: @test5_asan(
+; CHECK-NEXT:    [[Y:%.*]] = select i1 [[C:%.*]], ptr @X, ptr @X2
+; CHECK-NEXT:    [[Z:%.*]] = load i32, ptr [[Y]], align 4
+; CHECK-NEXT:    ret i32 [[Z]]
+;
+  %Y = select i1 %C, ptr @X, ptr @X2		; <ptr> [#uses=1]
+  %Z = load i32, ptr %Y		; <i32> [#uses=1]
+  ret i32 %Z
+}
+
 define i32 @load_gep_null_inbounds(i64 %X) {
 ; CHECK-LABEL: @load_gep_null_inbounds(
 ; CHECK-NEXT:    store i1 true, ptr poison, align 1

llvm/test/Transforms/InstCombine/multiple-uses-load-bitcast-select.ll

@@ -25,3 +25,27 @@ define void @PR35618(ptr %st1, ptr %st2) {
   ret void
 }
 
+define void @PR35618_asan(ptr %st1, ptr %st2) sanitize_address {
+; CHECK-LABEL: @PR35618_asan(
+; CHECK-NEXT:    [[Y1:%.*]] = alloca double, align 8
+; CHECK-NEXT:    [[Z1:%.*]] = alloca double, align 8
+; CHECK-NEXT:    [[LD1:%.*]] = load double, ptr [[Y1]], align 8
+; CHECK-NEXT:    [[LD2:%.*]] = load double, ptr [[Z1]], align 8
+; CHECK-NEXT:    [[TMP:%.*]] = fcmp olt double [[LD1]], [[LD2]]
+; CHECK-NEXT:    [[TMP12_V:%.*]] = select i1 [[TMP]], double [[LD1]], double [[LD2]]
+; CHECK-NEXT:    store double [[TMP12_V]], ptr [[ST1:%.*]], align 8
+; CHECK-NEXT:    store double [[TMP12_V]], ptr [[ST2:%.*]], align 8
+; CHECK-NEXT:    ret void
+;
+  %y1 = alloca double
+  %z1 = alloca double
+  %ld1 = load double, ptr %y1
+  %ld2 = load double, ptr %z1
+  %tmp = fcmp olt double %ld1, %ld2
+  %sel = select i1 %tmp, ptr %y1, ptr %z1
+  %tmp12 = load i64, ptr %sel
+  store i64 %tmp12, ptr %st1
+  store i64 %tmp12, ptr %st2
+  ret void
+}
+

llvm/test/Transforms/InstCombine/ptr-replace-alloca.ll

@@ -427,6 +427,23 @@ entry:
   ret i8 %load
 }
 
+define i8 @select_diff_addrspace_remove_alloca_asan(i1 %cond, ptr %p) sanitize_address {
+; CHECK-LABEL: @select_diff_addrspace_remove_alloca_asan(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[GEP2:%.*]] = select i1 [[COND:%.*]], ptr addrspace(1) getelementptr inbounds (i8, ptr addrspace(1) @g2, i64 4), ptr addrspace(1) getelementptr inbounds (i8, ptr addrspace(1) @g2, i64 6)
+; CHECK-NEXT:    [[LOAD:%.*]] = load i8, ptr addrspace(1) [[GEP2]], align 1
+; CHECK-NEXT:    ret i8 [[LOAD]]
+;
+entry:
+  %alloca = alloca [32 x i8]
+  call void @llvm.memcpy.p0.p1.i64(ptr %alloca, ptr addrspace(1) @g2, i64 32, i1 false)
+  %gep = getelementptr inbounds [32 x i8], ptr %alloca, i32 0, i32 2
+  %sel = select i1 %cond, ptr %alloca, ptr %gep
+  %gep2 = getelementptr inbounds i8, ptr %sel, i64 4
+  %load = load i8, ptr %gep2
+  ret i8 %load
+}
+
 declare i8 @readonly_callee(ptr readonly nocapture)
 
 ; FIXME: This should be able to fold to call i8 @readonly_callee(ptr nonnull @g1)

llvm/test/Transforms/InstCombine/select-load.ll

@@ -0,0 +1,101 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -passes=instcombine -S < %s | FileCheck %s
+
+target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-i128:128-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-grtev4-linux-gnu"
+
+define i32 @test_plain(i1 %f) {
+; CHECK-LABEL: @test_plain(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[A:%.*]] = alloca i32, align 8
+; CHECK-NEXT:    [[B:%.*]] = alloca i32, align 8
+; CHECK-NEXT:    [[A_VAL:%.*]] = load i32, ptr [[A]], align 8
+; CHECK-NEXT:    [[B_VAL:%.*]] = load i32, ptr [[B]], align 8
+; CHECK-NEXT:    [[L:%.*]] = select i1 [[F:%.*]], i32 [[A_VAL]], i32 [[B_VAL]]
+; CHECK-NEXT:    ret i32 [[L]]
+;
+entry:
+  %a = alloca i32, align 8
+  %b = alloca i32, align 8
+  %sel = select i1 %f, ptr %a, ptr %b
+  %l = load i32, ptr %sel, align 8
+  ret i32 %l
+}
+
+; Don't speculate as the condition may control which memory is valid from
+; sanitizer perspective.
+define i32 @test_asan(i1 %f) sanitize_address {
+; CHECK-LABEL: @test_asan(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[A:%.*]] = alloca i32, align 8
+; CHECK-NEXT:    [[B:%.*]] = alloca i32, align 8
+; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[F:%.*]], ptr [[A]], ptr [[B]]
+; CHECK-NEXT:    [[L:%.*]] = load i32, ptr [[SEL]], align 8
+; CHECK-NEXT:    ret i32 [[L]]
+;
+entry:
+  %a = alloca i32, align 8
+  %b = alloca i32, align 8
+  %sel = select i1 %f, ptr %a, ptr %b
+  %l = load i32, ptr %sel, align 8
+  ret i32 %l
+}
+
+
+; Don't speculate as the condition may control which memory is valid from
+; sanitizer perspective.
+define i32 @test_hwasan(i1 %f) sanitize_hwaddress {
+; CHECK-LABEL: @test_hwasan(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[A:%.*]] = alloca i32, align 8
+; CHECK-NEXT:    [[B:%.*]] = alloca i32, align 8
+; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[F:%.*]], ptr [[A]], ptr [[B]]
+; CHECK-NEXT:    [[L:%.*]] = load i32, ptr [[SEL]], align 8
+; CHECK-NEXT:    ret i32 [[L]]
+;
+entry:
+  %a = alloca i32, align 8
+  %b = alloca i32, align 8
+  %sel = select i1 %f, ptr %a, ptr %b
+  %l = load i32, ptr %sel, align 8
+  ret i32 %l
+}
+
+; Don't speculate as the condition may control which memory is valid from
+; sanitizer perspective.
+define i32 @test_tsan(i1 %f) sanitize_thread {
+; CHECK-LABEL: @test_tsan(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[A:%.*]] = alloca i32, align 8
+; CHECK-NEXT:    [[B:%.*]] = alloca i32, align 8
+; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[F:%.*]], ptr [[A]], ptr [[B]]
+; CHECK-NEXT:    [[L:%.*]] = load i32, ptr [[SEL]], align 8
+; CHECK-NEXT:    ret i32 [[L]]
+;
+entry:
+  %a = alloca i32, align 8
+  %b = alloca i32, align 8
+  %sel = select i1 %f, ptr %a, ptr %b
+  %l = load i32, ptr %sel, align 8
+  ret i32 %l
+}
+
+; Msan just propagates shadow, even if speculated load accesses uninitialized
+; value, instrumentation will select shadow of the desired value anyway.
+define i32 @test_msan(i1 %f) sanitize_memory {
+; CHECK-LABEL: @test_msan(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[A:%.*]] = alloca i32, align 8
+; CHECK-NEXT:    [[B:%.*]] = alloca i32, align 8
+; CHECK-NEXT:    [[A_VAL:%.*]] = load i32, ptr [[A]], align 8
+; CHECK-NEXT:    [[B_VAL:%.*]] = load i32, ptr [[B]], align 8
+; CHECK-NEXT:    [[L:%.*]] = select i1 [[F:%.*]], i32 [[A_VAL]], i32 [[B_VAL]]
+; CHECK-NEXT:    ret i32 [[L]]
+;
+entry:
+  %a = alloca i32, align 8
+  %b = alloca i32, align 8
+  %sel = select i1 %f, ptr %a, ptr %b
+  %l = load i32, ptr %sel, align 8
+  ret i32 %l
+}

llvm/test/Transforms/InstCombine/strnlen-2.ll

@@ -38,6 +38,21 @@ define i64 @fold_strnlen_s3_s5_1(i1 %C) {
   ret i64 %len
 }
 
+; FIXME: Constants should be allowed for this optimization.
+define i64 @fold_strnlen_s3_s5_1_asan(i1 %C) sanitize_address {
+; CHECK-LABEL: @fold_strnlen_s3_s5_1_asan(
+; CHECK-NEXT:    [[PTR:%.*]] = select i1 [[C:%.*]], ptr @s3, ptr @s6
+; CHECK-NEXT:    [[STRNLEN_CHAR0:%.*]] = load i8, ptr [[PTR]], align 1
+; CHECK-NEXT:    [[STRNLEN_CHAR0CMP:%.*]] = icmp ne i8 [[STRNLEN_CHAR0]], 0
+; CHECK-NEXT:    [[LEN:%.*]] = zext i1 [[STRNLEN_CHAR0CMP]] to i64
+; CHECK-NEXT:    ret i64 [[LEN]]
+;
+  %ptr = select i1 %C, ptr @s3, ptr @s6
+
+  %len = call i64 @strnlen(ptr %ptr, i64 1)
+  ret i64 %len
+}
+
 
 ; Fold strnlen (C ? s3 : s5, 3) to 3.
 

llvm/test/Transforms/SROA/phi-and-select.ll

@@ -344,6 +344,40 @@ entry:
   ret i32 %loaded
 }
 
+; We should not unconditionally load with sanitizers.
+define i32 @test9_asan(i32 %b, ptr %ptr) sanitize_address {
+; Same as @test8 but for a select rather than a PHI node.
+;
+; CHECK-PRESERVE-CFG-LABEL: @test9_asan(
+; CHECK-PRESERVE-CFG-NEXT:  entry:
+; CHECK-PRESERVE-CFG-NEXT:    [[F:%.*]] = alloca float, align 4
+; CHECK-PRESERVE-CFG-NEXT:    store i32 0, ptr [[PTR:%.*]], align 4
+; CHECK-PRESERVE-CFG-NEXT:    [[TEST:%.*]] = icmp ne i32 [[B:%.*]], 0
+; CHECK-PRESERVE-CFG-NEXT:    [[SELECT:%.*]] = select i1 [[TEST]], ptr [[F]], ptr [[PTR]]
+; CHECK-PRESERVE-CFG-NEXT:    [[LOADED:%.*]] = load i32, ptr [[SELECT]], align 4
+; CHECK-PRESERVE-CFG-NEXT:    ret i32 [[LOADED]]
+;
+; CHECK-MODIFY-CFG-LABEL: @test9_asan(
+; CHECK-MODIFY-CFG-NEXT:  entry:
+; CHECK-MODIFY-CFG-NEXT:    store i32 0, ptr [[PTR:%.*]], align 4
+; CHECK-MODIFY-CFG-NEXT:    [[TEST:%.*]] = icmp ne i32 [[B:%.*]], 0
+; CHECK-MODIFY-CFG-NEXT:    [[LOADED_ELSE_VAL:%.*]] = load i32, ptr [[PTR]], align 4
+; CHECK-MODIFY-CFG-NEXT:    br i1 [[TEST]], label [[ENTRY_THEN:%.*]], label [[ENTRY_CONT:%.*]]
+; CHECK-MODIFY-CFG:       entry.then:
+; CHECK-MODIFY-CFG-NEXT:    br label [[ENTRY_CONT]]
+; CHECK-MODIFY-CFG:       entry.cont:
+; CHECK-MODIFY-CFG-NEXT:    [[LOADED:%.*]] = phi i32 [ undef, [[ENTRY_THEN]] ], [ [[LOADED_ELSE_VAL]], [[ENTRY:%.*]] ]
+; CHECK-MODIFY-CFG-NEXT:    ret i32 [[LOADED]]
+;
+entry:
+  %f = alloca float
+  store i32 0, ptr %ptr
+  %test = icmp ne i32 %b, 0
+  %select = select i1 %test, ptr %f, ptr %ptr
+  %loaded = load i32, ptr %select, align 4
+  ret i32 %loaded
+}
+
 define float @test10(i32 %b, ptr %ptr) {
 ; Don't try to promote allocas which are not elligible for it even after
 ; rewriting due to the necessity of inserting bitcasts when speculating a PHI

llvm/test/Transforms/SROA/phi-with-duplicate-pred.ll

@@ -52,6 +52,52 @@ cleanup7:                                         ; preds = %cleanup
   ret void
 }
 
+define void @f2_hwasan(i1 %c1) sanitize_hwaddress {
+; CHECK-LABEL: @f2_hwasan(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[E:%.*]] = alloca i16, align 1
+; CHECK-NEXT:    br i1 [[C1:%.*]], label [[IF_THEN:%.*]], label [[IF_ELSE:%.*]]
+; CHECK:       if.then:
+; CHECK-NEXT:    br label [[CLEANUP:%.*]]
+; CHECK:       cleanup:
+; CHECK-NEXT:    switch i32 2, label [[CLEANUP7:%.*]] [
+; CHECK-NEXT:      i32 0, label [[LBL1:%.*]]
+; CHECK-NEXT:      i32 2, label [[LBL1]]
+; CHECK-NEXT:    ]
+; CHECK:       if.else:
+; CHECK-NEXT:    br label [[LBL1]]
+; CHECK:       lbl1:
+; CHECK-NEXT:    [[G_0:%.*]] = phi ptr [ @a, [[CLEANUP]] ], [ @a, [[CLEANUP]] ], [ [[E]], [[IF_ELSE]] ]
+; CHECK-NEXT:    [[TMP0:%.*]] = load i16, ptr [[G_0]], align 1
+; CHECK-NEXT:    unreachable
+; CHECK:       cleanup7:
+; CHECK-NEXT:    ret void
+;
+entry:
+  %e = alloca i16, align 1
+  br i1 %c1, label %if.then, label %if.else
+
+if.then:                                          ; preds = %entry
+  br label %cleanup
+
+cleanup:                                          ; preds = %if.then
+  switch i32 2, label %cleanup7 [
+  i32 0, label %lbl1
+  i32 2, label %lbl1
+  ]
+
+if.else:                                          ; preds = %entry
+  br label %lbl1
+
+lbl1:                                             ; preds = %if.else, %cleanup, %cleanup
+  %g.0 = phi ptr [ @a, %cleanup ], [ @a, %cleanup ], [ %e, %if.else ]
+  %0 = load i16, ptr %g.0, align 1
+  unreachable
+
+cleanup7:                                         ; preds = %cleanup
+  ret void
+}
+
 define void @f3(i1 %c1) {
 ; CHECK-LABEL: @f3(
 ; CHECK-NEXT:  entry: