2026-07-10 10:38:57 +02:00

362 lines
12 KiB
Python

from ScEpTIC.AST.elements.instructions.memory_operations import LoadOperation, StoreOperation
from ScEpTIC.AST.elements.instructions.other_operations import CallOperation
from ScEpTIC.AST.elements.types import BaseType
from ScEpTIC.AST.misc.virtual_memory_enum import VirtualMemoryEnum
from ScEpTIC.AST.transformations.ratchet.memory_tags_identifier import MemoryTagsIdentifier
class CheckpointPlacer:
def __init__(self, functions, functions_with_outside_frame_accesses, checkpoint_function_name):
self.functions = functions
self.functions_with_outside_frame_accesses = []
self.checkpoint_function_name = checkpoint_function_name
self.function_body_slices = {}
# adjust for LLVM IR
for function_name in functions_with_outside_frame_accesses:
name = f"@{function_name}"
self.functions_with_outside_frame_accesses.append(name)
def place_checkpoints(self):
"""
Places checkpoints in each function
"""
self._identify_memory_tags()
self._place_checkpoints_for_function_calls()
self._remove_duplicate_checkpoints()
self._split_function_bodies()
self._place_checkpoints()
self._synchronize_slices_to_body()
self._remove_duplicate_checkpoints()
self._set_nvm_default_memory()
def print_slices(self):
"""
Prints all the slices
"""
for name, slices in self.function_body_slices.items():
print(name)
for slice_index, slice in enumerate(slices):
print(f" --- Slice {slice_index} ---")
for index, instruction in enumerate(slice):
print(f" {index}: {instruction}")
print()
print()
def _set_nvm_default_memory(self):
"""
Sets each instruction to target NVM
"""
for function in self.functions.values():
for instruction in function.body:
if "virtual_memory_target" in dir(instruction):
instruction.virtual_memory_target = VirtualMemoryEnum.NON_VOLATILE
def _synchronize_slices_to_body(self):
"""
Updates the functions body with the modified slices
"""
for name, slices in self.function_body_slices.items():
instructions = []
for slice in slices:
for instruction in slice:
instructions.append(instruction)
self.functions[name].update_body(instructions)
def _place_checkpoints(self):
"""
Places checkpoints in function slices
"""
for slices in self.function_body_slices.values():
for slice in slices:
self._place_checkpoints_in_slice(slice)
def _identify_memory_tags(self):
"""
Identifies the memory tag of each instruction
"""
memory_tags = MemoryTagsIdentifier(self.functions)
memory_tags.resolve_memory_tags()
def _split_function_bodies(self):
"""
Splits the function body in slices
"""
for name, function in self.functions.items():
# Skip builtins processing
if function.is_builtin or function.is_input:
continue
self.function_body_slices[name] = []
slice = []
for instruction in function.body:
slice.append(instruction)
# slice end reached -> add to slices and start a new slice
if isinstance(instruction, CallOperation) and instruction.name == self.checkpoint_function_name:
self.function_body_slices[name].append(slice)
slice = []
# account for last slice
if len(slice) > 0:
self.function_body_slices[name].append(slice)
def _identify_possible_checkpoints(self, slice):
"""
Scans a slice and identifies all the pairs of store/load, which makes idempotent sections.
1. LOAD
2. LOAD
3. STORE
4. STORE
5. LOAD
Lines 3-4 makes an idempotent section, so we need a checkpoint between line 2 and 3
"""
checkpoint_ranges = {}
reads = {}
writes = {}
# scan in reverse order
for index, instruction in reversed(list(enumerate(slice))):
if isinstance(instruction, LoadOperation):
memory_tag = instruction.memory_tag
# if there is a write after the read -> checkpoint needed
if memory_tag in writes:
write = writes[memory_tag]
del writes[memory_tag]
if memory_tag not in checkpoint_ranges:
checkpoint_ranges[memory_tag] = []
data = {"locations": set(range(index, write)), "memory_tag": memory_tag, "covered": False}
checkpoint_ranges[memory_tag].append(data)
if isinstance(instruction, StoreOperation):
memory_tag = instruction.memory_tag
# only last identified store needed
writes[memory_tag] = index
return checkpoint_ranges
def _identify_ordered_possible_checkpoints(self, slice):
"""
Returns an ordered list of all the possible checkpoint locations
"""
possible_checkpoints = self._identify_possible_checkpoints(slice)
checkpoint_locations = {}
for data in possible_checkpoints.values():
for element in data:
index = len(element["locations"])
if index not in checkpoint_locations:
checkpoint_locations[index] = []
checkpoint_locations[index].append(element)
ordered_locations = []
for _, possible_locations in sorted(checkpoint_locations.items()):
for possible_location in possible_locations:
ordered_locations.append(possible_location)
return ordered_locations
def _place_checkpoints_in_slice(self, slice):
"""
Places checkpoints to create idempotent code sections in the provided slice of code
"""
checkpoints = self._identify_checkpoints(slice)
placed_checkpoints = 0
for checkpoint_index in checkpoints:
index = checkpoint_index + placed_checkpoints
# instruction at index
instruction = slice[index]
# build checkpoint
new_checkpoint = self._create_checkpoint()
new_checkpoint.basic_block_id = instruction.basic_block_id
# update labels
new_checkpoint.label = instruction.label
instruction.label = None
slice.insert(index, new_checkpoint)
placed_checkpoints += 1
def _identify_checkpoints(self, slice):
"""
Returns a list that contains the indexes where to place checkpoints.
The list is ordered.
Approximation: uses a map colouring-like algorithm.
Assumption valid for our benchmarks, as non-idempotent sections cannot be avoided inside loops
"""
ordered_locations = self._identify_ordered_possible_checkpoints(slice)
checkpoints = []
for location in ordered_locations:
# skip if already covered
if location["covered"]:
continue
candidates = {}
# tries each possible checkpoint placement
for candidate_checkpoint in location["locations"]:
candidates[candidate_checkpoint] = []
# check which locations the current candidate covers
for next_location in ordered_locations:
# skip previous locations and already covered locations
# Note: the current location is instead considered
if next_location["covered"]:
continue
if candidate_checkpoint in next_location["locations"]:
candidates[candidate_checkpoint].append(next_location)
# order possible checkpoint placements by covered locations
ordered_candidates = {}
# parse in order of checkpoint location
for candidate_checkpoint, data in sorted(candidates.items()):
index = len(data)
if index not in ordered_candidates:
ordered_candidates[index] = []
ordered_candidates[index].append(candidate_checkpoint)
# get the elements that cover the highest number of locations
max_covered_locations = max(ordered_candidates.keys())
# get the lowest location (instruction index), that is, where a lower number of live registers are (LOAD)
# Note: placing a checkpoint before a STORE requires to save two registers: target and value
target_checkpoint = ordered_candidates[max_covered_locations][0]
checkpoints.append(target_checkpoint)
# mark as covered the covered locations
for location in ordered_locations:
if target_checkpoint in location["locations"]:
location["covered"] = True
return sorted(checkpoints)
def _place_checkpoints_for_function_calls(self):
"""
Places a checkpoint before and after function calls that have outside frame accesses
"""
for function in self.functions.values():
# Skip builtins processing
if function.is_builtin or function.is_input:
continue
target_calls = []
for index, instruction in enumerate(function.body):
if isinstance(instruction, CallOperation) and instruction.name in self.functions_with_outside_frame_accesses:
target_calls.append(index)
new_checkpoints = 0
for call_index in sorted(target_calls):
index = call_index + new_checkpoints
target_call = function.body[index]
# create checkpoints
checkpoint_after = self._create_checkpoint()
checkpoint_before = self._create_checkpoint()
# set basic block id
checkpoint_after.basic_block_id = target_call.basic_block_id
checkpoint_before.basic_block_id = target_call.basic_block_id
# adjust label
checkpoint_before.label = target_call.label
target_call.label = None
# add checkpoint
function.body.insert(index, checkpoint_before)
# new index = index+1 -> add after -> index + 2
function.body.insert(index+2, checkpoint_after)
new_checkpoints += 2
def _remove_duplicate_checkpoints(self):
"""
Removes duplicate checkpoints
"""
for function in self.functions.values():
# Skip builtins processing
if function.is_builtin or function.is_input:
continue
to_remove = []
for index, instruction in enumerate(function.body[:-1]):
if isinstance(instruction, CallOperation) and instruction.name == self.checkpoint_function_name:
next_instruction = function.body[index+1]
if isinstance(next_instruction, CallOperation) and next_instruction.name == self.checkpoint_function_name:
to_remove.append(index)
removed = 0
for checkpoint_index in sorted(to_remove):
index = checkpoint_index - removed
checkpoint = function.body[index]
# sets the index on the next checkpoint
if checkpoint.label is not None:
function.body[index+1].label = checkpoint.label
del function.body[index]
removed += 1
def _create_checkpoint(self):
"""
Creates and returns a custom checkpoint function call
"""
return CallOperation(None, self.checkpoint_function_name, BaseType("int", 32), [], [], {})