Skip to content

PVMCode caching #266

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
arjanz merged 2 commits into from
Jul 6, 2026
Merged

PVMCode caching #266

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
359cf99
PVMCode caching in both numba and marshalling function
arjanz Jun 23, 2026
8cb0bbf
- code cleanup
matthijsb Jun 25, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
219 changes: 138 additions & 81 deletions pyjamaz/pvm/interpreters/numba/interpreter_numba_jit.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -65,6 +65,9 @@
from pyjamaz.pvm.interpreters.numba.memory import PVMMemory
from pyjamaz.pvm.basic_block import detect_basic_blocks

_CODE_METADATA_CACHE: dict[int, dict] = {}
_CODE_METADATA_CACHE_LIMIT = 64


@njit(uint32(
uint64[::1], # reg
Expand Down Expand Up @@ -436,8 +439,10 @@ def invoke_native(
skip_len = I64(initial_skip_len)
inst_nr = U32(inst_start)

# Copy registers
reg = registers_in.copy()
# Copy registers into the reusable output buffer and use it as the local register file.
reg = registers_out
for i in range(len(registers_in)):
reg[i] = registers_in[i]

timing_enabled = False
start_time = 0.0 # Note: only used when timing_enabled == True to measure time per opcode
Expand Down Expand Up @@ -1567,26 +1572,26 @@ def alloc_memory(
def __init__(self, program: PVMProgram, logger=None):

self.name = program.name
self.reg:npt.NDArray[U64] = np.zeros(13, dtype=U64)
self.inst_nr:U32 = U32(0)
self.pc:U32 = U32(0)
self.opcode:int = 0
self.reg: npt.NDArray[U64] = np.zeros(13, dtype=U64)
self.inst_nr: U32 = U32(0)
self.pc: U32 = U32(0)
self.opcode: int = 0
self.skip_len: int = 0
self.gas:I64 = I64(0)
self.code:npt.NDArray[U8] = np.array(1, dtype=U8)
self.gas: I64 = I64(0)
self.code: npt.NDArray[U8] = np.array(1, dtype=U8)
self.code_size: U64 = U64(0)
self.code_length: int = 0
self.jump_table = []
self.basic_block_starts_set = set()

self.inst_bitmask: List[bool] = []
self.inst_pos: Dict[int,int] = {0: 0}
self.inst_pos: Dict[int, int] = {0: 0}
self.inst_arg_len: List[int] = []
self.basic_block_start_mask = np.array([], dtype=np.uint8)

self.mem:PVMMemory = None
self.status:int = ExitReason.resume.value
self.exit_value:int = None
self.mem: PVMMemory = None
self.status: int = ExitReason.resume.value
self.exit_value: int = None

# Initialize memory operation lookups
self._init_mem_ops_lookup()
Expand All @@ -1598,20 +1603,124 @@ def __init__(self, program: PVMProgram, logger=None):
self.mem_section_starts = np.array([], dtype=U32)
self.mem_section_ends = np.array([], dtype=U32)
self.mem_section_size = np.array([], dtype=U32)
self._linked_memory_signature = None

self._mem_addr: int = -1

self.pc = U32(0)
self.gas = I64(0)
self.name = program.name
original_code = np.array(program.code.code, dtype=U8)
self.code_length = len(original_code)
self.code = np.empty(self.code_length + 1, dtype=U8)
self.code[:self.code_length] = original_code
self.code[self.code_length] = U8(op_trap)
self.code_size: U64 = U64(len(self.code))
self.mem = program.memory
self.jump_table = [x.value for x in program.code.jump_table]
code_cache_key = id(program.code)
cached_code = _CODE_METADATA_CACHE.get(code_cache_key)

if cached_code is None:
original_code = np.array(program.code.code, dtype=U8)
code_length = len(original_code)
code = np.empty(code_length + 1, dtype=U8)
code[:code_length] = original_code
code[code_length] = U8(op_trap)
jump_table = [x.value for x in program.code.jump_table]
inst_bitmask = program.code.opcode_bitmask

#Create lookups for byte_pos -> instruction_nr and instruction_nr->instruction_length
inst_pos = {0: 0}
inst_arg_len = []

inst_nr = 0
inst_bitmask_idx = 1
if len(inst_bitmask) == 1:
inst_arg_len.append(0)
else:
# Parse instruction bitmask and create a opcode offset and instruction length lookup
while inst_bitmask_idx < len(inst_bitmask):
inst_args = 0
is_opcode = False
while not is_opcode:
is_opcode = inst_bitmask[inst_bitmask_idx]
if not is_opcode:
inst_args += 1
inst_bitmask_idx += 1
if inst_bitmask_idx > len(inst_bitmask) - 1:
is_opcode = True

# GP-0.7.2-eq:A.20 (l)
inst_arg_len.append(inst_args)
inst_nr += 1
# Note: only add to inst_pos if this position has an opcode in the bitmask
if inst_bitmask_idx - 1 < len(inst_bitmask) and inst_bitmask[inst_bitmask_idx - 1]:
inst_pos[inst_bitmask_idx - 1] = inst_nr

inst_pos[code_length] = len(inst_arg_len)
inst_arg_len.append(0)
basic_block_starts_set = detect_basic_blocks(
code=bytes(code),
code_length=code_length,
inst_pos=inst_pos,
inst_arg_len=inst_arg_len,
)

inst_pos_keys = np.array(list(inst_pos.keys()), dtype=np.int32)
inst_pos_vals = np.array(list(inst_pos.values()), dtype=np.int32)
inst_arg_len_array = np.array(inst_arg_len, dtype=np.int32)
#opcode_scheme_array = _opcode_scheme_array()
self.opcode_scheme_array = np.full(256, 255, dtype=np.int32)
for opcode, scheme in OpcodeScheme.items():
self.opcode_scheme_array[opcode] = scheme.value

if inst_pos_keys.size > 0:
max_pc = int(max(inst_pos_keys))
size = max_pc + 1
else:
size = int(len(code))
pc_to_inst_index = np.full(size, -1, dtype=np.int32)
for k, v in zip(inst_pos_keys, inst_pos_vals):
idx = int(k)
if 0 <= idx < size:
pc_to_inst_index[idx] = int(v)
basic_block_start_mask = np.zeros(size, dtype=np.uint8)
for start in basic_block_starts_set:
idx = int(start)
if 0 <= idx < size:
basic_block_start_mask[idx] = np.uint8(1)

cached_code = {
"code_object": program.code,
"code": code,
"code_length": code_length,
"code_size": U64(len(code)),
"jump_table": jump_table,
"jump_table_array": np.array(jump_table, dtype=np.int32),
"inst_bitmask": inst_bitmask,
"inst_pos": inst_pos,
"inst_arg_len": inst_arg_len,
"basic_block_starts_set": basic_block_starts_set,
"inst_pos_keys": inst_pos_keys,
"inst_pos_vals": inst_pos_vals,
"inst_arg_len_array": inst_arg_len_array,
"opcode_scheme_array": self.opcode_scheme_array,
"pc_to_inst_index": pc_to_inst_index,
"basic_block_start_mask": basic_block_start_mask,
}
if len(_CODE_METADATA_CACHE) >= _CODE_METADATA_CACHE_LIMIT:
_CODE_METADATA_CACHE.pop(next(iter(_CODE_METADATA_CACHE)))
_CODE_METADATA_CACHE[code_cache_key] = cached_code

self.code = cached_code["code"]
self.code_length = cached_code["code_length"]
self.code_size = cached_code["code_size"]
self.jump_table = cached_code["jump_table"]
self.inst_bitmask = cached_code["inst_bitmask"]
self.inst_pos = cached_code["inst_pos"]
self.inst_arg_len = cached_code["inst_arg_len"]
self.basic_block_starts_set = cached_code["basic_block_starts_set"]
self.inst_pos_keys = cached_code["inst_pos_keys"]
self.inst_pos_vals = cached_code["inst_pos_vals"]
self.inst_arg_len_array = cached_code["inst_arg_len_array"]
self.opcode_scheme_array = cached_code["opcode_scheme_array"]
self.pc_to_inst_index = cached_code["pc_to_inst_index"]
self.basic_block_start_mask = cached_code["basic_block_start_mask"]
self.jump_table_array = cached_code["jump_table_array"]

# Initialize memory sections from the PVMMemory object (just reference where possible)
self._link_memory(program.memory)
Expand All @@ -1621,19 +1730,6 @@ def __init__(self, program: PVMProgram, logger=None):

self.status = ExitReason.resume.value

self.inst_bitmask: List[bool] = program.code.opcode_bitmask
self.inst_pos: Dict[int,int] = {0: 0}
self.inst_arg_len: List[int] = []
self.create_instruction_lookup()
self.inst_pos[self.code_length] = len(self.inst_arg_len)
self.inst_arg_len.append(0)
self.basic_block_starts_set = detect_basic_blocks(
code=bytes(self.code),
code_length=self.code_length,
inst_pos=self.inst_pos,
inst_arg_len=self.inst_arg_len,
)

# Create typed List of opcode names (max opcode value is 230)
# Initialize with "UNKNOWN" for all opcodes
self.opcode_names = List()
Expand All @@ -1655,23 +1751,25 @@ def __init__(self, program: PVMProgram, logger=None):
self.opcode_names[opcode_val] = str(opcode_name)

# Note: native (jit) caches, which we sync this back to the python side
self._prepare_jit_data()
self._jit_mem_cache_dirty = True
self._jit_section_starts_cache = None
self._jit_section_ends_cache = None
self._jit_section_arrays_cache = None
self._jit_section_access_cache = None
self._jit_acl_bitmaps_cache = None

self.jump_table_array = np.array(self.jump_table, dtype=np.int32)
self._registers_out = np.zeros(13, dtype=np.uint64)
self._state_out = np.zeros(7, dtype=np.int64)
self._heap_grew_out = np.zeros(1, dtype=np.int64)

# Prepare heap info (for sbrk)
current_heap_end = self.mem_section_ends[1] if len(self.mem_section_ends) > 1 else 0
self.heap_info = np.array([
current_heap_end, # current heap end
self.next_heap_section_start(),
MEM_WRITABLE # writable permission value
], dtype=np.uint64)
self.heap_info = np.array(
[
current_heap_end, # current heap end
self.next_heap_section_start(),
MEM_WRITABLE # writable permission value
], dtype=np.uint64
)


def get_heap_capacity(self, memory: PVMMemory, heap_section: MemorySection) -> int:
Expand Down Expand Up @@ -1853,47 +1951,6 @@ def addr_is_covered(addr: int) -> bool:
self._jit_acl_bitmaps_cache = None


def create_instruction_lookup(self):
"""
Create lookups for byte_pos -> instruction_nr and instruction_nr->instruction_length
"""
self.inst_pos = {0: 0}
self.inst_arg_len = []

inst_nr = 0
inst_bitmask = self.inst_bitmask
inst_bitmask_idx = 1

# Note: In the exceptional case we only have 1 instruction (trap or fallthrough), we add it manually and be done
if len(inst_bitmask) == 1:
self.inst_arg_len.append(0)
return

# Parse instruction bitmask and create a opcode offset and instruction length lookup
while inst_bitmask_idx < len(inst_bitmask):
inst_args = 0

is_opcode = False

while not is_opcode:

is_opcode = inst_bitmask[inst_bitmask_idx]
if not is_opcode:
inst_args += 1

inst_bitmask_idx += 1

if inst_bitmask_idx > len(inst_bitmask) - 1:
is_opcode = True

# GP-0.7.2-eq:A.20 (l)
self.inst_arg_len.append(inst_args)
inst_nr += 1
# Note: only add to inst_pos if this position has an opcode in the bitmask
if inst_bitmask_idx - 1 < len(inst_bitmask) and inst_bitmask[inst_bitmask_idx - 1]:
self.inst_pos[inst_bitmask_idx - 1] = inst_nr


def _prepare_jit_data(self):
# Convert dictionaries to arrays for JIT access
if self.inst_pos:
Expand Down
15 changes: 14 additions & 1 deletion pyjamaz/pvm/types.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
import logging

from abc import ABC, abstractmethod
from functools import lru_cache
from math import ceil
from dataclasses import dataclass, field
from typing import List, Union, Type, T, Optional, Sequence
Expand Down Expand Up @@ -32,6 +33,7 @@ def page_size(bytes: int) -> int:
_MAX_PAGE_IDX = (ADDR_MOD // PAGE_SIZE) - 1
_PAGE_CACHE_LIMIT = 16
_ZERO_PAGE = bytes(PAGE_SIZE)
_PVM_CODE_CACHE_LIMIT = 64


@dataclass
Expand Down Expand Up @@ -196,6 +198,12 @@ def from_jam_bytes(cls, scale_bytes: JamBytes, strict_decoding=True) -> 'PVMCode
opcode_bitmask=opcode_bitmask,
)

@classmethod
def from_bytes_cached(cls, data) -> "PVMCode":
if not isinstance(data, bytes):
data = bytes(data)
return _pvm_code_from_bytes_cached(data)

def to_codec_type(self) -> JamCodecType:
codec_def = self.to_codec_def()
# Change definition according to data
Expand All @@ -218,6 +226,11 @@ def serialize(self) -> List[int]:
return [b for b in self.to_jam_bytes().to_bytes()]


@lru_cache(maxsize=_PVM_CODE_CACHE_LIMIT)
def _pvm_code_from_bytes_cached(data: bytes) -> PVMCode:
return PVMCode.from_jam_bytes(JamBytes(data))


@dataclass
class PVMProgram(Serializable):
"""
Expand Down Expand Up @@ -338,7 +351,7 @@ def from_serialized_bytes(cls, serialized_program: bytes, argument_contents: byt
) <= 2**32:

instance = cls(
code=PVMCode.from_jam_bytes(JamBytes(pvm_code)),
code=PVMCode.from_bytes_cached(pvm_code),
registers=cls.init_registers(argument_contents),
memory=cls.init_memory(pvm_rom_contents, pvm_heap_contents, argument_contents, heap_mem_pages, stack_mem_size),
name=name
Expand Down
Loading