Implement kron runtime

Author: david-pl

src/bloqade/pyqrack/squin/op.py

@@ -6,6 +6,7 @@
 from bloqade.pyqrack.base import PyQrackInterpreter
 
 from .runtime import (
+    KronRuntime,
     MultRuntime,
     ControlRuntime,
     IdentityRuntime,
@@ -19,11 +20,13 @@
 @op.dialect.register(key="pyqrack")
 class PyQrackMethods(interp.MethodTable):
 
-    # @interp.impl(op.stmts.Kron)
-    # def kron(
-    #     self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: op.stmts.Kron
-    # ):
-    #     is_unitary: bool = info.attribute(default=False)
+    @interp.impl(op.stmts.Kron)
+    def kron(
+        self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: op.stmts.Kron
+    ):
+        lhs = frame.get(stmt.lhs)
+        rhs = frame.get(stmt.rhs)
+        return (KronRuntime(lhs, rhs),)
 
     @interp.impl(op.stmts.Mult)
     def mult(

src/bloqade/pyqrack/squin/runtime.py

@@ -56,3 +56,15 @@ class MultRuntime(OperatorRuntimeABC):
     def apply(self, *qubits: PyQrackQubit) -> None:
         self.rhs.apply(*qubits)
         self.lhs.apply(*qubits)
+
+
+@dataclass
+class KronRuntime(OperatorRuntimeABC):
+    lhs: OperatorRuntimeABC
+    rhs: OperatorRuntimeABC
+
+    def apply(self, *qubits: PyQrackQubit) -> None:
+        assert len(qubits) == 2
+        qbit1, qbit2 = qubits
+        self.lhs.apply(qbit1)
+        self.rhs.apply(qbit2)

test/pyqrack/test_squin.py

@@ -118,9 +118,25 @@ def main():
     assert result == [0]
 
 
+def test_kron():
+    @squin.kernel
+    def main():
+        q = squin.qubit.new(2)
+        x = squin.op.x()
+        k = squin.op.kron(x, x)
+        squin.qubit.apply(k, q)
+        return squin.qubit.measure(q)
+
+    target = PyQrack(2)
+    result = target.run(main)
+
+    assert result == [1, 1]
+
+
 # TODO: remove
 # test_qubit()
 # test_x()
 # test_basic_ops("x")
 # test_cx()
 # test_mult()
+# test_kron()