From 4ead266ca7ff81c77f82b59a857cae35136e550a Mon Sep 17 00:00:00 2001 From: Vitaly Buka Date: Thu, 22 Feb 2024 16:17:26 -0800 Subject: [PATCH 1/6] [ubsan][pgo] Pass to remove ubsan checks based on profile data UBSAN checks can be too expensive to be used in release binaries. However not all code affect performace in the same way. Removing small number of checks in hot code we can performance loss, preserving most of the checks. --- .../Instrumentation/RemoveTrapsPass.h | 32 ++ llvm/lib/Passes/PassBuilder.cpp | 1 + llvm/lib/Passes/PassRegistry.def | 1 + .../Transforms/Instrumentation/CMakeLists.txt | 1 + .../Instrumentation/RemoveTrapsPass.cpp | 95 +++++ .../Transforms/RemoveTraps/remove-traps.ll | 397 ++++++++++++++++++ 6 files changed, 527 insertions(+) create mode 100644 llvm/include/llvm/Transforms/Instrumentation/RemoveTrapsPass.h create mode 100644 llvm/lib/Transforms/Instrumentation/RemoveTrapsPass.cpp create mode 100644 llvm/test/Transforms/RemoveTraps/remove-traps.ll diff --git a/llvm/include/llvm/Transforms/Instrumentation/RemoveTrapsPass.h b/llvm/include/llvm/Transforms/Instrumentation/RemoveTrapsPass.h new file mode 100644 index 0000000000000..58f6bbcec5dc9 --- /dev/null +++ b/llvm/include/llvm/Transforms/Instrumentation/RemoveTrapsPass.h @@ -0,0 +1,32 @@ +//===- RemoveTrapsPass.h ----------------------------------------*- C++ -*-===// +// +// 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 +// +//===----------------------------------------------------------------------===// +/// \file +/// This file provides the interface for the pass responsible for removing +/// expensive ubsan checks. +/// +//===----------------------------------------------------------------------===// + +#ifndef LLVM_TRANSFORMS_INSTRUMENTATION_UBSANOPTIMIZATIONPASS_H +#define LLVM_TRANSFORMS_INSTRUMENTATION_UBSANOPTIMIZATIONPASS_H + +#include "llvm/IR/Function.h" +#include "llvm/IR/PassManager.h" +#include "llvm/Pass.h" + +namespace llvm { + +// This pass is responsible for removing optional traps, like llvm.ubsantrap +// from the hot code. +class RemoveTrapsPass : public PassInfoMixin { +public: + PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM); +}; + +} // namespace llvm + +#endif diff --git a/llvm/lib/Passes/PassBuilder.cpp b/llvm/lib/Passes/PassBuilder.cpp index e0bc57f8bf72f..4d1eb10d2d41c 100644 --- a/llvm/lib/Passes/PassBuilder.cpp +++ b/llvm/lib/Passes/PassBuilder.cpp @@ -177,6 +177,7 @@ #include "llvm/Transforms/Instrumentation/PGOForceFunctionAttrs.h" #include "llvm/Transforms/Instrumentation/PGOInstrumentation.h" #include "llvm/Transforms/Instrumentation/PoisonChecking.h" +#include "llvm/Transforms/Instrumentation/RemoveTrapsPass.h" #include "llvm/Transforms/Instrumentation/SanitizerBinaryMetadata.h" #include "llvm/Transforms/Instrumentation/SanitizerCoverage.h" #include "llvm/Transforms/Instrumentation/ThreadSanitizer.h" diff --git a/llvm/lib/Passes/PassRegistry.def b/llvm/lib/Passes/PassRegistry.def index a345e8d72d939..41f16d0915bf2 100644 --- a/llvm/lib/Passes/PassRegistry.def +++ b/llvm/lib/Passes/PassRegistry.def @@ -422,6 +422,7 @@ FUNCTION_PASS("print", UniformityInfoPrinterPass(dbgs())) FUNCTION_PASS("reassociate", ReassociatePass()) FUNCTION_PASS("redundant-dbg-inst-elim", RedundantDbgInstEliminationPass()) FUNCTION_PASS("reg2mem", RegToMemPass()) +FUNCTION_PASS("remove-traps", RemoveTrapsPass()) FUNCTION_PASS("safe-stack", SafeStackPass(TM)) FUNCTION_PASS("scalarize-masked-mem-intrin", ScalarizeMaskedMemIntrinPass()) FUNCTION_PASS("scalarizer", ScalarizerPass()) diff --git a/llvm/lib/Transforms/Instrumentation/CMakeLists.txt b/llvm/lib/Transforms/Instrumentation/CMakeLists.txt index ee9aa73ff0340..b23a6ed1f0841 100644 --- a/llvm/lib/Transforms/Instrumentation/CMakeLists.txt +++ b/llvm/lib/Transforms/Instrumentation/CMakeLists.txt @@ -17,6 +17,7 @@ add_llvm_component_library(LLVMInstrumentation PGOInstrumentation.cpp PGOMemOPSizeOpt.cpp PoisonChecking.cpp + RemoveTrapsPass.cpp SanitizerCoverage.cpp SanitizerBinaryMetadata.cpp ValueProfileCollector.cpp diff --git a/llvm/lib/Transforms/Instrumentation/RemoveTrapsPass.cpp b/llvm/lib/Transforms/Instrumentation/RemoveTrapsPass.cpp new file mode 100644 index 0000000000000..7a7c604741e92 --- /dev/null +++ b/llvm/lib/Transforms/Instrumentation/RemoveTrapsPass.cpp @@ -0,0 +1,95 @@ +//===- RemoveTrapsPass.cpp --------------------------------------*- C++ -*-===// +// +// 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 +// +//===----------------------------------------------------------------------===// +// +// +//===----------------------------------------------------------------------===// + +#include "llvm/Transforms/Instrumentation/RemoveTrapsPass.h" + +#include "llvm/ADT/SmallVector.h" +#include "llvm/ADT/Statistic.h" +#include "llvm/Analysis/ProfileSummaryInfo.h" +#include "llvm/IR/Instructions.h" +#include "llvm/IR/IntrinsicInst.h" +#include "llvm/IR/Intrinsics.h" +#include + +using namespace llvm; + +#define DEBUG_TYPE "remove-traps" + +static constexpr unsigned MaxRandomRate = 1000; + +static cl::opt HotPercentileCutoff( + "remove-traps-percentile-cutoff-hot", cl::init(0), + cl::desc("Alternative hot percentile cuttoff. By default " + "`-profile-summary-cutoff-hot` is used.")); +static cl::opt RandomRate( + "remove-traps-random-rate", cl::init(0.0), + cl::desc( + "Probability to use for pseudorandom unconditional checks removal.")); + +STATISTIC(NumChecksTotal, "Number of checks"); +STATISTIC(NumChecksRemoved, "Number of removed checks"); + +static SmallVector +removeUbsanTraps(Function &F, FunctionAnalysisManager &FAM, + ProfileSummaryInfo *PSI) { + SmallVector Remove; + + if (F.isDeclaration()) + return {}; + + auto &BFI = FAM.getResult(F); + + int BBCounter = 0; + for (BasicBlock &BB : F) { + for (Instruction &I : BB) { + IntrinsicInst *II = dyn_cast(&I); + if (!II) + continue; + auto ID = II->getIntrinsicID(); + if (ID != Intrinsic::ubsantrap) + continue; + ++NumChecksTotal; + + bool IsHot = false; + if (PSI) { + uint64_t Count = 0; + for (const auto *PR : predecessors(&BB)) + Count += BFI.getBlockProfileCount(PR).value_or(0); + + IsHot = HotPercentileCutoff.getNumOccurrences() + ? PSI->isHotCountNthPercentile(HotPercentileCutoff, Count) + : PSI->isHotCount(Count); + } + + if ((IsHot) || ((F.getGUID() + BBCounter++) % MaxRandomRate) < + RandomRate * RandomRate) { + Remove.push_back(II); + ++NumChecksRemoved; + } + } + } + return Remove; +} + +PreservedAnalyses RemoveTrapsPass::run(Function &F, + FunctionAnalysisManager &AM) { + if (F.isDeclaration()) + return PreservedAnalyses::all(); + auto &MAMProxy = AM.getResult(F); + ProfileSummaryInfo *PSI = + MAMProxy.getCachedResult(*F.getParent()); + + auto Remove = removeUbsanTraps(F, AM, PSI); + for (auto *I : Remove) + I->eraseFromParent(); + + return Remove.empty() ? PreservedAnalyses::all() : PreservedAnalyses::none(); +} diff --git a/llvm/test/Transforms/RemoveTraps/remove-traps.ll b/llvm/test/Transforms/RemoveTraps/remove-traps.ll new file mode 100644 index 0000000000000..af9cdab2f2920 --- /dev/null +++ b/llvm/test/Transforms/RemoveTraps/remove-traps.ll @@ -0,0 +1,397 @@ +; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 4 +; RUN: opt < %s -passes='function(remove-traps)' -S | FileCheck %s --check-prefixes=NOPROFILE +; RUN: opt < %s -passes='function(remove-traps)' -remove-traps-random-rate=999999 -S | FileCheck %s --check-prefixes=ALL +; RUN: opt < %s -passes='require,function(remove-traps)' -S | FileCheck %s --check-prefixes=HOT +; RUN: opt < %s -passes='require,function(remove-traps)' -remove-traps-percentile-cutoff-hot=700000 -S | FileCheck %s --check-prefixes=HOT70 + +target triple = "x86_64-pc-linux-gnu" + +declare void @llvm.ubsantrap(i8 immarg) + +define dso_local noundef i32 @simple(ptr noundef readonly %0) { +; NOPROFILE-LABEL: define dso_local noundef i32 @simple( +; NOPROFILE-SAME: ptr noundef readonly [[TMP0:%.*]]) { +; NOPROFILE-NEXT: [[TMP2:%.*]] = icmp eq ptr [[TMP0]], null +; NOPROFILE-NEXT: br i1 [[TMP2]], label [[TMP3:%.*]], label [[TMP4:%.*]] +; NOPROFILE: 3: +; NOPROFILE-NEXT: tail call void @llvm.ubsantrap(i8 22) +; NOPROFILE-NEXT: unreachable +; NOPROFILE: 4: +; NOPROFILE-NEXT: [[TMP5:%.*]] = load i32, ptr [[TMP0]], align 4 +; NOPROFILE-NEXT: ret i32 [[TMP5]] +; +; ALL-LABEL: define dso_local noundef i32 @simple( +; ALL-SAME: ptr noundef readonly [[TMP0:%.*]]) { +; ALL-NEXT: [[TMP2:%.*]] = icmp eq ptr [[TMP0]], null +; ALL-NEXT: br i1 [[TMP2]], label [[TMP3:%.*]], label [[TMP4:%.*]] +; ALL: 3: +; ALL-NEXT: unreachable +; ALL: 4: +; ALL-NEXT: [[TMP5:%.*]] = load i32, ptr [[TMP0]], align 4 +; ALL-NEXT: ret i32 [[TMP5]] +; +; HOT-LABEL: define dso_local noundef i32 @simple( +; HOT-SAME: ptr noundef readonly [[TMP0:%.*]]) { +; HOT-NEXT: [[TMP2:%.*]] = icmp eq ptr [[TMP0]], null +; HOT-NEXT: br i1 [[TMP2]], label [[TMP3:%.*]], label [[TMP4:%.*]] +; HOT: 3: +; HOT-NEXT: tail call void @llvm.ubsantrap(i8 22) +; HOT-NEXT: unreachable +; HOT: 4: +; HOT-NEXT: [[TMP5:%.*]] = load i32, ptr [[TMP0]], align 4 +; HOT-NEXT: ret i32 [[TMP5]] +; +; HOT70-LABEL: define dso_local noundef i32 @simple( +; HOT70-SAME: ptr noundef readonly [[TMP0:%.*]]) { +; HOT70-NEXT: [[TMP2:%.*]] = icmp eq ptr [[TMP0]], null +; HOT70-NEXT: br i1 [[TMP2]], label [[TMP3:%.*]], label [[TMP4:%.*]] +; HOT70: 3: +; HOT70-NEXT: tail call void @llvm.ubsantrap(i8 22) +; HOT70-NEXT: unreachable +; HOT70: 4: +; HOT70-NEXT: [[TMP5:%.*]] = load i32, ptr [[TMP0]], align 4 +; HOT70-NEXT: ret i32 [[TMP5]] +; + %2 = icmp eq ptr %0, null + br i1 %2, label %3, label %4 + +3: + tail call void @llvm.ubsantrap(i8 22) + unreachable + +4: + %5 = load i32, ptr %0, align 4 + ret i32 %5 +} + + +define dso_local noundef i32 @hot(ptr noundef readonly %0) !prof !36 { +; NOPROFILE-LABEL: define dso_local noundef i32 @hot( +; NOPROFILE-SAME: ptr noundef readonly [[TMP0:%.*]]) !prof [[PROF16:![0-9]+]] { +; NOPROFILE-NEXT: [[TMP2:%.*]] = icmp eq ptr [[TMP0]], null +; NOPROFILE-NEXT: br i1 [[TMP2]], label [[TMP3:%.*]], label [[TMP4:%.*]] +; NOPROFILE: 3: +; NOPROFILE-NEXT: tail call void @llvm.ubsantrap(i8 22) +; NOPROFILE-NEXT: unreachable +; NOPROFILE: 4: +; NOPROFILE-NEXT: [[TMP5:%.*]] = load i32, ptr [[TMP0]], align 4 +; NOPROFILE-NEXT: ret i32 [[TMP5]] +; +; ALL-LABEL: define dso_local noundef i32 @hot( +; ALL-SAME: ptr noundef readonly [[TMP0:%.*]]) !prof [[PROF16:![0-9]+]] { +; ALL-NEXT: [[TMP2:%.*]] = icmp eq ptr [[TMP0]], null +; ALL-NEXT: br i1 [[TMP2]], label [[TMP3:%.*]], label [[TMP4:%.*]] +; ALL: 3: +; ALL-NEXT: unreachable +; ALL: 4: +; ALL-NEXT: [[TMP5:%.*]] = load i32, ptr [[TMP0]], align 4 +; ALL-NEXT: ret i32 [[TMP5]] +; +; HOT-LABEL: define dso_local noundef i32 @hot( +; HOT-SAME: ptr noundef readonly [[TMP0:%.*]]) !prof [[PROF16:![0-9]+]] { +; HOT-NEXT: [[TMP2:%.*]] = icmp eq ptr [[TMP0]], null +; HOT-NEXT: br i1 [[TMP2]], label [[TMP3:%.*]], label [[TMP4:%.*]] +; HOT: 3: +; HOT-NEXT: unreachable +; HOT: 4: +; HOT-NEXT: [[TMP5:%.*]] = load i32, ptr [[TMP0]], align 4 +; HOT-NEXT: ret i32 [[TMP5]] +; +; HOT70-LABEL: define dso_local noundef i32 @hot( +; HOT70-SAME: ptr noundef readonly [[TMP0:%.*]]) !prof [[PROF16:![0-9]+]] { +; HOT70-NEXT: [[TMP2:%.*]] = icmp eq ptr [[TMP0]], null +; HOT70-NEXT: br i1 [[TMP2]], label [[TMP3:%.*]], label [[TMP4:%.*]] +; HOT70: 3: +; HOT70-NEXT: tail call void @llvm.ubsantrap(i8 22) +; HOT70-NEXT: unreachable +; HOT70: 4: +; HOT70-NEXT: [[TMP5:%.*]] = load i32, ptr [[TMP0]], align 4 +; HOT70-NEXT: ret i32 [[TMP5]] +; + %2 = icmp eq ptr %0, null + br i1 %2, label %3, label %4 + +3: + tail call void @llvm.ubsantrap(i8 22) + unreachable + +4: + %5 = load i32, ptr %0, align 4 + ret i32 %5 +} + +define dso_local noundef i32 @veryHot(ptr noundef readonly %0) !prof !39 { +; NOPROFILE-LABEL: define dso_local noundef i32 @veryHot( +; NOPROFILE-SAME: ptr noundef readonly [[TMP0:%.*]]) !prof [[PROF17:![0-9]+]] { +; NOPROFILE-NEXT: [[TMP2:%.*]] = icmp eq ptr [[TMP0]], null +; NOPROFILE-NEXT: br i1 [[TMP2]], label [[TMP3:%.*]], label [[TMP4:%.*]] +; NOPROFILE: 3: +; NOPROFILE-NEXT: tail call void @llvm.ubsantrap(i8 22) +; NOPROFILE-NEXT: unreachable +; NOPROFILE: 4: +; NOPROFILE-NEXT: [[TMP5:%.*]] = load i32, ptr [[TMP0]], align 4 +; NOPROFILE-NEXT: ret i32 [[TMP5]] +; +; ALL-LABEL: define dso_local noundef i32 @veryHot( +; ALL-SAME: ptr noundef readonly [[TMP0:%.*]]) !prof [[PROF17:![0-9]+]] { +; ALL-NEXT: [[TMP2:%.*]] = icmp eq ptr [[TMP0]], null +; ALL-NEXT: br i1 [[TMP2]], label [[TMP3:%.*]], label [[TMP4:%.*]] +; ALL: 3: +; ALL-NEXT: unreachable +; ALL: 4: +; ALL-NEXT: [[TMP5:%.*]] = load i32, ptr [[TMP0]], align 4 +; ALL-NEXT: ret i32 [[TMP5]] +; +; HOT-LABEL: define dso_local noundef i32 @veryHot( +; HOT-SAME: ptr noundef readonly [[TMP0:%.*]]) !prof [[PROF17:![0-9]+]] { +; HOT-NEXT: [[TMP2:%.*]] = icmp eq ptr [[TMP0]], null +; HOT-NEXT: br i1 [[TMP2]], label [[TMP3:%.*]], label [[TMP4:%.*]] +; HOT: 3: +; HOT-NEXT: unreachable +; HOT: 4: +; HOT-NEXT: [[TMP5:%.*]] = load i32, ptr [[TMP0]], align 4 +; HOT-NEXT: ret i32 [[TMP5]] +; +; HOT70-LABEL: define dso_local noundef i32 @veryHot( +; HOT70-SAME: ptr noundef readonly [[TMP0:%.*]]) !prof [[PROF17:![0-9]+]] { +; HOT70-NEXT: [[TMP2:%.*]] = icmp eq ptr [[TMP0]], null +; HOT70-NEXT: br i1 [[TMP2]], label [[TMP3:%.*]], label [[TMP4:%.*]] +; HOT70: 3: +; HOT70-NEXT: unreachable +; HOT70: 4: +; HOT70-NEXT: [[TMP5:%.*]] = load i32, ptr [[TMP0]], align 4 +; HOT70-NEXT: ret i32 [[TMP5]] +; + %2 = icmp eq ptr %0, null + br i1 %2, label %3, label %4 + +3: + tail call void @llvm.ubsantrap(i8 22) + unreachable + +4: + %5 = load i32, ptr %0, align 4 + ret i32 %5 +} + + +define dso_local noundef i32 @branchColdFnHot(i32 noundef %0, ptr noundef readonly %1) !prof !39 { +; NOPROFILE-LABEL: define dso_local noundef i32 @branchColdFnHot( +; NOPROFILE-SAME: i32 noundef [[TMP0:%.*]], ptr noundef readonly [[TMP1:%.*]]) !prof [[PROF17]] { +; NOPROFILE-NEXT: [[TMP3:%.*]] = icmp eq i32 [[TMP0]], 0 +; NOPROFILE-NEXT: br i1 [[TMP3]], label [[TMP9:%.*]], label [[TMP4:%.*]], !prof [[PROF18:![0-9]+]] +; NOPROFILE: 4: +; NOPROFILE-NEXT: [[TMP5:%.*]] = icmp eq ptr [[TMP1]], null +; NOPROFILE-NEXT: br i1 [[TMP5]], label [[TMP6:%.*]], label [[TMP7:%.*]] +; NOPROFILE: 6: +; NOPROFILE-NEXT: tail call void @llvm.ubsantrap(i8 22) +; NOPROFILE-NEXT: unreachable +; NOPROFILE: 7: +; NOPROFILE-NEXT: [[TMP8:%.*]] = load i32, ptr [[TMP1]], align 4 +; NOPROFILE-NEXT: br label [[TMP9]] +; NOPROFILE: 9: +; NOPROFILE-NEXT: [[TMP10:%.*]] = phi i32 [ [[TMP8]], [[TMP7]] ], [ 0, [[TMP2:%.*]] ] +; NOPROFILE-NEXT: ret i32 [[TMP10]] +; +; ALL-LABEL: define dso_local noundef i32 @branchColdFnHot( +; ALL-SAME: i32 noundef [[TMP0:%.*]], ptr noundef readonly [[TMP1:%.*]]) !prof [[PROF17]] { +; ALL-NEXT: [[TMP3:%.*]] = icmp eq i32 [[TMP0]], 0 +; ALL-NEXT: br i1 [[TMP3]], label [[TMP9:%.*]], label [[TMP4:%.*]], !prof [[PROF18:![0-9]+]] +; ALL: 4: +; ALL-NEXT: [[TMP5:%.*]] = icmp eq ptr [[TMP1]], null +; ALL-NEXT: br i1 [[TMP5]], label [[TMP6:%.*]], label [[TMP7:%.*]] +; ALL: 6: +; ALL-NEXT: unreachable +; ALL: 7: +; ALL-NEXT: [[TMP8:%.*]] = load i32, ptr [[TMP1]], align 4 +; ALL-NEXT: br label [[TMP9]] +; ALL: 9: +; ALL-NEXT: [[TMP10:%.*]] = phi i32 [ [[TMP8]], [[TMP7]] ], [ 0, [[TMP2:%.*]] ] +; ALL-NEXT: ret i32 [[TMP10]] +; +; HOT-LABEL: define dso_local noundef i32 @branchColdFnHot( +; HOT-SAME: i32 noundef [[TMP0:%.*]], ptr noundef readonly [[TMP1:%.*]]) !prof [[PROF17]] { +; HOT-NEXT: [[TMP3:%.*]] = icmp eq i32 [[TMP0]], 0 +; HOT-NEXT: br i1 [[TMP3]], label [[TMP9:%.*]], label [[TMP4:%.*]], !prof [[PROF18:![0-9]+]] +; HOT: 4: +; HOT-NEXT: [[TMP5:%.*]] = icmp eq ptr [[TMP1]], null +; HOT-NEXT: br i1 [[TMP5]], label [[TMP6:%.*]], label [[TMP7:%.*]] +; HOT: 6: +; HOT-NEXT: tail call void @llvm.ubsantrap(i8 22) +; HOT-NEXT: unreachable +; HOT: 7: +; HOT-NEXT: [[TMP8:%.*]] = load i32, ptr [[TMP1]], align 4 +; HOT-NEXT: br label [[TMP9]] +; HOT: 9: +; HOT-NEXT: [[TMP10:%.*]] = phi i32 [ [[TMP8]], [[TMP7]] ], [ 0, [[TMP2:%.*]] ] +; HOT-NEXT: ret i32 [[TMP10]] +; +; HOT70-LABEL: define dso_local noundef i32 @branchColdFnHot( +; HOT70-SAME: i32 noundef [[TMP0:%.*]], ptr noundef readonly [[TMP1:%.*]]) !prof [[PROF17]] { +; HOT70-NEXT: [[TMP3:%.*]] = icmp eq i32 [[TMP0]], 0 +; HOT70-NEXT: br i1 [[TMP3]], label [[TMP9:%.*]], label [[TMP4:%.*]], !prof [[PROF18:![0-9]+]] +; HOT70: 4: +; HOT70-NEXT: [[TMP5:%.*]] = icmp eq ptr [[TMP1]], null +; HOT70-NEXT: br i1 [[TMP5]], label [[TMP6:%.*]], label [[TMP7:%.*]] +; HOT70: 6: +; HOT70-NEXT: tail call void @llvm.ubsantrap(i8 22) +; HOT70-NEXT: unreachable +; HOT70: 7: +; HOT70-NEXT: [[TMP8:%.*]] = load i32, ptr [[TMP1]], align 4 +; HOT70-NEXT: br label [[TMP9]] +; HOT70: 9: +; HOT70-NEXT: [[TMP10:%.*]] = phi i32 [ [[TMP8]], [[TMP7]] ], [ 0, [[TMP2:%.*]] ] +; HOT70-NEXT: ret i32 [[TMP10]] +; + %3 = icmp eq i32 %0, 0 + br i1 %3, label %9, label %4, !prof !38 + +4: + %5 = icmp eq ptr %1, null + br i1 %5, label %6, label %7 + +6: + tail call void @llvm.ubsantrap(i8 22) #2 + unreachable + +7: + %8 = load i32, ptr %1, align 4 + br label %9 + +9: + %10 = phi i32 [ %8, %7 ], [ 0, %2 ] + ret i32 %10 +} + +define dso_local noundef i32 @branchHotFnCold(i32 noundef %0, ptr noundef readonly %1) !prof !36 { +; NOPROFILE-LABEL: define dso_local noundef i32 @branchHotFnCold( +; NOPROFILE-SAME: i32 noundef [[TMP0:%.*]], ptr noundef readonly [[TMP1:%.*]]) !prof [[PROF16]] { +; NOPROFILE-NEXT: [[TMP3:%.*]] = icmp eq i32 [[TMP0]], 0 +; NOPROFILE-NEXT: br i1 [[TMP3]], label [[TMP9:%.*]], label [[TMP4:%.*]], !prof [[PROF19:![0-9]+]] +; NOPROFILE: 4: +; NOPROFILE-NEXT: [[TMP5:%.*]] = icmp eq ptr [[TMP1]], null +; NOPROFILE-NEXT: br i1 [[TMP5]], label [[TMP6:%.*]], label [[TMP7:%.*]] +; NOPROFILE: 6: +; NOPROFILE-NEXT: tail call void @llvm.ubsantrap(i8 22) +; NOPROFILE-NEXT: unreachable +; NOPROFILE: 7: +; NOPROFILE-NEXT: [[TMP8:%.*]] = load i32, ptr [[TMP1]], align 4 +; NOPROFILE-NEXT: br label [[TMP9]] +; NOPROFILE: 9: +; NOPROFILE-NEXT: [[TMP10:%.*]] = phi i32 [ [[TMP8]], [[TMP7]] ], [ 0, [[TMP2:%.*]] ] +; NOPROFILE-NEXT: ret i32 [[TMP10]] +; +; ALL-LABEL: define dso_local noundef i32 @branchHotFnCold( +; ALL-SAME: i32 noundef [[TMP0:%.*]], ptr noundef readonly [[TMP1:%.*]]) !prof [[PROF16]] { +; ALL-NEXT: [[TMP3:%.*]] = icmp eq i32 [[TMP0]], 0 +; ALL-NEXT: br i1 [[TMP3]], label [[TMP9:%.*]], label [[TMP4:%.*]], !prof [[PROF19:![0-9]+]] +; ALL: 4: +; ALL-NEXT: [[TMP5:%.*]] = icmp eq ptr [[TMP1]], null +; ALL-NEXT: br i1 [[TMP5]], label [[TMP6:%.*]], label [[TMP7:%.*]] +; ALL: 6: +; ALL-NEXT: unreachable +; ALL: 7: +; ALL-NEXT: [[TMP8:%.*]] = load i32, ptr [[TMP1]], align 4 +; ALL-NEXT: br label [[TMP9]] +; ALL: 9: +; ALL-NEXT: [[TMP10:%.*]] = phi i32 [ [[TMP8]], [[TMP7]] ], [ 0, [[TMP2:%.*]] ] +; ALL-NEXT: ret i32 [[TMP10]] +; +; HOT-LABEL: define dso_local noundef i32 @branchHotFnCold( +; HOT-SAME: i32 noundef [[TMP0:%.*]], ptr noundef readonly [[TMP1:%.*]]) !prof [[PROF16]] { +; HOT-NEXT: [[TMP3:%.*]] = icmp eq i32 [[TMP0]], 0 +; HOT-NEXT: br i1 [[TMP3]], label [[TMP9:%.*]], label [[TMP4:%.*]], !prof [[PROF19:![0-9]+]] +; HOT: 4: +; HOT-NEXT: [[TMP5:%.*]] = icmp eq ptr [[TMP1]], null +; HOT-NEXT: br i1 [[TMP5]], label [[TMP6:%.*]], label [[TMP7:%.*]] +; HOT: 6: +; HOT-NEXT: unreachable +; HOT: 7: +; HOT-NEXT: [[TMP8:%.*]] = load i32, ptr [[TMP1]], align 4 +; HOT-NEXT: br label [[TMP9]] +; HOT: 9: +; HOT-NEXT: [[TMP10:%.*]] = phi i32 [ [[TMP8]], [[TMP7]] ], [ 0, [[TMP2:%.*]] ] +; HOT-NEXT: ret i32 [[TMP10]] +; +; HOT70-LABEL: define dso_local noundef i32 @branchHotFnCold( +; HOT70-SAME: i32 noundef [[TMP0:%.*]], ptr noundef readonly [[TMP1:%.*]]) !prof [[PROF16]] { +; HOT70-NEXT: [[TMP3:%.*]] = icmp eq i32 [[TMP0]], 0 +; HOT70-NEXT: br i1 [[TMP3]], label [[TMP9:%.*]], label [[TMP4:%.*]], !prof [[PROF19:![0-9]+]] +; HOT70: 4: +; HOT70-NEXT: [[TMP5:%.*]] = icmp eq ptr [[TMP1]], null +; HOT70-NEXT: br i1 [[TMP5]], label [[TMP6:%.*]], label [[TMP7:%.*]] +; HOT70: 6: +; HOT70-NEXT: tail call void @llvm.ubsantrap(i8 22) +; HOT70-NEXT: unreachable +; HOT70: 7: +; HOT70-NEXT: [[TMP8:%.*]] = load i32, ptr [[TMP1]], align 4 +; HOT70-NEXT: br label [[TMP9]] +; HOT70: 9: +; HOT70-NEXT: [[TMP10:%.*]] = phi i32 [ [[TMP8]], [[TMP7]] ], [ 0, [[TMP2:%.*]] ] +; HOT70-NEXT: ret i32 [[TMP10]] +; + %3 = icmp eq i32 %0, 0 + br i1 %3, label %9, label %4, !prof !37 + +4: + %5 = icmp eq ptr %1, null + br i1 %5, label %6, label %7 + +6: + tail call void @llvm.ubsantrap(i8 22) #2 + unreachable + +7: + %8 = load i32, ptr %1, align 4 + br label %9 + +9: + %10 = phi i32 [ %8, %7 ], [ 0, %2 ] + ret i32 %10 +} + +!llvm.module.flags = !{!6} +!6 = !{i32 1, !"ProfileSummary", !7} +!7 = !{!8, !9, !10, !11, !12, !13, !14, !17} +!8 = !{!"ProfileFormat", !"InstrProf"} +!9 = !{!"TotalCount", i64 30000} +!10 = !{!"MaxCount", i64 10000} +!11 = !{!"MaxInternalCount", i64 10000} +!12 = !{!"MaxFunctionCount", i64 10000} +!13 = !{!"NumCounts", i64 3} +!14 = !{!"NumFunctions", i64 5} +!17 = !{!"DetailedSummary", !18} +!18 = !{!19, !29, !30, !32, !34} +!19 = !{i32 10000, i64 10000, i32 3} +!29 = !{i32 950000, i64 5000, i32 3} +!30 = !{i32 990000, i64 500, i32 4} +!32 = !{i32 999900, i64 250, i32 4} +!34 = !{i32 999999, i64 1, i32 6} + +!36 = !{!"function_entry_count", i64 1000} +!39 = !{!"function_entry_count", i64 7000} + +!37 = !{!"branch_weights", i32 1, i32 1000} +!38 = !{!"branch_weights", i32 1000, i32 1} + +;. +; NOPROFILE: [[PROF16]] = !{!"function_entry_count", i64 1000} +; NOPROFILE: [[PROF17]] = !{!"function_entry_count", i64 7000} +; NOPROFILE: [[PROF18]] = !{!"branch_weights", i32 1000, i32 1} +; NOPROFILE: [[PROF19]] = !{!"branch_weights", i32 1, i32 1000} +;. +; ALL: [[PROF16]] = !{!"function_entry_count", i64 1000} +; ALL: [[PROF17]] = !{!"function_entry_count", i64 7000} +; ALL: [[PROF18]] = !{!"branch_weights", i32 1000, i32 1} +; ALL: [[PROF19]] = !{!"branch_weights", i32 1, i32 1000} +;. +; HOT: [[PROF16]] = !{!"function_entry_count", i64 1000} +; HOT: [[PROF17]] = !{!"function_entry_count", i64 7000} +; HOT: [[PROF18]] = !{!"branch_weights", i32 1000, i32 1} +; HOT: [[PROF19]] = !{!"branch_weights", i32 1, i32 1000} +;. +; HOT70: [[PROF16]] = !{!"function_entry_count", i64 1000} +; HOT70: [[PROF17]] = !{!"function_entry_count", i64 7000} +; HOT70: [[PROF18]] = !{!"branch_weights", i32 1000, i32 1} +; HOT70: [[PROF19]] = !{!"branch_weights", i32 1, i32 1000} +;. From a920b29aede06a293d274d4078ddf866d787baa4 Mon Sep 17 00:00:00 2001 From: Vitaly Buka Date: Mon, 4 Mar 2024 17:45:38 -0800 Subject: [PATCH 2/6] Update --- .../Instrumentation/RemoveTrapsPass.cpp | 57 ++++++++++--------- .../Transforms/RemoveTraps/remove-traps.ll | 2 +- 2 files changed, 32 insertions(+), 27 deletions(-) diff --git a/llvm/lib/Transforms/Instrumentation/RemoveTrapsPass.cpp b/llvm/lib/Transforms/Instrumentation/RemoveTrapsPass.cpp index 7a7c604741e92..52dba021b3c7f 100644 --- a/llvm/lib/Transforms/Instrumentation/RemoveTrapsPass.cpp +++ b/llvm/lib/Transforms/Instrumentation/RemoveTrapsPass.cpp @@ -5,9 +5,6 @@ // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception // //===----------------------------------------------------------------------===// -// -// -//===----------------------------------------------------------------------===// #include "llvm/Transforms/Instrumentation/RemoveTrapsPass.h" @@ -17,37 +14,31 @@ #include "llvm/IR/Instructions.h" #include "llvm/IR/IntrinsicInst.h" #include "llvm/IR/Intrinsics.h" -#include +#include "llvm/Support/RandomNumberGenerator.h" +#include +#include using namespace llvm; #define DEBUG_TYPE "remove-traps" -static constexpr unsigned MaxRandomRate = 1000; - static cl::opt HotPercentileCutoff( "remove-traps-percentile-cutoff-hot", cl::init(0), cl::desc("Alternative hot percentile cuttoff. By default " "`-profile-summary-cutoff-hot` is used.")); + static cl::opt RandomRate( "remove-traps-random-rate", cl::init(0.0), - cl::desc( - "Probability to use for pseudorandom unconditional checks removal.")); + cl::desc("Probability of unconditional pseudorandom checks removal.")); STATISTIC(NumChecksTotal, "Number of checks"); STATISTIC(NumChecksRemoved, "Number of removed checks"); -static SmallVector -removeUbsanTraps(Function &F, FunctionAnalysisManager &FAM, - ProfileSummaryInfo *PSI) { +static bool removeUbsanTraps(Function &F, const BlockFrequencyInfo &BFI, + const ProfileSummaryInfo *PSI) { SmallVector Remove; + std::unique_ptr Rng; - if (F.isDeclaration()) - return {}; - - auto &BFI = FAM.getResult(F); - - int BBCounter = 0; for (BasicBlock &BB : F) { for (Instruction &I : BB) { IntrinsicInst *II = dyn_cast(&I); @@ -65,18 +56,34 @@ removeUbsanTraps(Function &F, FunctionAnalysisManager &FAM, Count += BFI.getBlockProfileCount(PR).value_or(0); IsHot = HotPercentileCutoff.getNumOccurrences() - ? PSI->isHotCountNthPercentile(HotPercentileCutoff, Count) + ? (HotPercentileCutoff > 0 && + PSI->isHotCountNthPercentile(HotPercentileCutoff, Count)) : PSI->isHotCount(Count); } - if ((IsHot) || ((F.getGUID() + BBCounter++) % MaxRandomRate) < - RandomRate * RandomRate) { + auto ShouldRemove = [&]() { + if (IsHot) + return true; + if (!Rng) { + if (!RandomRate.getNumOccurrences()) + return false; + Rng = F.getParent()->createRNG(F.getName()); + } + std::bernoulli_distribution D(RandomRate); + return D(*Rng); + }; + + if (ShouldRemove()) { Remove.push_back(II); ++NumChecksRemoved; } } } - return Remove; + + for (auto *I : Remove) + I->eraseFromParent(); + + return !Remove.empty(); } PreservedAnalyses RemoveTrapsPass::run(Function &F, @@ -86,10 +93,8 @@ PreservedAnalyses RemoveTrapsPass::run(Function &F, auto &MAMProxy = AM.getResult(F); ProfileSummaryInfo *PSI = MAMProxy.getCachedResult(*F.getParent()); + BlockFrequencyInfo &BFI = AM.getResult(F); - auto Remove = removeUbsanTraps(F, AM, PSI); - for (auto *I : Remove) - I->eraseFromParent(); - - return Remove.empty() ? PreservedAnalyses::all() : PreservedAnalyses::none(); + return removeUbsanTraps(F, BFI, PSI) ? PreservedAnalyses::none() + : PreservedAnalyses::all(); } diff --git a/llvm/test/Transforms/RemoveTraps/remove-traps.ll b/llvm/test/Transforms/RemoveTraps/remove-traps.ll index af9cdab2f2920..71549e7d9b412 100644 --- a/llvm/test/Transforms/RemoveTraps/remove-traps.ll +++ b/llvm/test/Transforms/RemoveTraps/remove-traps.ll @@ -1,6 +1,6 @@ ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 4 ; RUN: opt < %s -passes='function(remove-traps)' -S | FileCheck %s --check-prefixes=NOPROFILE -; RUN: opt < %s -passes='function(remove-traps)' -remove-traps-random-rate=999999 -S | FileCheck %s --check-prefixes=ALL +; RUN: opt < %s -passes='function(remove-traps)' -remove-traps-random-rate=1 -S | FileCheck %s --check-prefixes=ALL ; RUN: opt < %s -passes='require,function(remove-traps)' -S | FileCheck %s --check-prefixes=HOT ; RUN: opt < %s -passes='require,function(remove-traps)' -remove-traps-percentile-cutoff-hot=700000 -S | FileCheck %s --check-prefixes=HOT70 From 71eae68d288d1540585454b41d172607b787c970 Mon Sep 17 00:00:00 2001 From: Vitaly Buka Date: Mon, 4 Mar 2024 17:53:10 -0800 Subject: [PATCH 3/6] Update --- .../Instrumentation/RemoveTrapsPass.cpp | 68 ++++++++++--------- 1 file changed, 37 insertions(+), 31 deletions(-) diff --git a/llvm/lib/Transforms/Instrumentation/RemoveTrapsPass.cpp b/llvm/lib/Transforms/Instrumentation/RemoveTrapsPass.cpp index 52dba021b3c7f..7a35134c8b65d 100644 --- a/llvm/lib/Transforms/Instrumentation/RemoveTrapsPass.cpp +++ b/llvm/lib/Transforms/Instrumentation/RemoveTrapsPass.cpp @@ -39,48 +39,54 @@ static bool removeUbsanTraps(Function &F, const BlockFrequencyInfo &BFI, SmallVector Remove; std::unique_ptr Rng; + auto ShouldRemove = [&](bool IsHot) { + if (IsHot && !RandomRate.getNumOccurrences()) + return true; + if (!Rng) { + if (!RandomRate.getNumOccurrences()) + return false; + Rng = F.getParent()->createRNG(F.getName()); + } + std::bernoulli_distribution D(RandomRate); + return D(*Rng); + }; + for (BasicBlock &BB : F) { for (Instruction &I : BB) { IntrinsicInst *II = dyn_cast(&I); if (!II) continue; auto ID = II->getIntrinsicID(); - if (ID != Intrinsic::ubsantrap) - continue; - ++NumChecksTotal; - - bool IsHot = false; - if (PSI) { - uint64_t Count = 0; - for (const auto *PR : predecessors(&BB)) - Count += BFI.getBlockProfileCount(PR).value_or(0); - - IsHot = HotPercentileCutoff.getNumOccurrences() - ? (HotPercentileCutoff > 0 && - PSI->isHotCountNthPercentile(HotPercentileCutoff, Count)) - : PSI->isHotCount(Count); - } - - auto ShouldRemove = [&]() { - if (IsHot) - return true; - if (!Rng) { - if (!RandomRate.getNumOccurrences()) - return false; - Rng = F.getParent()->createRNG(F.getName()); + switch (ID) { + case Intrinsic::ubsantrap: { + ++NumChecksTotal; + + bool IsHot = false; + if (PSI) { + uint64_t Count = 0; + for (const auto *PR : predecessors(&BB)) + Count += BFI.getBlockProfileCount(PR).value_or(0); + + IsHot = + HotPercentileCutoff.getNumOccurrences() + ? (HotPercentileCutoff > 0 && + PSI->isHotCountNthPercentile(HotPercentileCutoff, Count)) + : PSI->isHotCount(Count); } - std::bernoulli_distribution D(RandomRate); - return D(*Rng); - }; - if (ShouldRemove()) { - Remove.push_back(II); - ++NumChecksRemoved; + if (ShouldRemove(IsHot)) { + Remove.push_back(II); + ++NumChecksRemoved; + } + break; + } + default: + break; } } } - for (auto *I : Remove) + for (IntrinsicInst *I : Remove) I->eraseFromParent(); return !Remove.empty(); @@ -96,5 +102,5 @@ PreservedAnalyses RemoveTrapsPass::run(Function &F, BlockFrequencyInfo &BFI = AM.getResult(F); return removeUbsanTraps(F, BFI, PSI) ? PreservedAnalyses::none() - : PreservedAnalyses::all(); + : PreservedAnalyses::all(); } From d2ed11b872af6d72d52132bbb575ea42768395b0 Mon Sep 17 00:00:00 2001 From: Vitaly Buka Date: Wed, 6 Mar 2024 09:51:49 -0800 Subject: [PATCH 4/6] Simplify ShouldRemove --- llvm/lib/Transforms/Instrumentation/RemoveTrapsPass.cpp | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/llvm/lib/Transforms/Instrumentation/RemoveTrapsPass.cpp b/llvm/lib/Transforms/Instrumentation/RemoveTrapsPass.cpp index 7a35134c8b65d..6e909ae85ca05 100644 --- a/llvm/lib/Transforms/Instrumentation/RemoveTrapsPass.cpp +++ b/llvm/lib/Transforms/Instrumentation/RemoveTrapsPass.cpp @@ -40,13 +40,10 @@ static bool removeUbsanTraps(Function &F, const BlockFrequencyInfo &BFI, std::unique_ptr Rng; auto ShouldRemove = [&](bool IsHot) { - if (IsHot && !RandomRate.getNumOccurrences()) - return true; - if (!Rng) { - if (!RandomRate.getNumOccurrences()) - return false; + if (!RandomRate.getNumOccurrences()) + return IsHot; + if (!Rng) Rng = F.getParent()->createRNG(F.getName()); - } std::bernoulli_distribution D(RandomRate); return D(*Rng); }; From e011a2800c4e5ebc5239afd463869dff082d9889 Mon Sep 17 00:00:00 2001 From: Vitaly Buka Date: Wed, 6 Mar 2024 09:57:39 -0800 Subject: [PATCH 5/6] Add range to prob --- llvm/lib/Transforms/Instrumentation/RemoveTrapsPass.cpp | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/llvm/lib/Transforms/Instrumentation/RemoveTrapsPass.cpp b/llvm/lib/Transforms/Instrumentation/RemoveTrapsPass.cpp index 6e909ae85ca05..de25c22e76bdb 100644 --- a/llvm/lib/Transforms/Instrumentation/RemoveTrapsPass.cpp +++ b/llvm/lib/Transforms/Instrumentation/RemoveTrapsPass.cpp @@ -27,9 +27,10 @@ static cl::opt HotPercentileCutoff( cl::desc("Alternative hot percentile cuttoff. By default " "`-profile-summary-cutoff-hot` is used.")); -static cl::opt RandomRate( - "remove-traps-random-rate", cl::init(0.0), - cl::desc("Probability of unconditional pseudorandom checks removal.")); +static cl::opt + RandomRate("remove-traps-random-rate", cl::init(0.0), + cl::desc("Probability value in the range [0.0, 1.0] of " + "unconditional pseudorandom checks removal.")); STATISTIC(NumChecksTotal, "Number of checks"); STATISTIC(NumChecksRemoved, "Number of removed checks"); From c5460c5a5fb857aaca893cf462ba29ba32aaca0e Mon Sep 17 00:00:00 2001 From: Vitaly Buka Date: Wed, 6 Mar 2024 10:00:07 -0800 Subject: [PATCH 6/6] comment --- llvm/lib/Transforms/Instrumentation/RemoveTrapsPass.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/llvm/lib/Transforms/Instrumentation/RemoveTrapsPass.cpp b/llvm/lib/Transforms/Instrumentation/RemoveTrapsPass.cpp index de25c22e76bdb..d87f7482a21d2 100644 --- a/llvm/lib/Transforms/Instrumentation/RemoveTrapsPass.cpp +++ b/llvm/lib/Transforms/Instrumentation/RemoveTrapsPass.cpp @@ -30,7 +30,7 @@ static cl::opt HotPercentileCutoff( static cl::opt RandomRate("remove-traps-random-rate", cl::init(0.0), cl::desc("Probability value in the range [0.0, 1.0] of " - "unconditional pseudorandom checks removal.")); + "unconditional pseudo-random checks removal.")); STATISTIC(NumChecksTotal, "Number of checks"); STATISTIC(NumChecksRemoved, "Number of removed checks");