Skip to content

Commit

Permalink
Merge pull request rust-lang#105 from ptersilie/nomoreentryblockcalls
Browse files Browse the repository at this point in the history
Add pass disallowing calls in entry blocks.
  • Loading branch information
ltratt committed Nov 21, 2023
2 parents 881577b + 10bf29b commit 668b215
Show file tree
Hide file tree
Showing 6 changed files with 169 additions and 8 deletions.
10 changes: 10 additions & 0 deletions llvm/include/llvm/Transforms/Yk/NoCallsInEntryBlocks.h
Original file line number Diff line number Diff line change
@@ -0,0 +1,10 @@
#ifndef LLVM_TRANSFORMS_YK_NOCALLSINENTRYBLOCKS_H
#define LLVM_TRANSFORMS_YK_NOCALLSINENTRYBLOCKS_H

#include "llvm/Pass.h"

namespace llvm {
ModulePass *createYkNoCallsInEntryBlocksPass();
} // namespace llvm

#endif
20 changes: 17 additions & 3 deletions llvm/lib/CodeGen/StackMaps.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -305,8 +305,22 @@ StackMaps::parseOperand(MachineInstr::const_mop_iterator MOI,
// each one of them. Note, that the stackmap may track either of %rbx or
// %rcx, resulting in different ways below to retrieve the mappings.
int ExtraReg = 0;
Register R = MOI->getReg();
if (MOI->isReg()) {
Register R = MOI->getReg();
if (MOI->isKill()) {
// There's no point in tracking a killed register. Instead, try and
// find a copy of the register and use that to increase the likelyhood
// of us tracking that value.
// YKFIXME: This is likely only a temporary fix. In the future we might
// want to allow stackmaps to track arbitrarily many locations per live
// variable. Currently, we can only track two.
for (auto I : SpillOffsets) {
if (I.second == R) {
R = I.first;
break;
}
}
}
if (SpillOffsets.count(R) > 0) {
RHS = SpillOffsets[R];
assert(SpillOffsets[R] != 0);
Expand All @@ -333,9 +347,9 @@ StackMaps::parseOperand(MachineInstr::const_mop_iterator MOI,
}
}

unsigned DwarfRegNum = getDwarfRegNum(MOI->getReg(), TRI);
unsigned DwarfRegNum = getDwarfRegNum(R, TRI);
unsigned LLVMRegNum = *TRI->getLLVMRegNum(DwarfRegNum, false);
unsigned SubRegIdx = TRI->getSubRegIndex(LLVMRegNum, MOI->getReg());
unsigned SubRegIdx = TRI->getSubRegIndex(LLVMRegNum, R);
if (SubRegIdx)
Offset = TRI->getSubRegIdxOffset(SubRegIdx);

Expand Down
14 changes: 14 additions & 0 deletions llvm/lib/CodeGen/TargetPassConfig.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -53,6 +53,7 @@
#include "llvm/Transforms/Yk/Linkage.h"
#include "llvm/Transforms/Yk/ShadowStack.h"
#include "llvm/Transforms/Yk/Stackmaps.h"
#include "llvm/Transforms/Yk/NoCallsInEntryBlocks.h"
#include <cassert>
#include <optional>
#include <string>
Expand Down Expand Up @@ -282,6 +283,10 @@ static cl::opt<bool>
YkInsertStackMaps("yk-insert-stackmaps", cl::init(false), cl::NotHidden,
cl::desc("Insert stackmaps for JIT deoptimisation"));

static cl::opt<bool>
YkNoCallsInEntryBlocks("yk-no-calls-in-entryblocks", cl::init(false), cl::NotHidden,
cl::desc("Ensure there are no calls in the entryblock."));

/// Allow standard passes to be disabled by command line options. This supports
/// simple binary flags that either suppress the pass or do nothing.
/// i.e. -disable-mypass=false has no effect.
Expand Down Expand Up @@ -1126,6 +1131,15 @@ bool TargetPassConfig::addISelPasses() {
if (YkBlockDisambiguate)
addPass(createYkBlockDisambiguatePass());

if (YkNoCallsInEntryBlocks) {
// Make sure this pass runs before the shadowstack pass, so that we don't
// split the entry block after that pass inserted a `malloc` into `main`.
// This would otherwise result in allocas being moved outside the entry
// block which makes them dynamic allocas, and results in stackmaps not
// being able to record the functions stack size.
addPass(createYkNoCallsInEntryBlocksPass());
}

if (YkShadowStack) {
addPass(createYkShadowStackPass());
}
Expand Down
25 changes: 20 additions & 5 deletions llvm/lib/Target/X86/X86AsmPrinter.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -65,7 +65,21 @@ void clearRhs(Register Reg, std::map<Register, int64_t> &SpillMap) {
auto I = SpillMap.begin();
while (I != SpillMap.end()) {
if (I->second == Reg) {
I = SpillMap.erase(I);
// If there's a mapping A => B, where B has been reassigned and has a
// mapping B => C, then transitively apply the mapping so that A => C.
// This makes sure we don't remove mappings (especially to stack slots)
// when a register is reassigned. Example:
// store $r13 in [$rbp - 8]
// $rcx = $r13
// $r13 = $rbx
// Upon assigning `$r13 = $rbx`, we apply `$r13`'s previous mapping to
// `$rcx` so that `SpillMap[$rcx] = -8`.
if(SpillMap.count(Reg) > 0) {
SpillMap[I->first] = SpillMap[Reg];
++I;
} else {
I = SpillMap.erase(I);
}
} else {
++I;
}
Expand Down Expand Up @@ -99,13 +113,14 @@ void processInstructions(
const MachineOperand Rhs = Instr.getOperand(1);
assert(Lhs.isReg() && "Is register.");
assert(Rhs.isReg() && "Is register.");
// Reassigning a new value to LHS means any mappings to Lhs are now void
// and need to be removed. We need to do this before updating the
// mapping, so the transitive property of the SpillMap isn't violated
// (see `clearRhs` for more info).
clearRhs(Lhs.getReg(), SpillMap);
SpillMap[Lhs.getReg()] = Rhs.getReg();
// YKFIXME: If the `mov` instruction has a killed-flag, remove the
// register from the map.

// Reassigning a new value to Lhs means any mappings to Lhs are now void
// and need to be removed.
clearRhs(Lhs.getReg(), SpillMap);
continue;
}

Expand Down
1 change: 1 addition & 0 deletions llvm/lib/Transforms/Yk/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ add_llvm_component_library(LLVMYkPasses
LivenessAnalysis.cpp
StackMaps.cpp
ShadowStack.cpp
NoCallsInEntryBlocks.cpp

DEPENDS
intrinsics_gen
Expand Down
107 changes: 107 additions & 0 deletions llvm/lib/Transforms/Yk/NoCallsInEntryBlocks.cpp
Original file line number Diff line number Diff line change
@@ -0,0 +1,107 @@
// NoCallsInEntryBlocks.cpp - Ensure there are no calls in the first block.
//
// This pass splits up the entry block so that there are no calls inside it
// (one exception being the malloc call inserted by the shadowstack pass which
// will be applied later). This allows us to detect call backs from external
// functions more easily, which in turn simplifies our JITModBuilder
// implementation and fixes the soundness issue that existed in the old
// StackAdjust approach.
//
// As a simplified example, imagine this function:
//
// void foo(int argc) {
// if (argc == 0) {
// extfunc(foo)
// }
// }
//
// where `extfunc` may or may not call back to `foo`. Previously, this might
// result in a trace like
//
// foo.0
// extfunc
// foo.0
// ...
//
// Upon seeing the second `foo.0` we are unable to determine whether we are
// observing a callback to `foo` or are returning from `extfunc` which hasn't
// called foo. When running this pass we force `extfunc` not to be in the entry
// block, resulting in a trace like
//
// foo.0
// foo.1
// extfunc
// foo.1
// ...
//
// Here it is clear that `extfunc` has not called `foo` as we are seeing block
// `foo.1` next. If `extfunc` had called `foo`, the next block would be
// `foo.0`.

#include "llvm/Transforms/Yk/NoCallsInEntryBlocks.h"
#include "llvm/IR/BasicBlock.h"
#include "llvm/IR/Function.h"
#include "llvm/IR/IRBuilder.h"
#include "llvm/IR/Instructions.h"
#include "llvm/IR/Module.h"
#include "llvm/IR/Verifier.h"
#include "llvm/InitializePasses.h"
#include "llvm/Pass.h"
#include "llvm/Transforms/Yk/LivenessAnalysis.h"

#define DEBUG_TYPE "yk-no-more-entryblock-calls"

using namespace llvm;

namespace llvm {
void initializeYkNoCallsInEntryBlocksPass(PassRegistry &);
} // namespace llvm

namespace {

class YkNoCallsInEntryBlocks : public ModulePass {
public:
static char ID;
YkNoCallsInEntryBlocks() : ModulePass(ID) {
initializeYkNoCallsInEntryBlocksPass(*PassRegistry::getPassRegistry());
}

bool runOnModule(Module &M) override {
LLVMContext &Context = M.getContext();

// split block at first function call
std::set<Instruction *> Calls;
for (Function &F : M) {
if (F.empty()) // skip declarations.
continue;
BasicBlock &BB = F.getEntryBlock();
auto CurrFuncName = F.getName();
for (Instruction &I : BB) {
if (isa<CallInst>(I)) {
Function *F = cast<CallInst>(I).getCalledFunction();
BB.splitBasicBlock(&I);
break;
}
}
}

#ifndef NDEBUG
// Our pass runs after LLVM normally does its verify pass. In debug builds
// we run it again to check that our pass is generating valid IR.
if (verifyModule(M, &errs())) {
Context.emitError("Stackmap insertion pass generated invalid IR!");
return false;
}
#endif
return true;
}
};
} // namespace

char YkNoCallsInEntryBlocks::ID = 0;
INITIALIZE_PASS(YkNoCallsInEntryBlocks, DEBUG_TYPE, "no more entry block calls",
false, false)

ModulePass *llvm::createYkNoCallsInEntryBlocksPass() {
return new YkNoCallsInEntryBlocks();
}

0 comments on commit 668b215

Please sign in to comment.