206 lines
6.8 KiB
Python
206 lines
6.8 KiB
Python
from ScEpTIC.AST.elements.instructions.memory_operations import AllocaOperation
|
|
from ScEpTIC.AST.elements.instructions.other_operations import CallOperation
|
|
|
|
|
|
class RegisterSavingOptimization:
|
|
"""
|
|
Identifies the minimum set of registes that each checkpoint need to save
|
|
"""
|
|
|
|
def __init__(self, checkpoint_placer):
|
|
self.checkpoint_placer = checkpoint_placer
|
|
|
|
|
|
def set_checkpoints_registers(self):
|
|
"""
|
|
Identifies and sets the registers that each checkpoints need to save
|
|
"""
|
|
|
|
self.checkpoint_placer._split_function_bodies()
|
|
|
|
for name in self.checkpoint_placer.function_body_slices.keys():
|
|
self._reset_registers_info(name)
|
|
self._set_checkpoints_general_purpose_regs(name)
|
|
self._set_checkpoints_esp_reg(name)
|
|
|
|
self.checkpoint_placer._synchronize_slices_to_body()
|
|
|
|
|
|
def _reset_registers_info(self, name):
|
|
"""
|
|
Resets the register-saving information associated to each checkpoint
|
|
"""
|
|
slices = self.checkpoint_placer.function_body_slices[name]
|
|
|
|
for slice in slices:
|
|
checkpoints = self._get_checkpoints(slice)
|
|
|
|
for checkpoint in checkpoints:
|
|
checkpoint.checkpoint_save_pc = True
|
|
checkpoint.checkpoint_save_esp = False
|
|
checkpoint.checkpoint_save_regs = set()
|
|
|
|
|
|
def _get_checkpoints(self, slice):
|
|
"""
|
|
Returns a list of all the checkpoints included in the slices
|
|
"""
|
|
instructions = []
|
|
|
|
for instruction in slice:
|
|
if isinstance(instruction, CallOperation) and instruction.name == self.checkpoint_placer.checkpoint_function_name:
|
|
instructions.append(instruction)
|
|
|
|
return instructions
|
|
|
|
|
|
def _set_checkpoints_general_purpose_regs(self, name):
|
|
"""
|
|
Sets the general purpose registers that each checkpoint in the function need to save
|
|
"""
|
|
|
|
# account for function's argument registers
|
|
n_args = len(self.checkpoint_placer.functions[name].arguments)
|
|
|
|
slices = self.checkpoint_placer.function_body_slices[name]
|
|
alloca_regs = self._get_alloca_regs(slices, n_args)
|
|
|
|
alive_regs = set()
|
|
|
|
# analyze slices in reverse order to identify general purpose registers to save
|
|
for slice in reversed(slices):
|
|
to_save = set()
|
|
|
|
# retrieve regs infos
|
|
regs_external_uses, regs_defs = self._get_regs_first_use_def(slice, alloca_regs)
|
|
|
|
# add live set to registers to save
|
|
for reg in alive_regs:
|
|
to_save.add(reg)
|
|
|
|
# if a defined reg is in alive_regs -> remove it from alive_regs
|
|
for reg in regs_defs:
|
|
if reg in alive_regs:
|
|
alive_regs.remove(reg)
|
|
|
|
# add to alive_regs the used regs
|
|
for reg in regs_external_uses:
|
|
alive_regs.add(reg)
|
|
|
|
self._set_checkpoints_save_regs(slice, to_save)
|
|
|
|
|
|
def _get_alloca_regs(self, slices, n_args):
|
|
"""
|
|
Returns the registers that are target of alloca instructions
|
|
"""
|
|
alloca_regs = set()
|
|
|
|
# accounts for registers that contain arguments
|
|
for i in range(n_args):
|
|
alloca_regs.add(f"%{i}")
|
|
|
|
# order does not matter
|
|
for slice in slices:
|
|
for instruction in slice:
|
|
if isinstance(instruction, AllocaOperation):
|
|
alloca_regs.add(instruction.target.value)
|
|
|
|
return alloca_regs
|
|
|
|
|
|
def _get_regs_first_use_def(self, slice, alloca_regs):
|
|
"""
|
|
Returns two lists:
|
|
- registers whose use happens before their defs
|
|
- registers defined
|
|
"""
|
|
reg_first_uses = set()
|
|
reg_defs = set()
|
|
|
|
for instruction in slice:
|
|
for reg in instruction.get_uses():
|
|
# reg not defined + ignore alloca regs
|
|
if reg not in reg_defs and reg not in alloca_regs:
|
|
reg_first_uses.add(reg)
|
|
|
|
for reg in instruction.get_defs():
|
|
# ignore alloca regs
|
|
if reg not in alloca_regs:
|
|
reg_defs.add(reg)
|
|
|
|
return reg_first_uses, reg_defs
|
|
|
|
|
|
def _set_checkpoints_save_regs(self, slice, to_save):
|
|
"""
|
|
Sets the checkpoints contained in the slice to save the given registers
|
|
|
|
Assumption (safe) -> no reg is used after a function call
|
|
This is a safe assumption for the benchmarks we are considering, as we consider LLVM IR virtual regs
|
|
"""
|
|
checkpoints = self._get_checkpoints(slice)
|
|
|
|
for checkpoint in checkpoints:
|
|
checkpoint.checkpoint_save_pc = True
|
|
|
|
for reg in to_save:
|
|
checkpoint.checkpoint_save_regs.add(reg)
|
|
|
|
if len(to_save) > 0:
|
|
checkpoint.checkpoint_save_esp = True
|
|
|
|
|
|
def _set_checkpoints_esp_reg(self, name):
|
|
"""
|
|
Sets which checkpoint in the function need to save the stack pointer
|
|
"""
|
|
slices = self.checkpoint_placer.function_body_slices[name]
|
|
|
|
# all checkpoints need to save stack pointer
|
|
#for slice in slices:
|
|
# self._set_checkpoint_save_esp(slice)
|
|
|
|
# first slice always need to save stack pointer
|
|
self._set_checkpoint_save_esp(slices[0])
|
|
|
|
# a checkpoint at the end of a slice need to save the stack pointer only if the slice calls a function with a checkpoint
|
|
for slice in slices[1:]:
|
|
for instruction in slice:
|
|
# call to a function that is not a checkpoint
|
|
if isinstance(instruction, CallOperation) and instruction.name != self.checkpoint_placer.checkpoint_function_name:
|
|
# the function contains a checkpoint, that is, saves the stack pointer
|
|
if self._function_has_checkpoint(instruction.name):
|
|
self._set_checkpoint_save_esp(slice)
|
|
# go to next slice
|
|
break
|
|
|
|
|
|
def _set_checkpoint_save_esp(self, slice):
|
|
"""
|
|
Sets the checkpoints contained in the given slice to save the stack pointer
|
|
"""
|
|
checkpoints = self._get_checkpoints(slice)
|
|
|
|
for checkpoint in checkpoints:
|
|
checkpoint.checkpoint_save_esp = True
|
|
|
|
|
|
def _function_has_checkpoint(self, name):
|
|
"""
|
|
Returns if a given function contains a checkpoint
|
|
"""
|
|
|
|
# builtin or input function
|
|
if name not in self.checkpoint_placer.function_body_slices:
|
|
return False
|
|
|
|
slices = self.checkpoint_placer.function_body_slices[name]
|
|
|
|
for slice in slices:
|
|
for instruction in slice:
|
|
if isinstance(instruction, CallOperation) and instruction.name == self.checkpoint_placer.checkpoint_function_name:
|
|
return True
|
|
|
|
return False
|