fix: preserve SSE state in x86_64 bthread context switch

src/bthread/context.cpp

@@ -347,20 +347,40 @@ __asm (
 "    pushq  %r14  \n"
 "    pushq  %r13  \n"
 "    pushq  %r12  \n"
-"    leaq  -0x8(%rsp), %rsp\n"
+"    leaq  -0xa8(%rsp), %rsp\n"
+"    movups %xmm6,  0x00(%rsp)\n"
+"    movups %xmm7,  0x10(%rsp)\n"
+"    movups %xmm8,  0x20(%rsp)\n"
+"    movups %xmm9,  0x30(%rsp)\n"
+"    movups %xmm10, 0x40(%rsp)\n"
+"    movups %xmm11, 0x50(%rsp)\n"
+"    movups %xmm12, 0x60(%rsp)\n"
+"    movups %xmm13, 0x70(%rsp)\n"
+"    movups %xmm14, 0x80(%rsp)\n"
+"    movups %xmm15, 0x90(%rsp)\n"
 "    cmp  $0, %rcx\n"
 "    je  1f\n"
-"    stmxcsr  (%rsp)\n"
-"    fnstcw   0x4(%rsp)\n"
+"    stmxcsr  0xa0(%rsp)\n"
+"    fnstcw   0xa4(%rsp)\n"
 "1:\n"
 "    movq  %rsp, (%rdi)\n"
 "    movq  %rsi, %rsp\n"
+"    movups 0x00(%rsp), %xmm6\n"
+"    movups 0x10(%rsp), %xmm7\n"
+"    movups 0x20(%rsp), %xmm8\n"
+"    movups 0x30(%rsp), %xmm9\n"
+"    movups 0x40(%rsp), %xmm10\n"
+"    movups 0x50(%rsp), %xmm11\n"
+"    movups 0x60(%rsp), %xmm12\n"
+"    movups 0x70(%rsp), %xmm13\n"
+"    movups 0x80(%rsp), %xmm14\n"
+"    movups 0x90(%rsp), %xmm15\n"
 "    cmp  $0, %rcx\n"
 "    je  2f\n"
-"    ldmxcsr  (%rsp)\n"
-"    fldcw  0x4(%rsp)\n"
+"    ldmxcsr  0xa0(%rsp)\n"
+"    fldcw    0xa4(%rsp)\n"
 "2:\n"
-"    leaq  0x8(%rsp), %rsp\n"
+"    leaq  0xa8(%rsp), %rsp\n"
 "    popq  %r12  \n"
 "    popq  %r13  \n"
 "    popq  %r14  \n"
@@ -387,12 +407,12 @@ __asm (
 "bthread_make_fcontext:\n"
 "    movq  %rdi, %rax\n"
 "    andq  $-16, %rax\n"
-"    leaq  -0x48(%rax), %rax\n"
-"    movq  %rdx, 0x38(%rax)\n"
-"    stmxcsr  (%rax)\n"
-"    fnstcw   0x4(%rax)\n"
+"    leaq  -0xe8(%rax), %rax\n"
+"    movq  %rdx, 0xd8(%rax)\n"
+"    stmxcsr  0xa0(%rax)\n"
+"    fnstcw   0xa4(%rax)\n"
 "    leaq  finish(%rip), %rcx\n"
-"    movq  %rcx, 0x40(%rax)\n"
+"    movq  %rcx, 0xe0(%rax)\n"
 "    ret \n"
 "finish:\n"
 "    xorq  %rdi, %rdi\n"

src/bthread/task_tracer.cpp

@@ -35,6 +35,23 @@
 
 namespace bthread {
 
+namespace {
+
+// These slots must match the Linux x86_64 frame built by
+// bthread_jump_fcontext in context.cpp. XMM6-XMM15 add 10 128-bit values,
+// or 20 uintptr_t slots, ahead of the previously saved registers.
+constexpr size_t kSavedXmmRegisterCount = 10;
+constexpr size_t kSavedXmmRegisterBytes = 16;
+constexpr size_t kSavedXmmRegisterSlots =
+    kSavedXmmRegisterCount * kSavedXmmRegisterBytes / sizeof(uintptr_t);
+constexpr size_t kRbpContextSlot = 6 + kSavedXmmRegisterSlots;
+constexpr size_t kRipContextSlot = 7 + kSavedXmmRegisterSlots;
+#if UNW_VERSION_MAJOR >= 1 && UNW_VERSION_MINOR >= 7
+constexpr size_t kRspContextSlot = 8 + kSavedXmmRegisterSlots;
+#endif
+
+}  // namespace
+
 DEFINE_uint32(signal_trace_timeout_ms, 50, "Timeout for signal trace in ms");
 BUTIL_VALIDATE_GFLAG(signal_trace_timeout_ms, butil::PositiveInteger<uint32_t>);
 // Note that SIGURG handler may be registered by some library such as cgo
@@ -280,16 +297,16 @@ unw_cursor_t TaskTracer::MakeCursor(bthread_fcontext_t fcontext) {
 
     // Only need RBP, RIP, RSP on x86_64.
     // The base pointer (RBP).
-    if (unw_set_reg(&cursor, UNW_X86_64_RBP, regs[6]) != 0) {
+    if (unw_set_reg(&cursor, UNW_X86_64_RBP, regs[kRbpContextSlot]) != 0) {
         LOG(ERROR) << "Fail to set RBP";
     }
     // The instruction pointer (RIP).
-    if (unw_set_reg(&cursor, UNW_REG_IP, regs[7]) != 0) {
+    if (unw_set_reg(&cursor, UNW_REG_IP, regs[kRipContextSlot]) != 0) {
         LOG(ERROR) << "Fail to set RIP";
     }
 #if UNW_VERSION_MAJOR >= 1 && UNW_VERSION_MINOR >= 7
     // The stack pointer (RSP).
-    if (unw_set_reg(&cursor, UNW_REG_SP, regs[8]) != 0) {
+    if (unw_set_reg(&cursor, UNW_REG_SP, regs[kRspContextSlot]) != 0) {
         LOG(ERROR) << "Fail to set RSP";
     }
 #endif

test/bthread_context_unittest.cpp

@@ -0,0 +1,112 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <gtest/gtest.h>
+
+#include <cstdint>
+#include <cstdlib>
+#include <cstring>
+
+#include "bthread/context.h"
+
+#if defined(BTHREAD_CONTEXT_PLATFORM_linux_x86_64)
+
+namespace {
+
+constexpr size_t kStackSize = 8192;
+constexpr size_t kXmmRegisterCount = 10;
+
+struct XmmState {
+    uint8_t registers[kXmmRegisterCount][16];
+};
+
+bthread_fcontext_t g_main_context;
+bthread_fcontext_t g_test_context;
+XmmState g_clobber_state;
+
+extern "C" intptr_t jump_with_xmm_state(
+    bthread_fcontext_t* old_context, bthread_fcontext_t new_context,
+    intptr_t value, const XmmState* input, XmmState* output);
+
+// Keep the register setup, context switch, and register capture in one
+// assembly function so the compiler cannot use an XMM register between them.
+__asm__(
+".text\n"
+".p2align 4,,15\n"
+".type jump_with_xmm_state,@function\n"
+"jump_with_xmm_state:\n"
+"    movups 0x00(%rcx), %xmm6\n"
+"    movups 0x10(%rcx), %xmm7\n"
+"    movups 0x20(%rcx), %xmm8\n"
+"    movups 0x30(%rcx), %xmm9\n"
+"    movups 0x40(%rcx), %xmm10\n"
+"    movups 0x50(%rcx), %xmm11\n"
+"    movups 0x60(%rcx), %xmm12\n"
+"    movups 0x70(%rcx), %xmm13\n"
+"    movups 0x80(%rcx), %xmm14\n"
+"    movups 0x90(%rcx), %xmm15\n"
+"    pushq %r8\n"
+"    xorl %ecx, %ecx\n"
+"    call bthread_jump_fcontext\n"
+"    popq %r8\n"
+"    movups %xmm6, 0x00(%r8)\n"
+"    movups %xmm7, 0x10(%r8)\n"
+"    movups %xmm8, 0x20(%r8)\n"
+"    movups %xmm9, 0x30(%r8)\n"
+"    movups %xmm10, 0x40(%r8)\n"
+"    movups %xmm11, 0x50(%r8)\n"
+"    movups %xmm12, 0x60(%r8)\n"
+"    movups %xmm13, 0x70(%r8)\n"
+"    movups %xmm14, 0x80(%r8)\n"
+"    movups %xmm15, 0x90(%r8)\n"
+"    ret\n"
+".size jump_with_xmm_state,.-jump_with_xmm_state\n");
+
+void overwrite_xmm_and_return(intptr_t) {
+    XmmState ignored = {};
+    jump_with_xmm_state(&g_test_context, g_main_context, 0,
+                        &g_clobber_state, &ignored);
+}
+
+TEST(BthreadContextTest, preserves_xmm6_through_xmm15) {
+    XmmState expected;
+    XmmState actual = {};
+    for (size_t reg = 0; reg < kXmmRegisterCount; ++reg) {
+        for (size_t byte = 0; byte < 16; ++byte) {
+            expected.registers[reg][byte] =
+                static_cast<uint8_t>(reg * 16 + byte);
+            g_clobber_state.registers[reg][byte] =
+                static_cast<uint8_t>(0xff - reg * 16 - byte);
+        }
+    }
+
+    void* stack = std::malloc(kStackSize);
+    ASSERT_NE(nullptr, stack);
+    g_main_context = nullptr;
+    g_test_context = bthread_make_fcontext(
+        static_cast<char*>(stack) + kStackSize, kStackSize,
+        overwrite_xmm_and_return);
+
+    jump_with_xmm_state(&g_main_context, g_test_context, 0, &expected, &actual);
+
+    EXPECT_EQ(0, std::memcmp(&expected, &actual, sizeof(expected)));
+    std::free(stack);
+}
+
+}  // namespace
+
+#endif  // BTHREAD_CONTEXT_PLATFORM_linux_x86_64