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), [], [], {})