[LowerAllowCheck] Add llvm.allow.sanitize.* intrinsics

llvm/include/llvm/IR/Intrinsics.td

@@ -1884,6 +1884,20 @@ def int_allow_runtime_check : DefaultAttrsIntrinsic<[llvm_i1_ty], [llvm_metadata
     [IntrInaccessibleMemOnly, NoUndef<RetIndex>]>,
     ClangBuiltin<"__builtin_allow_runtime_check">;
 
+// Return true if the specific sanitizer is enabled for the function.
+def int_allow_sanitize_address
+    : DefaultAttrsIntrinsic<[llvm_i1_ty], [],
+                            [IntrInaccessibleMemOnly, NoUndef<RetIndex>]>;
+def int_allow_sanitize_thread
+    : DefaultAttrsIntrinsic<[llvm_i1_ty], [],
+                            [IntrInaccessibleMemOnly, NoUndef<RetIndex>]>;
+def int_allow_sanitize_memory
+    : DefaultAttrsIntrinsic<[llvm_i1_ty], [],
+                            [IntrInaccessibleMemOnly, NoUndef<RetIndex>]>;
+def int_allow_sanitize_hwaddress
+    : DefaultAttrsIntrinsic<[llvm_i1_ty], [],
+                            [IntrInaccessibleMemOnly, NoUndef<RetIndex>]>;
+
 // Support for dynamic deoptimization (or de-specialization)
 def int_experimental_deoptimize : Intrinsic<[llvm_any_ty], [llvm_vararg_ty],
                                             [Throws]>;

llvm/include/llvm/Transforms/Instrumentation/LowerAllowCheckPass.h

@@ -34,6 +34,8 @@ class LowerAllowCheckPass : public PassInfoMixin<LowerAllowCheckPass> {
       : Opts(std::move(Opts)) {};
   LLVM_ABI PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM);
 
+  static bool isRequired() { return true; }
+
   LLVM_ABI static bool IsRequested();
   LLVM_ABI void
   printPipeline(raw_ostream &OS,

llvm/lib/Transforms/Instrumentation/LowerAllowCheckPass.cpp

@@ -76,6 +76,7 @@ static bool lowerAllowChecks(Function &F, const BlockFrequencyInfo &BFI,
                              const ProfileSummaryInfo *PSI,
                              OptimizationRemarkEmitter &ORE,
                              const LowerAllowCheckPass::Options &Opts) {
+  // List of intrinsics and the constant value they should be lowered to.
   SmallVector<std::pair<IntrinsicInst *, bool>, 16> ReplaceWithValue;
   std::unique_ptr<RandomNumberGenerator> Rng;
 
@@ -123,26 +124,41 @@ static bool lowerAllowChecks(Function &F, const BlockFrequencyInfo &BFI,
     switch (ID) {
     case Intrinsic::allow_ubsan_check:
     case Intrinsic::allow_runtime_check: {
-      ++NumChecksTotal;
-
       bool ToRemove = ShouldRemove(II);
 
       ReplaceWithValue.push_back({
           II,
-          ToRemove,
+          !ToRemove,
       });
-      if (ToRemove)
-        ++NumChecksRemoved;
       emitRemark(II, ORE, ToRemove);
       break;
     }
+    case Intrinsic::allow_sanitize_address:
+      ReplaceWithValue.push_back(
+          {II, F.hasFnAttribute(Attribute::SanitizeAddress)});
+      break;
+    case Intrinsic::allow_sanitize_thread:
+      ReplaceWithValue.push_back(
+          {II, F.hasFnAttribute(Attribute::SanitizeThread)});
+      break;
+    case Intrinsic::allow_sanitize_memory:
+      ReplaceWithValue.push_back(
+          {II, F.hasFnAttribute(Attribute::SanitizeMemory)});
+      break;
+    case Intrinsic::allow_sanitize_hwaddress:
+      ReplaceWithValue.push_back(
+          {II, F.hasFnAttribute(Attribute::SanitizeHWAddress)});
+      break;
     default:
       break;
     }
   }
 
   for (auto [I, V] : ReplaceWithValue) {
-    I->replaceAllUsesWith(ConstantInt::getBool(I->getType(), !V));
+    ++NumChecksTotal;
+    if (!V) // If the final value is false, the check is considered removed
+      ++NumChecksRemoved;
+    I->replaceAllUsesWith(ConstantInt::getBool(I->getType(), V));
     I->eraseFromParent();
   }
 

llvm/test/Transforms/LowerAllowCheck/sanitize-check.ll

@@ -0,0 +1,70 @@
+; RUN: opt < %s -passes=lower-allow-check -S | FileCheck %s
+; RUN: opt < %s -passes=lower-allow-check -lower-allow-check-random-rate=0 -S | FileCheck %s
+
+declare i1 @llvm.allow.sanitize.address()
+declare i1 @llvm.allow.sanitize.thread()
+declare i1 @llvm.allow.sanitize.memory()
+declare i1 @llvm.allow.sanitize.hwaddress()
+
+define i1 @test_address() sanitize_address {
+; CHECK-LABEL: @test_address(
+; CHECK-NEXT:    ret i1 true
+  %1 = call i1 @llvm.allow.sanitize.address()
+  ret i1 %1
+}
+
+define i1 @test_no_sanitize_address() {
+; CHECK-LABEL: @test_no_sanitize_address(
+; CHECK-NEXT:    ret i1 false
+  %1 = call i1 @llvm.allow.sanitize.address()
+  ret i1 %1
+}
+
+define i1 @test_address_but_no_thread() sanitize_address {
+; CHECK-LABEL: @test_address_but_no_thread(
+; CHECK-NEXT:    ret i1 false
+  %1 = call i1 @llvm.allow.sanitize.thread()
+  ret i1 %1
+}
+
+define i1 @test_thread() sanitize_thread {
+; CHECK-LABEL: @test_thread(
+; CHECK-NEXT:    ret i1 true
+  %1 = call i1 @llvm.allow.sanitize.thread()
+  ret i1 %1
+}
+
+define i1 @test_no_sanitize_thread() {
+; CHECK-LABEL: @test_no_sanitize_thread(
+; CHECK-NEXT:    ret i1 false
+  %1 = call i1 @llvm.allow.sanitize.thread()
+  ret i1 %1
+}
+
+define i1 @test_memory() sanitize_memory {
+; CHECK-LABEL: @test_memory(
+; CHECK-NEXT:    ret i1 true
+  %1 = call i1 @llvm.allow.sanitize.memory()
+  ret i1 %1
+}
+
+define i1 @test_no_sanitize_memory() {
+; CHECK-LABEL: @test_no_sanitize_memory(
+; CHECK-NEXT:    ret i1 false
+  %1 = call i1 @llvm.allow.sanitize.memory()
+  ret i1 %1
+}
+
+define i1 @test_hwaddress() sanitize_hwaddress {
+; CHECK-LABEL: @test_hwaddress(
+; CHECK-NEXT:    ret i1 true
+  %1 = call i1 @llvm.allow.sanitize.hwaddress()
+  ret i1 %1
+}
+
+define i1 @test_no_sanitize_hwaddress() {
+; CHECK-LABEL: @test_no_sanitize_hwaddress(
+; CHECK-NEXT:    ret i1 false
+  %1 = call i1 @llvm.allow.sanitize.hwaddress()
+  ret i1 %1
+}