Adding WASM support

.pre-commit-config.yaml

@@ -7,6 +7,7 @@ repos:
       - id: check-toml
       - id: check-yaml
       - id: check-added-large-files
+        args: ['--maxkb=2000']
       # Python-specific
       - id: check-ast
       - id: check-docstring-first

pyproject.toml

@@ -56,7 +56,7 @@ pythonpath = [
 log_cli = true
 log_cli_level = "INFO"
 filterwarnings = ["ignore:::lark.s*"]
-log_format = "%(asctime)s.%(msecs)03d %(levelname)s %(message)s"
+log_format = "%(asctime)s.%(msecs)03d %(levelname)s %(name)s:%(lineno)s %(message)s"
 log_date_format = "%Y-%m-%d %H:%M:%S"
 
 [tool.setuptools_scm]

pytket/phir/cli.py

@@ -8,21 +8,18 @@
 
 # mypy: disable-error-code="misc"
 
-from argparse import ArgumentParser
+from argparse import ArgumentParser, BooleanOptionalAction
 from importlib.metadata import version
 
 from pecos.engines.hybrid_engine import HybridEngine  # type:ignore [import-not-found]
 
-from phir.model import PHIRModel
 from pytket.qasm.qasm import (
     circuit_from_qasm,
-    circuit_from_qasm_str,
-    circuit_to_qasm_str,
+    circuit_from_qasm_wasm,
 )
 
 from .api import pytket_to_phir
 from .qtm_machine import QtmMachine
-from .rebasing.rebaser import rebase_to_qtm_machine
 
 
 def main() -> None:
@@ -34,6 +31,12 @@ def main() -> None:
     parser.add_argument(
         "qasm_files", nargs="+", default=None, help="One or more QASM files to emulate"
     )
+    parser.add_argument(
+        "-w",
+        "--wasm-file-path",
+        default=None,
+        help="Optional WASM file path for use by the QASM programs",
+    )
     parser.add_argument(
         "-m",
         "--machine",
@@ -48,6 +51,7 @@ def main() -> None:
         default="0",
         help="TKET optimization level, 0 by default",
     )
+    parser.add_argument("--verbose", action=BooleanOptionalAction)
     parser.add_argument(
         "-v",
         "--version",
@@ -57,19 +61,24 @@ def main() -> None:
     args = parser.parse_args()
 
     for file in args.qasm_files:
-        print(f"Processing {file}")  # noqa: T201
-        c = circuit_from_qasm(file)
-        tket_opt_level = int(args.tk)
-        rc = rebase_to_qtm_machine(c, args.machine, tket_opt_level)
-        qasm = circuit_to_qasm_str(rc, header="hqslib1")
-        circ = circuit_from_qasm_str(qasm)
+        print(f"Processing {file}")
+        circuit = None
+        if args.wasm_file_path:
+            print(f"Including WASM from file {args.wasm_file_path}")
+            circuit = circuit_from_qasm_wasm(file, args.wasm_file_path)
+        else:
+            circuit = circuit_from_qasm(file)
 
         match args.machine:
             case "H1-1":
                 machine = QtmMachine.H1_1
             case "H1-2":
                 machine = QtmMachine.H1_2
-        phir = pytket_to_phir(circ, machine)
-        PHIRModel.model_validate_json(phir)
 
-        HybridEngine(qsim="state-vector").run(program=phir, shots=10)
+        phir = pytket_to_phir(circuit, machine, int(args.tket_opt_level))
+        if args.verbose:
+            print("\nPHIR to be simulated:")
+            print(phir)
+
+        print("\nPECOS results:")
+        print(HybridEngine(qsim="state-vector").run(program=phir, shots=10))

pytket/phir/phirgen.py

@@ -27,6 +27,8 @@
 
 logger = logging.getLogger(__name__)
 
+JsonDict = dict[str, Any]
+
 PHIR_HEADER: dict[str, Any] = {
     "format": "PHIR/JSON",
     "version": "0.1.0",
@@ -76,9 +78,7 @@ def arg_to_bit(arg: "UnitID") -> Bit:
     return [arg.reg_name, arg.index[0]]
 
 
-def assign_cop(
-    lhs: list[Var] | list[Bit], rhs: "Sequence[Var | int]"
-) -> dict[str, Any]:
+def assign_cop(lhs: list[Var] | list[Bit], rhs: "Sequence[Var | int]") -> JsonDict:
     """PHIR for classical assign operation."""
     return {
         "cop": "=",
@@ -87,16 +87,16 @@ def assign_cop(
     }
 
 
-def convert_subcmd(op: tk.Op, cmd: tk.Command) -> dict[str, Any]:
-    """Return PHIR dict give op and its arguments."""
+def convert_subcmd(op: tk.Op, cmd: tk.Command) -> JsonDict:
+    """Return PHIR dict given op and its arguments."""
     if op.is_gate():
         try:
             gate = tket_gate_to_phir[op.type]
         except KeyError:
             logging.exception("Gate %s unsupported by PHIR", op.get_name())
             raise
         angles = (op.params, "pi") if op.params else None
-        qop: dict[str, Any]
+        qop: JsonDict
         match gate:
             case "Measure":
                 qop = {
@@ -148,24 +148,27 @@ def convert_subcmd(op: tk.Op, cmd: tk.Command) -> dict[str, Any]:
                 [arg_to_bit(cmd.args[i]) for i in range(len(cmd.args) // 2)],
             )
 
-        case _:
+        case tk.WASMOp():
+            return create_wasm_op(cmd, op)
+
+        case m:
             # TODO(kartik): NYI
             # https://github.com/CQCL/pytket-phir/issues/25
-            raise NotImplementedError
+            raise NotImplementedError(m)
 
 
-def append_cmd(cmd: tk.Command, ops: list[dict[str, Any]]) -> None:
+def append_cmd(cmd: tk.Command, ops: list[JsonDict]) -> None:
     """Convert a pytket command to a PHIR command and append to `ops`.
 
     Args:
         cmd: pytket command obtained from pytket-phir
         ops: the list of ops to append to
     """
-    ops.append({"//": str(cmd)})
+    ops.append({"//": make_comment_text(cmd, cmd.op)})
     if cmd.op.is_gate():
         ops.append(convert_subcmd(cmd.op, cmd))
     else:
-        op: dict[str, Any] | None = None
+        op: JsonDict | None = None
         match cmd.op:
             case tk.BarrierOp():
                 # TODO(kartik): confirm with Ciaran/spec
@@ -188,7 +191,7 @@ def append_cmd(cmd: tk.Command, ops: list[dict[str, Any]]) -> None:
                 }
 
             case tk.RangePredicateOp():  # where the condition is a range
-                cond: dict[str, Any]
+                cond: JsonDict
                 match cmd.op.lower, cmd.op.upper:
                     case l, u if l == u:
                         cond = {
@@ -250,7 +253,7 @@ def append_cmd(cmd: tk.Command, ops: list[dict[str, Any]]) -> None:
                     "args": [arg["name"] for arg in exp.to_dict()["args"]],
                 }
 
-            case tk.ClassicalEvalOp():
+            case tk.ClassicalEvalOp() | tk.WASMOp():
                 op = convert_subcmd(cmd.op, cmd)
 
             case m:
@@ -259,6 +262,57 @@ def append_cmd(cmd: tk.Command, ops: list[dict[str, Any]]) -> None:
             ops.append(op)
 
 
+def create_wasm_op(cmd: tk.Command, wasm_op: tk.WASMOp) -> JsonDict:
+    """Creates a PHIR operation for a WASM command."""
+    args, returns = extract_wasm_args_and_returns(cmd, wasm_op)
+    op = {
+        "cop": "ffcall",
+        "function": wasm_op.func_name,
+        "args": args,
+        "metadata": {
+            "ff_object": f"WASM module uid: {wasm_op.wasm_uid}",
+        },
+    }
+    if cmd.bits:
+        op["returns"] = returns
+
+    return op
+
+
+def extract_wasm_args_and_returns(
+    command: tk.Command, op: tk.WASMOp
+) -> tuple[list[str], list[str]]:
+    """Extract the wasm args and return values as whole register names."""
+    # This slice removes the extra `_w` cregs (wires) that are not part of the
+    # circuit, and the output args which are appended after the input args
+    slice_index = op.num_w + sum(op.output_widths)
+    only_args = command.args[:-slice_index]
+    return (
+        dedupe_bits_to_registers(only_args),
+        dedupe_bits_to_registers(command.bits),
+    )
+
+
+def dedupe_bits_to_registers(bits: "Sequence[UnitID]") -> list[str]:
+    """Dedupes a list of bits to their registers, keeping order intact."""
+    return list(dict.fromkeys([bit.reg_name for bit in bits]))
+
+
+def make_comment_text(command: tk.Command, op: tk.Op) -> str:
+    """Converts a command + op to the PHIR comment spec."""
+    match op:
+        case tk.Conditional():
+            conditional_text = str(command)
+            cleaned = conditional_text[: conditional_text.find("THEN") + 4]
+            return f"{cleaned} {make_comment_text(command, op.op)}"
+
+        case tk.WASMOp():
+            args, returns = extract_wasm_args_and_returns(command, op)
+            return f"WASM function={op.func_name} args={args} returns={returns}"
+        case _:
+            return str(command)
+
+
 def get_decls(qbits: set["Qubit"], cbits: set["tkBit"]) -> list[dict[str, str | int]]:
     """Format the qvar and cvar define PHIR elements."""
     # TODO(kartik): this may not always be accurate
@@ -291,6 +345,7 @@ def get_decls(qbits: set["Qubit"], cbits: set["tkBit"]) -> list[dict[str, str |
             "size": dim,
         }
         for cvar, dim in cvar_dim.items()
+        if cvar != "_w"
     ]
 
     return decls

pytket/phir/sharding/sharder.py

@@ -14,7 +14,7 @@
 
 from .shard import Shard
 
-NOT_IMPLEMENTED_OP_TYPES = [OpType.CircBox, OpType.WASM]
+NOT_IMPLEMENTED_OP_TYPES = [OpType.CircBox]
 
 SHARD_TRIGGER_OP_TYPES = [
     OpType.Measure,
@@ -25,6 +25,7 @@
     OpType.RangePredicate,
     OpType.ExplicitPredicate,
     OpType.CopyBits,
+    OpType.WASM,
 ]
 
 logger = logging.getLogger(__name__)
@@ -185,7 +186,9 @@ def _resolve_shard_dependencies(
 
         for bit_read in bits_read:
             if bit_read in self._bit_written_by:
-                logger.debug("...adding shard dep %s -> RAW")
+                logger.debug(
+                    "...adding shard dep %s -> RAW", self._bit_written_by[bit_read]
+                )
                 depends_upon.add(self._bit_written_by[bit_read])
 
         for bit_written in bits_written:

ruff.toml

@@ -81,6 +81,9 @@ ignore = [
     "PLR2004",  # Magic constants
     "PLR6301",  # Method * could be a function, class method, or static method
   ]
+"pytket/phir/cli.py" = [
+    "T201"
+]
 
 [pydocstyle]
 convention = "google"

tests/data/qasm/bell.qasm

@@ -0,0 +1,10 @@
+OPENQASM 2.0;
+include "hqslib1.inc";
+
+qreg q[2];
+creg m[2];
+
+h q[0];
+CX q[0], q[1];
+
+measure q -> m;

tests/test_api.py

@@ -8,12 +8,14 @@
 
 import json
 import logging
+from pathlib import Path
 
 import pytest
 
 from pytket.circuit import Bit, Circuit
 from pytket.phir.api import pytket_to_phir, qasm_to_phir
 from pytket.phir.qtm_machine import QtmMachine
+from pytket.wasm.wasm import WasmFileHandler
 
 from .test_utils import QasmFile, get_qasm_as_circuit
 
@@ -85,3 +87,23 @@ def test_qasm_to_phir(self) -> None:
         """
 
         assert qasm_to_phir(qasm, QtmMachine.H1_1)
+
+    def test_pytket_with_wasm(self) -> None:
+        this_dir = Path(Path(__file__).resolve()).parent
+        w = WasmFileHandler(f"{this_dir}/data/wasm/testfile.wasm")
+
+        c = Circuit(6, 6)
+        c0 = c.add_c_register("c0", 3)
+        c1 = c.add_c_register("c1", 4)
+        c2 = c.add_c_register("c2", 5)
+
+        c.add_wasm_to_reg("multi", w, [c0, c1], [c2])
+        c.add_wasm_to_reg("add_one", w, [c2], [c2])
+        c.add_wasm_to_reg("no_return", w, [c2], [])
+        c.add_wasm_to_reg("no_parameters", w, [], [c2])
+
+        c.add_wasm_to_reg("add_one", w, [c0], [c0], condition=c1[0])
+
+        phir = pytket_to_phir(c, QtmMachine.H1_1)
+
+        assert phir