[XRay][compiler-rt+llvm] Update XRay register stashing semantics

[dberris]

Summary

This change expands the amount of registers stashed by the entry and
__xray_CustomEvent trampolines.

We've found that since the __xray_CustomEvent trampoline calls can show up in
situations where the scratch registers are being used, and since we don't
typically want to affect the code-gen around the disabled
__xray_customevent(...) intrinsic calls, that we need to save and restore the
state of even the scratch registers in the handling of these custom events.

[dberris]

Re-title and update the definition of the customevent intrinsics to signal side-effects.

[dberris]

Fix minor stack adjustment

[dberris]

Adding @dblaikie to the review for the CFI directives (please feel free to nominate someone else who might know better).

The problem I'm having in particular here is the computation in the ALIGNED_CALL_RAX macro, where we're doing stack re-alignment just before the call to the custom event handler function (line 77). Working offline with a teammate here in Sydney, we're prototyping some potential ways of doing this correctly but using .cfi_escape ... and providing raw DWARF opcodes. Ideally I wouldn't have to go through such lengths, but without the correct CFI directives and/or DWARF information, debugging XRay instrumentation handlers and the trampolines becomes really hard (since the CFI directives would probably be wrong).

Before I go uploading a work-in-progress implementation here, is there an example somewhere of how to potentially resolve this in assembler?

One sticking point is the fact that there seems to be no way to get the current CFI offset from the assembler using either a directive or a special macro -- at least not one I can find from the assembler manuals.

Pointers would be most appreciated.

[davide]

I think in this context CFI stands for Call Frame information rather than Control Flow Integrity.

[davide]

Probably Eric or Adrian can review this.

[dberris]

Rebase and back-out the CFI details. We can follow-up on those as the semantics changes are more important at this time.

[dberris]

Actually back out the CFI-specific changes.

[dberris]

• fixup: only stash r10 and r11 when calling __xray_CustomEvent
• fixup: mark that no registers are preserved by the call to the intrinsic

[dberris]

This is now ready for a look @dblaikie -- or if you know someone else who might know more about the lowering magic involved in SelectionDAG for X86 in particular, a redirect would be most appreciated too.

[dblaikie]

Yeah, can't say I know much of anything about this. @echristo might, or might be able to point to someone who would be suitable.

[dberris]

Thanks @dblaikie -- @echristo, when you have time, would really appreciate some thoughts on the following:

• The register stashing strategy for the custom event intrinsic -- whether it should stash all, some, and/or maybe do something else.

• In the SelectionDAG implementation, attempt to make the call to the intrinsic as non-disruptive to optimisations and register allocation as possible, or whether to go the current state and just signal that the intrinsic may very well use all the registers anyway.

I have a slight preference to incur the cost of stashing all/most of the registers in the trampoline when XRay is enabled and we're doing custom event handling, rather than affect the codegen around the call to the intrisic. But I understand that there might be a middleground that I haven't quite explored yet (or still can't begin to express properly).

Thoughts?

Adding in @chandlerc as well in case, for more feedback.

[echristo]

A bit late, but I have some concerns here.

I think the problem is that you were clobbering a bunch of registers that weren't "available" for you to clobber, but if you set this up as a normal calling convention they'll be saved - which I don't think you want to do. How about just saving all of the registers explicitly in your trampoline? If not, then you might just want to make it a normal call and save/restore over it? Given the goals of the work I'd probably lean toward the former. I think if you clobber all of the registers in call lowering you're just making the surrounding code worse?

Thoughts?

-eric

[echristo]

Also the register saves strike me as a somewhat odd optimization - you could also have just thrown this into SAVE_REGISTERS I think?

-eric

[dberris]

In D40894#975323, @echristo wrote:

A bit late, but I have some concerns here.

I think the problem is that you were clobbering a bunch of registers that weren't "available" for you to clobber, but if you set this up as a normal calling convention they'll be saved - which I don't think you want to do. How about just saving all of the registers explicitly in your trampoline? If not, then you might just want to make it a normal call and save/restore over it? Given the goals of the work I'd probably lean toward the former. I think if you clobber all of the registers in call lowering you're just making the surrounding code worse?

Thoughts?

I agree that we should try not to disturb the surrounding code when we lower the custom event instrumentation points. I'll update it to do all register stashing/restoration in the trampoline.

In D40894#975407, @echristo wrote:

Also the register saves strike me as a somewhat odd optimization - you could also have just thrown this into SAVE_REGISTERS I think?

Good point. Yes, it's much cleaner that way, thanks!

[dberris]

• fixup: Stash all registers in trampolines

[dberris]

• fixup: Remove unnecessary comments

[echristo]

llvm/include/llvm/CodeGen/TargetLowering.h

2524

AFAICT you don't use TRI in the implementation?

[dberris]

• fixup: Remove unused variable in API

llvm/include/llvm/CodeGen/TargetLowering.h

2524

Good catch, yes -- removed.

[dberris]

Re-title to reflect the changes better.

[dberris]

Ready for another look now @echristo -- thanks in advance!

[echristo]

One inline comment then LGTM.

-eric

compiler-rt/lib/xray/xray_trampoline_x86_64.S

244

Remove.

[dberris]

Done, thanks @echristo!

compiler-rt/lib/xray/xray_trampoline_x86_64.S

ASM_SYMBOL(__xray_FunctionEntry):
   SAVE_REGISTERS
 
   // This load has to be atomic, it's concurrent with __xray_patch().
   // On x86/amd64, a simple (type-aligned) MOV instruction is enough.
   movq  ASM_SYMBOL(_ZN6__xray19XRayPatchedFunctionE)(%rip), %rax
   testq %rax, %rax
   je  .Ltmp0
 
-  // The patched function prolog puts its xray_instr_map index into %r10d.
+  // The patched function prologue puts its xray_instr_map index into %r10d.
   movl  %r10d, %edi
   xor %esi,%esi
   ALIGNED_CALL_RAX
 
 .Ltmp0:
   RESTORE_REGISTERS
   retq
   ASM_SIZE(__xray_FunctionEntry)
//===----------------------------------------------------------------------===//
   .global ASM_SYMBOL(__xray_CustomEvent)
   .align 16, 0x90
   ASM_TYPE_FUNCTION(__xray_CustomEvent)
 ASM_SYMBOL(__xray_CustomEvent):
   CFI_STARTPROC
   SAVE_REGISTERS
 
   // We take two arguments to this trampoline, which should be in rdi and rsi
-  // already. We also make sure that we stash %rax because we use that register
-  // to call the logging handler.
+  // already.
   movq ASM_SYMBOL(_ZN6__xray22XRayPatchedCustomEventE)(%rip), %rax
   testq %rax,%rax
   je .LcustomEventCleanup
 
+  // Because calls to __xray_CustomEvent can occur in the middle of a function
+  // which may already be using scratch registers, we preserve more information
+  // in this trampoline than when we are handling entry/exit events.
+  subq $16, %rsp
+  CFI_ADJUST_CFA_OFFSET(16)
+  movq %r10, 8(%rsp)
+  movq %r11, 0(%rsp)
+
   ALIGNED_CALL_RAX
 
+  movq 8(%rsp), %r10
+  movq 0(%rsp), %r11
+  addq $16, %rsp
+  CFI_ADJUST_CFA_OFFSET(-16)
+
 .LcustomEventCleanup:
+
   RESTORE_REGISTERS
   retq
   ASM_SIZE(__xray_CustomEvent)
   CFI_ENDPROC
 
 NO_EXEC_STACK_DIRECTIVE

llvm/include/llvm/CodeGen/TargetLowering.h

protected:
   /// register class are all legal.
   bool isLegalRC(const TargetRegisterInfo &TRI,
                  const TargetRegisterClass &RC) const;
 
   /// Replace/modify any TargetFrameIndex operands with a targte-dependent
   /// sequence of memory operands that is recognized by PrologEpilogInserter.
   MachineBasicBlock *emitPatchPoint(MachineInstr &MI,
                                     MachineBasicBlock *MBB) const;
+
+  /// Replace/modify the XRay custom event operands to provide the correct
+  /// register mask (so that register allocation can deal with the fact that the
+  /// XRay custom event handlers may clobber all the registers).
+  MachineBasicBlock *emitXRayCustomEvent(MachineInstr &MI,
+                                         MachineBasicBlock *MBB,
+                                         const TargetRegisterInfo *TRI) const;
 };
 
 /// This class defines information used to lower LLVM code to legal SelectionDAG
 /// operators that the target instruction selector can accept natively.
 ///
 /// This class also defines callbacks that targets must implement to lower
 /// target-specific constructs to SelectionDAG operators.
 class TargetLowering : public TargetLoweringBase {

llvm/include/llvm/IR/Intrinsics.td

 def int_load_relative: Intrinsic<[llvm_ptr_ty], [llvm_ptr_ty, llvm_anyint_ty],
                                  [IntrReadMem, IntrArgMemOnly]>;
 
 // Xray intrinsics
 //===----------------------------------------------------------------------===//
 // Custom event logging for x-ray.
 // Takes a pointer to a string and the length of the string.
 def int_xray_customevent : Intrinsic<[], [llvm_ptr_ty, llvm_i32_ty],
-                                     [NoCapture<0>, ReadOnly<0>, IntrWriteMem]>;
+                                     [NoCapture<0>, ReadOnly<0>, IntrWriteMem, IntrHasSideEffects]>;
 //===----------------------------------------------------------------------===//
 
 //===------ Memory intrinsics with element-wise atomicity guarantees ------===//
 //
 
 // @llvm.memcpy.element.unordered.atomic.*(dest, src, length, elementsize)
 def int_memcpy_element_unordered_atomic
     : Intrinsic<[],

llvm/lib/CodeGen/SelectionDAG/FastISel.cpp

Ops.push_back(MachineOperand::CreateReg(getRegForValue(I->getArgOperand(0)),
                                           /*IsDef=*/false));
   Ops.push_back(MachineOperand::CreateReg(getRegForValue(I->getArgOperand(1)),
                                           /*IsDef=*/false));
   MachineInstrBuilder MIB =
       BuildMI(*FuncInfo.MBB, FuncInfo.InsertPt, DbgLoc,
               TII.get(TargetOpcode::PATCHABLE_EVENT_CALL));
   for (auto &MO : Ops)
     MIB.add(MO);
+
+  // We want to make sure that we mark all registers as clobbered for this
+  // function call. This is so that the codegen around the customevent call will
+  // treat this function call as something that will potentially use all the
+  // registers.
+  MIB.addRegMask(TRI.getNoPreservedMask());
+
   // Insert the Patchable Event Call instruction, that gets lowered properly.
   return true;
 }
 
 
 /// Returns an AttributeList representing the attributes applied to the return
 /// value of the given call.
 static AttributeList getReturnAttrs(FastISel::CallLoweringInfo &CLI) {

llvm/lib/CodeGen/TargetLoweringBase.cpp

for(unsigned OperIdx = 0; OperIdx != MI->getNumOperands(); ++OperIdx) {
     MBB->insert(MachineBasicBlock::iterator(MI), MIB);
     OperIdx += (MIB->getNumOperands() - MI->getNumOperands()) - 1;
     MI->eraseFromParent();
     MI = MIB;
   }
   return MBB;
 }
 
+MachineBasicBlock *
+TargetLoweringBase::emitXRayCustomEvent(MachineInstr &MI,
+                                        MachineBasicBlock *MBB,
+                                        const TargetRegisterInfo *TRI) const {
+  assert(MI.getOpcode() == TargetOpcode::PATCHABLE_EVENT_CALL &&
+         "Called emitXRayCustomEvent on the wrong MI!");
+  auto &MF = *MI.getMF();
+  auto MIB = BuildMI(MF, MI.getDebugLoc(), MI.getDesc());
+  for (unsigned OpIdx = 0; OpIdx != MI.getNumOperands(); ++OpIdx)
+    MIB.add(MI.getOperand(OpIdx));
+
+  // We need to make sure that all registers are considered clobbered after this
+  // call. There will be *no* safe registers after this call.
+  MIB.addRegMask(TRI->getNoPreservedMask());
+  MBB->insert(MachineBasicBlock::iterator(MI), MIB);
+  MI.eraseFromParent();
+  return MBB;
+}
+
 /// findRepresentativeClass - Return the largest legal super-reg register class
 /// of the register class for the specified type and its associated "cost".
 // This function is in TargetLowering because it uses RegClassForVT which would
 // need to be moved to TargetRegisterInfo and would necessitate moving
 // isTypeLegal over as well - a massive change that would just require
 // TargetLowering having a TargetRegisterInfo class member that it would use.
 std::pair<const TargetRegisterClass *, uint8_t>
 TargetLoweringBase::findRepresentativeClass(const TargetRegisterInfo *TRI,

llvm/lib/Target/X86/X86ISelLowering.cpp

case TargetOpcode::STATEPOINT:
     // this point in the process.  We diverge later.
     return emitPatchPoint(MI, BB);
 
   case TargetOpcode::STACKMAP:
   case TargetOpcode::PATCHPOINT:
     return emitPatchPoint(MI, BB);
 
   case TargetOpcode::PATCHABLE_EVENT_CALL:
-    // Do nothing here, handle in xray instrumentation pass.
-    return BB;
+    return emitXRayCustomEvent(MI, BB, Subtarget.getRegisterInfo());
 
   case X86::LCMPXCHG8B: {
     const X86RegisterInfo *TRI = Subtarget.getRegisterInfo();
     // In addition to 4 E[ABCD] registers implied by encoding, CMPXCHG8B
     // requires a memory operand. If it happens that current architecture is
     // i686 and for current function we need a base pointer
     // - which is ESI for i686 - register allocator would not be able to
     // allocate registers for an address in form of X(%reg, %reg, Y)