feat(verify): add Call trampolines for arm64, riscv64, loong64

Assisted-by: MiMo V2.5 Pro
This commit is contained in:
2026-08-21 00:35:21 +02:00
parent 1160c96427
commit 181d8e508c
14 changed files with 643 additions and 27 deletions
+1 -1
View File
@@ -1,7 +1,7 @@
// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
// SPDX-License-Identifier: BSD-3-Clause
//go:build !amd64
//go:build !amd64 && !arm64 && !riscv64 && !loong64
package verify
+57
View File
@@ -0,0 +1,57 @@
// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
// SPDX-License-Identifier: BSD-3-Clause
//go:build arm64
package verify
import (
"encoding/binary"
"fmt"
"reflect"
"syscall"
"unsafe"
)
// enterJIT switches to the prepared stack and branches to fn.
//
//go:nosplit
func enterJIT(fn uintptr, stack uintptr)
// leaveJIT restores the Go stack after a JIT function returns.
//
//go:nosplit
func leaveJIT()
// leaveJITAddr is the machine address of leaveJIT.
var leaveJITAddr uintptr
func init() {
leaveJITAddr = reflect.ValueOf(leaveJIT).Pointer()
}
const stackPad = 64
// Call invokes the assembled function at fnAddr with the given ABI0 argument
// block. On arm64, the ABI0 convention places arguments on the stack starting
// at [SP+8] (same as amd64 ABI0 in this project).
func Call(fnAddr uintptr, args []byte) ([]byte, error) {
stackSize := stackPad + 8 + len(args) + 64
stackMem, err := syscall.Mmap(-1, 0, stackSize,
syscall.PROT_READ|syscall.PROT_WRITE, syscall.MAP_PRIVATE|syscall.MAP_ANON)
if err != nil {
return nil, fmt.Errorf("verify: stack mmap: %w", err)
}
defer syscall.Munmap(stackMem)
retOff := stackPad
binary.LittleEndian.PutUint64(stackMem[retOff:retOff+8], uint64(leaveJITAddr))
copy(stackMem[retOff+8:], args)
stackBase := uintptr(unsafe.Pointer(&stackMem[retOff]))
enterJIT(fnAddr, stackBase)
out := make([]byte, len(args))
copy(out, stackMem[retOff+8:retOff+8+len(args)])
return out, nil
}
+50
View File
@@ -0,0 +1,50 @@
// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
// SPDX-License-Identifier: BSD-3-Clause
//go:build loong64
package verify
import (
"encoding/binary"
"fmt"
"reflect"
"syscall"
"unsafe"
)
//go:nosplit
func enterJIT(fn uintptr, stack uintptr)
//go:nosplit
func leaveJIT()
var leaveJITAddr uintptr
func init() {
leaveJITAddr = reflect.ValueOf(leaveJIT).Pointer()
}
const stackPad = 64
// Call invokes the assembled function at fnAddr with the given ABI0 argument block.
func Call(fnAddr uintptr, args []byte) ([]byte, error) {
stackSize := stackPad + 8 + len(args) + 64
stackMem, err := syscall.Mmap(-1, 0, stackSize,
syscall.PROT_READ|syscall.PROT_WRITE, syscall.MAP_PRIVATE|syscall.MAP_ANON)
if err != nil {
return nil, fmt.Errorf("verify: stack mmap: %w", err)
}
defer syscall.Munmap(stackMem)
retOff := stackPad
binary.LittleEndian.PutUint64(stackMem[retOff:retOff+8], uint64(leaveJITAddr))
copy(stackMem[retOff+8:], args)
stackBase := uintptr(unsafe.Pointer(&stackMem[retOff]))
enterJIT(fnAddr, stackBase)
out := make([]byte, len(args))
copy(out, stackMem[retOff+8:retOff+8+len(args)])
return out, nil
}
+1 -1
View File
@@ -1,7 +1,7 @@
// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
// SPDX-License-Identifier: BSD-3-Clause
//go:build !amd64
//go:build !amd64 && !arm64 && !riscv64 && !loong64
package verify
+50
View File
@@ -0,0 +1,50 @@
// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
// SPDX-License-Identifier: BSD-3-Clause
//go:build riscv64
package verify
import (
"encoding/binary"
"fmt"
"reflect"
"syscall"
"unsafe"
)
//go:nosplit
func enterJIT(fn uintptr, stack uintptr)
//go:nosplit
func leaveJIT()
var leaveJITAddr uintptr
func init() {
leaveJITAddr = reflect.ValueOf(leaveJIT).Pointer()
}
const stackPad = 64
// Call invokes the assembled function at fnAddr with the given ABI0 argument block.
func Call(fnAddr uintptr, args []byte) ([]byte, error) {
stackSize := stackPad + 8 + len(args) + 64
stackMem, err := syscall.Mmap(-1, 0, stackSize,
syscall.PROT_READ|syscall.PROT_WRITE, syscall.MAP_PRIVATE|syscall.MAP_ANON)
if err != nil {
return nil, fmt.Errorf("verify: stack mmap: %w", err)
}
defer syscall.Munmap(stackMem)
retOff := stackPad
binary.LittleEndian.PutUint64(stackMem[retOff:retOff+8], uint64(leaveJITAddr))
copy(stackMem[retOff+8:], args)
stackBase := uintptr(unsafe.Pointer(&stackMem[retOff]))
enterJIT(fnAddr, stackBase)
out := make([]byte, len(args))
copy(out, stackMem[retOff+8:retOff+8+len(args)])
return out, nil
}
+36
View File
@@ -0,0 +1,36 @@
// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
// SPDX-License-Identifier: BSD-3-Clause
#include "textflag.h"
// ABI0 JIT trampoline for arm64.
//
// enterJIT saves the Go SP and LR, switches to the prepared stack, and
// branches to the JIT function. When the function RETs (BR LR), control
// lands in leaveJIT, which restores the Go stack and returns.
//
// The prepared stack layout:
// [0] leaveJIT address (loaded into LR before jumping to fn)
// [8] ABI0 argument area
// func enterJIT(fn uintptr, stack uintptr)
TEXT ·enterJIT(SB), NOSPLIT, $0-16
MOVD fn+0(FP), R0 // target function address
MOVD R30, savedLR(SB) // save link register
MOVD R3, savedSP(SB) // save Go stack pointer
MOVD stack+8(FP), R3 // load prepared stack pointer
MOVD 0(R3), R30 // load leaveJIT address into LR
ADD $8, R3, R3 // advance past return address
MOVD R3, RSP // switch to prepared stack
JMP (R0) // branch to JIT function
// func leaveJIT()
TEXT ·leaveJIT(SB), NOSPLIT, $0-0
MOVD savedSP(SB), R3 // restore Go stack pointer
MOVD R3, RSP
MOVD savedLR(SB), R30 // restore link register
RET // return to Go caller
// Package-level storage for saved registers.
GLOBL savedLR(SB), NOPTR, $8
GLOBL savedSP(SB), NOPTR, $8
+31
View File
@@ -0,0 +1,31 @@
// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
// SPDX-License-Identifier: BSD-3-Clause
#include "textflag.h"
// ABI0 JIT trampoline for LoongArch 64.
//
// enterJIT saves Go SP and RA (R1), switches to the prepared stack, and
// jumps to the JIT function. When the function RETs (JIRL zero, ra, 0),
// control lands in leaveJIT.
// func enterJIT(fn uintptr, stack uintptr)
TEXT ·enterJIT(SB), NOSPLIT, $0-16
MOVV fn+0(FP), R4 // target function address (A0)
MOVV R1, savedRA(SB) // save return address (RA)
MOVV R3, savedSP(SB) // save Go stack pointer (SP)
MOVV stack+8(FP), R5 // load prepared stack pointer (A1)
MOVV 0(R5), R1 // load leaveJIT address into RA
ADDV $8, R5, R5 // advance past return address
MOVV R5, R3 // switch to prepared stack (SP)
JIRL R0, R4, 0 // jump to JIT function
// func leaveJIT()
TEXT ·leaveJIT(SB), NOSPLIT, $0-0
MOVV savedSP(SB), R5 // restore Go stack pointer
MOVV R5, R3 // restore SP
MOVV savedRA(SB), R1 // restore return address
JIRL R0, R1, 0 // return to Go caller
GLOBL savedRA(SB), NOPTR, $8
GLOBL savedSP(SB), NOPTR, $8
+31
View File
@@ -0,0 +1,31 @@
// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
// SPDX-License-Identifier: BSD-3-Clause
#include "textflag.h"
// ABI0 JIT trampoline for RISC-V 64.
//
// enterJIT saves Go SP and RA, switches to the prepared stack, and jumps
// to the JIT function. When the function RETs (JALR zero, ra, 0), control
// lands in leaveJIT.
// func enterJIT(fn uintptr, stack uintptr)
TEXT ·enterJIT(SB), NOSPLIT, $0-16
MOV fn+0(FP), X5 // target function address (T0)
MOV X1, savedRA(SB) // save return address
MOV X2, savedSP(SB) // save Go stack pointer
MOV stack+8(FP), X6 // load prepared stack pointer (T1)
LD 0(X6), X1 // load leaveJIT address into RA
ADD $8, X6, X6 // advance past return address
MOV X6, X2 // switch to prepared stack (SP)
JALR X0, 0(X5) // jump to JIT function
// func leaveJIT()
TEXT ·leaveJIT(SB), NOSPLIT, $0-0
MOV savedSP(SB), X6 // restore Go stack pointer
MOV X6, X2 // restore SP
MOV savedRA(SB), X1 // restore return address
JALR X0, 0(X1) // return to Go caller
GLOBL savedRA(SB), NOPTR, $8
GLOBL savedSP(SB), NOPTR, $8