| #!/usr/bin/env python3 |
| |
| from __future__ import annotations |
| from dataclasses import dataclass |
| from instructions import * |
| from typing import Any, Iterable, Callable, Optional, Tuple, List, Dict |
| import argparse |
| import fileinput |
| import inspect |
| import re |
| import sys |
| |
| RE_EXPECTS = re.compile(r"^([0-9]+,)*[0-9]+$") |
| |
| |
| # Parse the SPIR-V instructions. Some instructions are ignored because |
| # not required to simulate this module. |
| # Instructions are to be implemented in instructions.py |
| def parseInstruction(i): |
| IGNORED = set( |
| [ |
| "OpCapability", |
| "OpMemoryModel", |
| "OpExecutionMode", |
| "OpExtension", |
| "OpSource", |
| "OpTypeInt", |
| "OpTypeStruct", |
| "OpTypeFloat", |
| "OpTypeBool", |
| "OpTypeVoid", |
| "OpTypeFunction", |
| "OpTypePointer", |
| "OpTypeArray", |
| ] |
| ) |
| if i.opcode() in IGNORED: |
| return None |
| |
| try: |
| Type = getattr(sys.modules["instructions"], i.opcode()) |
| except AttributeError: |
| raise RuntimeError(f"Unsupported instruction {i}") |
| if not inspect.isclass(Type): |
| raise RuntimeError( |
| f"{i} instruction definition is not a class. Did you used 'def' instead of 'class'?" |
| ) |
| return Type(i.line) |
| |
| |
| # Split a list of instructions into pieces. Pieces are delimited by instructions of the type splitType. |
| # The delimiter is the first instruction of the next piece. |
| # This function returns no empty pieces: |
| # - if 2 subsequent delimiters will mean 2 pieces. One with only the first delimiter, and the second |
| # with the delimiter and following instructions. |
| # - if the first instruction is a delimiter, the first piece will begin with this delimiter. |
| def splitInstructions( |
| splitType: type, instructions: Iterable[Instruction] |
| ) -> List[List[Instruction]]: |
| blocks: List[List[Instruction]] = [[]] |
| for instruction in instructions: |
| if isinstance(instruction, splitType) and len(blocks[-1]) > 0: |
| blocks.append([]) |
| blocks[-1].append(instruction) |
| return blocks |
| |
| |
| # Defines a BasicBlock in the simulator. |
| # Begins at an OpLabel, and ends with a control-flow instruction. |
| class BasicBlock: |
| def __init__(self, instructions) -> None: |
| assert isinstance(instructions[0], OpLabel) |
| # The name of the basic block, which is the register of the leading |
| # OpLabel. |
| self._name = instructions[0].output_register() |
| # The list of instructions belonging to this block. |
| self._instructions = instructions[1:] |
| |
| # Returns the name of this basic block. |
| def name(self): |
| return self._name |
| |
| # Returns the instruction at index in this basic block. |
| def __getitem__(self, index: int) -> Instruction: |
| return self._instructions[index] |
| |
| # Returns the number of instructions in this basic block, excluding the |
| # leading OpLabel. |
| def __len__(self): |
| return len(self._instructions) |
| |
| def dump(self): |
| print(f" {self._name}:") |
| for instruction in self._instructions: |
| print(f" {instruction}") |
| |
| |
| # Defines a Function in the simulator. |
| class Function: |
| def __init__(self, instructions) -> None: |
| assert isinstance(instructions[0], OpFunction) |
| # The name of the function (name of the register returned by OpFunction). |
| self._name: str = instructions[0].output_register() |
| # The list of basic blocks that belongs to this function. |
| self._basic_blocks: List[BasicBlock] = [] |
| # The variables local to this function. |
| self._variables: List[OpVariable] = [ |
| x for x in instructions if isinstance(x, OpVariable) |
| ] |
| |
| assert isinstance(instructions[-1], OpFunctionEnd) |
| body = filter(lambda x: not isinstance(x, OpVariable), instructions[1:-1]) |
| for block in splitInstructions(OpLabel, body): |
| self._basic_blocks.append(BasicBlock(block)) |
| |
| # Returns the name of this function. |
| def name(self) -> str: |
| return self._name |
| |
| # Returns the basic block at index in this function. |
| def __getitem__(self, index: int) -> BasicBlock: |
| return self._basic_blocks[index] |
| |
| # Returns the index of the basic block with the given name if found, |
| # -1 otherwise. |
| def get_bb_index(self, name) -> int: |
| for i in range(len(self._basic_blocks)): |
| if self._basic_blocks[i].name() == name: |
| return i |
| return -1 |
| |
| def dump(self): |
| print(" Variables:") |
| for var in self._variables: |
| print(f" {var}") |
| print(" Blocks:") |
| for bb in self._basic_blocks: |
| bb.dump() |
| |
| |
| # Represents an instruction pointer in the simulator. |
| @dataclass |
| class InstructionPointer: |
| # The current function the IP points to. |
| function: Function |
| # The basic block index in function IP points to. |
| basic_block: int |
| # The instruction in basic_block IP points to. |
| instruction_index: int |
| |
| def __str__(self): |
| bb = self.function[self.basic_block] |
| i = bb[self.instruction_index] |
| return f"{bb.name()}:{self.instruction_index} in {self.function.name()} | {i}" |
| |
| def __hash__(self): |
| return hash((self.function.name(), self.basic_block, self.instruction_index)) |
| |
| # Returns the basic block IP points to. |
| def bb(self) -> BasicBlock: |
| return self.function[self.basic_block] |
| |
| # Returns the instruction IP points to. |
| def instruction(self): |
| return self.function[self.basic_block][self.instruction_index] |
| |
| # Increment IP by 1. This only works inside a basic-block boundary. |
| # Incrementing IP when at the boundary of a basic block will fail. |
| def __add__(self, value: int): |
| bb = self.function[self.basic_block] |
| assert len(bb) > self.instruction_index + value |
| return InstructionPointer( |
| self.function, self.basic_block, self.instruction_index + value |
| ) |
| |
| |
| # Defines a Lane in this simulator. |
| class Lane: |
| # The registers known by this lane. |
| _registers: Dict[str, Any] |
| # The current IP of this lane. |
| _ip: Optional[InstructionPointer] |
| # If this lane running. |
| _running: bool |
| # The wave this lane belongs to. |
| _wave: Wave |
| # The callstack of this lane. Each tuple represents 1 call. |
| # The first element is the IP the function will return to. |
| # The second element is the callback to call to store the return value |
| # into the correct register. |
| _callstack: List[Tuple[InstructionPointer, Callable[[Any], None]]] |
| |
| _previous_bb: Optional[BasicBlock] |
| _current_bb: Optional[BasicBlock] |
| |
| def __init__(self, wave: Wave, tid: int) -> None: |
| self._registers = dict() |
| self._ip = None |
| self._running = True |
| self._wave = wave |
| self._callstack = [] |
| |
| # The index of this lane in the wave. |
| self._tid = tid |
| # The last BB this lane was executing into. |
| self._previous_bb = None |
| # The current BB this lane is executing into. |
| self._current_bb = None |
| |
| # Returns the lane/thread ID of this lane in its wave. |
| def tid(self) -> int: |
| return self._tid |
| |
| # Returns true is this lane if the first by index in the current active tangle. |
| def is_first_active_lane(self) -> bool: |
| return self._tid == self._wave.get_first_active_lane_index() |
| |
| # Broadcast value into the registers of all active lanes. |
| def broadcast_register(self, register: str, value: Any) -> None: |
| self._wave.broadcast_register(register, value) |
| |
| # Returns the IP this lane is currently at. |
| def ip(self) -> InstructionPointer: |
| assert self._ip is not None |
| return self._ip |
| |
| # Returns true if this lane is running, false otherwise. |
| # Running means not dead. An inactive lane is running. |
| def running(self) -> bool: |
| return self._running |
| |
| # Set the register at "name" to "value" in this lane. |
| def set_register(self, name: str, value: Any) -> None: |
| self._registers[name] = value |
| |
| # Get the value in register "name" in this lane. |
| # If allow_undef is true, fetching an unknown register won't fail. |
| def get_register(self, name: str, allow_undef: bool = False) -> Optional[Any]: |
| if allow_undef and name not in self._registers: |
| return None |
| return self._registers[name] |
| |
| def set_ip(self, ip: InstructionPointer) -> None: |
| if ip.bb() != self._current_bb: |
| self._previous_bb = self._current_bb |
| self._current_bb = ip.bb() |
| self._ip = ip |
| |
| def get_previous_bb_name(self): |
| return self._previous_bb.name() |
| |
| def handle_convergence_header(self, instruction): |
| self._wave.handle_convergence_header(self, instruction) |
| |
| def do_call(self, ip, output_register): |
| return_ip = None if self._ip is None else self._ip + 1 |
| self._callstack.append( |
| (return_ip, lambda value: self.set_register(output_register, value)) |
| ) |
| self.set_ip(ip) |
| |
| def do_return(self, value): |
| ip, callback = self._callstack[-1] |
| self._callstack.pop() |
| |
| callback(value) |
| if len(self._callstack) == 0: |
| self._running = False |
| else: |
| self.set_ip(ip) |
| |
| |
| # Represents the SPIR-V module in the simulator. |
| class Module: |
| _functions: Dict[str, Function] |
| _prolog: List[Instruction] |
| _globals: List[Instruction] |
| _name2reg: Dict[str, str] |
| _reg2name: Dict[str, str] |
| |
| def __init__(self, instructions) -> None: |
| chunks = splitInstructions(OpFunction, instructions) |
| |
| # The instructions located outside of all functions. |
| self._prolog = chunks[0] |
| # The functions in this module. |
| self._functions = {} |
| # Global variables in this module. |
| self._globals = [ |
| x |
| for x in instructions |
| if isinstance(x, OpVariable) or issubclass(type(x), OpConstant) |
| ] |
| |
| # Helper dictionaries to get real names of registers, or registers by names. |
| self._name2reg = {} |
| self._reg2name = {} |
| for instruction in instructions: |
| if isinstance(instruction, OpName): |
| name = instruction.name() |
| reg = instruction.decoratedRegister() |
| self._name2reg[name] = reg |
| self._reg2name[reg] = name |
| |
| for chunk in chunks[1:]: |
| function = Function(chunk) |
| assert function.name() not in self._functions |
| self._functions[function.name()] = function |
| |
| # Returns the register matching "name" if any, None otherwise. |
| # This assumes names are unique. |
| def getRegisterFromName(self, name): |
| if name in self._name2reg: |
| return self._name2reg[name] |
| return None |
| |
| # Returns the name given to "register" if any, None otherwise. |
| def getNameFromRegister(self, register): |
| if register in self._reg2name: |
| return self._reg2name[register] |
| return None |
| |
| # Initialize the module before wave execution begins. |
| # See Instruction::static_execution for more details. |
| def initialize(self, lane): |
| for instruction in self._globals: |
| instruction.static_execution(lane) |
| |
| # Initialize builtins |
| for instruction in self._prolog: |
| if isinstance(instruction, OpDecorate): |
| instruction.static_execution(lane) |
| |
| def execute_one_instruction(self, lane: Lane, ip: InstructionPointer) -> None: |
| ip.instruction().runtime_execution(self, lane) |
| |
| # Returns the first valid IP for the function defined by the given register. |
| # Calling this with a register not returned by OpFunction is illegal. |
| def get_function_entry(self, register: str) -> InstructionPointer: |
| if register not in self._functions: |
| raise RuntimeError(f"Function defining {register} not found.") |
| return InstructionPointer(self._functions[register], 0, 0) |
| |
| # Returns the first valid IP for the basic block defined by register. |
| # Calling this with a register not returned by an OpLabel is illegal. |
| def get_bb_entry(self, register: str) -> InstructionPointer: |
| for name, function in self._functions.items(): |
| index = function.get_bb_index(register) |
| if index != -1: |
| return InstructionPointer(function, index, 0) |
| raise RuntimeError(f"Instruction defining {register} not found.") |
| |
| # Returns the list of function names in this module. |
| # If an OpName exists for this function, returns the pretty name, else |
| # returns the register name. |
| def get_function_names(self): |
| return [self.getNameFromRegister(reg) for reg, func in self._functions.items()] |
| |
| # Returns the global variables defined in this module. |
| def variables(self) -> Iterable: |
| return [x.output_register() for x in self._globals] |
| |
| def dump(self, function_name: Optional[str] = None): |
| print("Module:") |
| print(" globals:") |
| for instruction in self._globals: |
| print(f" {instruction}") |
| |
| if function_name is None: |
| print(" functions:") |
| for register, function in self._functions.items(): |
| name = self.getNameFromRegister(register) |
| print(f" Function {register} ({name})") |
| function.dump() |
| return |
| |
| register = self.getRegisterFromName(function_name) |
| print(f" function {register} ({function_name}):") |
| if register is not None: |
| self._functions[register].dump() |
| else: |
| print(f" error: cannot find function.") |
| |
| |
| # Defines a convergence requirement for the simulation: |
| # A list of lanes impacted by a merge and possibly the associated |
| # continue target. |
| @dataclass |
| class ConvergenceRequirement: |
| mergeTarget: InstructionPointer |
| continueTarget: Optional[InstructionPointer] |
| impactedLanes: set[int] |
| |
| |
| Task = Dict[InstructionPointer, List[Lane]] |
| |
| |
| # Defines a Lane group/Wave in the simulator. |
| class Wave: |
| # The module this wave will execute. |
| _module: Module |
| # The lanes this wave will be composed of. |
| _lanes: List[Lane] |
| # The instructions scheduled for execution. |
| _tasks: Task |
| # The actual requirements to comply with when executing instructions. |
| # E.g: the set of lanes required to merge before executing the merge block. |
| _convergence_requirements: List[ConvergenceRequirement] |
| # The indices of the active lanes for the current executing instruction. |
| _active_lane_indices: set[int] |
| |
| def __init__(self, module, wave_size: int) -> None: |
| assert wave_size > 0 |
| self._module = module |
| self._lanes = [] |
| |
| for i in range(wave_size): |
| self._lanes.append(Lane(self, i)) |
| |
| self._tasks = {} |
| self._convergence_requirements = [] |
| # The indices of the active lanes for the current executing instruction. |
| self._active_lane_indices = set() |
| |
| # Returns True if the given IP can be executed for the given list of lanes. |
| def _is_task_candidate(self, ip: InstructionPointer, lanes: List[Lane]): |
| merged_lanes: set[int] = set() |
| for lane in self._lanes: |
| if not lane.running(): |
| merged_lanes.add(lane.tid()) |
| |
| for requirement in self._convergence_requirements: |
| # This task is not executing a merge or continue target. |
| # Adding all lanes at those points into the ignore list. |
| if requirement.mergeTarget != ip and requirement.continueTarget != ip: |
| for tid in requirement.impactedLanes: |
| if self._lanes[tid].ip() == requirement.mergeTarget: |
| merged_lanes.add(tid) |
| if self._lanes[tid].ip() == requirement.continueTarget: |
| merged_lanes.add(tid) |
| continue |
| |
| # This task is executing the current requirement continue/merge |
| # target. |
| for tid in requirement.impactedLanes: |
| lane = self._lanes[tid] |
| if not lane.running(): |
| continue |
| |
| if lane.tid() in merged_lanes: |
| continue |
| |
| if ip == requirement.mergeTarget: |
| if lane.ip() != requirement.mergeTarget: |
| return False |
| else: |
| if ( |
| lane.ip() != requirement.mergeTarget |
| and lane.ip() != requirement.continueTarget |
| ): |
| return False |
| return True |
| |
| # Returns the next task we can schedule. This must always return a task. |
| # Calling this when all lanes are dead is invalid. |
| def _get_next_runnable_task(self) -> Tuple[InstructionPointer, List[Lane]]: |
| candidate = None |
| for ip, lanes in self._tasks.items(): |
| if len(lanes) == 0: |
| continue |
| if self._is_task_candidate(ip, lanes): |
| candidate = ip |
| break |
| |
| if candidate: |
| lanes = self._tasks[candidate] |
| del self._tasks[ip] |
| return (candidate, lanes) |
| raise RuntimeError("No task to execute. Deadlock?") |
| |
| # Handle an encountered merge instruction for the given lane. |
| def handle_convergence_header(self, lane: Lane, instruction: MergeInstruction): |
| mergeTarget = self._module.get_bb_entry(instruction.merge_location()) |
| for requirement in self._convergence_requirements: |
| if requirement.mergeTarget == mergeTarget: |
| requirement.impactedLanes.add(lane.tid()) |
| return |
| |
| continueTarget = None |
| if instruction.continue_location(): |
| continueTarget = self._module.get_bb_entry(instruction.continue_location()) |
| requirement = ConvergenceRequirement( |
| mergeTarget, continueTarget, set([lane.tid()]) |
| ) |
| self._convergence_requirements.append(requirement) |
| |
| # Returns true if some instructions are scheduled for execution. |
| def _has_tasks(self) -> bool: |
| return len(self._tasks) > 0 |
| |
| # Returns the index of the first active lane right now. |
| def get_first_active_lane_index(self) -> int: |
| return min(self._active_lane_indices) |
| |
| # Broadcast the given value to all active lane registers. |
| def broadcast_register(self, register: str, value: Any) -> None: |
| for tid in self._active_lane_indices: |
| self._lanes[tid].set_register(register, value) |
| |
| # Returns the entrypoint of the function associated with 'name'. |
| # Calling this function with an invalid name is illegal. |
| def _get_function_entry_from_name(self, name: str) -> InstructionPointer: |
| register = self._module.getRegisterFromName(name) |
| assert register is not None |
| return self._module.get_function_entry(register) |
| |
| # Run the wave on the function 'function_name' until all lanes are dead. |
| # If verbose is True, execution trace is printed. |
| # Returns the value returned by the function for each lane. |
| def run(self, function_name: str, verbose: bool = False) -> List[Any]: |
| for t in self._lanes: |
| self._module.initialize(t) |
| |
| entry_ip = self._get_function_entry_from_name(function_name) |
| assert entry_ip is not None |
| for t in self._lanes: |
| t.do_call(entry_ip, "__shader_output__") |
| |
| self._tasks[self._lanes[0].ip()] = self._lanes |
| while self._has_tasks(): |
| ip, lanes = self._get_next_runnable_task() |
| self._active_lane_indices = set([x.tid() for x in lanes]) |
| if verbose: |
| print( |
| f"Executing with lanes {self._active_lane_indices}: {ip.instruction()}" |
| ) |
| |
| for lane in lanes: |
| self._module.execute_one_instruction(lane, ip) |
| if not lane.running(): |
| continue |
| |
| if lane.ip() in self._tasks: |
| self._tasks[lane.ip()].append(lane) |
| else: |
| self._tasks[lane.ip()] = [lane] |
| |
| if verbose and ip.instruction().has_output_register(): |
| register = ip.instruction().output_register() |
| print( |
| f" {register:3} = {[ x.get_register(register, allow_undef=True) for x in lanes ]}" |
| ) |
| |
| output = [] |
| for lane in self._lanes: |
| output.append(lane.get_register("__shader_output__")) |
| return output |
| |
| def dump_register(self, register: str) -> None: |
| for lane in self._lanes: |
| print( |
| f" Lane {lane.tid():2} | {register:3} = {lane.get_register(register)}" |
| ) |
| |
| |
| parser = argparse.ArgumentParser( |
| description="simulator", formatter_class=argparse.ArgumentDefaultsHelpFormatter |
| ) |
| parser.add_argument( |
| "-i", "--input", help="Text SPIR-V to read from", required=False, default="-" |
| ) |
| parser.add_argument("-f", "--function", help="Function to execute") |
| parser.add_argument("-w", "--wave", help="Wave size", default=32, required=False) |
| parser.add_argument( |
| "-e", |
| "--expects", |
| help="Expected results per lanes, expects a list of values. Ex: '1, 2, 3'.", |
| ) |
| parser.add_argument("-v", "--verbose", help="verbose", action="store_true") |
| args = parser.parse_args() |
| |
| |
| def load_instructions(filename: str): |
| if filename is None: |
| return [] |
| |
| if filename.strip() != "-": |
| try: |
| with open(filename, "r") as f: |
| lines = f.read().split("\n") |
| except Exception: # (FileNotFoundError, PermissionError): |
| return [] |
| else: |
| lines = sys.stdin.readlines() |
| |
| # Remove leading/trailing whitespaces. |
| lines = [x.strip() for x in lines] |
| # Strip comments. |
| lines = [x for x in filter(lambda x: len(x) != 0 and x[0] != ";", lines)] |
| |
| instructions = [] |
| for i in [Instruction(x) for x in lines]: |
| out = parseInstruction(i) |
| if out != None: |
| instructions.append(out) |
| return instructions |
| |
| |
| def main(): |
| if args.expects is None or not RE_EXPECTS.match(args.expects): |
| print("Invalid format for --expects/-e flag.", file=sys.stderr) |
| sys.exit(1) |
| if args.function is None: |
| print("Invalid format for --function/-f flag.", file=sys.stderr) |
| sys.exit(1) |
| try: |
| int(args.wave) |
| except ValueError: |
| print("Invalid format for --wave/-w flag.", file=sys.stderr) |
| sys.exit(1) |
| |
| expected_results = [int(x.strip()) for x in args.expects.split(",")] |
| wave_size = int(args.wave) |
| if len(expected_results) != wave_size: |
| print("Wave size != expected result array size", file=sys.stderr) |
| sys.exit(1) |
| |
| instructions = load_instructions(args.input) |
| if len(instructions) == 0: |
| print("Invalid input. Expected a text SPIR-V module.") |
| sys.exit(1) |
| |
| module = Module(instructions) |
| if args.verbose: |
| module.dump() |
| module.dump(args.function) |
| |
| function_names = module.get_function_names() |
| if args.function not in function_names: |
| print( |
| f"'{args.function}' function not found. Known functions are:", |
| file=sys.stderr, |
| ) |
| for name in function_names: |
| print(f" - {name}", file=sys.stderr) |
| sys.exit(1) |
| |
| wave = Wave(module, wave_size) |
| results = wave.run(args.function, verbose=args.verbose) |
| |
| if expected_results != results: |
| print("Expected != Observed", file=sys.stderr) |
| print(f"{expected_results} != {results}", file=sys.stderr) |
| sys.exit(1) |
| sys.exit(0) |
| |
| |
| main() |
