#!/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()
Codebase Browser
llvm