1//===-- CopyProf.cpp ------------------------------------------------------===//
2//
3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4// See https://llvm.org/LICENSE.txt for license information.
5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6//
7//===----------------------------------------------------------------------===//
8///
9/// This file implements the LLVM IR instrumentation passes for CopyProf.
10/// It adds enter/exit callbacks to C++ special member functions, and
11/// instruments store instructions.
12///
13/// The basic idea of the CopyProf algorithm works like this:
14/// An object copy Y is made from original object X. The shadow memory
15/// corresponding to (and owned by) Y is marked as "copied". Any subsequent
16/// memory store to the memory corresponding to Y marks the shadow memory as
17/// "modified". When Y is destroyed and all of its corresponding shadow memory
18/// is marked as "copied", the object is reported as an unnecessary copy.
19///
20//===----------------------------------------------------------------------===//
21
22#include "llvm/Transforms/Instrumentation/CopyProf.h"
23
24#include "llvm/ADT/SmallVector.h"
25#include "llvm/ADT/StringExtras.h"
26#include "llvm/IR/Attributes.h"
27#include "llvm/IR/DerivedTypes.h"
28#include "llvm/IR/Function.h"
29#include "llvm/IR/IRBuilder.h"
30#include "llvm/IR/Instruction.h"
31#include "llvm/IR/Instructions.h"
32#include "llvm/IR/Module.h"
33#include "llvm/IR/PassManager.h"
34#include "llvm/Support/Casting.h"
35#include "llvm/Transforms/Utils/Instrumentation.h"
36#include "llvm/Transforms/Utils/ModuleUtils.h"
37#include <array>
38#include <cstddef>
39#include <cstdint>
40
41// TODO: Convert CopyProfPass and CopyProfStoresPass to module passes so that
42// the runtime callbacks can be cached, thus avoiding repetitive symbol table
43// lookups.
44
45using namespace llvm;
46
47// Names for the module c'tor to initialize the runtime, and the runtime
48// initialization function itself.
49constexpr StringRef CopyProfModuleCtorName = "copyprof.module_ctor";
50constexpr StringRef CopyProfInitName = "__copyprof_init_once";
51
52// Runtime callback function names.
53constexpr StringRef CopyProfCtorEnterCallbackName =
54 "__copyprof_ctor_enter_callback";
55constexpr StringRef CopyProfCtorExitCallbackName =
56 "__copyprof_ctor_exit_callback";
57constexpr StringRef CopyProfCopyCtorEnterCallbackName =
58 "__copyprof_copy_ctor_enter_callback";
59constexpr StringRef CopyProfCopyCtorExitCallbackName =
60 "__copyprof_copy_ctor_exit_callback";
61constexpr StringRef CopyProfCopyAssignOpEnterCallbackName =
62 "__copyprof_copy_assign_op_enter_callback";
63constexpr StringRef CopyProfCopyAssignOpExitCallbackName =
64 "__copyprof_copy_assign_op_exit_callback";
65constexpr StringRef CopyProfDtorEnterCallbackName =
66 "__copyprof_dtor_enter_callback";
67constexpr StringRef CopyProfDtorExitCallbackName =
68 "__copyprof_dtor_exit_callback";
69constexpr StringRef CopyProfStoreCallbackName = "__copyprof_store_callback";
70
71// Attribute strings used by the frontend to mark special member functions.
72constexpr StringRef CopyProfCtorAttr = "copyprof-ctor";
73constexpr StringRef CopyProfCopyCtorAttr = "copyprof-copy-ctor";
74constexpr StringRef CopyProfCopyAssignAttr = "copyprof-copy-assign-op";
75constexpr StringRef CopyProfDtorAttr = "copyprof-dtor";
76
77static bool insertModuleCtor(Module &M) {
78 bool Modified = false;
79 getOrCreateSanitizerCtorAndInitFunctions(
80 M, CtorName: CopyProfModuleCtorName, InitName: CopyProfInitName,
81 /*InitArgTypes=*/{},
82 /*InitArgs=*/{}, FunctionsCreatedCallback: [&](Function *Ctor, FunctionCallee) {
83 // Mark the ctor so it's never instrumented itself.
84 Ctor->addFnAttr(Kind: Attribute::DisableSanitizerInstrumentation);
85 appendToGlobalCtors(M, F: Ctor, Priority: 0);
86 Modified = true;
87 });
88 return Modified;
89}
90
91static bool isCopyProfCandidate(const Function &F) {
92 // Must not instrument functions that are explicitly disallowed for
93 // instrumentation, or naked functions.
94 if (F.isDeclaration() ||
95 F.hasFnAttribute(Kind: Attribute::DisableSanitizerInstrumentation) ||
96 F.hasFnAttribute(Kind: Attribute::Naked))
97 return false;
98
99 if (!F.hasFnAttribute(Kind: CopyProfCtorAttr) &&
100 !F.hasFnAttribute(Kind: CopyProfCopyCtorAttr) &&
101 !F.hasFnAttribute(Kind: CopyProfCopyAssignAttr) &&
102 !F.hasFnAttribute(Kind: CopyProfDtorAttr))
103 return false;
104
105 // Don't instrument a function at all if it ends in a tail call.
106 // Alternatively, the exit callback could be placed before the tail call, but
107 // that would risk missing observable side-effects needed by CopyProf to infer
108 // memory ownership (potentially leading to false positive reports).
109 // For example, if the tail would deallocate memory then CopyProf would be
110 // unable to inspect that memory and the object could be misclassified as
111 // having been unnecessarily copied. Skipping functions ending in musttail
112 // calls therefore favors false negatives over false positives.
113 for (const BasicBlock &BB : F)
114 if (BB.getTerminatingMustTailCall())
115 return false;
116
117 return true;
118}
119
120static bool isCopyProfStoresCandidate(const Function &F) {
121 return !F.isDeclaration() &&
122 !F.hasFnAttribute(Kind: Attribute::DisableSanitizerInstrumentation) &&
123 !F.hasFnAttribute(Kind: Attribute::Naked);
124}
125
126// Returns the object size in bytes that was stored in the given function
127// attribute during parsing in the frontend.
128static size_t getAttrValueAsInt(const Function &F, StringRef Attr) {
129 size_t IntValue = 0;
130 [[maybe_unused]] bool Success =
131 to_integer<size_t>(S: F.getFnAttribute(Kind: Attr).getValueAsString(), Num&: IntValue,
132 /*Base=*/10);
133 assert(Success &&
134 "Unable to parse object size from CopyProf function attribute value.");
135 return IntValue;
136}
137
138namespace {
139
140// Instruments special member functions to call into the CopyProf runtime.
141class CopyProf {
142public:
143 explicit CopyProf(Module &M);
144 bool instrumentFunction(Function &F);
145
146private:
147 void insertCallback(Function &F, size_t ObjSize, unsigned NumArgs,
148 FunctionCallee Callback, FunctionCallee ExitCallback);
149
150 Type *IntPtrTy;
151 FunctionCallee CtorEnterCallback;
152 FunctionCallee CtorExitCallback;
153 FunctionCallee CopyCtorEnterCallback;
154 FunctionCallee CopyCtorExitCallback;
155 FunctionCallee CopyAssignOpEnterCallback;
156 FunctionCallee CopyAssignOpExitCallback;
157 FunctionCallee DtorEnterCallback;
158 FunctionCallee DtorExitCallback;
159};
160
161// Late-stage pass that instruments store instructions after all optimizations
162// have run (to avoid instrumenting stores that would be eliminated).
163class CopyProfStores {
164public:
165 explicit CopyProfStores(Module &M);
166 bool instrumentFunction(Function &F);
167
168private:
169 Type *IntPtrTy;
170 FunctionCallee StoreCallback;
171};
172
173} // namespace
174
175CopyProf::CopyProf(Module &M) {
176 LLVMContext &Ctx = M.getContext();
177 IRBuilder<> IRB(M);
178 IntPtrTy = IRB.getIntPtrTy(DL: M.getDataLayout());
179 Type *PtrTy = IRB.getPtrTy();
180 Type *VoidTy = IRB.getVoidTy();
181 // CopyProf callbacks never throw exceptions.
182 AttributeList Attr;
183 Attr = Attr.addFnAttribute(C&: Ctx, Kind: Attribute::NoUnwind);
184 CtorEnterCallback = M.getOrInsertFunction(Name: CopyProfCtorEnterCallbackName, AttributeList: Attr,
185 RetTy: VoidTy, Args: PtrTy, Args: IntPtrTy);
186 CtorExitCallback = M.getOrInsertFunction(Name: CopyProfCtorExitCallbackName, AttributeList: Attr,
187 RetTy: VoidTy, Args: PtrTy, Args: IntPtrTy);
188 CopyCtorEnterCallback = M.getOrInsertFunction(
189 Name: CopyProfCopyCtorEnterCallbackName, AttributeList: Attr, RetTy: VoidTy, Args: PtrTy, Args: PtrTy, Args: IntPtrTy);
190 CopyCtorExitCallback = M.getOrInsertFunction(
191 Name: CopyProfCopyCtorExitCallbackName, AttributeList: Attr, RetTy: VoidTy, Args: PtrTy, Args: PtrTy, Args: IntPtrTy);
192 CopyAssignOpEnterCallback =
193 M.getOrInsertFunction(Name: CopyProfCopyAssignOpEnterCallbackName, AttributeList: Attr, RetTy: VoidTy,
194 Args: PtrTy, Args: PtrTy, Args: IntPtrTy);
195 CopyAssignOpExitCallback =
196 M.getOrInsertFunction(Name: CopyProfCopyAssignOpExitCallbackName, AttributeList: Attr, RetTy: VoidTy,
197 Args: PtrTy, Args: PtrTy, Args: IntPtrTy);
198 DtorEnterCallback = M.getOrInsertFunction(Name: CopyProfDtorEnterCallbackName, AttributeList: Attr,
199 RetTy: VoidTy, Args: PtrTy, Args: IntPtrTy);
200 DtorExitCallback = M.getOrInsertFunction(Name: CopyProfDtorExitCallbackName, AttributeList: Attr,
201 RetTy: VoidTy, Args: PtrTy, Args: IntPtrTy);
202}
203
204bool CopyProf::instrumentFunction(Function &F) {
205 bool Modified = true;
206 if (F.hasFnAttribute(Kind: CopyProfCtorAttr))
207 insertCallback(F, ObjSize: getAttrValueAsInt(F, Attr: CopyProfCtorAttr), /*NumArgs=*/1,
208 Callback: CtorEnterCallback, ExitCallback: CtorExitCallback);
209 else if (F.hasFnAttribute(Kind: CopyProfCopyCtorAttr))
210 insertCallback(F, ObjSize: getAttrValueAsInt(F, Attr: CopyProfCopyCtorAttr), /*NumArgs=*/2,
211 Callback: CopyCtorEnterCallback, ExitCallback: CopyCtorExitCallback);
212 else if (F.hasFnAttribute(Kind: CopyProfCopyAssignAttr))
213 insertCallback(F, ObjSize: getAttrValueAsInt(F, Attr: CopyProfCopyAssignAttr),
214 /*NumArgs=*/2, Callback: CopyAssignOpEnterCallback,
215 ExitCallback: CopyAssignOpExitCallback);
216 else if (F.hasFnAttribute(Kind: CopyProfDtorAttr))
217 insertCallback(F, ObjSize: getAttrValueAsInt(F, Attr: CopyProfDtorAttr), /*NumArgs=*/1,
218 Callback: DtorEnterCallback, ExitCallback: DtorExitCallback);
219 else
220 Modified = false;
221
222 return Modified;
223}
224
225void CopyProf::insertCallback(Function &F, size_t ObjSize, unsigned NumArgs,
226 FunctionCallee EntryCallback,
227 FunctionCallee ExitCallback) {
228 auto InsertCallback = [IntPtrTy = IntPtrTy, ObjSize,
229 NumArgs](Function &F, InstrumentationIRBuilder &&IRB,
230 FunctionCallee Callback) {
231 SmallVector<Value *, 3> Args;
232 // `this` is always the first argument to a special member function, but
233 // copy c'tor / copy assignment operator will have the other `this` ptr
234 // passed as their second argument.
235 assert(NumArgs == 1 || NumArgs == 2);
236 for (unsigned I = 0; I < NumArgs; ++I)
237 Args.push_back(Elt: F.getArg(i: I));
238 // The last argument to the callback is the static size of the object
239 // pointed at by `this`.
240 Args.push_back(Elt: ConstantInt::get(Ty: IntPtrTy, V: ObjSize));
241 IRB.CreateCall(Callee: Callback, Args);
242 };
243
244 InsertCallback(
245 F,
246 InstrumentationIRBuilder{F.getEntryBlock().getFirstNonPHIOrDbgOrAlloca()},
247 EntryCallback);
248 for (BasicBlock &BB : F) {
249 Instruction *Term = BB.getTerminator();
250 if (isa<ReturnInst>(Val: Term) || isa<ResumeInst>(Val: Term))
251 InsertCallback(F, InstrumentationIRBuilder{Term}, ExitCallback);
252 }
253}
254
255CopyProfStores::CopyProfStores(Module &M) {
256 LLVMContext &Ctx = M.getContext();
257 IRBuilder<> IRB(M);
258 IntPtrTy = IRB.getIntPtrTy(DL: M.getDataLayout());
259 Type *PtrTy = IRB.getPtrTy();
260 Type *VoidTy = IRB.getVoidTy();
261 // CopyProf callbacks never throw exceptions.
262 AttributeList Attr;
263 Attr = Attr.addFnAttribute(C&: Ctx, Kind: Attribute::NoUnwind);
264 StoreCallback = M.getOrInsertFunction(Name: CopyProfStoreCallbackName, AttributeList: Attr, RetTy: VoidTy,
265 Args: PtrTy, Args: IntPtrTy);
266}
267
268bool CopyProfStores::instrumentFunction(Function &F) {
269 // TODO: Handle all types of memory stores (memory intrinsics, masked store
270 // intrinsics, AtomicRMW, and AtomicCmpXchg).
271 // TODO: Skip stores to alloca if only made of fundamental types, arrays
272 // thereof and (possibly) class types that are trivial and aggregate.
273 const DataLayout &DL = F.getDataLayout();
274 SmallVector<StoreInst *, 16> ToInstrument;
275 for (BasicBlock &BB : F) {
276 for (Instruction &I : BB) {
277 if (auto *SI = dyn_cast<StoreInst>(Val: &I);
278 SI != nullptr && SI->getPointerAddressSpace() == 0 &&
279 !SI->hasMetadata(KindID: LLVMContext::MD_nosanitize) &&
280 // Scalable vector stores have no compile-time-constant size so skip
281 // them.
282 !DL.getTypeStoreSize(Ty: SI->getValueOperand()->getType()).isScalable())
283 ToInstrument.push_back(Elt: SI);
284 }
285 }
286 if (ToInstrument.empty())
287 return false;
288
289 for (StoreInst *SI : ToInstrument) {
290 uint64_t StoredSize =
291 DL.getTypeStoreSize(Ty: SI->getValueOperand()->getType()).getFixedValue();
292 InstrumentationIRBuilder IRB(SI);
293 std::array<Value *, 2> Args = {SI->getPointerOperand(),
294 ConstantInt::get(Ty: IntPtrTy, V: StoredSize)};
295 IRB.CreateCall(Callee: StoreCallback, Args);
296 }
297 return true;
298}
299
300PreservedAnalyses CopyProfPass::run(Function &F, FunctionAnalysisManager &) {
301 if (!isCopyProfCandidate(F))
302 return PreservedAnalyses::all();
303 CopyProf Impl(*F.getParent());
304 if (!Impl.instrumentFunction(F))
305 return PreservedAnalyses::all();
306 PreservedAnalyses PA;
307 PA.preserveSet<CFGAnalyses>();
308 return PA;
309}
310
311PreservedAnalyses ModuleCopyProfPass::run(Module &M, ModuleAnalysisManager &) {
312 return insertModuleCtor(M) ? PreservedAnalyses::none()
313 : PreservedAnalyses::all();
314}
315
316PreservedAnalyses CopyProfStoresPass::run(Function &F,
317 FunctionAnalysisManager &) {
318 if (!isCopyProfStoresCandidate(F))
319 return PreservedAnalyses::all();
320 CopyProfStores Impl(*F.getParent());
321 if (!Impl.instrumentFunction(F))
322 return PreservedAnalyses::all();
323 PreservedAnalyses PA;
324 PA.preserveSet<CFGAnalyses>();
325 return PA;
326}
327