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";
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(Ctx);
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(),
247 F.getEntryBlock().getFirstNonPHIOrDbgOrAlloca()},
248 EntryCallback);
249 for (BasicBlock &BB : F) {
250 Instruction *Term = BB.getTerminator();
251 if (isa<ReturnInst>(Val: Term) || isa<ResumeInst>(Val: Term))
252 InsertCallback(F, InstrumentationIRBuilder{Term}, ExitCallback);
253 }
254}
255
256CopyProfStores::CopyProfStores(Module &M) {
257 LLVMContext &Ctx = M.getContext();
258 IRBuilder<> IRB(Ctx);
259 IntPtrTy = IRB.getIntPtrTy(DL: M.getDataLayout());
260 Type *PtrTy = IRB.getPtrTy();
261 Type *VoidTy = IRB.getVoidTy();
262 // CopyProf callbacks never throw exceptions.
263 AttributeList Attr;
264 Attr = Attr.addFnAttribute(C&: Ctx, Kind: Attribute::NoUnwind);
265 StoreCallback = M.getOrInsertFunction(Name: CopyProfStoreCallbackName, AttributeList: Attr, RetTy: VoidTy,
266 Args: PtrTy, Args: IntPtrTy);
267}
268
269bool CopyProfStores::instrumentFunction(Function &F) {
270 // TODO: Handle all types of memory stores (memory intrinsics, masked store
271 // intrinsics, AtomicRMW, and AtomicCmpXchg).
272 // TODO: Skip stores to alloca if only made of fundamental types, arrays
273 // thereof and (possibly) class types that are trivial and aggregate.
274 const DataLayout &DL = F.getParent()->getDataLayout();
275 SmallVector<StoreInst *, 16> ToInstrument;
276 for (BasicBlock &BB : F) {
277 for (Instruction &I : BB) {
278 if (auto *SI = dyn_cast<StoreInst>(Val: &I);
279 SI != nullptr && SI->getPointerAddressSpace() == 0 &&
280 !SI->hasMetadata(KindID: LLVMContext::MD_nosanitize) &&
281 // Scalable vector stores have no compile-time-constant size so skip
282 // them.
283 !DL.getTypeStoreSize(Ty: SI->getValueOperand()->getType()).isScalable())
284 ToInstrument.push_back(Elt: SI);
285 }
286 }
287 if (ToInstrument.empty())
288 return false;
289
290 for (StoreInst *SI : ToInstrument) {
291 uint64_t StoredSize =
292 DL.getTypeStoreSize(Ty: SI->getValueOperand()->getType()).getFixedValue();
293 InstrumentationIRBuilder IRB(SI);
294 std::array<Value *, 2> Args = {SI->getPointerOperand(),
295 ConstantInt::get(Ty: IntPtrTy, V: StoredSize)};
296 IRB.CreateCall(Callee: StoreCallback, Args);
297 }
298 return true;
299}
300
301PreservedAnalyses CopyProfPass::run(Function &F, FunctionAnalysisManager &) {
302 if (!isCopyProfCandidate(F))
303 return PreservedAnalyses::all();
304 CopyProf Impl(*F.getParent());
305 if (!Impl.instrumentFunction(F))
306 return PreservedAnalyses::all();
307 PreservedAnalyses PA;
308 PA.preserveSet<CFGAnalyses>();
309 return PA;
310}
311
312PreservedAnalyses ModuleCopyProfPass::run(Module &M, ModuleAnalysisManager &) {
313 return insertModuleCtor(M) ? PreservedAnalyses::none()
314 : PreservedAnalyses::all();
315}
316
317PreservedAnalyses CopyProfStoresPass::run(Function &F,
318 FunctionAnalysisManager &) {
319 if (!isCopyProfStoresCandidate(F))
320 return PreservedAnalyses::all();
321 CopyProfStores Impl(*F.getParent());
322 if (!Impl.instrumentFunction(F))
323 return PreservedAnalyses::all();
324 PreservedAnalyses PA;
325 PA.preserveSet<CFGAnalyses>();
326 return PA;
327}
328