compiled-knowledge 4.0.0a20__cp313-cp313-macosx_11_0_arm64.whl

ck/circuit_compiler/interpret_compiler.py

@@ -0,0 +1,223 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Sequence, Optional, Dict, List, Tuple, Callable
+
+import numpy as np
+import ctypes as ct
+
+from ..circuit import Circuit, CircuitNode, VarNode, OpNode, ADD, MUL
+from ..program.raw_program import RawProgram, RawProgramFunction
+from ..utils.iter_extras import multiply, first
+from ..utils.np_extras import NDArrayNumeric, DTypeNumeric
+from .support.circuit_analyser import CircuitAnalysis, analyze_circuit
+from .support.input_vars import InputVars, InferVars, infer_input_vars
+
+# index to a value array
+_VARS = 0
+_CONSTS = 1
+_TMPS = 2
+_RESULT = 3
+
+
+def compile_circuit(
+        *result: CircuitNode,
+        input_vars: InputVars = InferVars.ALL,
+        circuit: Optional[Circuit] = None,
+        dtype: DTypeNumeric = np.double,
+) -> InterpreterRawProgram:
+    """
+    Make a RawProgram that interprets the given circuit.
+
+    Args:
+        *result: result nodes nominating the results of the returned program.
+        input_vars: How to determine the input variables.
+        circuit: optionally explicitly specify the Circuit.
+        dtype: the numpy DType to use for the raw program.
+
+    Returns:
+        a raw program.
+
+    Raises:
+        ValueError: if the circuit is unknown, but it is needed.
+        ValueError: if not all nodes are from the same circuit.
+    """
+    in_vars: Sequence[VarNode] = infer_input_vars(circuit, result, input_vars)
+    analysis: CircuitAnalysis = analyze_circuit(in_vars, result)
+    instructions: List[_Instruction]
+    np_consts: NDArrayNumeric
+    instructions, np_consts = _make_instructions(analysis, dtype)
+
+    return InterpreterRawProgram(
+        in_vars=in_vars,
+        result=result,
+        op_nodes=analysis.op_nodes,
+        dtype=dtype,
+        instructions=instructions,
+        np_consts=np_consts,
+    )
+
+
+class InterpreterRawProgram(RawProgram):
+    def __init__(
+            self,
+            in_vars: Sequence[VarNode],
+            result: Sequence[CircuitNode],
+            op_nodes: Sequence[OpNode],
+            dtype: DTypeNumeric,
+            instructions: List[_Instruction],
+            np_consts: NDArrayNumeric,
+    ):
+        self.instructions = instructions
+        self.np_consts = np_consts
+
+        function = _make_function(
+            instructions=instructions,
+            np_consts=np_consts,
+        )
+
+        super().__init__(
+            function=function,
+            dtype=dtype,
+            number_of_vars=len(in_vars),
+            number_of_tmps=len(op_nodes),
+            number_of_results=len(result),
+            var_indices=tuple(var.idx for var in in_vars),
+        )
+
+    def __getstate__(self):
+        """
+        Support for pickle.
+        """
+        return {
+            'dtype': self.dtype,
+            'number_of_vars': self.number_of_vars,
+            'number_of_tmps': self.number_of_tmps,
+            'number_of_results': self.number_of_results,
+            'var_indices': self.var_indices,
+            #
+            'instructions': self.instructions,
+            'np_consts': self.np_consts,
+        }
+
+    def __setstate__(self, state):
+        """
+        Support for pickle.
+        """
+        self.dtype = state['dtype']
+        self.number_of_vars = state['number_of_vars']
+        self.number_of_tmps = state['number_of_tmps']
+        self.number_of_results = state['number_of_results']
+        self.var_indices = state['var_indices']
+        #
+        self.instructions = state['instructions']
+        self.np_consts = state['np_consts']
+
+        self.function = _make_function(
+            instructions=self.instructions,
+            np_consts=self.np_consts,
+        )
+
+
+def _make_instructions(
+        analysis: CircuitAnalysis,
+        dtype: DTypeNumeric,
+) -> Tuple[Sequence[_Instruction], NDArrayNumeric]:
+
+    # Store const values in a numpy array
+    node_to_const_idx: Dict[int, int] = {
+        id(node): i
+        for i, node in enumerate(analysis.const_nodes)
+    }
+    np_consts: NDArrayNumeric = np.zeros(len(node_to_const_idx), dtype=dtype)
+    for i, node in enumerate(analysis.const_nodes):
+        np_consts[i] = node.value
+
+    # Where to get input values for each possible node.
+    node_to_element: Dict[int, _ElementID] = {}
+    # const nodes
+    for node_id, const_idx in node_to_const_idx.items():
+        node_to_element[node_id] = _ElementID(_CONSTS, const_idx)
+    # var nodes
+    for i, var_node in enumerate(analysis.var_nodes):
+        if var_node.is_const():
+            node_to_element[id(var_node)] = node_to_element[id(var_node.const)]
+        else:
+            node_to_element[id(var_node)] = _ElementID(_VARS, i)
+    # op nodes
+    for node_id, tmp_index in analysis.op_to_tmp.items():
+        node_to_element[node_id] = _ElementID(_TMPS, tmp_index)
+    for node_id, tmp_index in analysis.op_to_result.items():
+        node_to_element[node_id] = _ElementID(_RESULT, tmp_index)
+
+    # Build instructions
+    instructions: List[_Instruction] = []
+
+    op_node: OpNode
+    for op_node in analysis.op_nodes:
+        dest: _ElementID = node_to_element[id(op_node)]
+        args: List[_ElementID] = [
+            node_to_element[id(arg)]
+            for arg in op_node.args
+        ]
+        if op_node.symbol == MUL:
+            operation = multiply
+        elif op_node.symbol == ADD:
+            operation = sum
+        else:
+            assert False, 'symbol not understood'
+
+        instructions.append(_Instruction(operation, args, dest))
+
+    # Add any copy operations, i.e., result nodes that are not op nodes
+    for i, node in enumerate(analysis.result_nodes):
+        if not isinstance(node, OpNode):
+            source: _ElementID = node_to_element[id(node)]
+            instructions.append(_Instruction(first, [source], _ElementID(_RESULT, i)))
+
+    return instructions, np_consts
+
+
+def _make_function(
+        instructions: List[_Instruction],
+        np_consts: NDArrayNumeric,
+) -> RawProgramFunction:
+    """
+    Make a RawProgram function that executes the given instructions.
+    """
+
+    # RawProgramFunction = Callable[[ct.POINTER, ct.POINTER, ct.POINTER], None]
+    def raw_program_function(vars_in: ct.POINTER, tmps: ct.POINTER, result_out: ct.POINTER) -> None:
+        nonlocal np_consts
+        nonlocal instructions
+
+        arrays: List[ct.POINTER] = [None, None, None, None]
+        arrays[_VARS] = vars_in
+        arrays[_TMPS] = tmps
+        arrays[_RESULT] = result_out
+        arrays[_CONSTS] = np_consts
+
+        def get_value(_element: _ElementID):
+            return arrays[_element.array][_element.index]
+
+        instruction: _Instruction
+        element: _ElementID
+        for instruction in instructions:
+            value = instruction.operation(get_value(element) for element in instruction.args)
+            dest: _ElementID = instruction.dest
+            arrays[dest.array][dest.index] = value
+
+    return raw_program_function
+
+
+@dataclass
+class _ElementID:
+    array: int  # VARS, TMPS, CONSTS, RESULT
+    index: int  # index into the array
+
+
+@dataclass
+class _Instruction:
+    operation: Callable
+    args: Sequence[_ElementID]
+    dest: _ElementID

ck/circuit_compiler/llvm_compiler.py

@@ -0,0 +1,388 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from enum import Enum
+from typing import Sequence, Optional, Tuple, Dict, Protocol
+
+import llvmlite.binding as llvm
+import llvmlite.ir as ir
+
+from .support.circuit_analyser import CircuitAnalysis, analyze_circuit
+from .support.input_vars import InputVars, InferVars, infer_input_vars
+from .support.llvm_ir_function import IRFunction, DataType, TypeInfo, compile_llvm_program, LLVMRawProgram, IrBOp
+from ..circuit import Circuit, VarNode, CircuitNode, OpNode, MUL, ADD, ConstNode
+from ..program.raw_program import RawProgramFunction
+
+
+class Flavour(Enum):
+    STACK = 0  # No working temporary memory requested - all on stack.
+    TMPS = 1  # Working temporary memory used for op node calculations.
+    FUNCS = 2  # Working temporary memory used for op node calculations, one sub-function per op-node.
+
+
+DEFAULT_TYPE_INFO: TypeInfo = DataType.FLOAT_64.value
+DEFAULT_FLAVOUR: Flavour = Flavour.TMPS
+
+
+def compile_circuit(
+        *result: CircuitNode,
+        input_vars: InputVars = InferVars.ALL,
+        circuit: Optional[Circuit] = None,
+        data_type: DataType | TypeInfo = DEFAULT_TYPE_INFO,
+        flavour: Flavour = DEFAULT_FLAVOUR,
+        keep_llvm_program: bool = True,
+        opt: int = 2,
+) -> LLVMRawProgram:
+    """
+    Compile the given circuit using LLVM.
+
+    This creates an LLVM program where each circuit op node is converted to
+    one or more LLVM binary op machine code instructions. For large circuits
+    this results in a large LLVM program which can be slow to compile.
+
+    This compiler produces a RawProgram that _does_ use client managed working memory.
+
+    Conforms to the CircuitCompiler protocol.
+
+    Args:
+        *result: result nodes nominating the results of the returned program.
+        input_vars: How to determine the input variables.
+        circuit: optionally explicitly specify the Circuit.
+        data_type: What data type to use for arithmetic calculations. Either a DataType member or TypeInfo.
+        keep_llvm_program: if true, the LLVM program will be kept. This is required for picking.
+        flavour: what flavour of LLVM program to construct.
+        opt:The optimization level to use by LLVM MC JIT.
+
+    Returns:
+        a raw program.
+
+    Raises:
+        ValueError: if the circuit is unknown, but it is needed.
+        ValueError: if not all nodes are from the same circuit.
+        ValueError: if the program data type could not be interpreted.
+    """
+    in_vars: Sequence[VarNode] = infer_input_vars(circuit, result, input_vars)
+    var_indices: Sequence[int] = tuple(var.idx for var in in_vars)
+
+    # Get the type info
+    type_info: TypeInfo
+    if isinstance(data_type, DataType):
+        type_info = data_type.value
+    elif isinstance(data_type, TypeInfo):
+        type_info = data_type
+    else:
+        raise ValueError(f'could not interpret program data type: {data_type!r}')
+
+    # Compile the circuit to an LLVM module representing a RawProgramFunction
+    llvm_program: str
+    number_of_tmps: int
+    llvm_program, number_of_tmps = _make_llvm_program(in_vars, result, type_info, flavour)
+
+    # Compile the LLVM program to a native executable
+    engine: llvm.ExecutionEngine
+    function: RawProgramFunction
+    engine, function = compile_llvm_program(llvm_program, dtype=type_info.dtype, opt=opt)
+
+    return LLVMRawProgram(
+        function=function,
+        dtype=type_info.dtype,
+        number_of_vars=len(var_indices),
+        number_of_tmps=number_of_tmps,
+        number_of_results=len(result),
+        var_indices=var_indices,
+        llvm_program=llvm_program if keep_llvm_program else None,
+        engine=engine,
+        opt=opt,
+    )
+
+
+def _make_llvm_program(
+        in_vars: Sequence[VarNode],
+        result: Sequence[CircuitNode],
+        type_info: TypeInfo,
+        flavour: Flavour,
+) -> Tuple[str, int]:
+    """
+    Returns:
+        (llvm_program, number_of_tmps)
+    """
+    llvm_function = IRFunction(type_info)
+
+    builder = llvm_function.builder
+    type_info = llvm_function.type_info
+    function = llvm_function.function
+
+    analysis: CircuitAnalysis = analyze_circuit(in_vars, result)
+
+    function_builder: _FunctionBuilder
+    if flavour == Flavour.STACK:
+        function_builder = _FunctionBuilderStack(
+            builder=builder,
+            analysis=analysis,
+            llvm_type=type_info.llvm_type,
+            llvm_idx_type=ir.IntType(32),
+            in_args=function.args[0],
+            out_args=function.args[2],
+            ir_cache={},
+        )
+    elif flavour == Flavour.TMPS:
+        function_builder = _FunctionBuilderTmps(
+            builder=builder,
+            analysis=analysis,
+            llvm_type=type_info.llvm_type,
+            llvm_idx_type=ir.IntType(32),
+            in_args=function.args[0],
+            tmp_args=function.args[1],
+            out_args=function.args[2],
+        )
+    elif flavour == Flavour.FUNCS:
+        function_builder = _FunctionBuilderFuncs(
+            builder=builder,
+            analysis=analysis,
+            llvm_type=type_info.llvm_type,
+            llvm_idx_type=ir.IntType(32),
+            in_args=function.args[0],
+            tmp_args=function.args[1],
+            out_args=function.args[2],
+        )
+    else:
+        raise ValueError(f'unknown LLVM program flavour: {flavour!r}')
+
+    # Add a calculation for each op node
+    for op_node in analysis.op_nodes:
+        if op_node.symbol == ADD:
+            op: IrBOp = type_info.add
+        elif op_node.symbol == MUL:
+            op: IrBOp = type_info.mul
+        else:
+            raise RuntimeError(f'unknown op node: {op_node.symbol!r}')
+        function_builder.process_op_node(op_node, op)
+
+    # Copy any non-op node values to the results
+    for idx, node in enumerate(result):
+        if not isinstance(node, OpNode):
+            value: ir.Value = function_builder.value(node)
+            function_builder.store_result(value, idx)
+
+    # Return from the function
+    builder.ret_void()
+
+    return llvm_function.llvm_program(), function_builder.number_of_tmps()
+
+
+class _FunctionBuilder(Protocol):
+    def process_op_node(self, op_node: OpNode, op: IrBOp) -> None:
+        ...
+
+    def value(self, node: CircuitNode) -> ir.Value:
+        ...
+
+    def store_result(self, value: ir.Value, idx: int) -> None:
+        ...
+
+    def number_of_tmps(self) -> int:
+        ...
+
+
+@dataclass
+class _FunctionBuilderTmps(_FunctionBuilder):
+    """
+    A function builder that puts op node calculations into the temporary working memory.
+    """
+    builder: ir.IRBuilder
+    analysis: CircuitAnalysis
+    llvm_type: ir.Type
+    llvm_idx_type: ir.Type
+    in_args: ir.Value
+    tmp_args: ir.Value
+    out_args: ir.Value
+
+    def number_of_tmps(self) -> int:
+        return len(self.analysis.op_to_tmp)
+
+    def process_op_node(self, op_node: OpNode, op: IrBOp) -> None:
+        value: ir.Value = self.value(op_node.args[0])
+        for arg in op_node.args[1:]:
+            next_value: ir.Value = self.value(arg)
+            value = op(self.builder, value, next_value)
+        self.store_calculation(value, op_node)
+
+    def value(self, node: CircuitNode) -> ir.Value:
+        """
+        Return an IR value for the given circuit node.
+        """
+        node_id: int = id(node)
+
+        # If it is a constant...
+        if isinstance(node, ConstNode):
+            return ir.Constant(self.llvm_type, node.value)
+
+        builder = self.builder
+
+        # If it is a var...
+        if isinstance(node, VarNode):
+            if node.is_const():
+                return ir.Constant(self.llvm_type, node.const.value)
+            else:
+                return builder.load(builder.gep(self.in_args, [ir.Constant(self.llvm_idx_type, node.idx)]))
+
+        analysis = self.analysis
+
+        # If it is an op _not_ in the results...
+        idx: Optional[int] = analysis.op_to_tmp.get(node_id)
+        if idx is not None:
+            return builder.load(builder.gep(self.tmp_args, [ir.Constant(self.llvm_idx_type, idx)]))
+
+        # If it is an op in the results...
+        idx: Optional[int] = analysis.op_to_result.get(node_id)
+        if idx is not None:
+            return builder.load(builder.gep(self.out_args, [ir.Constant(self.llvm_idx_type, idx)]))
+
+        assert False, 'not reached'
+
+    def store_calculation(self, value: ir.Value, op_node: OpNode) -> None:
+        """
+        Store the given IR value as a result for the given op node.
+        """
+        builder = self.builder
+        analysis = self.analysis
+        node_id: int = id(op_node)
+
+        # If it is an op _not_ in the results...
+        idx: Optional[int] = analysis.op_to_tmp.get(node_id)
+        if idx is not None:
+            ptr: ir.GEPInstr = builder.gep(self.tmp_args, [ir.Constant(self.llvm_idx_type, idx)])
+            builder.store(value, ptr)
+            return
+
+        # If it is an op in the results...
+        idx: Optional[int] = analysis.op_to_result.get(node_id)
+        if idx is not None:
+            ptr: ir.GEPInstr = builder.gep(self.out_args, [ir.Constant(self.llvm_idx_type, idx)])
+            builder.store(value, ptr)
+            return
+
+        assert False, 'not reached'
+
+    def store_result(self, value: ir.Value, idx: int) -> None:
+        """
+        Store the given IR value in the indexed result slot.
+        """
+        builder = self.builder
+        ptr: ir.GEPInstr = builder.gep(self.out_args, [ir.Constant(self.llvm_idx_type, idx)])
+        builder.store(value, ptr)
+
+
+class _FunctionBuilderFuncs(_FunctionBuilderTmps):
+    """
+    A function builder that puts op node calculations into the temporary working memory,
+    but each op node becomes its own sub-function.
+    """
+
+    def process_op_node(self, op_node: OpNode, op: IrBOp) -> None:
+        builder: ir.IRBuilder = self.builder
+        save_block = builder.block
+
+        sub_function_name: str = f'sub_{id(op_node)}'
+        function_type = builder.function.type.pointee
+        sub_function = ir.Function(builder.module, function_type, name=sub_function_name)
+        sub_function.attributes.add('noinline')  # alwaysinline, noinline
+        bb_entry = sub_function.append_basic_block(sub_function_name + '_entry')
+        self.builder.position_at_end(bb_entry)
+
+        value: ir.Value = self.value(op_node.args[0])
+        for arg in op_node.args[1:]:
+            next_value: ir.Value = self.value(arg)
+            value = op(self.builder, value, next_value)
+        self.store_calculation(value, op_node)
+
+        builder.ret_void()
+
+        # Restore builder to main function
+        builder.position_at_end(save_block)
+        builder.call(sub_function, [self.in_args, self.tmp_args, self.out_args])
+
+
+@dataclass
+class _FunctionBuilderStack(_FunctionBuilder):
+    """
+    A function builder that puts op node calculations onto the stack.
+    """
+    builder: ir.IRBuilder
+    analysis: CircuitAnalysis
+    llvm_type: ir.Type
+    llvm_idx_type: ir.Type
+    in_args: ir.Value
+    out_args: ir.Value
+    ir_cache: Dict[int, ir.Value]
+
+    def number_of_tmps(self) -> int:
+        return 0
+
+    def process_op_node(self, op_node: OpNode, op: IrBOp) -> None:
+        value: ir.Value = self.value(op_node.args[0])
+        for arg in op_node.args[1:]:
+            next_value: ir.Value = self.value(arg)
+            value = op(self.builder, value, next_value)
+        self.store_calculation(value, op_node)
+
+    def value(self, node: CircuitNode) -> ir.Value:
+        """
+        Return an IR value for the given circuit node.
+        """
+        node_id: int = id(node)
+
+        # First check if it is in the IR cache
+        cached: Optional[ir.Value] = self.ir_cache.get(node_id)
+        if cached is not None:
+            return cached
+
+        # If it is a constant...
+        if isinstance(node, ConstNode):
+            value = ir.Constant(self.llvm_type, node.value)
+            self.ir_cache[node_id] = value
+            return value
+
+        builder = self.builder
+
+        # If it is a var...
+        if isinstance(node, VarNode):
+            if node.is_const():
+                value = ir.Constant(self.llvm_type, node.const.value)
+            else:
+                value = builder.load(builder.gep(self.in_args, [ir.Constant(self.llvm_idx_type, node.idx)]))
+            self.ir_cache[node_id] = value
+            return value
+
+        # If it is an op in the results...
+        idx: Optional[int] = self.analysis.op_to_result.get(node_id)
+        if idx is not None:
+            return builder.load(builder.gep(self.out_args, [ir.Constant(self.llvm_idx_type, idx)]))
+
+        assert False, 'not reached'
+
+    def store_calculation(self, value: ir.Value, op_node: OpNode) -> None:
+        """
+        Store the given IR value as a result for the given op node.
+        """
+        node_id: int = id(op_node)
+
+        # If it is an op in the results...
+        idx: Optional[int] = self.analysis.op_to_result.get(node_id)
+        if idx is not None:
+            builder = self.builder
+            ptr: ir.GEPInstr = builder.gep(self.out_args, [ir.Constant(self.llvm_idx_type, idx)])
+            builder.store(value, ptr)
+            return
+
+        # Just put it in the ir_cache.
+        # This effectively forces the LLVM compiler to put it on the stack when registers run out.
+        self.ir_cache[node_id] = value
+
+    def store_result(self, value: ir.Value, idx: int) -> None:
+        """
+        Store the given IR value in the indexed result slot.
+        """
+        builder = self.builder
+        ptr: ir.GEPInstr = builder.gep(self.out_args, [ir.Constant(self.llvm_idx_type, idx)])
+        builder.store(value, ptr)