|
| 1 | +//===-- OffloadSanitizer.cpp - Offload sanitizer --------------------------===// |
| 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 | +//===----------------------------------------------------------------------===// |
| 10 | + |
| 11 | +#include "llvm/Transforms/Instrumentation/OffloadSanitizer.h" |
| 12 | + |
| 13 | +#include "llvm/ADT/SetVector.h" |
| 14 | +#include "llvm/ADT/SmallVector.h" |
| 15 | +#include "llvm/IR/DebugInfoMetadata.h" |
| 16 | +#include "llvm/IR/DerivedTypes.h" |
| 17 | +#include "llvm/IR/IRBuilder.h" |
| 18 | +#include "llvm/IR/InstIterator.h" |
| 19 | +#include "llvm/IR/Instructions.h" |
| 20 | +#include "llvm/IR/IntrinsicInst.h" |
| 21 | +#include "llvm/IR/Intrinsics.h" |
| 22 | +#include "llvm/IR/IntrinsicsAMDGPU.h" |
| 23 | +#include "llvm/IR/Module.h" |
| 24 | +#include "llvm/IR/Value.h" |
| 25 | +#include "llvm/Transforms/Utils/Cloning.h" |
| 26 | +#include "llvm/Transforms/Utils/ModuleUtils.h" |
| 27 | + |
| 28 | +using namespace llvm; |
| 29 | + |
| 30 | +#define DEBUG_TYPE "offload-sanitizer" |
| 31 | + |
| 32 | +namespace { |
| 33 | + |
| 34 | +class OffloadSanitizerImpl final { |
| 35 | +public: |
| 36 | + OffloadSanitizerImpl(Module &M, FunctionAnalysisManager &FAM) |
| 37 | + : M(M), FAM(FAM), Ctx(M.getContext()) {} |
| 38 | + |
| 39 | + bool instrument(); |
| 40 | + |
| 41 | +private: |
| 42 | + bool shouldInstrumentFunction(Function &Fn); |
| 43 | + bool instrumentFunction(Function &Fn); |
| 44 | + bool instrumentTrapInstructions(SmallVectorImpl<IntrinsicInst *> &TrapCalls); |
| 45 | + |
| 46 | + FunctionCallee getOrCreateFn(FunctionCallee &FC, StringRef Name, Type *RetTy, |
| 47 | + ArrayRef<Type *> ArgTys) { |
| 48 | + if (!FC) { |
| 49 | + auto *NewAllocationFnTy = FunctionType::get(RetTy, ArgTys, false); |
| 50 | + FC = M.getOrInsertFunction(Name, NewAllocationFnTy); |
| 51 | + } |
| 52 | + return FC; |
| 53 | + } |
| 54 | + |
| 55 | + /// void __offload_san_trap_info(Int64Ty); |
| 56 | + FunctionCallee TrapInfoFn; |
| 57 | + FunctionCallee getTrapInfoFn() { |
| 58 | + return getOrCreateFn(TrapInfoFn, "__offload_san_trap_info", VoidTy, |
| 59 | + {/*PC*/ Int64Ty}); |
| 60 | + } |
| 61 | + |
| 62 | + CallInst *createCall(IRBuilder<> &IRB, FunctionCallee Callee, |
| 63 | + ArrayRef<Value *> Args = std::nullopt, |
| 64 | + const Twine &Name = "") { |
| 65 | + Calls.push_back(IRB.CreateCall(Callee, Args, Name)); |
| 66 | + return Calls.back(); |
| 67 | + } |
| 68 | + SmallVector<CallInst *> Calls; |
| 69 | + |
| 70 | + Value *getPC(IRBuilder<> &IRB) { |
| 71 | + return IRB.CreateIntrinsic(Int64Ty, Intrinsic::amdgcn_s_getpc, {}, nullptr, |
| 72 | + "PC"); |
| 73 | + } |
| 74 | + |
| 75 | + Module &M; |
| 76 | + FunctionAnalysisManager &FAM; |
| 77 | + LLVMContext &Ctx; |
| 78 | + |
| 79 | + Type *VoidTy = Type::getVoidTy(Ctx); |
| 80 | + Type *IntptrTy = M.getDataLayout().getIntPtrType(Ctx); |
| 81 | + PointerType *PtrTy = PointerType::getUnqual(Ctx); |
| 82 | + IntegerType *Int8Ty = Type::getInt8Ty(Ctx); |
| 83 | + IntegerType *Int32Ty = Type::getInt32Ty(Ctx); |
| 84 | + IntegerType *Int64Ty = Type::getInt64Ty(Ctx); |
| 85 | + |
| 86 | + const DataLayout &DL = M.getDataLayout(); |
| 87 | +}; |
| 88 | + |
| 89 | +} // end anonymous namespace |
| 90 | + |
| 91 | +bool OffloadSanitizerImpl::shouldInstrumentFunction(Function &Fn) { |
| 92 | + if (Fn.isDeclaration()) |
| 93 | + return false; |
| 94 | + if (Fn.getName().contains("ompx") || Fn.getName().contains("__kmpc") || |
| 95 | + Fn.getName().starts_with("rpc_")) |
| 96 | + return false; |
| 97 | + return !Fn.hasFnAttribute(Attribute::DisableSanitizerInstrumentation); |
| 98 | +} |
| 99 | + |
| 100 | +bool OffloadSanitizerImpl::instrumentTrapInstructions( |
| 101 | + SmallVectorImpl<IntrinsicInst *> &TrapCalls) { |
| 102 | + bool Changed = false; |
| 103 | + for (auto *II : TrapCalls) { |
| 104 | + IRBuilder<> IRB(II); |
| 105 | + createCall(IRB, getTrapInfoFn(), {getPC(IRB)}); |
| 106 | + } |
| 107 | + return Changed; |
| 108 | +} |
| 109 | + |
| 110 | +bool OffloadSanitizerImpl::instrumentFunction(Function &Fn) { |
| 111 | + if (!shouldInstrumentFunction(Fn)) |
| 112 | + return false; |
| 113 | + |
| 114 | + SmallVector<IntrinsicInst *> TrapCalls; |
| 115 | + |
| 116 | + bool Changed = false; |
| 117 | + for (auto &I : instructions(Fn)) { |
| 118 | + switch (I.getOpcode()) { |
| 119 | + case Instruction::Call: { |
| 120 | + auto &CI = cast<CallInst>(I); |
| 121 | + if (auto *II = dyn_cast<IntrinsicInst>(&CI)) |
| 122 | + if (II->getIntrinsicID() == Intrinsic::trap) |
| 123 | + TrapCalls.push_back(II); |
| 124 | + break; |
| 125 | + } |
| 126 | + default: |
| 127 | + break; |
| 128 | + } |
| 129 | + } |
| 130 | + |
| 131 | + Changed |= instrumentTrapInstructions(TrapCalls); |
| 132 | + |
| 133 | + return Changed; |
| 134 | +} |
| 135 | + |
| 136 | +bool OffloadSanitizerImpl::instrument() { |
| 137 | + bool Changed = false; |
| 138 | + |
| 139 | + for (Function &Fn : M) |
| 140 | + Changed |= instrumentFunction(Fn); |
| 141 | + |
| 142 | + removeFromUsedLists(M, [&](Constant *C) { |
| 143 | + if (!C->getName().starts_with("__offload_san")) |
| 144 | + return false; |
| 145 | + return Changed = true; |
| 146 | + }); |
| 147 | + |
| 148 | + return Changed; |
| 149 | +} |
| 150 | + |
| 151 | +PreservedAnalyses OffloadSanitizerPass::run(Module &M, |
| 152 | + ModuleAnalysisManager &AM) { |
| 153 | + FunctionAnalysisManager &FAM = |
| 154 | + AM.getResult<FunctionAnalysisManagerModuleProxy>(M).getManager(); |
| 155 | + OffloadSanitizerImpl Impl(M, FAM); |
| 156 | + if (!Impl.instrument()) |
| 157 | + return PreservedAnalyses::all(); |
| 158 | + LLVM_DEBUG(M.dump()); |
| 159 | + return PreservedAnalyses::none(); |
| 160 | +} |
0 commit comments