Introduce a "patchable-function" function attribute

[sanjoy]

Summary

The "patchable-function" attribute can be used by an LLVM client to
influence LLVM's code generation in ways that makes the generated code
easily patchable at runtime (for instance, to redirect control).
Right now only one patchability scheme is supported,
"prologue-short-redirect", but this can be expanded in the future.

[rnk]

include/llvm/Target/TargetLowering.h

1805 ↗

(On Diff #53513)

This is about the prologue, so I would put this in TargetFrameLowering / X86FrameLowering, rather than growing the massive TargetLowering interface.

1807 ↗

(On Diff #53513)

This file is not consistent on this point, but this should be report_fatal_error, since we want to keep the message in release builds.

lib/Target/X86/X86ISelLowering.cpp

30642 ↗

(On Diff #53513)

I guess we are confident that modifying an instruction during its execution is not problematic.

[sanjoy]

lib/Target/X86/X86ISelLowering.cpp

30642 ↗

(On Diff #53513)

At this point I'm fairly sure that replacing an instruction with another instruction of the exact same size is okay in practice. What I'm less sure of is replacing an executing instruction with another one that is smaller (i.e. replace only the prefix of an instruction), which is what this patch does.

Unfortunately, I can't think of a way to determine if the second assertion is correct or not except by running a lot of code compiled with "patchable-prologue"="hotpatch-compact" on machines with high core counts (the patch as is passes some basic sanity checks). If more thorough testing uncovers issues, then we'll deal with them as they come.

Given what I just said, do you think it is a good idea to rename the attribute to "experimental-hotpatch-compact"?

[sanjoy]

• Address @rnk 's review

[rnk]

lgtm with the adjusted naming

include/llvm/Target/TargetFrameLowering.h

326 ↗

(On Diff #53601)

Why not "Kind" instead of "Flavor"? That's way more common across LLVM. Also, our enum naming convention would make this look like:

enum PatchablePrologueKind { PPF_HotpatchCompact, PPF_Unknown };

http://llvm.org/docs/CodingStandards.html#name-types-functions-variables-and-enumerators-properly

[rnk]

lib/Target/X86/X86ISelLowering.cpp

30642 ↗

(On Diff #53513)

Hit submit too soon...

I actually think you'll be OK here with the 2 byte alignment that you already have. No icache fetch is going to be able to observe any tearing. If we discover problems, we can nop-pad before subs.

I wouldn't add experimental here. All we need to guarantee is that there are two bytes to patch. Changing what we do for sub after the fact won't break any users.

[sanjoy]

Rename "patchable-prologue" to ""patchable-function" + what @rnk
suggested around enum names.

[echristo]

docs/LangRef.rst

1408

Does this need to be a hard coded attribute? Why not something similar to the floating point ones while we're still working out things? Avoids needing to worry about bitcode reading/writing.

1415

Perhaps a different name for it? hotpatch-compact isn't particularly enlightening without the description. Is the "compact" because it only handles the small code model? It might be best to talk about the option in an architecture neutral way and then explain the particular implementation in a cpu specific section below for it.

lib/CodeGen/PatchableFunction.cpp

45–48

Can merge all of this.

lib/Target/X86/X86FrameLowering.cpp

2924 ↗

(On Diff #53768)

Interesting. Is the idea here to avoid too much code growth? I'm assuming the performance of a pile of nops isn't that bad.

Also, are you just using the address of the symbol as the patchable address for the function?

[sanjoy]

docs/LangRef.rst

1408

Does this need to be a hard coded attribute?

Are you objecting to specifically documenting this attribute in the language reference? I don't mind that at all, given that means less work for me. :)

Avoids needing to worry about bitcode reading/writing.

If I understood you correctly, we don't have to worry about that here either, since this is a string attribute.

1415

"compact" as in "two bytes". I've tried to not mention any arch-specific details here (while avoid making things vague). Can you be more specific about how I can make this description less arch specific?

lib/Target/X86/X86FrameLowering.cpp

2924 ↗

(On Diff #53768)

Interesting. Is the idea here to avoid too much code growth? I'm assuming the performance of a pile of nops isn't that bad.

Yes, we want to avoid too much code growth -- we have to do this for every function. In an older scheme where we had a 5 byte nop in the function prologue unconditionally, we did see some performance impact on old amd64 chips.

Also, are you just using the address of the symbol as the patchable address for the function?

Yes. To redirect control away from foo, we basically patch &foo.

[dberris]

PS. First review in LLVM, please be gentle? :)

docs/LangRef.rst

1415

My thought here is something that's recognisable. Consider things like:

• compact-redirect-prologue
• compact-rewrite-prologue
• prologue-short-redirect
• short-prologue-redirect
• short-prologue-rewrite

If you intend to use "hotpatch" as a namespace of sorts (if there will be more later), something like:

• hotpatch-short-prologue
• hotpatch-prologue-small

lib/Target/X86/X86FrameLowering.cpp

2928–2930 ↗

(On Diff #53768)

Have you considered inserting a pseudo instruction that gets translated instead when emitting the assembler?

lib/Target/X86/X86FrameLowering.h

207–209 ↗

(On Diff #53768)

Is this intended to only handle prologues? Consider making this a single entry point, naming it something like makeFunctionPatchable(...). There may be other places where the patch-sleds could be inserted (before calls to functions, before entering loops, before returning, etc.) and it would be really great if this wasn't tied just to the prologue.

[sanjoy]

docs/LangRef.rst

1415

Thanks! I think I'll go with "prologue-short-redirect".

lib/CodeGen/PatchableFunction.cpp

45–48

Did not quite understand what you meant here. :)

lib/Target/X86/X86FrameLowering.cpp

2928–2930 ↗

(On Diff #53768)

That's a good idea, let me give that a try.

lib/Target/X86/X86FrameLowering.h

207–209 ↗

(On Diff #53768)

That's a good point, will do.

[echristo]

I don't have anything to add to dberris's comments. One reply inline.

[sanjoy]

• Changed to use a pseudo instruction PATCHABLE_OP, as per @dberris, the code looks a lot cleaner now!
• Renamed "prologue-hotpatch-compact" to "prologue-short-redirect"

There is some cleanup we can do after this, that I'll do separately
once this lands:

• Simplify StackMapShadowTracker
• Split out EmitNops so that the assert using OnlyOneNop can instead live in its caller.

[sanjoy]

• Remove unnecessary callback

[sanjoy]

Guard asserts-only work under #ifndef NDEBUG

[dberris]

lib/Target/X86/X86MCInstLower.cpp

837–838

I'm not sure this assertion makes sense here. I would have thought this assert should have been done in the calling code, that it doesn't ask for a single nop in the first place?

948–949

So in this branch, MinSize != 2 or Opcode != X86::PUSH64r.

Question: If you check instead whether the function where MI is included in had the correct type of patchable-function attribute, and see that the nops being added is less than 2, you can assert and say this is actually a bug in a higher implementation detail? i.e. this would be a bug in the insertion of this instruction.

This way you wouldn't need to touch EmitNops.

[dberris]

I'm not very familiar with Differential but is there a way for you to update the summary to more accurately describe what the patch is doing now?

[sanjoy]

(Just about to update the description).

lib/Target/X86/X86MCInstLower.cpp

837–838

That's the first point under the "There is some cleanup we can do after this" note I sent in with this update. :) Basically, I'd rather not make this already large patch any larger if it can be helped. NFC cleanups like these are easy to do once the hard stuff has been reviewed and checked in, IMO.

948–949

Question: If you check instead whether the function where MI is included in had the correct type of patchable-function attribute, and see that the nops being added is less than 2, you can assert and say this is actually a bug in a higher implementation detail? i.e. this would be a bug in the insertion of this instruction.

Do you mean something like:

assert(MinSize < 2 && !MF.getPatchableFnType() == "prologue-short-redirect");

I'd say that is an incorrect layering. It feels cleaner for MC to not have to know why the PATCHABLE_OP of a certain variety is present. It should just understand its end of the contract of how it needs to lower PATCHABLE_OP.

[dberris]

lib/Target/X86/X86MCInstLower.cpp

948–949

I'd say that is an incorrect layering. It feels cleaner for MC to not have to know why the PATCHABLE_OP of a certain variety is present. It should just understand its end of the contract of how it needs to lower PATCHABLE_OP.

Consider the true branch of this if-statement though -- it already feels like it's already bleeding details in based on what a short 'prologue-short-redirect' patchable function attribute is already expecting. I'd say we're already breaking some layering guidelines here. :D

[sanjoy]

lib/Target/X86/X86MCInstLower.cpp

948–949

I don't think those two are the same things. One is saying "in this specific case I know can do better" (and you can extend the logic later by not just handling push'es but also returns, for instance), and the other is saying "this is the only case I support".

I'm not opposed to having PATCHABLE_OP specifically only work with "prologue-short-redirect", but then I'd rather have it not have a minsize operand at all. Do you think that will be cleaner? (I could go either way on this -- no strong preferences).

Actually, now that I think about it, what I dislike about the current scheme (i.e. this patch) is that only the minsize == 2 case is tested, so from just a testing / code coverage POV, removing minsize sounds slightly better.

[dberris]

lib/Target/X86/X86MCInstLower.cpp

948–949

I think minsize still makes sense, but I'm thinking that there's two inputs here really:

1. That there are instructions to be placed here with a given minimum size. Note that this could be not just a single instruction, or could be treated just as a placeholder.
2. That the semantics of what kinds of instructions will be emitted would be based on the type of function patching is supported. In this case you're implementing 'prologue-short-redirect' which expects certain kinds of operations to be valid where a PATCHABLE_OP appears.

Which is why I think looking at the attribute that defines the semantics of the PATCHABLE_OP still makes sense at this level. It allows us to make decisions of what's valid or not valid based on the attribute provided on the function. And I think this is the correct layer to make that decision, because this is where we're actually generating the instructions.

Does that reasoning make sense?

[sanjoy]

lib/Target/X86/X86MCInstLower.cpp

948–949

I think minsize still makes sense, but I'm thinking that there's two
inputs here really:

That there are instructions to be placed here with a given minimum
size. Note that this could be not just a single instruction, or could
be treated just as a placeholder.

That the semantics of what kinds of instructions will be emitted would
be based on the type of function patching is supported. In this case
you're implementing 'prologue-short-redirect' which expects certain
kinds of operations to be valid where a PATCHABLE_OP appears.

Which is why I think looking at the attribute that defines the
semantics of the PATCHABLE_OP still makes sense at this level. It
allows us to make decisions of what's valid or not valid based on the
attribute provided on the function. And I think this is the correct

I've been thinking about this as: an instance of PATCHABLE_OP is
(should be) self contained with regards to what needs to be emitted
when MC encounters one. Right now, as you said, emitting a
PATCHABLE_OP involves ensuring two things:

1. There is *one* instruction of at least MinSize bytes at the place in the instruction stream the PATCHABLE_OP appeared at.
2. The set of instructions emitted as part of lowering the PATCHABLE_OP pseudo instruction is "equivalent" (i.e. has the same effect on the CPU state, modulo the delta by which the instruction pointer is advanced) to the instruction bundled with PATCHABLE_OP.

(I could be more explicit about this in comment over PATCHABLE_OP --
let me know if you think that will help)

This is all that a PATCHABLE_OP implies. Any optimizations that
happens on top of (1) and (2) are strictly optional, and cannot tread
beyond what is allowed by (1) and (2).

If I understand you correctly, you're saying PATCHABLE_OP is not
self contained, but interpreting MC needs to do when it sees a
PATCHABLE_OP depends on what attributes the containing function has.
I still don't think this is correct (assuming I haven't misrepresented
your position): unless we gain something by by coupling function
attributes with PATCHABLE_OP, I'd rather have these de-coupled
(i.e. have PATCHABLE_OP be the mechanism by which the
"patchable-function" policy is implemented). For instance, a
potential use case for PATCHABLE_OP is to make some call sites
patchable, and one way to do that is via call site attributes.
Without making PATCHABLE_OP self sufficient we'd have to resort to
walking back to the relevant (IR level) call instruction (which may be
difficult to locate), in addition to checking the function attributes.
What if after this we want to add a third source for PATCHABLE_OP
(an intrinsic, say)? The complexity of lowering PATCHABLE_OP,
unless it is self sufficient, will scale linearly with the number of
ways we can generate one.

layer to make that decision, because this is where we're actually
generating the instructions.

Does that reasoning make sense?

[dberris]

lib/Target/X86/X86MCInstLower.cpp

948–949

I think I originally misunderstood the purpose of PATCHABLE_OP -- I had been thinking it was a standalone pseudo-instruction which would be reduced the the op as a parameter if patching was not enabled on the function (or due to some other consideration, like the size of the instruction that proceeds it). In my head (and my current implementation for something similar I and echristo are working on) we just have a pure pseudo-instruction that expands in a context-sensitive manner.

As implemented, I think it's fine for the purpose of handling this specific attribute. I suppose later implementations of a different attribute can dispatch to the correct behaviour (and re-use/extend PATCHABLE_OP) appropriately.

docs/LangRef.rst

``optnone``
     the function as well, so the function is never inlined into any caller.
     Only functions with the ``alwaysinline`` attribute are valid
     candidates for inlining into the body of this function.
 ``optsize``
     This attribute suggests that optimization passes and code generator
     passes make choices that keep the code size of this function low,
     and otherwise do optimizations specifically to reduce code size as
     long as they do not significantly impact runtime performance.
+``"patchable-function"``
+    This attribute tells the code generator that the code
+    generated for this function needs to follow certain conventions that
+    make it possible for a runtime function to patch over it later.
+    The exact effect of this attribute depends on its string value,
+    for which there currently is one legal possiblity:
+
+     * ``"prologue-short-redirect"`` - This style of patchable
+       function is intended to support patching a function prologue to
+       redirect control away from the function in a thread safe
+       manner.  It guarantees that the first instruction of the
+       function will be large enough to accommodate a short jump
+       instruction, and will be sufficiently aligned to allow being
+       fully changed via an atomic compare-and-swap instruction.
+       While the first requirement can be satisfied by inserting large
+       enough NOP, LLVM can and will try to re-purpose an existing
+       instruction (i.e. one that would have to be emitted anyway) as
+       the patchable instruction larger than a short jump.
+
+       ``"prologue-short-redirect"`` is currently only supported on
+       x86-64.
+
+    This attribute by itself does not imply restrictions on
+    inter-procedural optimizations.  All of the semantic effects the
+    patching may have to be separately conveyed via the linkage type.
 ``readnone``
     On a function, this attribute indicates that the function computes its
     result (or decides to unwind an exception) based strictly on its arguments,
     without dereferencing any pointer arguments or otherwise accessing
     any mutable state (e.g. memory, control registers, etc) visible to
     caller functions. It does not write through any pointer arguments
     (including ``byval`` arguments) and never changes any state visible
     to callers. This means that it cannot unwind exceptions by calling

include/llvm/CodeGen/Passes.h

/// MachineDominanaceFrontier - This pass is a machine dominators analysis pass.
   extern char &OptimizePHIsID;
 
   /// StackSlotColoring - This pass performs stack slot coloring.
   extern char &StackSlotColoringID;
 
   /// \brief This pass lays out funclets contiguously.
   extern char &FuncletLayoutID;
 
+  /// \brief This pass implements the "patchable-function" attribute.
+  extern char &PatchableFunctionID;
+
   /// createStackProtectorPass - This pass adds stack protectors to functions.
   ///
   FunctionPass *createStackProtectorPass(const TargetMachine *TM);
 
   /// createMachineVerifierPass - This pass verifies cenerated machine code
   /// instructions for correctness.
   ///
   FunctionPass *createMachineVerifierPass(const std::string& Banner);

include/llvm/InitializePasses.h

 void initializeLoopDistributePass(PassRegistry&);
 void initializeSjLjEHPreparePass(PassRegistry&);
 void initializeDemandedBitsPass(PassRegistry&);
 void initializeFuncletLayoutPass(PassRegistry &);
 void initializeLoopLoadEliminationPass(PassRegistry&);
 void initializeFunctionImportPassPass(PassRegistry &);
 void initializeLoopVersioningPassPass(PassRegistry &);
 void initializeWholeProgramDevirtPass(PassRegistry &);
+void initializePatchableFunctionPass(PassRegistry &);
 }
 
 #endif

include/llvm/Target/Target.td

def LOCAL_ESCAPE : Instruction {
   let hasCtrlDep = 1;
 }
 def FAULTING_LOAD_OP : Instruction {
   let OutOperandList = (outs unknown:$dst);
   let InOperandList = (ins variable_ops);
   let usesCustomInserter = 1;
   let mayLoad = 1;
 }
+def PATCHABLE_OP : Instruction {
+  let OutOperandList = (outs unknown:$dst);
+  let InOperandList = (ins variable_ops);
+  let usesCustomInserter = 1;
+  let mayLoad = 1;
+  let mayStore = 1;
+  let hasSideEffects = 1;
+}
 
 // Generic opcodes used in GlobalISel.
 include "llvm/Target/GenericOpcodes.td"
 
 }
 
 //===----------------------------------------------------------------------===//
 // AsmParser - This class can be implemented by targets that wish to implement

include/llvm/Target/TargetOpcodes.def

 HANDLE_TARGET_OPCODE(LOCAL_ESCAPE, 21)
 
 /// Loading instruction that may page fault, bundled with associated
 /// information on how to handle such a page fault.  It is intended to support
 /// "zero cost" null checks in managed languages by allowing LLVM to fold
 /// comparisons into existing memory operations.
 HANDLE_TARGET_OPCODE(FAULTING_LOAD_OP, 22)
 
+/// Wraps a machine instruction to add patchability constraints.  An
+/// instruction wrapped in PATCHABLE_OP has to either have a minimum
+/// size or be preceded with a nop of that size.  The first operand is
+/// an immediate denoting the minimum size of the instruction, the
+/// second operand is an immediate denoting the opcode of the original
+/// instruction.  The rest of the operands are the operands of the
+/// original instruction.
+HANDLE_TARGET_OPCODE(PATCHABLE_OP, 23)
+
 /// The following generic opcodes are not supposed to appear after ISel.
 /// This is something we might want to relax, but for now, this is convenient
 /// to produce diagnostics.
 
 /// Generic ADD instruction. This is an integer add.
-HANDLE_TARGET_OPCODE(G_ADD, 23)
+HANDLE_TARGET_OPCODE(G_ADD, 24)
 HANDLE_TARGET_OPCODE_MARKER(PRE_ISEL_GENERIC_OPCODE_START, G_ADD)
 
 /// Generic BRANCH instruction. This is an unconditional branch.
-HANDLE_TARGET_OPCODE(G_BR, 24)
+HANDLE_TARGET_OPCODE(G_BR, 25)
 
 // TODO: Add more generic opcodes as we move along.
 
 /// Marker for the end of the generic opcode.
 /// This is used to check if an opcode is in the range of the
 /// generic opcodes.
 HANDLE_TARGET_OPCODE_MARKER(PRE_ISEL_GENERIC_OPCODE_END, G_BR)
 
 /// BUILTIN_OP_END - This must be the last enum value in this list.
 /// The target-specific post-isel opcode values start here.
 HANDLE_TARGET_OPCODE_MARKER(GENERIC_OP_END, PRE_ISEL_GENERIC_OPCODE_END)

lib/CodeGen/CMakeLists.txt

add_llvm_library(LLVMCodeGen
   MachinePostDominators.cpp
   MachineRegionInfo.cpp
   MachineRegisterInfo.cpp
   MachineScheduler.cpp
   MachineSink.cpp
   MachineSSAUpdater.cpp
   MachineTraceMetrics.cpp
   MachineVerifier.cpp
+  PatchableFunction.cpp
   MIRPrinter.cpp
   MIRPrintingPass.cpp
   OptimizePHIs.cpp
   ParallelCG.cpp
   Passes.cpp
   PeepholeOptimizer.cpp
   PHIElimination.cpp
   PHIEliminationUtils.cpp

lib/CodeGen/CodeGen.cpp

void llvm::initializeCodeGen(PassRegistry &Registry) {
   initializeMachineFunctionPrinterPassPass(Registry);
   initializeMachineLICMPass(Registry);
   initializeMachineLoopInfoPass(Registry);
   initializeMachineModuleInfoPass(Registry);
   initializeMachinePostDominatorTreePass(Registry);
   initializeMachineSchedulerPass(Registry);
   initializeMachineSinkingPass(Registry);
   initializeMachineVerifierPassPass(Registry);
+  initializePatchableFunctionPass(Registry);
   initializeOptimizePHIsPass(Registry);
   initializePEIPass(Registry);
   initializePHIEliminationPass(Registry);
   initializePeepholeOptimizerPass(Registry);
   initializePostMachineSchedulerPass(Registry);
   initializePostRASchedulerPass(Registry);
   initializeProcessImplicitDefsPass(Registry);
   initializeRegisterCoalescerPass(Registry);

lib/CodeGen/Passes.cpp

void TargetPassConfig::addMachinePasses() {
 
   addPreEmitPass();
 
   addPass(&FuncletLayoutID, false);
 
   addPass(&StackMapLivenessID, false);
   addPass(&LiveDebugValuesID, false);
 
+  addPass(&PatchableFunctionID, false);
+
   AddingMachinePasses = false;
 }
 
 /// Add passes that optimize machine instructions in SSA form.
 void TargetPassConfig::addMachineSSAOptimization() {
   // Pre-ra tail duplication.
   addPass(&EarlyTailDuplicateID);
 

lib/CodeGen/PatchableFunction.cpp

+//===-- PatchableFunction.cpp - Patchable prologues for LLVM -------------===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements edits function bodies in place to support the
+// "patchable-function" attribute.
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/CodeGen/Passes.h"
+#include "llvm/CodeGen/Analysis.h"
+#include "llvm/CodeGen/MachineFunction.h"
+#include "llvm/CodeGen/MachineFunctionPass.h"
+#include "llvm/CodeGen/MachineInstrBuilder.h"
+#include "llvm/Target/TargetFrameLowering.h"
+#include "llvm/Target/TargetInstrInfo.h"
+#include "llvm/Target/TargetSubtargetInfo.h"
+
+using namespace llvm;
+
+namespace {
+struct PatchableFunction : public MachineFunctionPass {
+  static char ID; // Pass identification, replacement for typeid
+  PatchableFunction() : MachineFunctionPass(ID) {
+    initializePatchableFunctionPass(*PassRegistry::getPassRegistry());
+  }
+
+  bool runOnMachineFunction(MachineFunction &F) override;
+   MachineFunctionProperties getRequiredProperties() const override {
+    return MachineFunctionProperties().set(
+        MachineFunctionProperties::Property::AllVRegsAllocated);
+  }
+};
+}
+
+bool PatchableFunction::runOnMachineFunction(MachineFunction &MF) {
+  if (!MF.getFunction()->hasFnAttribute("patchable-function"))
+    return false;
+
+#ifndef NDEBUG
+  Attribute PatchAttr = MF.getFunction()->getFnAttribute("patchable-function");
+  StringRef PatchType = PatchAttr.getValueAsString();
+  assert(PatchType == "prologue-short-redirect" && "Only possibility today!");
+#endif
+
+  auto &FirstMBB = *MF.begin();
+  auto &FirstMI = *FirstMBB.begin();
+
+  auto *TII = MF.getSubtarget().getInstrInfo();
+  auto MIB = BuildMI(FirstMBB, FirstMBB.begin(), FirstMI.getDebugLoc(),
+                     TII->get(TargetOpcode::PATCHABLE_OP))
+                 .addImm(2)
+                 .addImm(FirstMI.getOpcode());
+
+  for (auto &MO : FirstMI.operands())
+    MIB.addOperand(MO);
+
+  FirstMI.eraseFromParent();
+  MF.ensureAlignment(4);
+  return true;
+}
+
+char PatchableFunction::ID = 0;
+char &llvm::PatchableFunctionID = PatchableFunction::ID;
+INITIALIZE_PASS(PatchableFunction, "patchable-function", "", false, false)

lib/Target/X86/X86AsmPrinter.h

 namespace llvm {
 class MCStreamer;
 class MCSymbol;
 
 class LLVM_LIBRARY_VISIBILITY X86AsmPrinter : public AsmPrinter {
   const X86Subtarget *Subtarget;
   StackMaps SM;
   FaultMaps FM;
+  std::unique_ptr<MCCodeEmitter> CodeEmitter;
 
   // This utility class tracks the length of a stackmap instruction's 'shadow'.
   // It is used by the X86AsmPrinter to ensure that the stackmap shadow
   // invariants (i.e. no other stackmaps, patchpoints, or control flow within
   // the shadow) are met, while outputting a minimal number of NOPs for padding.
   //
   // To minimise the number of NOPs used, the shadow tracker counts the number
   // of instruction bytes output since the last stackmap. Only if there are too
   // few instruction bytes to cover the shadow are NOPs used for padding.
   class StackMapShadowTracker {
   public:
-    StackMapShadowTracker(TargetMachine &TM);
+    StackMapShadowTracker();
     ~StackMapShadowTracker();
     void startFunction(MachineFunction &MF);
-    void count(MCInst &Inst, const MCSubtargetInfo &STI);
+    void count(MCInst &Inst, const MCSubtargetInfo &STI,
+               MCCodeEmitter *CodeEmitter);
 
     // Called to signal the start of a shadow of RequiredSize bytes.
     void reset(unsigned RequiredSize) {
       RequiredShadowSize = RequiredSize;
       CurrentShadowSize = 0;
       InShadow = true;
     }
 
     // Called before every stackmap/patchpoint, and at the end of basic blocks,
     // to emit any necessary padding-NOPs.
     void emitShadowPadding(MCStreamer &OutStreamer, const MCSubtargetInfo &STI);
   private:
-    TargetMachine &TM;
     const MachineFunction *MF;
-    std::unique_ptr<MCCodeEmitter> CodeEmitter;
     bool InShadow;
 
     // RequiredShadowSize holds the length of the shadow specified in the most
     // recently encountered STACKMAP instruction.
     // CurrentShadowSize counts the number of bytes encoded since the most
     // recently encountered STACKMAP, stopping when that number is greater than
     // or equal to RequiredShadowSize.
     unsigned RequiredShadowSize, CurrentShadowSize;
   };
 
   StackMapShadowTracker SMShadowTracker;
 
   // All instructions emitted by the X86AsmPrinter should use this helper
   // method.
   //
   // This helper function invokes the SMShadowTracker on each instruction before
   // outputting it to the OutStream. This allows the shadow tracker to minimise
   // the number of NOPs used for stackmap padding.
   void EmitAndCountInstruction(MCInst &Inst);
   void LowerSTACKMAP(const MachineInstr &MI);
   void LowerPATCHPOINT(const MachineInstr &MI, X86MCInstLower &MCIL);
   void LowerSTATEPOINT(const MachineInstr &MI, X86MCInstLower &MCIL);
   void LowerFAULTING_LOAD_OP(const MachineInstr &MI, X86MCInstLower &MCIL);
+  void LowerPATCHABLE_OP(const MachineInstr &MI, X86MCInstLower &MCIL);
 
   void LowerTlsAddr(X86MCInstLower &MCInstLowering, const MachineInstr &MI);
 
  public:
    explicit X86AsmPrinter(TargetMachine &TM,
                           std::unique_ptr<MCStreamer> Streamer)
-       : AsmPrinter(TM, std::move(Streamer)), SM(*this), FM(*this),
-         SMShadowTracker(TM) {}
+       : AsmPrinter(TM, std::move(Streamer)), SM(*this), FM(*this) {}
 
   const char *getPassName() const override {
     return "X86 Assembly / Object Emitter";
   }
 
   const X86Subtarget &getSubtarget() const { return *Subtarget; }
 
   void EmitStartOfAsmFile(Module &M) override;

lib/Target/X86/X86AsmPrinter.cpp

 #include "llvm/CodeGen/MachineValueType.h"
 #include "llvm/CodeGen/TargetLoweringObjectFileImpl.h"
 #include "llvm/IR/DebugInfo.h"
 #include "llvm/IR/DerivedTypes.h"
 #include "llvm/IR/Mangler.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/Type.h"
 #include "llvm/MC/MCAsmInfo.h"
+#include "llvm/MC/MCCodeEmitter.h"
 #include "llvm/MC/MCContext.h"
 #include "llvm/MC/MCExpr.h"
 #include "llvm/MC/MCSectionCOFF.h"
 #include "llvm/MC/MCSectionMachO.h"
 #include "llvm/MC/MCStreamer.h"
 #include "llvm/MC/MCSymbol.h"
 #include "llvm/Support/COFF.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/TargetRegistry.h"
 using namespace llvm;
 
 //===----------------------------------------------------------------------===//
 // Primitive Helper Functions.
 //===----------------------------------------------------------------------===//
 
 /// runOnMachineFunction - Emit the function body.
 ///
 bool X86AsmPrinter::runOnMachineFunction(MachineFunction &MF) {
   Subtarget = &MF.getSubtarget<X86Subtarget>();
 
   SMShadowTracker.startFunction(MF);
+  CodeEmitter.reset(TM.getTarget().createMCCodeEmitter(
+      *MF.getSubtarget().getInstrInfo(), *MF.getSubtarget().getRegisterInfo(),
+      MF.getContext()));
 
   SetupMachineFunction(MF);
 
   if (Subtarget->isTargetCOFF()) {
     bool Intrn = MF.getFunction()->hasInternalLinkage();
     OutStreamer->BeginCOFFSymbolDef(CurrentFnSym);
     OutStreamer->EmitCOFFSymbolStorageClass(Intrn ? COFF::IMAGE_SYM_CLASS_STATIC
                                             : COFF::IMAGE_SYM_CLASS_EXTERNAL);

lib/Target/X86/X86MCInstLower.cpp

 #include "X86AsmPrinter.h"
 #include "X86RegisterInfo.h"
 #include "X86ShuffleDecodeConstantPool.h"
 #include "InstPrinter/X86ATTInstPrinter.h"
 #include "MCTargetDesc/X86BaseInfo.h"
 #include "Utils/X86ShuffleDecode.h"
 #include "llvm/ADT/Optional.h"
 #include "llvm/ADT/SmallString.h"
+#include "llvm/ADT/iterator_range.h"
 #include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/CodeGen/MachineConstantPool.h"
 #include "llvm/CodeGen/MachineOperand.h"
 #include "llvm/CodeGen/MachineModuleInfoImpls.h"
 #include "llvm/CodeGen/StackMaps.h"
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/GlobalValue.h"
 #include "llvm/IR/Mangler.h"
Mangler *getMang() const {
     return AsmPrinter.Mang;
   }
 };
 
 } // end anonymous namespace
 
 // Emit a minimal sequence of nops spanning NumBytes bytes.
 static void EmitNops(MCStreamer &OS, unsigned NumBytes, bool Is64Bit,
-                     const MCSubtargetInfo &STI);
+                     const MCSubtargetInfo &STI, bool OnlyOneNop = false);
 
 namespace llvm {
-   X86AsmPrinter::StackMapShadowTracker::StackMapShadowTracker(TargetMachine &TM)
-     : TM(TM), InShadow(false), RequiredShadowSize(0), CurrentShadowSize(0) {}
+   X86AsmPrinter::StackMapShadowTracker::StackMapShadowTracker()
+     : InShadow(false), RequiredShadowSize(0), CurrentShadowSize(0) {}
 
   X86AsmPrinter::StackMapShadowTracker::~StackMapShadowTracker() {}
 
-  void
-  X86AsmPrinter::StackMapShadowTracker::startFunction(MachineFunction &F) {
+  void X86AsmPrinter::StackMapShadowTracker::startFunction(MachineFunction &F) {
     MF = &F;
-    CodeEmitter.reset(TM.getTarget().createMCCodeEmitter(
-        *MF->getSubtarget().getInstrInfo(),
-        *MF->getSubtarget().getRegisterInfo(), MF->getContext()));
   }
 
   void X86AsmPrinter::StackMapShadowTracker::count(MCInst &Inst,
-                                                   const MCSubtargetInfo &STI) {
+                                                   const MCSubtargetInfo &STI,
+                                                   MCCodeEmitter *CodeEmitter) {
     if (InShadow) {
       SmallString<256> Code;
       SmallVector<MCFixup, 4> Fixups;
       raw_svector_ostream VecOS(Code);
       CodeEmitter->encodeInstruction(Inst, VecOS, Fixups, STI);
       CurrentShadowSize += Code.size();
       if (CurrentShadowSize >= RequiredShadowSize)
         InShadow = false; // The shadow is big enough. Stop counting.
     }
   }
 
   void X86AsmPrinter::StackMapShadowTracker::emitShadowPadding(
     MCStreamer &OutStreamer, const MCSubtargetInfo &STI) {
     if (InShadow && CurrentShadowSize < RequiredShadowSize) {
       InShadow = false;
       EmitNops(OutStreamer, RequiredShadowSize - CurrentShadowSize,
                MF->getSubtarget<X86Subtarget>().is64Bit(), STI);
     }
   }
 
   void X86AsmPrinter::EmitAndCountInstruction(MCInst &Inst) {
     OutStreamer->EmitInstruction(Inst, getSubtargetInfo());
-    SMShadowTracker.count(Inst, getSubtargetInfo());
+    SMShadowTracker.count(Inst, getSubtargetInfo(), CodeEmitter.get());
   }
 } // end llvm namespace
 
 X86MCInstLower::X86MCInstLower(const MachineFunction &mf,
                                X86AsmPrinter &asmprinter)
     : Ctx(mf.getContext()), MF(mf), TM(mf.getTarget()), MAI(*TM.getMCAsmInfo()),
       AsmPrinter(asmprinter) {}
 
MCSymbolRefExpr::create(tlsGetAddr,
                             context);
 
   EmitAndCountInstruction(MCInstBuilder(is64Bits ? X86::CALL64pcrel32
                                                  : X86::CALLpcrel32)
                             .addExpr(tlsRef));
 }
 
 /// \brief Emit the optimal amount of multi-byte nops on X86.
-static void EmitNops(MCStreamer &OS, unsigned NumBytes, bool Is64Bit, const MCSubtargetInfo &STI) {
+static void EmitNops(MCStreamer &OS, unsigned NumBytes, bool Is64Bit,
+                     const MCSubtargetInfo &STI, bool OnlyOneNop) {
   // This works only for 64bit. For 32bit we have to do additional checking if
   // the CPU supports multi-byte nops.
   assert(Is64Bit && "EmitNops only supports X86-64");
   while (NumBytes) {
     unsigned Opc, BaseReg, ScaleVal, IndexReg, Displacement, SegmentReg;
     Opc = IndexReg = Displacement = SegmentReg = 0;
     BaseReg = X86::RAX; ScaleVal = 1;
     switch (NumBytes) {
case X86::XCHG16ar:
       break;
     case X86::NOOPL:
     case X86::NOOPW:
       OS.EmitInstruction(MCInstBuilder(Opc).addReg(BaseReg)
                          .addImm(ScaleVal).addReg(IndexReg)
                          .addImm(Displacement).addReg(SegmentReg), STI);
       break;
     }
+
+    (void) OnlyOneNop;
+    assert((!OnlyOneNop || NumBytes == 0) &&
+           "Allowed only one nop instruction!");
   } // while (NumBytes)
 }
 
 void X86AsmPrinter::LowerSTATEPOINT(const MachineInstr &MI,
                                     X86MCInstLower &MCIL) {
   assert(Subtarget->is64Bit() && "Statepoint currently only supports X86-64");
 
   StatepointOpers SOpers(&MI);
for (auto I = MI.operands_begin() + LoadOperandsBeginIdx,
             E = MI.operands_end();
        I != E; ++I)
     if (auto MaybeOperand = MCIL.LowerMachineOperand(&MI, *I))
       LoadMI.addOperand(MaybeOperand.getValue());
 
   OutStreamer->EmitInstruction(LoadMI, getSubtargetInfo());
 }
 
+void X86AsmPrinter::LowerPATCHABLE_OP(const MachineInstr &MI,
+                                      X86MCInstLower &MCIL) {
+  // PATCHABLE_OP minsize, opcode, operands
+
+  unsigned MinSize = MI.getOperand(0).getImm();
+  unsigned Opcode = MI.getOperand(1).getImm();
+
+  MCInst MCI;
+  MCI.setOpcode(Opcode);
+  for (auto &MO : make_range(MI.operands_begin() + 2, MI.operands_end()))
+    if (auto MaybeOperand = MCIL.LowerMachineOperand(&MI, MO))
+      MCI.addOperand(MaybeOperand.getValue());
+
+  SmallString<256> Code;
+  SmallVector<MCFixup, 4> Fixups;
+  raw_svector_ostream VecOS(Code);
+  CodeEmitter->encodeInstruction(MCI, VecOS, Fixups, getSubtargetInfo());
+
+  if (Code.size() < MinSize) {
+    if (MinSize == 2 && Opcode == X86::PUSH64r) {
+      // This is an optimization that lets us get away without emitting a nop in
+      // many cases.
+      //
+      // NB! In some cases the encoding for PUSH64r (e.g. PUSH64r %R9) takes two
+      // bytes too, so the check on MinSize is important.
+      MCI.setOpcode(X86::PUSH64rmr);
+    } else {
+      EmitNops(*OutStreamer, MinSize, Subtarget->is64Bit(), getSubtargetInfo(),
+               /* OnlyOneNop = */ true);
+    }
+  }
+
+  OutStreamer->EmitInstruction(MCI, getSubtargetInfo());
+}
+
 // Lower a stackmap of the form:
 // <id>, <shadowBytes>, ...
 void X86AsmPrinter::LowerSTACKMAP(const MachineInstr &MI) {
   SMShadowTracker.emitShadowPadding(*OutStreamer, getSubtargetInfo());
   SM.recordStackMap(MI);
   unsigned NumShadowBytes = MI.getOperand(1).getImm();
   SMShadowTracker.reset(NumShadowBytes);
 }
case X86::ADD32ri: {
     return;
   }
   case TargetOpcode::STATEPOINT:
     return LowerSTATEPOINT(*MI, MCInstLowering);
 
   case TargetOpcode::FAULTING_LOAD_OP:
     return LowerFAULTING_LOAD_OP(*MI, MCInstLowering);
 
+  case TargetOpcode::PATCHABLE_OP:
+    return LowerPATCHABLE_OP(*MI, MCInstLowering);
+
   case TargetOpcode::STACKMAP:
     return LowerSTACKMAP(*MI);
 
   case TargetOpcode::PATCHPOINT:
     return LowerPATCHPOINT(*MI, MCInstLowering);
 
   case X86::MORESTACK_RET:
     EmitAndCountInstruction(MCInstBuilder(getRetOpcode(*Subtarget)));
#define CASE_ALL_MOV_RM()               \
   MCInstLowering.Lower(MI, TmpInst);
 
   // Stackmap shadows cannot include branch targets, so we can count the bytes
   // in a call towards the shadow, but must ensure that the no thread returns
   // in to the stackmap shadow.  The only way to achieve this is if the call
   // is at the end of the shadow.
   if (MI->isCall()) {
     // Count then size of the call towards the shadow
-    SMShadowTracker.count(TmpInst, getSubtargetInfo());
+    SMShadowTracker.count(TmpInst, getSubtargetInfo(), CodeEmitter.get());
     // Then flush the shadow so that we fill with nops before the call, not
     // after it.
     SMShadowTracker.emitShadowPadding(*OutStreamer, getSubtargetInfo());
     // Then emit the call
     OutStreamer->EmitInstruction(TmpInst, getSubtargetInfo());
     return;
   }
 
   EmitAndCountInstruction(TmpInst);
 }

test/CodeGen/X86/patchable-prologue.ll

+; RUN: llc -filetype=obj -o - -mtriple=x86_64-apple-macosx < %s | llvm-objdump -triple x86_64-apple-macosx -disassemble - | FileCheck %s
+; RUN: llc -mtriple=x86_64-apple-macosx < %s | FileCheck %s --check-prefix=CHECK-ALIGN
+
+declare void @callee(i64*)
+
+define void @f0() "patchable-function"="prologue-short-redirect" {
+; CHECK-LABEL: _f0:
+; CHECK-NEXT:  66 90  nop
+
+; CHECK-ALIGN:  .p2align  4, 0x90
+; CHECK-ALIGN: _f0:
+
+  ret void
+}
+
+define void @f1() "patchable-function"="prologue-short-redirect" "no-frame-pointer-elim"="true" {
+; CHECK-LABEL: _f1
+; CHECK-NEXT: ff f5   pushq %rbp
+
+; CHECK-ALIGN:  .p2align  4, 0x90
+; CHECK-ALIGN: _f1:
+  ret void
+}
+
+define void @f2() "patchable-function"="prologue-short-redirect" {
+; CHECK-LABEL: _f2
+; CHECK-NEXT: 48 81 ec a8 00 00 00  subq  $168, %rsp
+
+; CHECK-ALIGN:  .p2align  4, 0x90
+; CHECK-ALIGN: _f2:
+  %ptr = alloca i64, i32 20
+  call void @callee(i64* %ptr)
+  ret void
+}
+
+define void @f3() "patchable-function"="prologue-short-redirect" optsize {
+; CHECK-LABEL: _f3
+; CHECK-NEXT: 66 90   nop
+
+; CHECK-ALIGN:  .p2align  4, 0x90
+; CHECK-ALIGN: _f3:
+  ret void
+}

test/TableGen/trydecode-emission.td

def InstB : TestInstruction {
   let AsmString = "InstB";
   let DecoderMethod = "DecodeInstB";
   let hasCompleteDecoder = 0;
 }
 
 // CHECK:      /* 0 */       MCD::OPC_ExtractField, 4, 4,  // Inst{7-4} ...
 // CHECK-NEXT: /* 3 */       MCD::OPC_FilterValue, 0, 14, 0, // Skip to: 21
 // CHECK-NEXT: /* 7 */       MCD::OPC_CheckField, 2, 2, 0, 5, 0, // Skip to: 18
-// CHECK-NEXT: /* 13 */      MCD::OPC_TryDecode, 26, 0, 0, 0, // Opcode: InstB, skip to: 18
-// CHECK-NEXT: /* 18 */      MCD::OPC_Decode, 25, 1, // Opcode: InstA
+// CHECK-NEXT: /* 13 */      MCD::OPC_TryDecode, 27, 0, 0, 0, // Opcode: InstB, skip to: 18
+// CHECK-NEXT: /* 18 */      MCD::OPC_Decode, 26, 1, // Opcode: InstA
 // CHECK-NEXT: /* 21 */      MCD::OPC_Fail,
 
 // CHECK: if (DecodeInstB(MI, insn, Address, Decoder) == MCDisassembler::Fail) { DecodeComplete = false; return MCDisassembler::Fail; }

test/TableGen/trydecode-emission2.td

def InstB : TestInstruction {
   let hasCompleteDecoder = 0;
 }
 
 // CHECK:      /* 0 */       MCD::OPC_ExtractField, 2, 1,  // Inst{2} ...
 // CHECK-NEXT: /* 3 */       MCD::OPC_FilterValue, 0, 29, 0, // Skip to: 36
 // CHECK-NEXT: /* 7 */       MCD::OPC_ExtractField, 5, 3,  // Inst{7-5} ...
 // CHECK-NEXT: /* 10 */      MCD::OPC_FilterValue, 0, 22, 0, // Skip to: 36
 // CHECK-NEXT: /* 14 */      MCD::OPC_CheckField, 0, 2, 3, 5, 0, // Skip to: 25
-// CHECK-NEXT: /* 20 */      MCD::OPC_TryDecode, 26, 0, 0, 0, // Opcode: InstB, skip to: 25
+// CHECK-NEXT: /* 20 */      MCD::OPC_TryDecode, 27, 0, 0, 0, // Opcode: InstB, skip to: 25
 // CHECK-NEXT: /* 25 */      MCD::OPC_CheckField, 3, 2, 0, 5, 0, // Skip to: 36
-// CHECK-NEXT: /* 31 */      MCD::OPC_TryDecode, 25, 1, 0, 0, // Opcode: InstA, skip to: 36
+// CHECK-NEXT: /* 31 */      MCD::OPC_TryDecode, 26, 1, 0, 0, // Opcode: InstA, skip to: 36
 // CHECK-NEXT: /* 36 */      MCD::OPC_Fail,
 
 // CHECK: if (DecodeInstB(MI, insn, Address, Decoder) == MCDisassembler::Fail) { DecodeComplete = false; return MCDisassembler::Fail; }
 // CHECK: if (DecodeInstA(MI, insn, Address, Decoder) == MCDisassembler::Fail) { DecodeComplete = false; return MCDisassembler::Fail; }

test/TableGen/trydecode-emission3.td

def InstB : TestInstruction {
   let Inst{1-0} = op;
   let OutOperandList = (outs InstBOp:$op);
   let AsmString = "InstB";
 }
 
 // CHECK:      /* 0 */       MCD::OPC_ExtractField, 4, 4,  // Inst{7-4} ...
 // CHECK-NEXT: /* 3 */       MCD::OPC_FilterValue, 0, 14, 0, // Skip to: 21
 // CHECK-NEXT: /* 7 */       MCD::OPC_CheckField, 2, 2, 0, 5, 0, // Skip to: 18
-// CHECK-NEXT: /* 13 */      MCD::OPC_TryDecode, 26, 0, 0, 0, // Opcode: InstB, skip to: 18
-// CHECK-NEXT: /* 18 */      MCD::OPC_Decode, 25, 1, // Opcode: InstA
+// CHECK-NEXT: /* 13 */      MCD::OPC_TryDecode, 27, 0, 0, 0, // Opcode: InstB, skip to: 18
+// CHECK-NEXT: /* 18 */      MCD::OPC_Decode, 26, 1, // Opcode: InstA
 // CHECK-NEXT: /* 21 */      MCD::OPC_Fail,
 
 // CHECK: if (DecodeInstBOp(MI, tmp, Address, Decoder) == MCDisassembler::Fail) { DecodeComplete = false; return MCDisassembler::Fail; }