sceptic-biomem/ScEpTIC/emulator/state_retention_manager.py
2026-07-10 10:38:57 +02:00

450 lines
18 KiB
Python

import copy
import gc
import logging
from ScEpTIC.AST.elements.instructions.other_operations import CallOperation
from ScEpTIC.emulator.io.output import OutputManager
from ScEpTIC.exceptions import MemoryException
class StateRetentionManager:
"""
Manages checkpoint and restore operations
"""
def __init__(self, vmstate, config):
self.vmstate = vmstate
self.routine_names = {
'save': config.state_save_function_name,
'restore': config.state_restore_function_name,
}
self.physical_mmu = config.physical_state_retention
self.volatile_stack = vmstate.memory.memory_positions['stack'] == 'volatile'
# Ignore state save and restore routines
CallOperation.ignore.append(self.routine_names['save'])
CallOperation.ignore.append(self.routine_names['restore'])
self.use_checkpoints = config.type == 'checkpoint'
self.static_placement = config.state_save_strategy == 'static_placement' and config.type == 'checkpoint'
self.probe_energy_buffer = config.state_save_behaviour == 'probe' and config.type == 'checkpoint'
self.probe_with_adc = config.probe_with_adc and config.type == 'checkpoint'
self.hibernate_after_save = config.hibernate_after_save and config.type == 'checkpoint'
self.restore_register_file = config.restore_register_file
self.restore_stack = config.restore_stack
self.restore_heap = config.restore_heap
self.restore_volatile_gst = config.restore_volatile_gst
self.restore_non_volatile_gst = config.restore_non_volatile_gst
self.restore_environment = config.restore_environment
self.calculate_v_save_state = config.calculate_v_save_state
self.calculate_v_save_state_additional_cc = config.calculate_v_save_state_additional_cc
self.calculate_v_resume_additional_cc = config.calculate_v_resume_additional_cc
self.calculate_v_restore_min_cc = config.calculate_v_restore_min_cc
self.v_save_state = config.v_save_state
self.v_resume = config.v_resume
self.callbacks = {
'before_save': config.before_save_function_name,
'after_save': config.after_save_function_name,
'before_restore': config.before_restore_function_name,
'after_restore': config.after_restore_function_name,
}
self.reset()
def reset(self):
"""
Resets the saved dump and state
"""
self.state_dump = {'register_file': None, 'reg_saving_pool': None, 'stack': None, 'heap': None, 'volatile_gst': None, 'non_volatile_gst': None, 'stack_top_address': None, 'global_clock': 0, 'function_call_lookup': None, 'output': None}
self.state_saved = {'register_file': None, 'reg_saving_pool': None, 'stack': None, 'heap': None, 'volatile_gst': None, 'non_volatile_gst': None, 'stack_top_address': None, 'global_clock': 0, 'function_call_lookup': None, 'output': None}
def process_state_save_routines(self, functions, declarations):
"""
If dynamic checkpoint mechanisms: removes all references to the checkpoint and restore routines (if any).
If static checkpoint mechanisms: removes all references to restore routine and removes target from checkpoint routine calls.
NB: Checkpoint routine is considered just as a marker to verify when interrupt and test the execution.
"""
# remove declaration of state-save routines
if self.routine_names['save'] in declarations:
del declarations[self.routine_names['save']]
# remove declaration of state-restore routines
if self.routine_names['restore'] in declarations:
del declarations[self.routine_names['restore']]
# parse each function and remove restore routine (usually is at the start of the main)
for name in functions:
function = functions[name]
i = 0
while i < len(function.body):
instruction = function.body[i]
# if is a restore operation, remove it
if isinstance(instruction, CallOperation) and instruction.name == self.routine_names['restore']:
# preserve label of the instruction to be removed
if instruction.label is not None:
function.body[i+1].label = copy.deepcopy(instruction.label)
del function.body[i]
# if is a checkpoint operation
elif isinstance(instruction, CallOperation) and instruction.name == self.routine_names['save']:
# if state-save routines not statically placed, remove them
if not self.static_placement:
# preserve label
if instruction.label is not None:
function.body[i+1].label = copy.deepcopy(instruction.label)
del function.body[i]
# if statically placed, remove its target and continue
else:
function.body[i].target = None
i += 1
else:
i += 1
# update function labels
function.update_labels()
# remove from internal calls the state save and restore functions
# so they are ignored during register allocation
function._remove_from_calls(self.routine_names['save'])
function._remove_from_calls(self.routine_names['restore'])
def execute_state_save(self):
"""
Saves the state.
"""
logging.info('[CheckpointManager] Saving checkpoint.')
# clean
del self.state_saved['register_file']
del self.state_saved['reg_saving_pool']
del self.state_saved['stack']
del self.state_saved['heap']
del self.state_saved['volatile_gst']
del self.state_saved['non_volatile_gst']
del self.state_saved['stack_top_address']
del self.state_saved['global_clock']
del self.state_saved['function_call_lookup']
del self.state_saved['output']
gc.collect()
self._inject_function_call('before_save')
# save register file
if self.restore_register_file:
self.state_saved['register_file'] = self.vmstate.register_file.dump()
self.state_saved['stack_top_address'] = copy.deepcopy(self.vmstate.memory.stack.top_address)
self.state_saved['reg_saving_pool'] = copy.deepcopy(self.vmstate.reg_saving_pool)
else:
self.state_saved['register_file'] = None
self.state_saved['stack_top_address'] = None
self.state_saved['reg_saving_pool'] = None
# save stack
if self.restore_stack:
self.state_saved['stack'] = self.vmstate.memory.stack.dump()
else:
self.state_saved['stack'] = None
# save heap
if self.restore_heap:
self.state_saved['heap'] = self.vmstate.memory.heap.dump()
else:
self.state_saved['heap'] = None
# save volatile gst
if self.restore_volatile_gst:
volatile_gst = self.vmstate.memory.gst._get_gst_from_ram_name('volatile')
if volatile_gst is not None:
self.state_saved['volatile_gst'] = volatile_gst.dump()
else:
self.state_saved['volatile_gst'] = None
else:
self.state_saved['volatile_gst'] = None
# save non-volatile gst
if self.restore_non_volatile_gst:
non_volatile_gst = self.vmstate.memory.gst._get_gst_from_ram_name('non_volatile')
if non_volatile_gst is not None:
self.state_saved['non_volatile_gst'] = non_volatile_gst.dump()
else:
self.state_saved['non_volatile_gst'] = None
else:
self.state_saved['non_volatile_gst'] = None
self.state_saved['global_clock'] = self.vmstate.global_clock
# save function call lookup
self.state_saved['function_call_lookup'] = copy.deepcopy(self.vmstate.function_call_lookup)
# output
if self.restore_environment:
self.state_saved['output'] = OutputManager.dump()
else:
self.state_saved['output'] = None
if self.physical_mmu is not None:
self.physical_mmu.write(self.state_saved)
# stats
#TODO: new stats
self.vmstate.stats.checkpoint_executed()
self._inject_function_call('after_save')
def execute_state_restore(self):
"""
Restores the previously saved state.
"""
logging.info('[CheckpointManager] Restoring checkpoint.')
self._inject_function_call('before_restore')
if self.physical_mmu is not None:
self.state_saved = self.physical_mmu.read()
#self.physical_mmu.read(self.state_saved)
# restore stack
if self.state_saved['stack'] is not None:
self.vmstate.memory.stack.restore(self.state_saved['stack'])
# restore register saving pool used by register allocation to track registers
if self.state_saved['reg_saving_pool'] is not None:
self.vmstate.reg_saving_pool = copy.deepcopy(self.state_saved['reg_saving_pool'])
# restore register file
if self.state_saved['register_file'] is not None:
self.vmstate.register_file.restore(self.state_saved['register_file'])
# if stack in volatile memory, the saved state contains only the real portion of the stack
# so "real" memory zeroing is required.
if self.volatile_stack:
# Deallocate and remove memory cells that exceed stack_top_address
try:
self.vmstate.memory.stack._deallocate(self.state_saved['stack_top_address'], False, True)
except MemoryException as e:
print(e)
pass
# if stack not restored, restore ESP (is a relative address of the stack)
if self.state_saved['stack'] is None:
# Deallocate memory cells that exceed stack_top_address (just pops)
try:
self.vmstate.memory.stack.deallocate(self.state_saved['stack_top_address'], False, True)
except MemoryException as e:
print(e)
pass
# restore heap
if self.state_saved['heap'] is not None:
self.vmstate.memory.heap.restore(self.state_saved['heap'])
# restore voltaile gst
if self.state_saved['volatile_gst'] is not None:
volatile_gst = self.vmstate.memory.gst._get_gst_from_ram_name('volatile')
volatile_gst.restore(self.state_saved['volatile_gst'])
# restore non-volatile gst
if self.state_saved['non_volatile_gst'] is not None:
non_volatile_gst = self.vmstate.memory.gst._get_gst_from_ram_name('non_volatile')
non_volatile_gst.restore(self.state_saved['non_volatile_gst'])
# restore global_clock
if self.state_saved['global_clock'] > 0:
self.vmstate.global_clock = self.state_saved['global_clock']
# restore function call's stack lookup
if self.state_saved['function_call_lookup'] is not None:
self.vmstate.function_call_lookup = copy.deepcopy(self.state_saved['function_call_lookup'])
# restore environment state
if self.restore_environment:
OutputManager.restore(self.state_saved['output'])
# stats
#TODO: new stats
self.vmstate.stats.checkpoint_restored()
self._inject_function_call('after_restore')
def save_dump(self):
"""
Performs a dump of the vm state
"""
del self.state_dump['register_file']
del self.state_dump['reg_saving_pool']
del self.state_dump['stack']
del self.state_dump['heap']
del self.state_dump['volatile_gst']
del self.state_dump['non_volatile_gst']
del self.state_dump['global_clock']
del self.state_dump['function_call_lookup']
del self.state_dump['output']
gc.collect()
# save register file
self.state_dump['register_file'] = self.vmstate.register_file.dump()
self.state_dump['reg_saving_pool'] = copy.deepcopy(self.vmstate.reg_saving_pool)
# save stack
self.state_dump['stack'] = self.vmstate.memory.stack.dump()
# save heap
self.state_dump['heap'] = self.vmstate.memory.heap.dump()
# save volatile gst
volatile_gst = self.vmstate.memory.gst._get_gst_from_ram_name('volatile')
if volatile_gst is not None:
self.state_dump['volatile_gst'] = volatile_gst.dump()
else:
self.state_dump['volatile_gst'] = None
# save non-volatile gst
non_volatile_gst = self.vmstate.memory.gst._get_gst_from_ram_name('non_volatile')
if non_volatile_gst is not None:
self.state_dump['non_volatile_gst'] = non_volatile_gst.dump()
else:
self.state_dump['non_volatile_gst'] = None
self.state_dump['global_clock'] = self.vmstate.global_clock
# save function call lookup
self.state_dump['function_call_lookup'] = copy.deepcopy(self.vmstate.function_call_lookup)
# save environment state
self.state_dump['output'] = OutputManager.dump()
def restore_dump(self):
"""
Restores the previously taken dump, to make memory consistent again.
"""
logging.info('[CheckpointManager] Restoring memory state to a consistent one.')
# restore register file
self.vmstate.register_file.restore(self.state_dump['register_file'])
# restore register saving pool used by register allocation to track registers
self.vmstate.reg_saving_pool = copy.deepcopy(self.state_dump['reg_saving_pool'])
# restore stack
self.vmstate.memory.stack.restore(self.state_dump['stack'])
# restore heap
self.vmstate.memory.heap.restore(self.state_dump['heap'])
# restore volatile gst
volatile_gst = self.vmstate.memory.gst._get_gst_from_ram_name('volatile')
if volatile_gst is not None:
volatile_gst.restore(self.state_dump['volatile_gst'])
# restore volatile gst
non_volatile_gst = self.vmstate.memory.gst._get_gst_from_ram_name('non_volatile')
if non_volatile_gst is not None:
non_volatile_gst.restore(self.state_dump['non_volatile_gst'])
# restore global_clock, if its dump is present
self.vmstate.global_clock = self.state_dump['global_clock']
# restore function call's stack lookup
self.vmstate.function_call_lookup = copy.deepcopy(self.state_dump['function_call_lookup'])
# restore output table
OutputManager.restore(self.state_dump['output'])
# stats
# TODO: new stats
self.vmstate.stats.dump_restored()
def get_vm_state_diff(self):
"""
Compares the vm state with the vm dump taken during the checkpoint.
It returns a list containing the memory section names that differs.
"""
comparations = {
'register_file': self.vmstate.register_file == self.state_dump['register_file'],
'stack': self.vmstate.memory.stack == self.state_dump['stack'],
'heap': self.vmstate.memory.heap == self.state_dump['heap'],
'volatile_gst': self.vmstate.memory.gst._get_gst_from_ram_name('volatile') == self.state_dump['volatile_gst'],
'non_volatile_gst': self.vmstate.memory.gst._get_gst_from_ram_name('non_volatile') == self.state_dump['non_volatile_gst'],
'global_clock': self.vmstate.global_clock == self.state_dump['global_clock'],
'environment': OutputManager.output_table == self.state_dump['output']
}
for name in list(comparations.keys()):
# remove memory sections names which are not changed.
if comparations[name]:
del comparations[name]
return list(comparations.keys())
def get_vm_state_diff_elements(self):
"""
Returns a dictionary containing the differences for each memory component between the current state and the saved dump.
"""
if self.vmstate.global_clock != self.state_dump['global_clock']:
global_clock_diff = [{'element': 'global_clock', 'dump_value': self.state_dump['global_clock'], 'current_value': self.vmstate.global_clock}]
else:
global_clock_diff = []
diffs = {
'register_file': self.vmstate.register_file.diff(self.state_dump['register_file']),
'stack': self.vmstate.memory.stack.diff(self.state_dump['stack']),
'heap': self.vmstate.memory.heap.diff(self.state_dump['heap']),
'volatile_gst': self.vmstate.memory.gst._get_gst_from_ram_name('volatile').diff(self.state_dump['volatile_gst']),
'non_volatile_gst': self.vmstate.memory.gst._get_gst_from_ram_name('non_volatile').diff(self.state_dump['non_volatile_gst']),
'global_clock': global_clock_diff,
'environment': OutputManager.diff(self.state_dump['output'])
}
return diffs
def _inject_function_call(self, name):
"""
Inject a custom function call
"""
if name not in self.vmstate.custom_callbacks:
return
function_name = self.vmstate.custom_callbacks[name]['name']
# correct return operation(s) of the injected function
ret_ops = self.vmstate.custom_callbacks[name]['ret_ops']
pc = self.vmstate.register_file.pc.save()
# inject PC
self.vmstate.register_file.pc.update(function_name, 0)
while self.vmstate.current_instruction not in ret_ops:
self.vmstate.run_step()
# restore PC
self.vmstate.register_file.pc.restore(pc)