[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

2541 ↗

(On Diff #132020)

AFAICT you don't use TRI in the implementation?

[dberris]

• fixup: Remove unused variable in API

llvm/include/llvm/CodeGen/TargetLowering.h

2541 ↗

(On Diff #132020)

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

243

Remove.

[dberris]

Done, thanks @echristo!

compiler-rt/lib/xray/xray_trampoline_x86_64.S

 
 #include "../builtins/assembly.h"
 #include "../sanitizer_common/sanitizer_asm.h"
 
 
 
 .macro SAVE_REGISTERS
   subq $192, %rsp
-  CFI_DEF_CFA_OFFSET(200)
+  CFI_ADJUST_CFA_OFFSET(192)
   // At this point, the stack pointer should be aligned to an 8-byte boundary,
   // because any call instructions that come after this will add another 8
   // bytes and therefore align it to 16-bytes.
   movq %rbp, 184(%rsp)
   movupd  %xmm0, 168(%rsp)
   movupd  %xmm1, 152(%rsp)
   movupd  %xmm2, 136(%rsp)
   movupd  %xmm3, 120(%rsp)
.macro RESTORE_REGISTERS
   movq  48(%rsp), %rdi
   movq  40(%rsp), %rax
   movq  32(%rsp), %rdx
   movq  24(%rsp), %rsi
   movq  16(%rsp), %rcx
   movq  8(%rsp), %r8
   movq  0(%rsp), %r9
   addq  $192, %rsp
-  CFI_DEF_CFA_OFFSET(8)
+  CFI_ADJUST_CFA_OFFSET(-192)
 .endm
 
 .macro ALIGNED_CALL_RAX
   // Call the logging handler, after aligning the stack to a 16-byte boundary.
   // The approach we're taking here uses additional stack space to stash the
   // stack pointer twice before aligning the pointer to 16-bytes. If the stack
   // was 8-byte aligned, it will become 16-byte aligned -- when restoring the
   // pointer, we can always look -8 bytes from the current position to get
   // either of the values we've stashed in the first place.
   pushq %rsp
+  CFI_ADJUST_CFA_OFFSET(8)
   pushq (%rsp)
+  CFI_ADJUST_CFA_OFFSET(8)
   andq $-0x10, %rsp
   callq *%rax
   movq 8(%rsp), %rsp
+  CFI_ADJUST_CFA_OFFSET(-16)
 .endm
 
   .text
 #if !defined(__APPLE__)
   .section .text
 #else
   .section __TEXT,__text
 #endif
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)
   CFI_ENDPROC
 
 //===----------------------------------------------------------------------===//
 
   .globl ASM_SYMBOL(__xray_FunctionExit)
   .align 16, 0x90
   ASM_TYPE_FUNCTION(__xray_FunctionExit)
 ASM_SYMBOL(__xray_FunctionExit):
   CFI_STARTPROC
   // Save the important registers first. Since we're assuming that this
   // function is only jumped into, we only preserve the registers for
   // returning.
   subq  $56, %rsp
-  CFI_DEF_CFA_OFFSET(64)
+  CFI_ADJUST_CFA_OFFSET(56)
   movq  %rbp, 48(%rsp)
   movupd  %xmm0, 32(%rsp)
   movupd  %xmm1, 16(%rsp)
   movq  %rax, 8(%rsp)
   movq  %rdx, 0(%rsp)
   movq  ASM_SYMBOL(_ZN6__xray19XRayPatchedFunctionE)(%rip), %rax
   testq %rax,%rax
   je  .Ltmp2
 
   movl  %r10d, %edi
   movl  $1, %esi
   ALIGNED_CALL_RAX
 
 .Ltmp2:
   // Restore the important registers.
   movq  48(%rsp), %rbp
   movupd  32(%rsp), %xmm0
   movupd  16(%rsp), %xmm1
   movq  8(%rsp), %rax
   movq  0(%rsp), %rdx
   addq  $56, %rsp
-  CFI_DEF_CFA_OFFSET(8)
+  CFI_ADJUST_CFA_OFFSET(-56)
   retq
   ASM_SIZE(__xray_FunctionExit)
   CFI_ENDPROC
 
 //===----------------------------------------------------------------------===//
 
   .globl ASM_SYMBOL(__xray_FunctionTailExit)
   .align 16, 0x90
//===----------------------------------------------------------------------===//
 
   .global ASM_SYMBOL(__xray_CustomEvent)
   .align 16, 0x90
   ASM_TYPE_FUNCTION(__xray_CustomEvent)
 ASM_SYMBOL(__xray_CustomEvent):
   CFI_STARTPROC
   SAVE_REGISTERS
 
+  // 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)
+
   // 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.
   movq ASM_SYMBOL(_ZN6__xray22XRayPatchedCustomEventE)(%rip), %rax
   testq %rax,%rax
   je .LcustomEventCleanup
 
   ALIGNED_CALL_RAX
 
 .LcustomEventCleanup:
+
+  movq 8(%rsp), %r10
+  movq 0(%rsp), %r11
+  addq $24, %rsp
+  CFI_ADJUST_CFA_OFFSET(-24)
   RESTORE_REGISTERS
   retq
   ASM_SIZE(__xray_CustomEvent)
   CFI_ENDPROC
 
 NO_EXEC_STACK_DIRECTIVE

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<[],