mplang-nightly 0.1.dev158__py3-none-any.whl → 0.1.dev268__py3-none-any.whl

mplang/v2/backends/field_impl.py

@@ -0,0 +1,454 @@
+# Copyright 2025 Ant Group Co., Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Field Backend Implementation.
+
+Implements runtime execution logic for Field dialect primitives,
+including bindings to C++ kernels (libmplang_kernels.so) and
+NumPy fallbacks where appropriate.
+"""
+
+from __future__ import annotations
+
+import ctypes
+import os
+import threading
+
+# print("DEBUG: Importing field_impl.py")
+import jax.numpy as jnp
+import numpy as np
+
+from mplang.v2.backends.tensor_impl import TensorValue, _unwrap, _wrap
+from mplang.v2.dialects import field
+from mplang.v2.edsl.graph import Operation
+from mplang.v2.kernels import py_kernels
+from mplang.v2.runtime.interpreter import Interpreter
+
+# =============================================================================
+# Kernel Loading
+# =============================================================================
+
+# Load Kernel Library
+# In a real package, this path would be resolved robustly
+_KERNEL_LIB_PATH = os.path.join(
+    os.path.dirname(__file__), "..", "kernels", "libmplang_kernels.so"
+)
+_LIB = None
+_LIB_LOCK = threading.Lock()
+
+
+def _get_lib() -> ctypes.CDLL | None:
+    global _LIB
+    with _LIB_LOCK:
+        if _LIB is None:
+            try:
+                _LIB = ctypes.CDLL(_KERNEL_LIB_PATH)
+                # Define signatures
+                _LIB.gf128_mul.argtypes = [
+                    ctypes.POINTER(ctypes.c_uint64),
+                    ctypes.POINTER(ctypes.c_uint64),
+                    ctypes.POINTER(ctypes.c_uint64),
+                ]
+                _LIB.gf128_mul_batch.argtypes = [
+                    ctypes.POINTER(ctypes.c_uint64),
+                    ctypes.POINTER(ctypes.c_uint64),
+                    ctypes.POINTER(ctypes.c_uint64),
+                    ctypes.c_int64,
+                ]
+                _LIB.solve_okvs.argtypes = [
+                    ctypes.POINTER(ctypes.c_uint64),  # keys
+                    ctypes.POINTER(ctypes.c_uint64),  # values
+                    ctypes.POINTER(ctypes.c_uint64),  # output
+                    ctypes.c_uint64,  # n
+                    ctypes.c_uint64,  # m
+                    ctypes.POINTER(ctypes.c_uint64),  # seed
+                ]
+                _LIB.decode_okvs.argtypes = [
+                    ctypes.POINTER(ctypes.c_uint64),  # keys
+                    ctypes.POINTER(ctypes.c_uint64),  # storage
+                    ctypes.POINTER(ctypes.c_uint64),  # output
+                    ctypes.c_uint64,  # n
+                    ctypes.c_uint64,  # m
+                    ctypes.POINTER(ctypes.c_uint64),  # seed
+                ]
+                # Optimized Mega-Binning Versions
+                _LIB.solve_okvs_opt.argtypes = _LIB.solve_okvs.argtypes
+                _LIB.decode_okvs_opt.argtypes = _LIB.decode_okvs.argtypes
+
+                _LIB.aes_128_expand.argtypes = [
+                    ctypes.POINTER(ctypes.c_uint64),  # seeds
+                    ctypes.POINTER(ctypes.c_uint64),  # output
+                    ctypes.c_uint64,  # num_seeds
+                    ctypes.c_uint64,  # length
+                ]
+                _LIB.ldpc_encode.argtypes = [
+                    ctypes.POINTER(ctypes.c_uint64),  # message
+                    ctypes.POINTER(ctypes.c_uint64),  # indices
+                    ctypes.POINTER(ctypes.c_uint64),  # indptr
+                    ctypes.POINTER(ctypes.c_uint64),  # output
+                    ctypes.c_uint64,  # m
+                    ctypes.c_uint64,  # n
+                ]
+            except OSError:
+                print(f"WARNING: Could not load kernels from {_KERNEL_LIB_PATH}")
+    return _LIB
+
+
+# =============================================================================
+# Helper Implementations (C++ Wrappers)
+# =============================================================================
+
+
+def _gf128_mul_impl(a: np.ndarray, b: np.ndarray) -> np.ndarray:
+    # a, b are numpy arrays (uint64) usually (N, 2)
+
+    lib = _get_lib()
+    if lib is None:
+        # Use pure Python fallback
+        return py_kernels.gf128_mul_batch(a, b)
+
+    # Enforce contiguous C-order arrays (important for ctypes)
+    # Use ascontiguousarray to avoid copy if already contiguous
+    a_contig = np.ascontiguousarray(a, dtype=np.uint64)
+    b_contig = np.ascontiguousarray(b, dtype=np.uint64)
+    out = np.zeros_like(a_contig)
+
+    # Calculate number of elements
+    # Assumes last dim is 2.
+    # Total uint64 count / 2
+    n_elements = a_contig.size // 2
+
+    a_ptr = a_contig.ctypes.data_as(ctypes.POINTER(ctypes.c_uint64))
+    b_ptr = b_contig.ctypes.data_as(ctypes.POINTER(ctypes.c_uint64))
+    out_ptr = out.ctypes.data_as(ctypes.POINTER(ctypes.c_uint64))
+
+    lib.gf128_mul_batch(a_ptr, b_ptr, out_ptr, n_elements)
+
+    return out
+
+
+def _okvs_solve_opt_impl(
+    keys: np.ndarray, values: np.ndarray, m: int, seed: np.ndarray
+) -> np.ndarray:
+    lib = _get_lib()
+    if seed.ndim > 1:
+        seed = seed.flatten()
+
+    if lib is None:
+        # Fallback to standard (no python impl for opt)
+        return _okvs_solve_impl(keys, values, m, seed)
+
+    n = keys.shape[0]
+
+    # Heuristic: Mega-Binning is unstable < 200k.
+    if n < 200_000:
+        return _okvs_solve_impl(keys, values, m, seed)
+
+    # Heuristic: Mega-Binning requires higher expansion (epsilon ~ 1.35)
+    # If m/n is too tight, fallback to Naive (which works with 1.25)
+    if m / n < 1.32:
+        return _okvs_solve_impl(keys, values, m, seed)
+
+    keys_c = np.ascontiguousarray(keys, dtype=np.uint64)
+    values_c = np.ascontiguousarray(values, dtype=np.uint64)
+    seed_c = np.ascontiguousarray(seed, dtype=np.uint64)
+    output = np.zeros((m, 2), dtype=np.uint64)
+
+    lib.solve_okvs_opt(
+        keys_c.ctypes.data_as(ctypes.POINTER(ctypes.c_uint64)),
+        values_c.ctypes.data_as(ctypes.POINTER(ctypes.c_uint64)),
+        output.ctypes.data_as(ctypes.POINTER(ctypes.c_uint64)),
+        n,
+        m,
+        seed_c.ctypes.data_as(ctypes.POINTER(ctypes.c_uint64)),
+    )
+    return output
+
+
+def _okvs_decode_opt_impl(
+    keys: np.ndarray, storage: np.ndarray, m: int, seed: np.ndarray
+) -> np.ndarray:
+    lib = _get_lib()
+    if seed.ndim > 1:
+        seed = seed.flatten()
+
+    if lib is None:
+        return _okvs_decode_impl(keys, storage, m, seed)
+
+    n = keys.shape[0]
+
+    # Heuristic: Mega-Binning (1024 Bins) is unstable for small N due to variance.
+    # It requires ~1000 items/bin to be efficient and stable with epsilon=1.3.
+    # Threshold: 200,000 (approx 200 items/bin). Below this, Naive is fast enough (<50ms).
+    if n < 200_000:
+        return _okvs_decode_impl(keys, storage, m, seed)
+
+    if m / n < 1.32:
+        return _okvs_decode_impl(keys, storage, m, seed)
+
+    keys_c = np.ascontiguousarray(keys, dtype=np.uint64)
+    storage_c = np.ascontiguousarray(storage, dtype=np.uint64)
+    seed_c = np.ascontiguousarray(seed, dtype=np.uint64)
+    output = np.zeros((n, 2), dtype=np.uint64)
+
+    lib.decode_okvs_opt(
+        keys_c.ctypes.data_as(ctypes.POINTER(ctypes.c_uint64)),
+        storage_c.ctypes.data_as(ctypes.POINTER(ctypes.c_uint64)),
+        output.ctypes.data_as(ctypes.POINTER(ctypes.c_uint64)),
+        n,
+        m,
+        seed_c.ctypes.data_as(ctypes.POINTER(ctypes.c_uint64)),
+    )
+    return output
+
+
+def _okvs_solve_impl(
+    keys: np.ndarray, values: np.ndarray, m: int, seed: np.ndarray
+) -> np.ndarray:
+    lib = _get_lib()
+    # Ensure seed is flat tuple or array
+    if seed.ndim > 1:
+        seed = seed.flatten()
+    s_tuple = (int(seed[0]), int(seed[1]))
+
+    if lib is None:
+        # Use pure Python fallback
+        keys_flat = keys.flatten() if keys.ndim > 1 else keys
+        return py_kernels.okvs_solve(keys_flat, values, m, seed=s_tuple)
+
+    n = keys.shape[0]
+    keys_c = np.ascontiguousarray(keys, dtype=np.uint64)
+    values_c = np.ascontiguousarray(values, dtype=np.uint64)
+    seed_c = np.ascontiguousarray(seed, dtype=np.uint64)
+    output = np.zeros((m, 2), dtype=np.uint64)
+
+    lib.solve_okvs(
+        keys_c.ctypes.data_as(ctypes.POINTER(ctypes.c_uint64)),
+        values_c.ctypes.data_as(ctypes.POINTER(ctypes.c_uint64)),
+        output.ctypes.data_as(ctypes.POINTER(ctypes.c_uint64)),
+        n,
+        m,
+        seed_c.ctypes.data_as(ctypes.POINTER(ctypes.c_uint64)),
+    )
+    return output
+
+
+def _okvs_decode_impl(
+    keys: np.ndarray, storage: np.ndarray, m: int, seed: np.ndarray
+) -> np.ndarray:
+    lib = _get_lib()
+    # Ensure seed is flat tuple or array
+    if seed.ndim > 1:
+        seed = seed.flatten()
+    s_tuple = (int(seed[0]), int(seed[1]))
+
+    if lib is None:
+        # Use pure Python fallback
+        keys_flat = keys.flatten() if keys.ndim > 1 else keys
+        return py_kernels.okvs_decode(keys_flat, storage, m, seed=s_tuple)
+
+    n = keys.shape[0]
+    keys_c = np.ascontiguousarray(keys, dtype=np.uint64)
+    storage_c = np.ascontiguousarray(storage, dtype=np.uint64)
+    seed_c = np.ascontiguousarray(seed, dtype=np.uint64)
+    output = np.zeros((n, 2), dtype=np.uint64)
+
+    lib.decode_okvs(
+        keys_c.ctypes.data_as(ctypes.POINTER(ctypes.c_uint64)),
+        storage_c.ctypes.data_as(ctypes.POINTER(ctypes.c_uint64)),
+        output.ctypes.data_as(ctypes.POINTER(ctypes.c_uint64)),
+        n,
+        m,
+        seed_c.ctypes.data_as(ctypes.POINTER(ctypes.c_uint64)),
+    )
+    return output
+
+
+def ldpc_encode_impl(
+    message: np.ndarray, h_indices: np.ndarray, h_indptr: np.ndarray, m: int
+) -> np.ndarray:
+    lib = _get_lib()
+    if lib is None:
+        # Use pure Python fallback
+        h_idx_flat = h_indices.flatten() if h_indices.ndim > 1 else h_indices
+        h_ptr_flat = h_indptr.flatten() if h_indptr.ndim > 1 else h_indptr
+        return py_kernels.ldpc_encode(message, h_idx_flat, h_ptr_flat, m)
+
+    # Fast C++ Path
+    msg_c = np.ascontiguousarray(message, dtype=np.uint64)
+    idx_c = np.ascontiguousarray(h_indices, dtype=np.uint64)
+    ptr_c = np.ascontiguousarray(h_indptr, dtype=np.uint64)
+
+    output = np.zeros((m, 2), dtype=np.uint64)
+
+    # n is inferred from message length
+    n = message.shape[0]
+
+    lib.ldpc_encode(
+        msg_c.ctypes.data_as(ctypes.POINTER(ctypes.c_uint64)),
+        idx_c.ctypes.data_as(ctypes.POINTER(ctypes.c_uint64)),
+        ptr_c.ctypes.data_as(ctypes.POINTER(ctypes.c_uint64)),
+        output.ctypes.data_as(ctypes.POINTER(ctypes.c_uint64)),
+        m,
+        n,
+    )
+    return output
+
+
+# =============================================================================
+# Primitive Implementations
+# =============================================================================
+
+
+@field.ldpc_encode_p.def_impl
+def _ldpc_encode_impl_prim(
+    interpreter: Interpreter,
+    op: Operation,
+    message_val: TensorValue,
+    indices_val: TensorValue,
+    indptr_val: TensorValue,
+) -> TensorValue:
+    m = op.attrs["m"]
+    message = _unwrap(message_val)
+    indices = _unwrap(indices_val)
+    indptr = _unwrap(indptr_val)
+    res = ldpc_encode_impl(message, indices, indptr, m)
+    return _wrap(res)
+
+
+@field.aes_expand_p.def_impl
+def _aes_expand_impl_prim(
+    interpreter: Interpreter, op: Operation, seeds_val: TensorValue
+) -> TensorValue:
+    length = op.attrs["length"]
+    seeds = _unwrap(seeds_val)
+
+    # JAX PRG Fallback crashed. Switching to NumPy PRG.
+
+    # Check if bytes
+    if seeds.dtype == np.uint8 and seeds.shape[-1] == 16:
+        seeds = seeds.view(np.uint64)
+
+    if seeds.shape[-1] != 2:
+        seeds = seeds.reshape(-1, 2)
+
+    num_seeds = seeds.shape[0]
+    out_shape = (num_seeds, length, 2)
+    output = np.zeros(out_shape, dtype=np.uint64)
+
+    lib = _get_lib()
+    if lib is not None:
+        # Fast C++ Path
+        seeds_c = np.ascontiguousarray(seeds, dtype=np.uint64)
+
+        lib.aes_128_expand(
+            seeds_c.ctypes.data_as(ctypes.POINTER(ctypes.c_uint64)),
+            output.ctypes.data_as(ctypes.POINTER(ctypes.c_uint64)),
+            num_seeds,
+            length,
+        )
+    else:
+        # Slow Python Path (Fallback)
+        # Iterate and generate
+        for i in range(num_seeds):
+            # Seed from pair
+            s0 = int(seeds[i, 0])
+            s1 = int(seeds[i, 1])
+            seed_val = [s0, s1]
+
+            rng = np.random.default_rng(seed_val)
+            vals = rng.integers(
+                0, 0xFFFFFFFFFFFFFFFF, size=(length, 2), dtype=np.uint64
+            )
+            output[i] = vals
+
+    # Return as JAX array to keep downstream happy
+    res_jax = jnp.array(output)
+
+    return _wrap(res_jax)
+
+
+@field.mul_p.def_impl
+def _mul_impl(
+    interpreter: Interpreter, op: Operation, a_val: TensorValue, b_val: TensorValue
+) -> TensorValue:
+    a = a_val.unwrap()
+    b = b_val.unwrap()
+    res = _gf128_mul_impl(a, b)
+    return TensorValue(res)
+
+
+@field.solve_okvs_p.def_impl
+def _solve_okvs_impl(
+    interpreter: Interpreter,
+    op: Operation,
+    keys_val: TensorValue,
+    values_val: TensorValue,
+    seed_val: TensorValue,
+) -> TensorValue:
+    m = op.attrs["m"]
+    keys = _unwrap(keys_val)
+    values = _unwrap(values_val)
+    seed = _unwrap(seed_val)
+    res = _okvs_solve_impl(keys, values, m, seed)
+    return _wrap(res)
+
+
+@field.decode_okvs_p.def_impl
+def _decode_okvs_impl(
+    interpreter: Interpreter,
+    op: Operation,
+    keys_val: TensorValue,
+    store_val: TensorValue,
+    seed_val: TensorValue,
+) -> TensorValue:
+    keys = _unwrap(keys_val)
+    storage = _unwrap(store_val)
+    seed = _unwrap(seed_val)
+    m = storage.shape[0]
+    res = _okvs_decode_impl(keys, storage, m, seed)
+    return _wrap(res)
+    return _wrap(res)
+
+
+@field.solve_okvs_opt_p.def_impl
+def _solve_okvs_opt_impl_prim(
+    interpreter: Interpreter,
+    op: Operation,
+    keys_val: TensorValue,
+    values_val: TensorValue,
+    seed_val: TensorValue,
+) -> TensorValue:
+    m = op.attrs["m"]
+    keys = _unwrap(keys_val)
+    values = _unwrap(values_val)
+    seed = _unwrap(seed_val)
+    res = _okvs_solve_opt_impl(keys, values, m, seed)
+    return _wrap(res)
+
+
+@field.decode_okvs_opt_p.def_impl
+def _decode_okvs_opt_impl_prim(
+    interpreter: Interpreter,
+    op: Operation,
+    keys_val: TensorValue,
+    store_val: TensorValue,
+    seed_val: TensorValue,
+) -> TensorValue:
+    keys = _unwrap(keys_val)
+    storage = _unwrap(store_val)
+    seed = _unwrap(seed_val)
+    m = storage.shape[0]
+    res = _okvs_decode_opt_impl(keys, storage, m, seed)
+    return _wrap(res)

mplang/v2/backends/func_impl.py

@@ -0,0 +1,107 @@
+# Copyright 2025 Ant Group Co., Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Generic kernel implementations for the `func` dialect.
+
+Design: Function as Value
+- func.func impl: Returns a FunctionValue wrapping the traced graph.
+- func.call impl: Executes the function graph with provided arguments.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, ClassVar
+
+from mplang.v2.dialects.func import call_p, func_def_p
+from mplang.v2.edsl import serde
+from mplang.v2.edsl.graph import Graph, Operation
+from mplang.v2.runtime.value import Value
+
+if TYPE_CHECKING:
+    from typing import Self
+
+
+@serde.register_class
+class FunctionValue(Value):
+    """Runtime representation of a traced function.
+
+    This is a first-class runtime Value that wraps a Graph.
+    Produced by func.func, consumed by func.call.
+
+    Semantic rationale:
+        In the interpreter's computation model, Values are data that flow
+        between Operations. A function (Graph) is just another kind of data
+        that can be passed around, stored, and invoked - hence it's a Value.
+    """
+
+    _serde_kind: ClassVar[str] = "func.FunctionValue"
+
+    def __init__(self, graph: Graph, name: str = "anonymous") -> None:
+        self._graph = graph
+        self._name = name
+
+    @property
+    def graph(self) -> Graph:
+        return self._graph
+
+    @property
+    def name(self) -> str:
+        return self._name
+
+    def __repr__(self) -> str:
+        return f"FunctionValue({self._name!r}, ops={len(self._graph.operations)})"
+
+    def to_json(self) -> dict[str, Any]:
+        """Serialize function to JSON."""
+        return {
+            "graph": serde.to_json(self._graph),
+            "name": self._name,
+        }
+
+    @classmethod
+    def from_json(cls, data: dict[str, Any]) -> Self:
+        """Deserialize function from JSON."""
+        graph = serde.from_json(data["graph"])
+        return cls(graph=graph, name=data.get("name", "anonymous"))
+
+
+@func_def_p.def_impl
+def _func_def_impl(interpreter: Any, op: Operation, *args: Any) -> FunctionValue:
+    """Implementation of func.func: return a FunctionValue wrapping the body graph."""
+    if not op.regions:
+        raise ValueError("func.func operation missing body region")
+
+    name = op.attrs.get("sym_name", "anonymous")
+    return FunctionValue(graph=op.regions[0], name=name)
+
+
+@call_p.def_impl
+def _call_impl(
+    interpreter: Any, op: Operation, fn_obj: FunctionValue, *args: Any
+) -> Any:
+    """Implementation of func.call: execute the function graph.
+
+    Args:
+        interpreter: The interpreter instance.
+        op: The func.call operation.
+        fn_obj: The FunctionValue returned by func.func.
+        *args: Arguments to pass to the function.
+    """
+    if not isinstance(fn_obj, FunctionValue):
+        raise TypeError(f"func.call expects FunctionValue, got {type(fn_obj)}")
+
+    call_args = list(args)
+    result = interpreter.evaluate_graph(fn_obj.graph, call_args)
+    # Return single value or list based on graph outputs
+    return result[0] if len(fn_obj.graph.outputs) == 1 else result

mplang/v2/backends/phe_impl.py

@@ -0,0 +1,148 @@
+# Copyright 2025 Ant Group Co., Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""PHE Runtime Implementation using LightPHE."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, cast
+
+from lightphe import LightPHE
+from lightphe.models.Ciphertext import Ciphertext
+
+from mplang.v2.dialects import phe
+from mplang.v2.edsl.graph import Operation
+from mplang.v2.runtime.interpreter import Interpreter
+
+
+class PHEContext:
+    """Wraps LightPHE context."""
+
+    def __init__(self, algorithm_name: str = "Paillier", key_size: int = 2048):
+        # Normalize algorithm name (LightPHE expects capitalized names)
+        normalized_name = algorithm_name.capitalize()
+        self.cs = LightPHE(algorithm_name=normalized_name, key_size=key_size)
+
+    def encrypt(self, value: int) -> Ciphertext:
+        return self.cs.encrypt(value)
+
+    def decrypt(self, ct: Ciphertext) -> int:
+        return cast(int, self.cs.decrypt(ct))
+
+
+@dataclass
+class PHEEncoder:
+    """Simple fixed-point encoder."""
+
+    scale: float
+
+
+@dataclass
+class WrappedCiphertext:
+    ct: Ciphertext
+    ctx: PHEContext
+
+    def __add__(self, other: Any) -> WrappedCiphertext:
+        if isinstance(other, WrappedCiphertext):
+            # ct + ct
+            new_ct = self.ct + other.ct
+            return WrappedCiphertext(new_ct, self.ctx)
+        elif isinstance(other, int):
+            # ct + int -> ct + encrypt(int)
+            ct_other = self.ctx.encrypt(other)
+            new_ct = self.ct + ct_other
+            return WrappedCiphertext(new_ct, self.ctx)
+        return NotImplemented
+
+    def __mul__(self, other: Any) -> WrappedCiphertext:
+        if isinstance(other, int):
+            # ct * int
+            new_ct = self.ct * other
+            return WrappedCiphertext(new_ct, self.ctx)
+        return NotImplemented
+
+
+@phe.keygen_p.def_impl
+def keygen_impl(
+    interpreter: Interpreter, op: Operation, *args: Any
+) -> tuple[PHEContext, PHEContext]:
+    key_size = op.attrs.get("key_size", 2048)
+    scheme = op.attrs.get("scheme", "Paillier")
+
+    ctx = PHEContext(algorithm_name=scheme, key_size=key_size)
+
+    return ctx, ctx
+
+
+@phe.create_encoder_p.def_impl
+def create_encoder_impl(
+    interpreter: Interpreter, op: Operation, *args: Any
+) -> PHEEncoder:
+    fxp_bits = op.attrs.get("fxp_bits", 16)
+    scale = 2.0**fxp_bits
+    return PHEEncoder(scale=scale)
+
+
+@phe.encode_p.def_impl
+def encode_impl(
+    interpreter: Interpreter, op: Operation, value: float, encoder: PHEEncoder
+) -> int:
+    return int(value * encoder.scale)
+
+
+@phe.decode_p.def_impl
+def decode_impl(
+    interpreter: Interpreter, op: Operation, value: int, encoder: PHEEncoder
+) -> float:
+    return float(value) / encoder.scale
+
+
+@phe.encrypt_p.def_impl
+def encrypt_impl(
+    interpreter: Interpreter, op: Operation, value: int, pk: PHEContext
+) -> WrappedCiphertext:
+    ct = pk.encrypt(value)
+    return WrappedCiphertext(ct, pk)
+
+
+@phe.decrypt_p.def_impl
+def decrypt_impl(
+    interpreter: Interpreter, op: Operation, wct: WrappedCiphertext, sk: PHEContext
+) -> int:
+    return sk.decrypt(wct.ct)
+
+
+@phe.add_cc_p.def_impl
+def add_cc_impl(
+    interpreter: Interpreter,
+    op: Operation,
+    lhs: WrappedCiphertext,
+    rhs: WrappedCiphertext,
+) -> WrappedCiphertext:
+    return lhs + rhs
+
+
+@phe.add_cp_p.def_impl
+def add_cp_impl(
+    interpreter: Interpreter, op: Operation, lhs: WrappedCiphertext, rhs: int
+) -> WrappedCiphertext:
+    return lhs + rhs
+
+
+@phe.mul_cp_p.def_impl
+def mul_cp_impl(
+    interpreter: Interpreter, op: Operation, lhs: WrappedCiphertext, rhs: int
+) -> WrappedCiphertext:
+    return lhs * rhs