611 lines
25 KiB
Python
611 lines
25 KiB
Python
import copy
|
|
import logging
|
|
|
|
from ScEpTIC import tools
|
|
from ScEpTIC.AST.elements import value
|
|
from ScEpTIC.AST.elements.instructions import memory_operations
|
|
from ScEpTIC.AST.elements.instructions import other_operations
|
|
from ScEpTIC.AST.elements.metadata import Metadata
|
|
from ScEpTIC.AST.register_allocation.linear_scan.register_operations import SaveRegistersOperation, RestoreRegistersOperation
|
|
from ScEpTIC.AST.register_allocation.linear_scan.register_pool import RegisterPool
|
|
from ScEpTIC.exceptions import RegAllocException
|
|
from ScEpTIC.llvmir_parser.sections_parser import global_vars
|
|
|
|
|
|
class LinearScanRegisterAllocator:
|
|
"""
|
|
Implementation of linear scan register allocation.
|
|
It must be done for each function and this class refers to a single function body.
|
|
Implementation details can be found here:
|
|
https://www2.seas.gwu.edu/~hchoi/teaching/cs160d/linearscan.pdf
|
|
"""
|
|
|
|
def __init__(self, function, regs_number, reg_prefix = 'R', spill_prefix = '%spill_', spill_type = 'i32'):
|
|
self.function = function
|
|
self.code = function.body
|
|
|
|
self.virtual_regs = {}
|
|
self.intervals = []
|
|
|
|
self.active = []
|
|
|
|
self.register_pool = RegisterPool(regs_number, reg_prefix)
|
|
|
|
# prefix of spill registers
|
|
self.spill_prefix = spill_prefix
|
|
self.spill_count = 0
|
|
self.spill_dimension = global_vars.parse_type(spill_type)
|
|
self.reg_count = 0
|
|
|
|
# registers to be ignored due to data layout
|
|
self.ignores = []
|
|
|
|
# identifies the line of the last alloca
|
|
# it is used to put alloca operations in the same area of the code, as they are translated with a single ESP increment.
|
|
self.last_alloca = 0
|
|
|
|
# calculate latest alloca operation
|
|
for i in range(0, len(self.code)):
|
|
self.last_alloca = i
|
|
if not isinstance(self.code[i], memory_operations.AllocaOperation):
|
|
break
|
|
|
|
|
|
def run_register_allocation(self):
|
|
"""
|
|
Runs the actual register allocation.
|
|
"""
|
|
|
|
self.do_call_pre_processing()
|
|
self.do_liveness_analysis()
|
|
self.do_register_allocation()
|
|
|
|
|
|
@property
|
|
def spill_type(self):
|
|
"""
|
|
Returns an instance of the spill type, to be used in the operations.
|
|
"""
|
|
|
|
return copy.deepcopy(self.spill_dimension)
|
|
|
|
|
|
def _name_to_virtual_reg(self, name):
|
|
"""
|
|
Returns a Value representing a virtual register with the provided name
|
|
"""
|
|
|
|
return value.Value('virtual_reg', name, None)
|
|
|
|
|
|
def _create_virtual_reg(self):
|
|
"""
|
|
Creates a new virtual register.
|
|
"""
|
|
|
|
name = '{}{}'.format(self.spill_prefix, self.spill_count)
|
|
self.spill_count += 1
|
|
logging.debug('[RegisterAllocation] Creating new virtual register {}.'.format(name))
|
|
return self._name_to_virtual_reg(name)
|
|
|
|
|
|
def append_alloca_operation(self, alloca_type):
|
|
"""
|
|
Appends an alloca operation in the code at the end of the alloca section.
|
|
It returns the virtual registers contining the alloca address.
|
|
"""
|
|
logging.debug('[RegisterAllocation] Appending alloca({}) operation on top.'.format(len(alloca_type)))
|
|
|
|
spill_reg = self._create_virtual_reg()
|
|
target = copy.deepcopy(spill_reg)
|
|
|
|
# create the operation
|
|
alloca = memory_operations.AllocaOperation(target, alloca_type, 1, 1)
|
|
alloca.metadata = Metadata(None, None, None, [], [])
|
|
|
|
# insert it in the code
|
|
self.code.insert(self.last_alloca, alloca)
|
|
|
|
# update last_alloca, so to consider this added one.
|
|
self.last_alloca += 1
|
|
|
|
return spill_reg
|
|
|
|
|
|
def append_store_operation(self, position, target, value, target_type, label, metadata):
|
|
"""
|
|
Appends a store operation at the required position.
|
|
"""
|
|
|
|
logging.debug('[RegisterAllocation] Appending store {}, {} at {}.'.format(target.value, value.value, position))
|
|
|
|
target = copy.deepcopy(target)
|
|
value = copy.deepcopy(value)
|
|
|
|
value.type = copy.deepcopy(target_type)
|
|
|
|
# create the operation
|
|
store = memory_operations.StoreOperation(target, value, 1, False)
|
|
store.label = label
|
|
store.metadata = metadata
|
|
|
|
# insert in the code the operation
|
|
self.code.insert(position, store)
|
|
|
|
|
|
def append_load_operation(self, position, element, load_type, is_arg_of_function_call, metadata):
|
|
"""
|
|
Appends a load operation at the required position.
|
|
It returns the virtual register containing the loaded value.
|
|
"""
|
|
|
|
logging.debug('[RegisterAllocation] Appending load {} at {}.'.format(element.value, position))
|
|
|
|
spill_reg = self._create_virtual_reg()
|
|
target = copy.deepcopy(spill_reg)
|
|
element = copy.deepcopy(element)
|
|
|
|
# create the operation
|
|
load = memory_operations.LoadOperation(target, load_type, element, 1, False)
|
|
load.metadata = metadata
|
|
|
|
# eventually set if it loads a function call's argument
|
|
if is_arg_of_function_call:
|
|
load.is_arg_of_function_call = True
|
|
|
|
# insert in the code the operation
|
|
self.code.insert(position, load)
|
|
|
|
return spill_reg
|
|
|
|
|
|
def do_call_pre_processing(self):
|
|
"""
|
|
Applies call pre-processing, consisting in:
|
|
- marking load of arguments as is_arg_of_function_call
|
|
- creating alloca-store-load for immediate
|
|
- creating alloca-store-load for conversion and other operation s.t. their
|
|
results are passed as arguments of the function call.
|
|
"""
|
|
|
|
search_from = 0
|
|
search_to = 0
|
|
i = 0
|
|
|
|
# scan the code by line. while is used since the lines number may change
|
|
while i < len(self.code):
|
|
operation = self.code[i]
|
|
|
|
# process only call operations which are not part of the checkpoint mechanism.
|
|
if isinstance(operation, other_operations.CallOperation) and operation.name in self.function._calls:
|
|
targets = operation.get_uses()
|
|
search_to = i
|
|
j = search_from
|
|
|
|
# adjust label (if present)
|
|
append_label = operation.label
|
|
|
|
# search for instructions s.t. their target is in function arguments
|
|
while j < search_to:
|
|
op = self.code[j]
|
|
|
|
# target of operation is an argument of the call
|
|
if 'target' in op.__dict__ and op.target is not None and op.target.value in targets:
|
|
targets.remove(op.target.value)
|
|
|
|
# if is a load instruction, mark it
|
|
if isinstance(op, memory_operations.LoadOperation):
|
|
op.is_arg_of_function_call = True
|
|
# else generate alloca, store and load, as it would be done
|
|
# in a real scenario (value saved onto the stack to be passed)
|
|
# NB: load is marked as is_arg_of_function_call (so to emulate stack passing)
|
|
|
|
else:
|
|
# get type of argument
|
|
arg_type = operation.get_type_from_virtual_reg(op.target.value)
|
|
|
|
# append alloca and compensate for the added instruction
|
|
tmp_reg = self.append_alloca_operation(arg_type)
|
|
|
|
# append store before call and compensate for the added instruction
|
|
# search_to = call position
|
|
# nb: here is set the eventual label of the call operation.
|
|
self.append_store_operation(search_to+1, tmp_reg, op.target, arg_type, append_label, operation.metadata)
|
|
|
|
# set the lable to be appended to None
|
|
append_label = None
|
|
|
|
# append load before call and compensate for the added instruction
|
|
tmp_reg = self.append_load_operation(search_to+2, tmp_reg, arg_type, True, operation.metadata)
|
|
|
|
# update the new register name containing the argument value.
|
|
operation.replace_reg_name(op.target.value, tmp_reg.value)
|
|
|
|
# compensate for the 3 added instructions
|
|
i += 3
|
|
|
|
# update search_to value, to consider added instructions
|
|
search_to = i
|
|
|
|
# update code index
|
|
j += 1
|
|
|
|
# create alloca, store and load for immediate values, as it is done
|
|
# for operations on the above lines.
|
|
|
|
for arg in operation.function_args:
|
|
if arg.value_class == 'immediate':
|
|
|
|
# llvm builtins may have 1bit operations, but memory can only work
|
|
# with bytes, so modify the memory dimension to do not have problems.
|
|
# data integrity is preserved.
|
|
if arg.type is not None and arg.type.base_type.bits < 8:
|
|
arg.type.base_type.bits = 8
|
|
|
|
tmp_reg = self.append_alloca_operation(arg.type)
|
|
i += 1
|
|
|
|
# nb: here is set the eventual label of the call operation.
|
|
self.append_store_operation(i, tmp_reg, arg, arg.type, append_label, operation.metadata)
|
|
# set the lable to be appended to None
|
|
append_label = None
|
|
i += 1
|
|
|
|
tmp_reg = self.append_load_operation(i, tmp_reg, arg.type, True, operation.metadata)
|
|
i += 1
|
|
|
|
# update the arguments to be a virtual register instead of an immediate
|
|
arg.value_class = 'virtual_reg'
|
|
|
|
# set the virutal register name
|
|
arg.value = tmp_reg.value
|
|
|
|
# update search from. it is used to limit the area in which function arguments are searched.
|
|
# in llvm all arguments are reloaded if multiple calls happens from them, so there is no need
|
|
# to search arguments from the top (it is sufficient to go from the previous call to the current one)
|
|
search_from = i + 1
|
|
|
|
# refresh the operation label: if it is assigned to another operation, it will be set to false
|
|
# otherwise is unchanged.
|
|
operation.label = append_label
|
|
|
|
# update code index
|
|
i += 1
|
|
|
|
|
|
def do_call_post_processing(self):
|
|
"""
|
|
Applies function call post-processing, which consists in:
|
|
- Insertion of SaveRegistersOperation and RestoreRegistersOperation respectively before and after
|
|
a function call
|
|
- Updating the labels mappings (if code is appended, a label is mapped with a line which won't be exact)
|
|
"""
|
|
|
|
i = 0
|
|
|
|
# scan the code by line. while is used since the lines number may change
|
|
while i < len(self.code):
|
|
operation = self.code[i]
|
|
|
|
# process only call operations which are not part of the checkpoint mechanism.
|
|
if isinstance(operation, other_operations.CallOperation) and operation.name in self.function._calls:
|
|
|
|
# if recursive call, save used registers
|
|
if operation.name == self.function.name:
|
|
tick_count = self.function.reg_count
|
|
|
|
# else estimates the number of registers to be saved in case of a well-partitioned
|
|
# register allocation among functions.
|
|
# e.g. if main uses 2 registers and calls pow which uses 2 registers, a well-partitioned
|
|
# register allocation will give R0 and R1 to main, and R2 and R3 to pow
|
|
# with the result of no register saving needed
|
|
else:
|
|
# usage = reg used by current function + max reg used by the called function (and calls inside it)
|
|
# registers to be saved = usage - available
|
|
# if I have 10 registers, the function uses 2 registers and the called one uses 2 registers,
|
|
# no saving needed.
|
|
# max_reg_usage is calculated iteratively before running do_call_post_processing
|
|
tick_count = self.function.reg_count + operation.get_function_max_reg_usage() - self.register_pool.regs_number
|
|
|
|
# if result is < 0 means that no register should be saved.
|
|
if tick_count < 0:
|
|
tick_count = 0
|
|
|
|
# get target register name
|
|
if 'target' in operation.__dict__ and operation.target is not None:
|
|
target = operation.target.get_uses()[0]
|
|
else:
|
|
target = None
|
|
|
|
# create the saving and restoring operations
|
|
save_reg_op = SaveRegistersOperation(tick_count, target, self.register_pool)
|
|
save_reg_op.metadata = operation.metadata
|
|
|
|
# compensate label
|
|
save_reg_op.label = operation.label
|
|
operation.label = None
|
|
|
|
load_reg_op = RestoreRegistersOperation(save_reg_op)
|
|
load_reg_op.metadata = operation.metadata
|
|
|
|
# insert the save operation before the function call
|
|
self.code.insert(i, save_reg_op)
|
|
|
|
# 1 for compensating save_reg_op, 1 for finding the operation after the call
|
|
i += 2
|
|
|
|
# insert the restore operation after the function call
|
|
self.code.insert(i, load_reg_op)
|
|
|
|
i += 1
|
|
|
|
# update the function's label mappings.
|
|
self.function.update_labels()
|
|
|
|
|
|
def _init_intervals(self):
|
|
"""
|
|
Inits the intervals used for calculating the active interva of the registers for the linear scan.
|
|
"""
|
|
|
|
for _ in range(0, len(self.code)):
|
|
self.intervals.append(None)
|
|
|
|
|
|
def do_liveness_analysis(self):
|
|
"""
|
|
This method performs the liveness analysis, which finds the live interval for each register.
|
|
NB: LLVM IR is in Static Single Assignment (SSA) foarm, so for obtaining the liveness intervals
|
|
is sufficient to map the definition with its last use.
|
|
"""
|
|
# registers to be ignored, since will be considered as stack offsets/addresses
|
|
# which will be resolved by datalayout operation.
|
|
self.ignores = self.function.get_ignore()
|
|
|
|
# for each line, compute the uses and definitions of each register.
|
|
for i in range(0, len(self.code)):
|
|
self.ignores = self.ignores + self.code[i].get_ignore()
|
|
defs = tools.list_sanitize_from_list(self.code[i].get_defs(), self.ignores)
|
|
uses = tools.list_sanitize_from_list(self.code[i].get_uses(), self.ignores)
|
|
|
|
for reg in defs:
|
|
self.virtual_regs[reg] = {'uses': [], 'def': i, 'target': None}
|
|
|
|
for reg in uses:
|
|
self.virtual_regs[reg]['uses'].append(i)
|
|
|
|
self._init_intervals()
|
|
|
|
# find last use of each register and create interval dict
|
|
for reg_id, reg in self.virtual_regs.items():
|
|
# if no use, the last use will be in the same instruction in which the register
|
|
# is defined
|
|
if len(reg['uses']) == 0:
|
|
reg['uses'].append(reg['def'])
|
|
|
|
reg['last_use'] = max(reg['uses'])
|
|
|
|
# set register intervals
|
|
self.intervals[reg['def']] = {'reg': reg_id, 'last_use': reg['last_use']}
|
|
|
|
|
|
def _compensate_for_alloca(self):
|
|
"""
|
|
Updates the liveness analysis to account for the insertion of an alloca operations.
|
|
"""
|
|
|
|
logging.debug('[RegisterAllocation] Compensating for alloca operation.')
|
|
|
|
# alloca operations do not introduces any use/def and are appended on the top of the code
|
|
# so i can insert a None interval on top of the list (since all alloca operations are in the same region)
|
|
self.intervals.insert(0, None)
|
|
|
|
# increment uses of each interval by 1
|
|
for interval in self.intervals:
|
|
if interval is None:
|
|
continue
|
|
|
|
interval['last_use'] += 1
|
|
|
|
# increament uses and definitions of each virtual register by one
|
|
for virtual_reg in self.virtual_regs.values():
|
|
virtual_reg['def'] += 1
|
|
virtual_reg['uses'] = [x+1 for x in virtual_reg['uses']]
|
|
virtual_reg['last_use'] += 1
|
|
|
|
|
|
def _compensate_for_operation(self, instant):
|
|
"""
|
|
Updates the liveness analysis to account for the insertion of an operation which may
|
|
insert use/def and is inserted in a precise line of code.
|
|
instant specifies also the line of code
|
|
"""
|
|
logging.debug('[RegisterAllocation] Compensating for operation at {}.'.format(instant))
|
|
|
|
# insert the proper interval
|
|
self.intervals.insert(instant, None)
|
|
|
|
# increment data by 1 for subsequent intervals
|
|
for interval in self.intervals:
|
|
if interval is None:
|
|
continue
|
|
|
|
# if subsequent interval, it needs to be updated
|
|
if interval['last_use'] >= instant:
|
|
interval['last_use'] += 1
|
|
|
|
# get corresponding virtual register
|
|
virtual_reg = self.virtual_regs[interval['reg']]
|
|
|
|
# increment its definition instant if happens after the given instant
|
|
if virtual_reg['def'] >= instant:
|
|
virtual_reg['def'] += 1
|
|
|
|
# increment all the uses by 1 if they are after the given instant
|
|
virtual_reg['uses'] = [x+1 if x >= instant else x for x in virtual_reg['uses']]
|
|
# increment last use by 1, since it will be certainly after the given instant
|
|
virtual_reg['last_use'] += 1
|
|
|
|
|
|
def _assign_register(self, instant, new_reg_name):
|
|
"""
|
|
Assigns a physical register to a virtual register which is defined in a given instant.
|
|
"""
|
|
|
|
interval = self.intervals[instant]
|
|
reg_name = interval['reg']
|
|
logging.debug('[RegisterAllocation] Assigning register {} to virtual register {} at {}.'.format(new_reg_name, reg_name, instant))
|
|
|
|
# update each occurrence of the old register with the new register
|
|
for line in self.virtual_regs[reg_name]['uses']:
|
|
self.code[line].replace_reg_name(reg_name, new_reg_name)
|
|
|
|
self.code[instant].replace_reg_name(reg_name, new_reg_name)
|
|
self.virtual_regs[reg_name]['target'] = new_reg_name
|
|
|
|
# insert register into the active ones.
|
|
self.active.append(interval['reg'])
|
|
|
|
def do_register_allocation(self):
|
|
"""
|
|
Performs the linear scan register allocation over the given function.
|
|
"""
|
|
|
|
instant = 0
|
|
|
|
# scan all intervals and insert register spills / promotions
|
|
while instant < len(self.intervals):
|
|
interval = self.intervals[instant]
|
|
|
|
if interval is None:
|
|
instant += 1
|
|
continue
|
|
|
|
# expire old registers
|
|
self.expire_old(instant)
|
|
|
|
# set max reg count (used to estimate the tick_count)
|
|
self.reg_count = max(self.reg_count, len(self.active))
|
|
|
|
# if all registers used, need to spill.
|
|
if len(self.active) == self.register_pool.regs_number:
|
|
# spill at current instant
|
|
self.spill_at(instant)
|
|
# compensate for alloca and store
|
|
instant += 2
|
|
|
|
# else assign register
|
|
else:
|
|
new_reg_name = self.register_pool.get_reg()
|
|
self._assign_register(instant, new_reg_name)
|
|
|
|
# increment instant index
|
|
instant += 1
|
|
|
|
# set function regiser count.
|
|
self.function.reg_count = self.reg_count
|
|
|
|
|
|
def expire_old(self, instant):
|
|
"""
|
|
Updates the active registers by removing expired ones.
|
|
"""
|
|
|
|
logging.debug('[RegisterAllocation] Expiring registers at {}.'.format(instant))
|
|
|
|
# scan active registers to find if some of them can be freed.
|
|
for reg in self.active:
|
|
|
|
# if the last use is before current instant, the register can be
|
|
# considered as free.
|
|
if self.virtual_regs[reg]['last_use'] < instant:
|
|
|
|
# remove from active registers
|
|
self.active.remove(reg)
|
|
|
|
# free register
|
|
reg_name = self.virtual_regs[reg]['target']
|
|
self.register_pool.free_reg(reg_name)
|
|
|
|
def _get_fist_use_after_instant(self, reg, instant):
|
|
"""
|
|
Returns the first use of a register after a given instant.
|
|
"""
|
|
|
|
uses = self.virtual_regs[reg]['uses']
|
|
first_use = self.virtual_regs[reg]['last_use']
|
|
|
|
for i in uses:
|
|
if instant < i < first_use:
|
|
first_use = i
|
|
|
|
return first_use
|
|
|
|
def spill_at(self, instant):
|
|
"""
|
|
Selects a register to be spilled in the given instant and spills it.
|
|
"""
|
|
|
|
selected_spill = None
|
|
selected_reg = None
|
|
latest_use = -1
|
|
|
|
# scans each active register.
|
|
for reg in self.active:
|
|
reg_data = self.virtual_regs[reg]
|
|
|
|
# if register can be spilled (= not used in this instant nor defined)
|
|
if instant not in reg_data['uses'] and instant != reg_data['def']:
|
|
# get first use after current instant
|
|
first_use = self._get_fist_use_after_instant(reg, instant)
|
|
|
|
# update max usage informations, used for spilling selection
|
|
if first_use > latest_use:
|
|
selected_spill = reg_data['target']
|
|
selected_reg = reg
|
|
latest_use = first_use
|
|
|
|
# no register can be spilled, exit.
|
|
if selected_spill is None:
|
|
raise RegAllocException('Unable to select a register for spilling')
|
|
|
|
logging.debug('[RegisterAllocation] Spilling register {} at instant {}'.format(selected_spill, instant))
|
|
|
|
# append alloca operation, to account for stack space
|
|
tmp_reg = self.append_alloca_operation(self.spill_type)
|
|
|
|
# compensate for alloca
|
|
self._compensate_for_alloca()
|
|
instant += 1
|
|
|
|
# save register value onto the stack before current instant
|
|
value_reg = self._name_to_virtual_reg(selected_spill)
|
|
self.append_store_operation(instant, tmp_reg, value_reg, self.spill_type, None, self.code[instant].metadata)
|
|
|
|
# compensate for store
|
|
self._compensate_for_operation(instant)
|
|
instant += 1
|
|
|
|
# compensate for alloca and store. NB: use is after alloca and store
|
|
latest_use += 2
|
|
# create load before the first use of spilled register and
|
|
tmp_reg = self.append_load_operation(latest_use, tmp_reg, self.spill_type, False, self.code[instant].metadata)
|
|
self._compensate_for_operation(latest_use)
|
|
|
|
# retarget uses of spilled register to created load (new reg name, which for now is virtual and in
|
|
# a next step will be updated with a physical one)
|
|
uses = [x for x in self.virtual_regs[selected_reg]['uses'] if x >= latest_use]
|
|
for i in uses:
|
|
self.code[i].replace_reg_name(selected_spill, tmp_reg.value)
|
|
|
|
# create interval and register for load, so to assign a physical register to it in a next iteration.
|
|
reg_id = tmp_reg.value
|
|
last_use = max(uses)
|
|
virtual_reg = {'uses': uses, 'def': latest_use, 'target': None, 'last_use': last_use}
|
|
self.virtual_regs[reg_id] = virtual_reg
|
|
self.intervals[latest_use] = {'reg': reg_id, 'last_use': last_use}
|
|
|
|
# assign register
|
|
self.active.remove(selected_reg)
|
|
self._assign_register(instant, selected_spill)
|