Repository: bytecodealliance/wasmtime
Pull request #7978 (2024-02-22 11:55:18)

s390x: Fix TLS GD relocation order
When emitting a call to __tls_get_offset, the instruction needs to carry two relocations, a R_390_PLT32DBL targeting __tls_get_offset and a R_390_TLS_GDCALL targeting the TLS symbol.  Specifically, the system linker expects to see these two relocation in that order.

However, the cranelift backend currently emits the relocations in reverse order - this unfortunately causes the linker to corrupt the instruction sequence when performing a TLS relaxation.

To fix this in the backend, I need support in machinst common code to emit a relocation at some offset to the instruction about to be emitted (e.g. the relocation should target two bytes into the 6-byte instruction that will be emitted next).  I've added a new routine add_reloc_at_offset to that effect.  This also allowed to simplify some existing code in the backend.

<!--
Please make sure you include the following information:

- If this work has been discussed elsewhere, please include a link to that
  conversation. If it was discussed in an issue, just mention "issue #...".

- Explain why this change is needed. If the details are in an issue already,
  this can be brief.

Our development process is documented in the Wasmtime book:
https://docs.wasmtime.dev/contributing-development-process.html

Please ensure all communication follows the code of conduct:
https://github.com/bytecodealliance/wasmtime/blob/main/CODE_OF_CONDUCT.md
-->



Patch:
diff --git a/cranelift/codegen/src/isa/s390x/inst/emit.rs b/cranelift/codegen/src/isa/s390x/inst/emit.rs
index 7efb8b4df..8d9c72c70 100644
--- a/cranelift/codegen/src/isa/s390x/inst/emit.rs
+++ b/cranelift/codegen/src/isa/s390x/inst/emit.rs
@@ -218,15 +218,8 @@ pub fn mem_emit(
         &MemArg::Symbol {
             ref name, offset, ..
         } => {
-            let reloc = Reloc::S390xPCRel32Dbl;
-            put_with_reloc(
-                sink,
-                &enc_ril_b(opcode_ril.unwrap(), rd, 0),
-                2,
-                reloc,
-                name,
-                offset.into(),
-            );
+            sink.add_reloc_at_offset(2, Reloc::S390xPCRel32Dbl, &**name, (offset + 2).into());
+            put(sink, &enc_ril_b(opcode_ril.unwrap(), rd, 0));
         }
         _ => unreachable!(),
     }
@@ -1320,25 +1313,6 @@ fn put_with_trap(sink: &mut MachBuffer<Inst>, enc: &[u8], trap_code: TrapCode) {
     sink.put1(enc[len - 1]);
 }
 
-/// Emit encoding to sink, adding a relocation at byte offset.
-fn put_with_reloc(
-    sink: &mut MachBuffer<Inst>,
-    enc: &[u8],
-    offset: usize,
-    ri2_reloc: Reloc,
-    ri2_name: &ExternalName,
-    ri2_offset: i64,
-) {
-    let len = enc.len();
-    for i in 0..offset {
-        sink.put1(enc[i]);
-    }
-    sink.add_reloc(ri2_reloc, ri2_name, ri2_offset + offset as i64);
-    for i in offset..len {
-        sink.put1(enc[i]);
-    }
-}
-
 /// State carried between emissions of a sequence of instructions.
 #[derive(Default, Clone, Debug)]
 pub struct EmitState {
@@ -3505,6 +3479,13 @@ impl Inst {
             &Inst::Call { link, ref info } => {
                 debug_assert_eq!(link.to_reg(), gpr(14));
 
+                let opcode = 0xc05; // BRASL
+
+                // Add relocation for target function.  This has to be done *before*
+                // the S390xTlsGdCall relocation if any, to ensure linker relaxation
+                // works correctly.
+                sink.add_reloc_at_offset(2, Reloc::S390xPLTRel32Dbl, &info.dest, 2);
+
                 // Add relocation for TLS libcalls to enable linker optimizations.
                 match &info.tls_symbol {
                     None => {}
@@ -3514,19 +3495,10 @@ impl Inst {
                     _ => unreachable!(),
                 }
 
-                let opcode = 0xc05; // BRASL
-                let reloc = Reloc::S390xPLTRel32Dbl;
                 if let Some(s) = state.take_stack_map() {
                     sink.add_stack_map(StackMapExtent::UpcomingBytes(6), s);
                 }
-                put_with_reloc(
-                    sink,
-                    &enc_ril_b(opcode, link.to_reg(), 0),
-                    2,
-                    reloc,
-                    &info.dest,
-                    0,
-                );
+                put(sink, &enc_ril_b(opcode, link.to_reg(), 0));
                 if info.opcode.is_call() {
                     sink.add_call_site(info.opcode);
                 }
diff --git a/cranelift/codegen/src/machinst/buffer.rs b/cranelift/codegen/src/machinst/buffer.rs
index 901102cf7..57c3faa8d 100644
--- a/cranelift/codegen/src/machinst/buffer.rs
+++ b/cranelift/codegen/src/machinst/buffer.rs
@@ -1549,9 +1549,10 @@ impl<I: VCodeInst> MachBuffer<I> {
         }
     }
 
-    /// Add an external relocation at the current offset.
-    pub fn add_reloc<T: Into<RelocTarget> + Clone>(
+    /// Add an external relocation at the given offset from current offset.
+    pub fn add_reloc_at_offset<T: Into<RelocTarget> + Clone>(
         &mut self,
+        offset: CodeOffset,
         kind: Reloc,
         target: &T,
         addend: Addend,
@@ -1591,13 +1592,23 @@ impl<I: VCodeInst> MachBuffer<I> {
         // when a relocation can't otherwise be resolved later, so it shouldn't
         // actually result in any memory unsafety or anything like that.
         self.relocs.push(MachReloc {
-            offset: self.data.len() as CodeOffset,
+            offset: self.data.len() as CodeOffset + offset,
             kind,
             target,
             addend,
         });
     }
 
+    /// Add an external relocation at the current offset.
+    pub fn add_reloc<T: Into<RelocTarget> + Clone>(
+        &mut self,
+        kind: Reloc,
+        target: &T,
+        addend: Addend,
+    ) {
+        self.add_reloc_at_offset(0, kind, target, addend);
+    }
+
     /// Add a trap record at the current offset.
     pub fn add_trap(&mut self, code: TrapCode) {
         self.traps.push(MachTrap {
diff --git a/cranelift/codegen/src/machinst/mod.rs b/cranelift/codegen/src/machinst/mod.rs
index 3ebcf8b74..3e3034c9a 100644
--- a/cranelift/codegen/src/machinst/mod.rs
+++ b/cranelift/codegen/src/machinst/mod.rs
@@ -430,7 +430,7 @@ impl<T: CompilePhase> CompiledCodeBase<T> {
                 let end = i.address() + i.bytes().len() as u64;
                 let contains = |off| i.address() <= off && off < end;
 
-                if let Some(reloc) = relocs.iter().find(|reloc| contains(reloc.offset as u64)) {
+                for reloc in relocs.iter().filter(|reloc| contains(reloc.offset as u64)) {
                     write!(
                         buf,
                         " ; reloc_external {} {} {}",
diff --git a/winch/codegen/src/codegen/bounds.rs b/winch/codegen/src/codegen/bounds.rs
index b788da731..c44ce172f 100644
--- a/winch/codegen/src/codegen/bounds.rs
+++ b/winch/codegen/src/codegen/bounds.rs
@@ -3,7 +3,7 @@
 //! recommended when working on this area of Winch.
 use super::env::{HeapData, HeapStyle};
 use crate::{
-    abi::ABI,
+    abi::{vmctx, ABI},
     codegen::CodeGenContext,
     isa::reg::Reg,
     masm::{IntCmpKind, MacroAssembler, OperandSize, RegImm, TrapCode},
@@ -102,7 +102,7 @@ where
                 masm.load_ptr(addr, scratch);
                 scratch
             } else {
-                <M::ABI as ABI>::vmctx_reg()
+                vmctx!(M)
             };
             let addr = masm.address_at_reg(base, heap.current_length_offset);
             masm.load_ptr(addr, dst);
@@ -205,7 +205,7 @@ pub(crate) fn load_heap_addr_unchecked<M>(
     } else {
         // Else if the WebAssembly memory is defined in the current module,
         // simply use the `VMContext` as the base for subsequent operations.
-        <M::ABI as ABI>::vmctx_reg()
+        vmctx!(M)
     };
 
     // Load the base of the memory into the `addr` register.
diff --git a/winch/codegen/src/codegen/mod.rs b/winch/codegen/src/codegen/mod.rs
index 1438b2026..e867cc380 100644
--- a/winch/codegen/src/codegen/mod.rs
+++ b/winch/codegen/src/codegen/mod.rs
@@ -94,7 +94,7 @@ where
         // Stack overflow checks must occur during the function prologue to ensure that unwinding
         // will not assume they're user-handlable exceptions. As the `save_clobbers` call below
         // marks the end of the prologue for unwinding annotations, we make the stack check here.
-        self.masm.check_stack();
+        self.masm.check_stack(vmctx!(M));
 
         // We don't have any callee save registers in the winch calling convention, but
         // `save_clobbers` does some useful work for setting up unwinding state, and marks the end
diff --git a/winch/codegen/src/isa/aarch64/masm.rs b/winch/codegen/src/isa/aarch64/masm.rs
index 01fa6fd27..bc6fbe097 100644
--- a/winch/codegen/src/isa/aarch64/masm.rs
+++ b/winch/codegen/src/isa/aarch64/masm.rs
@@ -49,7 +49,7 @@ impl Masm for MacroAssembler {
         self.move_sp_to_shadow_sp();
     }
 
-    fn check_stack(&mut self) {
+    fn check_stack(&mut self, _vmctx: Reg) {
         // TODO: Implement when we have more complete assembler support.
     }
 
diff --git a/winch/codegen/src/isa/x64/masm.rs b/winch/codegen/src/isa/x64/masm.rs
index a1f42488d..0e6b02247 100644
--- a/winch/codegen/src/isa/x64/masm.rs
+++ b/winch/codegen/src/isa/x64/masm.rs
@@ -9,16 +9,16 @@ use crate::masm::{
     DivKind, ExtendKind, FloatCmpKind, Imm as I, IntCmpKind, MacroAssembler as Masm, OperandSize,
     RegImm, RemKind, RoundingMode, ShiftKind, TrapCode, TruncKind, TRUSTED_FLAGS, UNTRUSTED_FLAGS,
 };
-use crate::{
-    abi::ABI,
-    masm::{SPOffset, StackSlot},
-    stack::TypedReg,
-};
 use crate::{
     abi::{self, align_to, calculate_frame_adjustment, LocalSlot},
     codegen::{ptr_type_from_ptr_size, CodeGenContext, FuncEnv, HeapData, TableData},
     stack::Val,
 };
+use crate::{
+    abi::{vmctx, ABI},
+    masm::{SPOffset, StackSlot},
+    stack::TypedReg,
+};
 use crate::{
     isa::reg::{Reg, RegClass},
     masm::CalleeKind,
@@ -77,12 +77,12 @@ impl Masm for MacroAssembler {
             .mov_rr(stack_pointer, frame_pointer, OperandSize::S64);
     }
 
-    fn check_stack(&mut self) {
+    fn check_stack(&mut self, vmctx: Reg) {
         let ptr_size: u8 = self.ptr_size.bytes().try_into().unwrap();
         let scratch = regs::scratch();
 
         self.load_ptr(
-            self.address_at_vmctx(ptr_size.vmcontext_runtime_limits().into()),
+            self.address_at_reg(vmctx, ptr_size.vmcontext_runtime_limits().into()),
             scratch,
         );
 
@@ -110,7 +110,7 @@ impl Masm for MacroAssembler {
                     RegClass::Float => align_to(total, float_bytes) + float_bytes,
                     RegClass::Vector => unimplemented!(),
                 }),
-            16,
+            float_bytes,
         );
 
         // Emit unwind info.
@@ -148,6 +148,9 @@ impl Masm for MacroAssembler {
     }
 
     fn restore_clobbers(&mut self, clobbers: &[(Reg, OperandSize)]) {
+        let int_bytes: u32 = Self::ABI::word_bytes().try_into().unwrap();
+        let float_bytes = int_bytes * 2;
+
         let mut off = 0;
         for &(reg, size) in clobbers {
             // Align the current offset
@@ -161,7 +164,7 @@ impl Masm for MacroAssembler {
             off += size.bytes();
         }
 
-        self.free_stack(align_to(off, 16));
+        self.free_stack(align_to(off, float_bytes));
     }
 
     fn push(&mut self, reg: Reg, size: OperandSize) -> StackSlot {
@@ -240,7 +243,6 @@ impl Masm for MacroAssembler {
         table_data: &TableData,
         context: &mut CodeGenContext,
     ) -> Self::Address {
-        let vmctx = <Self::ABI as ABI>::vmctx_reg();
         let scratch = regs::scratch();
         let bound = context.any_gpr(self);
         let tmp = context.any_gpr(self);
@@ -258,7 +260,7 @@ impl Masm for MacroAssembler {
         } else {
             // Else, simply move the vmctx register into the addr register as
             // the base to calculate the table address.
-            self.asm.mov_rr(vmctx, ptr_base, self.ptr_size);
+            self.asm.mov_rr(vmctx!(Self), ptr_base, self.ptr_size);
         };
 
         // OOB check.
@@ -302,7 +304,6 @@ impl Masm for MacroAssembler {
     }
 
     fn table_size(&mut self, table_data: &TableData, context: &mut CodeGenContext) {
-        let vmctx = <Self::ABI as ABI>::vmctx_reg();
         let scratch = regs::scratch();
         let size = context.any_gpr(self);
 
@@ -314,7 +315,7 @@ impl Masm for MacroAssembler {
                 TRUSTED_FLAGS,
             );
         } else {
-            self.asm.mov_rr(vmctx, scratch, self.ptr_size);
+            self.asm.mov_rr(vmctx!(Self), scratch, self.ptr_size);
         };
 
         let size_addr = Address::offset(scratch, table_data.current_elems_offset);
@@ -331,7 +332,6 @@ impl Masm for MacroAssembler {
     fn memory_size(&mut self, heap_data: &HeapData, context: &mut CodeGenContext) {
         let size_reg = context.any_gpr(self);
         let scratch = regs::scratch();
-        let vmctx = <Self::ABI as ABI>::vmctx_reg();
 
         let base = if let Some(offset) = heap_data.import_from {
             self.asm.movzx_mr(
@@ -342,7 +342,7 @@ impl Masm for MacroAssembler {
             );
             scratch
         } else {
-            vmctx
+            vmctx!(Self)
         };
 
         let size_addr = Address::offset(base, heap_data.current_length_offset);
@@ -380,7 +380,7 @@ impl Masm for MacroAssembler {
     }
 
     fn address_at_vmctx(&self, offset: u32) -> Self::Address {
-        Address::offset(<Self::ABI as ABI>::vmctx_reg(), offset)
+        Address::offset(vmctx!(Self), offset)
     }
 
     fn store_ptr(&mut self, src: Reg, dst: Self::Address) {
diff --git a/winch/codegen/src/masm.rs b/winch/codegen/src/masm.rs
index 7aba29328..890d67902 100644
--- a/winch/codegen/src/masm.rs
+++ b/winch/codegen/src/masm.rs
@@ -488,7 +488,7 @@ pub(crate) trait MacroAssembler {
     }
 
     /// Emit a stack check.
-    fn check_stack(&mut self);
+    fn check_stack(&mut self, vmctx: Reg);
 
     /// Emit the function epilogue.
     fn epilogue(&mut self, locals_size: u32);
diff --git a/winch/codegen/src/trampoline.rs b/winch/codegen/src/trampoline.rs
index 37e2ae1a0..23169262e 100644
--- a/winch/codegen/src/trampoline.rs
+++ b/winch/codegen/src/trampoline.rs
@@ -98,11 +98,12 @@ where
             .map(|operand| RegImm::reg(operand.unwrap_reg()))
             .ok_or_else(|| anyhow!("Expected value pointer to be in a register"))?;
 
-        self.prologue_with_callee_saved();
-
         // Assign the caller and caller VMContext arguments.
         let (vmctx, caller_vmctx) = Self::callee_and_caller_vmctx(&array_sig.params)?;
         let (dst_callee_vmctx, dst_caller_vmctx) = Self::callee_and_caller_vmctx(&wasm_sig.params)?;
+
+        self.prologue_with_callee_saved(caller_vmctx);
+
         self.masm
             .mov(vmctx.into(), dst_callee_vmctx, self.pointer_type.into());
         self.masm.mov(
@@ -194,9 +195,9 @@ where
     pub fn emit_native_to_wasm(mut self, ty: &WasmFuncType, callee_index: FuncIndex) -> Result<()> {
         let native_sig = native_sig::<M::ABI>(&ty, &self.call_conv);
         let wasm_sig = wasm_sig::<M::ABI>(&ty);
-        let (vmctx, _) = Self::callee_and_caller_vmctx(&native_sig.params)?;
+        let (vmctx, caller_vmctx) = Self::callee_and_caller_vmctx(&native_sig.params)?;
 
-        self.prologue_with_callee_saved();
+        self.prologue_with_callee_saved(caller_vmctx);
 
         let vmctx_runtime_limits_addr = self.vmctx_runtime_limits_addr(vmctx);
         let ret_area = self.make_ret_area(&wasm_sig);
@@ -363,7 +364,7 @@ where
         let (vmctx, caller_vmctx) = Self::callee_and_caller_vmctx(&wasm_sig.params).unwrap();
         let vmctx_runtime_limits_addr = self.vmctx_runtime_limits_addr(caller_vmctx);
 
-        self.prologue();
+        self.prologue(caller_vmctx);
 
         // Save the FP and return address when exiting Wasm.
         // TODO: Once Winch supports comparison operators,
@@ -620,17 +621,17 @@ where
     }
 
     /// The trampoline's prologue.
-    fn prologue(&mut self) {
+    fn prologue(&mut self, vmctx: Reg) {
         self.masm.prologue();
-        // TODO: emit a stack check
+        self.masm.check_stack(vmctx);
         self.masm.save_clobbers(&[]);
     }
 
     /// Similar to [Trampoline::prologue], but saves
     /// callee-saved registers.
-    fn prologue_with_callee_saved(&mut self) {
+    fn prologue_with_callee_saved(&mut self, vmctx: Reg) {
         self.masm.prologue();
-        // TODO: emit a stack check
+        self.masm.check_stack(vmctx);
         // Save any callee-saved registers.
         self.masm.save_clobbers(&self.callee_saved_regs);
     }


Test patch:
diff --git a/cranelift/filetests/filetests/isa/s390x/tls_elf.clif b/cranelift/filetests/filetests/isa/s390x/tls_elf.clif
index 0436b892b..d0011a86f 100644
--- a/cranelift/filetests/filetests/isa/s390x/tls_elf.clif
+++ b/cranelift/filetests/filetests/isa/s390x/tls_elf.clif
@@ -37,7 +37,7 @@ block0(v0: i32):
 ;   .byte 0x00, 0x00
 ;   .byte 0x00, 0x00
 ;   lg %r2, 0(%r1)
-;   brasl %r14, 0x22 ; reloc_external TlsGdCall u1:0 0
+;   brasl %r14, 0x22 ; reloc_external PLTRel32Dbl %ElfTlsGetOffset 2 ; reloc_external TlsGdCall u1:0 0
 ;   ear %r3, %a0
 ;   sllg %r5, %r3, 0x20
 ;   ear %r5, %a1