feat(asm): extend arm64 encoder with atomics, bitfield, SIMD and more test kernels
Assisted-by: MiMo V2.5 Pro
This commit is contained in:
+202
-20
@@ -263,6 +263,31 @@ func encodeARM64Instr(instr *ast.Instr, pc int, offsets map[string]int, fi arm64
|
||||
return encodeARM64CRC32(mnem, enc.op, ops)
|
||||
}
|
||||
|
||||
// Exclusive load/store (LDXR, STXR, LDAXR, STLXR).
|
||||
if enc, ok := a64InstrTable[mnem]; ok && enc.format == a64FExcl {
|
||||
return encodeARM64Excl(mnem, enc.op, ops)
|
||||
}
|
||||
|
||||
// LSE atomics (LDADD, CAS, SWP).
|
||||
if enc, ok := a64InstrTable[mnem]; ok && enc.format == a64FLSE {
|
||||
return encodeARM64LSEAtom(mnem, enc.op, ops)
|
||||
}
|
||||
|
||||
// Bitfield/shift (ASR, LSL, LSR, ROR, BFI, BFXIL, SBFM, UBFM).
|
||||
if enc, ok := a64InstrTable[mnem]; ok && enc.format == a64FBitfield {
|
||||
return encodeARM64Bitfield(mnem, enc.op, ops)
|
||||
}
|
||||
|
||||
// EXTR.
|
||||
if enc, ok := a64InstrTable[mnem]; ok && enc.format == a64FEXTR {
|
||||
return encodeARM64Extr(mnem, enc.op, ops)
|
||||
}
|
||||
|
||||
// SIMD 3-operand (VADD, VSUB, VMUL).
|
||||
if enc, ok := a64InstrTable[mnem]; ok && enc.format == a64FSIMD3 {
|
||||
return encodeARM64SIMD3(mnem, enc.op, ops)
|
||||
}
|
||||
|
||||
return nil, fmt.Errorf("unsupported arm64 instruction %q", mnem)
|
||||
}
|
||||
|
||||
@@ -273,21 +298,28 @@ func encodeARM64Branch(mnem string, ops []*ast.Operand, pc int, offsets map[stri
|
||||
if len(ops) != 1 {
|
||||
return nil, fmt.Errorf("%s expects 1 operand, got %d", mnem, len(ops))
|
||||
}
|
||||
target := resolve(arm64Label(ops[0]))
|
||||
op := ops[0]
|
||||
|
||||
// External symbol reference: BL sym(SB).
|
||||
if link && op.Addr.Sym != nil && op.Addr.Sym.Pseudo == "SB" {
|
||||
// Emit BL with zero offset; the linker fills in the target.
|
||||
return a64wordLE(a64Branch(1, 0)), nil
|
||||
}
|
||||
|
||||
target := resolve(arm64Label(op))
|
||||
targetOff, ok := offsets[target]
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("undefined label %q", target)
|
||||
}
|
||||
// Branch offset in bytes, shifted right by 2 (instructions are 4-byte aligned).
|
||||
rel := (targetOff - pc) >> 2
|
||||
if rel < -(1<<25) || rel >= (1<<25) {
|
||||
return nil, fmt.Errorf("branch to %q too far (26-bit range)", target)
|
||||
}
|
||||
op := uint32(0) // B
|
||||
bop := uint32(0) // B
|
||||
if link {
|
||||
op = 1 // BL
|
||||
bop = 1 // BL
|
||||
}
|
||||
return a64wordLE(a64Branch(op, int32(rel))), nil
|
||||
return a64wordLE(a64Branch(bop, int32(rel))), nil
|
||||
}
|
||||
|
||||
// encodeARM64BranchCond encodes a conditional branch (B.cond) to a label.
|
||||
@@ -446,7 +478,7 @@ func encodeARM64Mov(instr *ast.Instr, mnem string, fi arm64FrameInfo, relocs *[]
|
||||
if rd < 0 {
|
||||
return nil, fmt.Errorf("%s $imm: invalid destination register", mnem)
|
||||
}
|
||||
return encodeARM64LoadImm(rd, immFromOperand(src), mnem)
|
||||
return encodeARM64LoadImm(rd, arm64Imm64(src), mnem)
|
||||
}
|
||||
|
||||
// Static symbol load/store via ADRP.
|
||||
@@ -496,11 +528,11 @@ func arm64MovSize(mnem string, ops []*ast.Operand, fi arm64FrameInfo) int {
|
||||
if src.Imm.Sym != nil && src.Imm.Sym.Pseudo == "SB" {
|
||||
return 8 // ADRP + ADD
|
||||
}
|
||||
v := immFromOperand(src)
|
||||
v := arm64Imm64(src)
|
||||
if v == 0 {
|
||||
return 4
|
||||
}
|
||||
if arm64Movcon(int64(v)) >= 0 || arm64Movcon(^int64(v)) >= 0 {
|
||||
if arm64Movcon(v) >= 0 || arm64Movcon(^v) >= 0 {
|
||||
return 4
|
||||
}
|
||||
return 8 // MOVZ + MOVK
|
||||
@@ -534,8 +566,8 @@ func arm64MovSize(mnem string, ops []*ast.Operand, fi arm64FrameInfo) int {
|
||||
|
||||
// encodeARM64LoadImm loads an immediate into a register, matching the
|
||||
// toolchain's MOVZ/MOVN/MOVK sequence.
|
||||
func encodeARM64LoadImm(rd int, v int32, mnem string) ([]byte, error) {
|
||||
d := int64(v)
|
||||
func encodeARM64LoadImm(rd int, v int64, mnem string) ([]byte, error) {
|
||||
d := v
|
||||
// For 32-bit MOVW, zero-extend.
|
||||
if mnem == "MOVW" || mnem == "MOVWU" {
|
||||
d = int64(uint32(v))
|
||||
@@ -555,26 +587,39 @@ func encodeARM64LoadImm(rd int, v int32, mnem string) ([]byte, error) {
|
||||
sf = 0
|
||||
}
|
||||
|
||||
// Try logical immediate (bitmask) encoding. The Go toolchain uses ORR
|
||||
// with a bitmask immediate for constants like $1, $-2, $0xFF, etc.
|
||||
// that can be represented as a repeating pattern of contiguous 1s.
|
||||
N, immr, imms, ok := arm64Bitmask(uint64(d), int(sf))
|
||||
if ok {
|
||||
// ORR Rd, XZR, #bitmask (logical immediate)
|
||||
return a64wordLE(sf<<31 | 1<<29 | 0x24<<23 | N<<22 | immr<<16 | imms<<10 | 31<<5 | uint32(rd)), nil
|
||||
// The Go toolchain classifies immediates:
|
||||
// - C_ABCON0 (0 < v ≤ 4095): bitmask first for positive values
|
||||
// - Negative values: MOVN first, then bitmask
|
||||
// - C_MOVCON (movcon-eligible, outside ABCON range): MOVZ/MOVN first
|
||||
tryBitmaskFirst := (d > 0 && d <= 0xFFF)
|
||||
|
||||
if tryBitmaskFirst {
|
||||
// Small immediate: try bitmask first (Go uses ORR for values like $1, $256).
|
||||
N, immr, imms, ok := arm64Bitmask(uint64(d), int(sf))
|
||||
if ok {
|
||||
return a64wordLE(sf<<31 | 1<<29 | 0x24<<23 | N<<22 | immr<<16 | imms<<10 | 31<<5 | uint32(rd)), nil
|
||||
}
|
||||
}
|
||||
|
||||
// Try MOVZ (single non-zero16-bit chunk).
|
||||
// Try MOVZ (single non-zero 16-bit chunk).
|
||||
s := arm64Movcon(d)
|
||||
if s >= 0 {
|
||||
return a64wordLE(a64MoveWide(sf, 2, uint32(s>>4), uint32((d>>uint(s))&0xFFFF), uint32(rd))), nil
|
||||
}
|
||||
// Try MOVN (single non-0xFFFF16-bit chunk of ^d).
|
||||
// Try MOVN (single non-0xFFFF 16-bit chunk of ^d).
|
||||
sn := arm64Movcon(^d)
|
||||
if sn >= 0 {
|
||||
return a64wordLE(a64MoveWide(sf, 0, uint32(sn>>4), uint32((^d>>uint(sn))&0xFFFF), uint32(rd))), nil
|
||||
}
|
||||
|
||||
// For values outside the bitmask-first range that are not movcon: try bitmask.
|
||||
if !tryBitmaskFirst {
|
||||
N, immr, imms, ok := arm64Bitmask(uint64(d), int(sf))
|
||||
if ok {
|
||||
return a64wordLE(sf<<31 | 1<<29 | 0x24<<23 | N<<22 | immr<<16 | imms<<10 | 31<<5 | uint32(rd)), nil
|
||||
}
|
||||
}
|
||||
|
||||
// Multi-instruction: MOVZ + MOVK for each non-zero16-bit chunk.
|
||||
var ws []uint32
|
||||
first := true
|
||||
@@ -659,7 +704,7 @@ func arm64Bitmask(v uint64, sf int) (N, immr, imms uint32, ok bool) {
|
||||
N = 0
|
||||
}
|
||||
imms = uint32((^(esize - 1))&0x3F) | uint32(ones-1)
|
||||
immr = uint32(r)
|
||||
immr = uint32((esize - r) % esize)
|
||||
return N, immr, imms, true
|
||||
}
|
||||
}
|
||||
@@ -823,6 +868,18 @@ func arm64Reg(op *ast.Operand) int {
|
||||
return arm64RegNum(operandRegName(op))
|
||||
}
|
||||
|
||||
// arm64Imm64 returns the full 64-bit immediate value of an operand.
|
||||
func arm64Imm64(op *ast.Operand) int64 {
|
||||
if op.Imm.HasVal {
|
||||
v := op.Imm.Val
|
||||
if op.Imm.Neg {
|
||||
v = -v
|
||||
}
|
||||
return v
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
// arm64MemWithFrame resolves a memory operand, translating FP/SP pseudo-
|
||||
// registers via the frame mapping.
|
||||
func arm64MemWithFrame(op *ast.Operand, fi arm64FrameInfo) (rn int, off int32) {
|
||||
@@ -1068,6 +1125,131 @@ func encodeARM64CRC32(mnem string, baseOp uint32, ops []*ast.Operand) ([]byte, e
|
||||
return a64wordLE(baseOp | uint32(rm)<<16 | uint32(rd)<<5 | uint32(rd)), nil
|
||||
}
|
||||
|
||||
// ---- Atomics encoding ----
|
||||
|
||||
// encodeARM64Excl encodes an exclusive load/store instruction.
|
||||
// LDXR (Rn), Rt → LDXR Rt, [Rn] (2 operands: mem, reg or reg, mem)
|
||||
// STXR Rs, (Rn), Rt → STXR Rs, Rt, [Rn] (3 operands: Rs, mem, Rt-status)
|
||||
func encodeARM64Excl(mnem string, baseOp uint32, ops []*ast.Operand) ([]byte, error) {
|
||||
// LDXR/STXR have different operand forms.
|
||||
isLoad := strings.HasPrefix(mnem, "LD")
|
||||
if isLoad {
|
||||
// LDXR (Rn), Rt → 2 operands: mem, reg
|
||||
if len(ops) != 2 {
|
||||
return nil, fmt.Errorf("%s expects 2 operands, got %d", mnem, len(ops))
|
||||
}
|
||||
rn, _ := arm64MemWithFrame(ops[0], arm64FrameInfo{})
|
||||
rt := arm64RegNum(operandRegName(ops[1]))
|
||||
if rn < 0 || rt < 0 {
|
||||
return nil, fmt.Errorf("invalid operand in %s", mnem)
|
||||
}
|
||||
return a64wordLE(baseOp | uint32(rn)<<5 | uint32(rt)), nil
|
||||
}
|
||||
// STXR Rs, (Rn), Rt → 3 operands: Rs, mem, Rt
|
||||
if len(ops) != 3 {
|
||||
return nil, fmt.Errorf("%s expects 3 operands, got %d", mnem, len(ops))
|
||||
}
|
||||
rs := arm64RegNum(operandRegName(ops[0]))
|
||||
rn, _ := arm64MemWithFrame(ops[1], arm64FrameInfo{})
|
||||
rt := arm64RegNum(operandRegName(ops[2]))
|
||||
if rs < 0 || rn < 0 || rt < 0 {
|
||||
return nil, fmt.Errorf("invalid operand in %s", mnem)
|
||||
}
|
||||
return a64wordLE(baseOp | uint32(rs)<<16 | uint32(rn)<<5 | uint32(rt)), nil
|
||||
}
|
||||
|
||||
// encodeARM64LSEAtom encodes an LSE atomic instruction (LDADD, CAS, SWP).
|
||||
// LDADD Rs, (Rn), Rt → 3 operands: Rs, mem, Rt
|
||||
// CAS Rs, (Rn), Rt → 3 operands: Rs, mem, Rt
|
||||
func encodeARM64LSEAtom(mnem string, baseOp uint32, ops []*ast.Operand) ([]byte, error) {
|
||||
if len(ops) != 3 {
|
||||
return nil, fmt.Errorf("%s expects 3 operands, got %d", mnem, len(ops))
|
||||
}
|
||||
rs := arm64RegNum(operandRegName(ops[0]))
|
||||
rn, _ := arm64MemWithFrame(ops[1], arm64FrameInfo{})
|
||||
rt := arm64RegNum(operandRegName(ops[2]))
|
||||
if rs < 0 || rn < 0 || rt < 0 {
|
||||
return nil, fmt.Errorf("invalid operand in %s", mnem)
|
||||
}
|
||||
return a64wordLE(baseOp | uint32(rs)<<16 | uint32(rn)<<5 | uint32(rt)), nil
|
||||
}
|
||||
|
||||
// ---- Bitfield/EXTR encoding ----
|
||||
|
||||
// encodeARM64Bitfield encodes a bitfield instruction.
|
||||
// ASR/LSL/LSR/ROR $shamt, Rn, Rd → 3 operands: $imm, Rn, Rd
|
||||
// BFI/BFXIL/SBFM/UBFM $immr, Rn, $imms, Rd → 4 operands
|
||||
func encodeARM64Bitfield(mnem string, baseOp uint32, ops []*ast.Operand) ([]byte, error) {
|
||||
isShift := mnem == "ASR" || mnem == "ASRW" || mnem == "LSL" || mnem == "LSLW" ||
|
||||
mnem == "LSR" || mnem == "LSRW" || mnem == "ROR" || mnem == "RORW"
|
||||
|
||||
if isShift {
|
||||
// ASR $shamt, Rn, Rd → SBFM with immr=shamt, imms=31/63
|
||||
if len(ops) != 3 {
|
||||
return nil, fmt.Errorf("%s expects 3 operands, got %d", mnem, len(ops))
|
||||
}
|
||||
shamt := int(immFromOperand(ops[0]))
|
||||
rn := arm64RegNum(operandRegName(ops[1]))
|
||||
rd := arm64RegNum(operandRegName(ops[2]))
|
||||
if rn < 0 || rd < 0 {
|
||||
return nil, fmt.Errorf("invalid register operand in %s", mnem)
|
||||
}
|
||||
// ASR: SBFM with immr=shamt, imms=31(32-bit) or 63(64-bit)
|
||||
is64 := mnem == "ASR"
|
||||
imms := 31
|
||||
if is64 {
|
||||
imms = 63
|
||||
}
|
||||
return a64wordLE(baseOp | uint32(shamt)<<16 | uint32(imms)<<10 | uint32(rn)<<5 | uint32(rd)), nil
|
||||
}
|
||||
|
||||
// BFI/BFXIL/SBFM/UBFM: 4 operands ($immr, Rn, $imms, Rd)
|
||||
if len(ops) != 4 {
|
||||
return nil, fmt.Errorf("%s expects 4 operands, got %d", mnem, len(ops))
|
||||
}
|
||||
immr := int(immFromOperand(ops[0]))
|
||||
rn := arm64RegNum(operandRegName(ops[1]))
|
||||
imms := int(immFromOperand(ops[2]))
|
||||
rd := arm64RegNum(operandRegName(ops[3]))
|
||||
if rn < 0 || rd < 0 {
|
||||
return nil, fmt.Errorf("invalid register operand in %s", mnem)
|
||||
}
|
||||
return a64wordLE(baseOp | uint32(immr)<<16 | uint32(imms)<<10 | uint32(rn)<<5 | uint32(rd)), nil
|
||||
}
|
||||
|
||||
// encodeARM64Extr encodes an EXTR instruction.
|
||||
// EXTR $lsb, Rm, Rn, Rd → 4 operands
|
||||
func encodeARM64Extr(mnem string, baseOp uint32, ops []*ast.Operand) ([]byte, error) {
|
||||
if len(ops) != 4 {
|
||||
return nil, fmt.Errorf("%s expects 4 operands, got %d", mnem, len(ops))
|
||||
}
|
||||
lsb := int(immFromOperand(ops[0]))
|
||||
rm := arm64RegNum(operandRegName(ops[1]))
|
||||
rn := arm64RegNum(operandRegName(ops[2]))
|
||||
rd := arm64RegNum(operandRegName(ops[3]))
|
||||
if rm < 0 || rn < 0 || rd < 0 {
|
||||
return nil, fmt.Errorf("invalid register operand in %s", mnem)
|
||||
}
|
||||
return a64wordLE(baseOp | uint32(rm)<<16 | uint32(lsb)<<10 | uint32(rn)<<5 | uint32(rd)), nil
|
||||
}
|
||||
|
||||
// ---- SIMD/NEON encoding ----
|
||||
|
||||
// encodeARM64SIMD3 encodes a SIMD 3-operand instruction.
|
||||
// VADD Vm, Vn, Vd → base | Rm<<16 | Rn<<5 | Rd (Q and size bits in base)
|
||||
func encodeARM64SIMD3(mnem string, baseOp uint32, ops []*ast.Operand) ([]byte, error) {
|
||||
if len(ops) != 3 {
|
||||
return nil, fmt.Errorf("%s expects 3 operands, got %d", mnem, len(ops))
|
||||
}
|
||||
rm := arm64RegNum(operandRegName(ops[0]))
|
||||
rn := arm64RegNum(operandRegName(ops[1]))
|
||||
rd := arm64RegNum(operandRegName(ops[2]))
|
||||
if rm < 0 || rn < 0 || rd < 0 {
|
||||
return nil, fmt.Errorf("invalid register operand in %s", mnem)
|
||||
}
|
||||
return a64wordLE(baseOp | uint32(rm)<<16 | uint32(rn)<<5 | uint32(rd)), nil
|
||||
}
|
||||
|
||||
// AssembleFileARM64 assembles every TEXT function of a parsed arm64 file
|
||||
// and lays out its static symbols (GLOBL/DATA) in a data section behind the
|
||||
// code. SB references in the code are encoded as ADRP pairs with zero
|
||||
|
||||
Reference in New Issue
Block a user