jitasm

A Python x86-64 JIT assembler.

I've tested on Windows and Linux. I have no macOS computer for testing, but theoretically it should also work on Intel macOS.

pip install jitasm

Example

The following example creates an int64_t max(int64_t a, int64_t b) function equivalent to:

int64_t max(int64_t a, int64_t b) {
    if (a >= b) return a;
    return b;
}

On Linux, the generated instructions are:

max:
    mov rax, rdi
    cmp rdi, rsi
    jge .done
    mov rax, rsi
.done:
    ret

And we provide a pseudo-instruction branch.cc to do CMP and Jcc:

max:
    mov rax, rdi
    bge rdi, rsi, .done
    mov rax, rsi
.done:
    ret

And we implement it in Python using jitasm:

from jitasm.x86_64 import *
from jitasm.utils import ccall

e = Emitter()
(e.label("max"),
    e.mov(RAX, RDI),
    e.bge(RDI, RSI, ".done"),
    e.mov(RAX, RSI),
 e.label(".done"),
    e.ret())
e.finalize()
max_fn_ptr = e.symbol("max")

assert ccall(max_fn_ptr, 3, 7)   == 7
assert ccall(max_fn_ptr, 42, -1) == 42

finalize() allocates executable memory and resolves labels. symbol() returns the address of a public label, and ctypes.CFUNCTYPE converts that address into a callable Python object. Call unmap() when the generated code is no longer needed. Functions created from symbol() must not be called after unmapping. The emitter can be finalized again to create a fresh mapping.

See test_qsort.py for a larger example on Linux.

Windows

Windows is using a different ABI, registers for function arguments are RCX, RDX, R8, R9:

from jitasm.x86_64 import *
from jitasm.utils import ccall

e = Emitter()

e.label('add_two')
e.mov(RAX, RCX)
e.add(RAX, RDX)
e.ret()

e.finalize()
assert ccall(e.symbol('add_two'), 20, 22) == 42

Save Registers on Stack Using with

with e.save_regs(rax, rbx, rcx, rdx):
    do_something()

will generate code like this:

push rax
push rbx
push rcx
push rdx

do_something

pop rdx
pop rcx
pop rbx
pop rax

Assembly Spec

`m8/m16/m32/m64/m128/m256/m*`: `[r64]`
                   | `[r64 + r64 * scale +/- simm32]` (scale = 1/2/4/8)
                   | `[r64 * scale +/- simm32]` (scale = 1/2/4/8)
                   | `[r64 +/- simm32]`
                   | `[rip + rel32]`
`rel32`: label 
`cond`: EQ | NE | LT | GT | LE | GE | LTU | GTU | GEU | LEU | P | NP | O | NO | S | NS

Implemented Instructions

Only an essential subset of the x86-64 instruction set with several pseudo-instructions is implemented.

Data directives

• DB: db int8...
• DW: dw int16...
• DD: dd int32... / dd float32... / dd (target_label, base_label)...
• DQ: dq int64... / dq float64...
• ASCII: ascii str
• ASCIZ: asciz str
• ALIGN: align bytes // DATA section zero-padding

Data movement and addressing

• MOV: mov r64, r64
• MOV: mov r64, imm64 // zero uses XOR r64, r64
• MOV: mov r64, m64
• MOV: mov m64, r64
• MOV: mov m32, r64 // low bits
• MOV: mov m16, r64 // low bits
• MOV: mov m8, r64 // low bits
• MOVZX: movzx r64, r8
• MOVZX: movzx r64, r16
• MOVZX: movzx r64, r32
• MOVZX: movzx r64, m8
• MOVZX: movzx r64, m16
• MOVZX: movzx r64, m32
• MOVSX: movsx r64, r8
• MOVSX: movsx r64, r16
• MOVSX: movsx r64, r32
• MOVSX: movsx r64, m8
• MOVSX: movsx r64, m16
• MOVSX: movsx r64, m32
• LEA: lea r64, m*

Integer arithmetic and bitwise operations

• ADD: add r64, r64 / add r64, simm32
• SUB: sub r64, r64 / sub r64, simm32
• BITAND: bitand r64, r64 / bitand r64, simm32
• BITOR: bitor r64, r64 / bitor r64, simm32
• XOR: xor r64, r64 / xor r64, simm32
• BITNOT: bitnot r64 // XOR r64, -1
• NEG: neg r64
• IMUL: imul r64, r64 / imul r64, simm32
• IDIV: idiv r64, r64 // quotient, remainder
• DIV: div r64, r64 // quotient, remainder
• SHL: shl r64, r64 / shl r64, uimm8
• SAR: sar r64, r64 / sar r64, uimm8
• SHR: shr r64, r64 / shr r64, uimm8
• ROR: ror r64, r64 / ror r64, uimm8
• ROL: rol r64, r64 / rol r64, uimm8

Scalar floating point

MOVSS, MOVSD, and scalar arithmetic select AVX when available and otherwise use SSE.

• MOVSS: movss xmm, xmm
• MOVSS: movss xmm, m32
• MOVSS: movss m32, xmm
• MOVSD: movsd xmm, xmm
• MOVSD: movsd xmm, m64
• MOVSD: movsd m64, xmm
• ADDSS: addss xmm, xmm
• SUBSS: subss xmm, xmm
• MULSS: mulss xmm, xmm
• DIVSS: divss xmm, xmm
• ADDSD: addsd xmm, xmm
• SUBSD: subsd xmm, xmm
• MULSD: mulsd xmm, xmm
• DIVSD: divsd xmm, xmm
• CVTSI2SS: cvtsi2ss xmm, r64
• CVTTSS2SI: cvttss2si r64, xmm
• CVTSI2SD: cvtsi2sd xmm, r64
• CVTTSD2SI: cvttsd2si r64, xmm
• ROUNDS: rounds xmm, xmm // single precision, round to nearest, ties to even
• CEILS: ceils xmm, xmm
• FLOORS: floors xmm, xmm
• TRUNCS: truncs xmm, xmm
• ROUNDD: roundd xmm, xmm // double precision, round to nearest, ties to even
• CEILD: ceild xmm, xmm
• FLOORD: floord xmm, xmm
• TRUNCD: truncd xmm, xmm

Comparisons and branches

Floating-point comparisons select AVX when available and otherwise use SSE.

• CMP: cmp r64, r64 / cmp r64, simm32
• UCOMISS: ucomiss xmm, xmm
• UCOMISD: ucomisd xmm, xmm
• JCC: jcc cond, rel32
• Jcc helpers (each takes rel32):
  • Equality: je/jeq/jz, jne/jnz
  • Unsigned: ja/jnbe/jgtu, jae/jnb/jnc/jgeu, jb/jnae/jc/jltu, jbe/jna/jleu
  • Signed: jg/jnle/jgt, jge/jnl, jl/jnge/jlt, jle/jng
  • Flags: jo/jno (overflow), js/jns (sign), jp/jpe/jnp/jpo (parity)
• SETCC: setcc cond, r8
• CMOVcc: cmoveq/cmovne/cmovgt/cmovge/cmovlt/cmovle r64, r64
• CMOVcc: cmovgtu/cmovgeu/cmovltu/cmovleu/cmovp r64, r64
• BRANCH: branch cond, r64, r64, rel32
• BRANCH: branch cond, r64, simm32, rel32
• BRANCHS: branchs cond, xmm, xmm, rel32 // beqs/bnes/bgts/bges/blts/bles
• BRANCHD: branchd cond, xmm, xmm, rel32 // beqd/bned/bgtd/bged/bltd/bled
• CSET: cset cond, r64, r64, r8
• CSET: cset cond, r64, simm32, r8
• CSETS: csets cond, xmm, xmm, r8 // seteqs/setnes/setgts/setges/setlts/setles
• CSETD: csetd cond, xmm, xmm, r8 // seteqd/setned/setgtd/setged/setltd/setled

Stack, calls, and control flow

• PUSH: push r64
• POP: pop r64
• BEGIN: begin // PUSH RBP + MOV RBP, RSP
• END: end // MOV RSP, RBP + POP RBP + RET
• CALL: call rel32
• CALL: call r64
• JMP: jmp rel32
• JMP: jmp r64
• RET: ret

System

• CPUID: cpuid // query processor identification and features

AVX

• VMOVAPS: vmovaps xmm, xmm / vmovaps ymm, ymm
• VMOVAPS: vmovaps xmm, m128 / vmovaps m128, xmm
• VMOVAPS: vmovaps ymm, m256 / vmovaps m256, ymm
• VMOVUPS: vmovups xmm, xmm / vmovups ymm, ymm
• VMOVUPS: vmovups xmm, m128 / vmovups m128, xmm
• VMOVUPS: vmovups ymm, m256 / vmovups m256, ymm
• VADDPS: vaddps xmm, xmm, xmm / vaddps ymm, ymm, ymm
• VSUBPS: vsubps xmm, xmm, xmm / vsubps ymm, ymm, ymm
• VMULPS: vmulps xmm, xmm, xmm / vmulps ymm, ymm, ymm
• VDIVPS: vdivps xmm, xmm, xmm / vdivps ymm, ymm, ymm
• VADDSUBPS: vaddsubps xmm, xmm, xmm / vaddsubps ymm, ymm, ymm
• VSQRTPS: vsqrtps xmm, xmm / vsqrtps ymm, ymm
• VMAXPS: vmaxps xmm, xmm, xmm / vmaxps ymm, ymm, ymm
• VMINPS: vminps xmm, xmm, xmm / vminps ymm, ymm, ymm
• VANDPS: vandps xmm, xmm, xmm / vandps ymm, ymm, ymm
• VANDNPS: vandnps xmm, xmm, xmm / vandnps ymm, ymm, ymm
• VORPS: vorps xmm, xmm, xmm / vorps ymm, ymm, ymm
• VXORPS: vxorps xmm, xmm, xmm / vxorps ymm, ymm, ymm
• VHADDPS: vhaddps xmm, xmm, xmm / vhaddps ymm, ymm, ymm
• VHSUBPS: vhsubps xmm, xmm, xmm / vhsubps ymm, ymm, ymm
• VDPPS: vdpps xmm, xmm, xmm, input_mask, output_mask
• VDPPS: vdpps ymm, ymm, ymm, input_mask, output_mask
• VRCPPS: vrcpps xmm, xmm / vrcpps ymm, ymm
• VRSQRTPS: vrsqrtps xmm, xmm / vrsqrtps ymm, ymm
• VROUNDPS: vroundps xmm, xmm / vroundps ymm, ymm // round to nearest, ties to even
• VFLOORPS: vfloorps xmm, xmm / vfloorps ymm, ymm
• VCEILPS: vceilps xmm, xmm / vceilps ymm, ymm
• VTRUNCPS: vtruncps xmm, xmm / vtruncps ymm, ymm
• VCMPPS: vcmpps xmm, xmm, xmm, predicate / vcmpps ymm, ymm, ymm, predicate
• VCMPPS helpers: veqps/vltps/vleps/vunordps/vneps/vnltps/vnleps/vordps/vgtps/vgeps
• VBLENDPS: vblendps xmm, xmm, xmm, mask / vblendps ymm, ymm, ymm, mask
• VSHUFPS: vshufps xmm, xmm, xmm, imm / vshufps ymm, ymm, ymm, imm
• VPERMILPS: vpermilps xmm, xmm, imm / vpermilps ymm, ymm, imm
• VPERMILPS: vpermilps xmm, xmm, xmm / vpermilps ymm, ymm, ymm
• VPERM2F128: vperm2f128 ymm, ymm, ymm, sel, zmask // sel: two lanes (0-3); zmask: 1 zeroes a lane
• VUNPCKLPS: vunpcklps xmm, xmm, xmm / vunpcklps ymm, ymm, ymm
• VUNPCKHPS: vunpckhps xmm, xmm, xmm / vunpckhps ymm, ymm, ymm
• VINSERTPS: vinsertps xmm, xmm, xmm, count_dst, count_src, zero_mask
• VMASKMOVPS: vmaskmovps xmm, xmm, m128 / vmaskmovps m128, xmm, xmm // mask is the middle operand
• VMASKMOVPS: vmaskmovps ymm, ymm, m256 / vmaskmovps m256, ymm, ymm
• VMOVMSKPS: vmovmskps r32, xmm / vmovmskps r32, ymm
• VMOVSHDUP: vmovshdup xmm, xmm / vmovshdup ymm, ymm
• VMOVSLDUP: vmovsldup xmm, xmm / vmovsldup ymm, ymm
• VEXTRACTPS: vextractps r32, xmm, imm // imm: element index 0-3
• VEXTRACTF128: vextractf128 xmm, ymm, imm // imm: lane 0-1
• VINSERTF128: vinsertf128 ymm, ymm, xmm, imm // imm: lane 0-1
• VBROADCASTSS: vbroadcastss xmm, m32 / vbroadcastss ymm, m32
• VBROADCASTSS: vbroadcastss xmm, xmm / vbroadcastss ymm, xmm // AVX2
• VPTEST: vptest xmm, xmm / vptest ymm, ymm // set ZF and CF from packed bitwise tests
• VZEROUPPER: vzeroupper // clear bits 128–255 of all YMM registers