from typing import Optional, List
# Base class for an instruction. To implement a basic instruction that doesn't
# impact the control-flow, create a new class inheriting from this.
class Instruction:
# Contains the name of the output register, if any.
_result: Optional[str]
# Contains the instruction opcode.
_opcode: str
# Contains all the instruction operands, except result and opcode.
_operands: List[str]
def __init__(self, line: str):
self.line = line
tokens = line.split()
if len(tokens) > 1 and tokens[1] == "=":
self._result = tokens[0]
self._opcode = tokens[2]
self._operands = tokens[3:] if len(tokens) > 2 else []
else:
self._result = None
self._opcode = tokens[0]
self._operands = tokens[1:] if len(tokens) > 1 else []
def __str__(self):
if self._result is None:
return f" {self._opcode} {self._operands}"
return f"{self._result:3} = {self._opcode} {self._operands}"
# Returns the instruction opcode.
def opcode(self) -> str:
return self._opcode
# Returns the instruction operands.
def operands(self) -> List[str]:
return self._operands
# Returns the instruction output register. Calling this function is
# only allowed if has_output_register() is true.
def output_register(self) -> str:
assert self._result is not None
return self._result
# Returns true if this function has an output register. False otherwise.
def has_output_register(self) -> bool:
return self._result is not None
# This function is used to initialize state related to this instruction
# before module execution begins. For example, global Input variables
# can use this to store the lane ID into the register.
def static_execution(self, lane):
pass
# This function is called everytime this instruction is executed by a
# tangle. This function should not be directly overriden, instead see
# _impl and _advance_ip.
def runtime_execution(self, module, lane):
self._impl(module, lane)
self._advance_ip(module, lane)
# This function needs to be overriden if your instruction can be executed.
# It implements the logic of the instruction.
# 'Static' instructions like OpConstant should not override this since
# they are not supposed to be executed at runtime.
def _impl(self, module, lane):
raise RuntimeError(f"Unimplemented instruction {self}")
# By default, IP is incremented to point to the next instruction.
# If the instruction modifies IP (like OpBranch), this must be overridden.
def _advance_ip(self, module, lane):
lane.set_ip(lane.ip() + 1)
# Those are parsed, but never executed.
class OpEntryPoint(Instruction):
pass
class OpFunction(Instruction):
pass
class OpFunctionEnd(Instruction):
pass
class OpLabel(Instruction):
pass
class OpVariable(Instruction):
pass
class OpName(Instruction):
def name(self) -> str:
return self._operands[1][1:-1]
def decoratedRegister(self) -> str:
return self._operands[0]
# The only decoration we use if the BuiltIn one to initialize the values.
class OpDecorate(Instruction):
def static_execution(self, lane):
if self._operands[1] == "LinkageAttributes":
return
assert (
self._operands[1] == "BuiltIn"
and self._operands[2] == "SubgroupLocalInvocationId"
)
lane.set_register(self._operands[0], lane.tid())
# Constants
class OpConstant(Instruction):
def static_execution(self, lane):
lane.set_register(self._result, int(self._operands[1]))
class OpConstantTrue(OpConstant):
def static_execution(self, lane):
lane.set_register(self._result, True)
class OpConstantFalse(OpConstant):
def static_execution(self, lane):
lane.set_register(self._result, False)
class OpConstantComposite(OpConstant):
def static_execution(self, lane):
result = []
for op in self._operands[1:]:
result.append(lane.get_register(op))
lane.set_register(self._result, result)
# Control flow instructions
class OpFunctionCall(Instruction):
def _impl(self, module, lane):
pass
def _advance_ip(self, module, lane):
entry = module.get_function_entry(self._operands[1])
lane.do_call(entry, self._result)
class OpReturn(Instruction):
def _impl(self, module, lane):
pass
def _advance_ip(self, module, lane):
lane.do_return(None)
class OpReturnValue(Instruction):
def _impl(self, module, lane):
pass
def _advance_ip(self, module, lane):
lane.do_return(lane.get_register(self._operands[0]))
class OpBranch(Instruction):
def _impl(self, module, lane):
pass
def _advance_ip(self, module, lane):
lane.set_ip(module.get_bb_entry(self._operands[0]))
pass
class OpBranchConditional(Instruction):
def _impl(self, module, lane):
pass
def _advance_ip(self, module, lane):
condition = lane.get_register(self._operands[0])
if condition:
lane.set_ip(module.get_bb_entry(self._operands[1]))
else:
lane.set_ip(module.get_bb_entry(self._operands[2]))
class OpSwitch(Instruction):
def _impl(self, module, lane):
pass
def _advance_ip(self, module, lane):
value = lane.get_register(self._operands[0])
default_label = self._operands[1]
i = 2
while i < len(self._operands):
imm = int(self._operands[i])
label = self._operands[i + 1]
if value == imm:
lane.set_ip(module.get_bb_entry(label))
return
i += 2
lane.set_ip(module.get_bb_entry(default_label))
class OpUnreachable(Instruction):
def _impl(self, module, lane):
raise RuntimeError("This instruction should never be executed.")
# Convergence instructions
class MergeInstruction(Instruction):
def merge_location(self):
return self._operands[0]
def continue_location(self):
return None if len(self._operands) < 3 else self._operands[1]
def _impl(self, module, lane):
lane.handle_convergence_header(self)
class OpLoopMerge(MergeInstruction):
pass
class OpSelectionMerge(MergeInstruction):
pass
# Other instructions
class OpBitcast(Instruction):
def _impl(self, module, lane):
# TODO: find out the type from the defining instruction.
# This can only work for DXC.
if self._operands[0] == "%int":
lane.set_register(self._result, int(lane.get_register(self._operands[1])))
else:
raise RuntimeError("Unsupported OpBitcast operand")
class OpAccessChain(Instruction):
def _impl(self, module, lane):
# Python dynamic types allows me to simplify. As long as the SPIR-V
# is legal, this should be fine.
# Note: SPIR-V structs are stored as tuples
value = lane.get_register(self._operands[1])
for operand in self._operands[2:]:
value = value[lane.get_register(operand)]
lane.set_register(self._result, value)
class OpCompositeConstruct(Instruction):
def _impl(self, module, lane):
output = []
for op in self._operands[1:]:
output.append(lane.get_register(op))
lane.set_register(self._result, output)
class OpCompositeExtract(Instruction):
def _impl(self, module, lane):
value = lane.get_register(self._operands[1])
output = value
for op in self._operands[2:]:
output = output[int(op)]
lane.set_register(self._result, output)
class OpStore(Instruction):
def _impl(self, module, lane):
lane.set_register(self._operands[0], lane.get_register(self._operands[1]))
class OpLoad(Instruction):
def _impl(self, module, lane):
lane.set_register(self._result, lane.get_register(self._operands[1]))
class OpIAdd(Instruction):
def _impl(self, module, lane):
LHS = lane.get_register(self._operands[1])
RHS = lane.get_register(self._operands[2])
lane.set_register(self._result, LHS + RHS)
class OpISub(Instruction):
def _impl(self, module, lane):
LHS = lane.get_register(self._operands[1])
RHS = lane.get_register(self._operands[2])
lane.set_register(self._result, LHS - RHS)
class OpIMul(Instruction):
def _impl(self, module, lane):
LHS = lane.get_register(self._operands[1])
RHS = lane.get_register(self._operands[2])
lane.set_register(self._result, LHS * RHS)
class OpLogicalNot(Instruction):
def _impl(self, module, lane):
LHS = lane.get_register(self._operands[1])
lane.set_register(self._result, not LHS)
class _LessThan(Instruction):
def _impl(self, module, lane):
LHS = lane.get_register(self._operands[1])
RHS = lane.get_register(self._operands[2])
lane.set_register(self._result, LHS < RHS)
class _GreaterThan(Instruction):
def _impl(self, module, lane):
LHS = lane.get_register(self._operands[1])
RHS = lane.get_register(self._operands[2])
lane.set_register(self._result, LHS > RHS)
class OpSLessThan(_LessThan):
pass
class OpULessThan(_LessThan):
pass
class OpSGreaterThan(_GreaterThan):
pass
class OpUGreaterThan(_GreaterThan):
pass
class OpIEqual(Instruction):
def _impl(self, module, lane):
LHS = lane.get_register(self._operands[1])
RHS = lane.get_register(self._operands[2])
lane.set_register(self._result, LHS == RHS)
class OpINotEqual(Instruction):
def _impl(self, module, lane):
LHS = lane.get_register(self._operands[1])
RHS = lane.get_register(self._operands[2])
lane.set_register(self._result, LHS != RHS)
class OpPhi(Instruction):
def _impl(self, module, lane):
previousBBName = lane.get_previous_bb_name()
i = 1
while i < len(self._operands):
label = self._operands[i + 1]
if label == previousBBName:
lane.set_register(self._result, lane.get_register(self._operands[i]))
return
i += 2
raise RuntimeError("previousBB not in the OpPhi _operands")
class OpSelect(Instruction):
def _impl(self, module, lane):
condition = lane.get_register(self._operands[1])
value = lane.get_register(self._operands[2 if condition else 3])
lane.set_register(self._result, value)
# Wave intrinsics
class OpGroupNonUniformBroadcastFirst(Instruction):
def _impl(self, module, lane):
assert lane.get_register(self._operands[1]) == 3
if lane.is_first_active_lane():
lane.broadcast_register(self._result, lane.get_register(self._operands[2]))
class OpGroupNonUniformElect(Instruction):
def _impl(self, module, lane):
lane.set_register(self._result, lane.is_first_active_lane())
Codebase Browser
llvm