Scale runtime

Author: david-pl

src/bloqade/pyqrack/squin/op.py

@@ -1,21 +1,18 @@
 from kirin import interp
 
 from bloqade.squin import op
-
-# from bloqade.pyqrack.reg import QubitState, PyQrackQubit
 from bloqade.pyqrack.base import PyQrackInterpreter
 
 from .runtime import (
     KronRuntime,
     MultRuntime,
+    ScaleRuntime,
     ControlRuntime,
     IdentityRuntime,
     OperatorRuntime,
     ProjectorRuntime,
 )
 
-# from kirin.dialects import ilist
-
 
 @op.dialect.register(key="pyqrack")
 class PyQrackMethods(interp.MethodTable):
@@ -44,14 +41,13 @@ def mult(
     #     op: ir.SSAValue = info.argument(OpType)
     #     result: ir.ResultValue = info.result(OpType)
 
-    # @interp.impl(op.stmts.Scale)
-    # def scale(
-    #     self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: op.stmts.Scale
-    # ):
-    #     is_unitary: bool = info.attribute(default=False)
-    #     op: ir.SSAValue = info.argument(OpType)
-    #     factor: ir.SSAValue = info.argument(Complex)
-    #     result: ir.ResultValue = info.result(OpType)
+    @interp.impl(op.stmts.Scale)
+    def scale(
+        self, interp: PyQrackInterpreter, frame: interp.Frame, stmt: op.stmts.Scale
+    ):
+        op = frame.get(stmt.op)
+        factor = frame.get(stmt.factor)
+        return (ScaleRuntime(op, factor),)
 
     @interp.impl(op.stmts.Control)
     def control(

src/bloqade/pyqrack/squin/runtime.py

@@ -16,7 +16,7 @@ class OperatorRuntime(OperatorRuntimeABC):
     method_name: str
 
     def apply(self, *qubits: PyQrackQubit) -> None:
-        getattr(qubits[-1].sim_reg, self.method_name)(qubits[-1].addr)
+        getattr(qubits[0].sim_reg, self.method_name)(qubits[0].addr)
 
 
 @dataclass
@@ -64,7 +64,21 @@ class KronRuntime(OperatorRuntimeABC):
     rhs: OperatorRuntimeABC
 
     def apply(self, *qubits: PyQrackQubit) -> None:
-        assert len(qubits) == 2
-        qbit1, qbit2 = qubits
-        self.lhs.apply(qbit1)
-        self.rhs.apply(qbit2)
+        self.lhs.apply(qubits[0])
+        self.rhs.apply(qubits[1])
+
+
+@dataclass
+class ScaleRuntime(OperatorRuntimeABC):
+    op: OperatorRuntimeABC
+    factor: complex
+
+    def apply(self, *qubits: PyQrackQubit) -> None:
+        target = qubits[0]
+        self.op.apply(target)
+
+        # NOTE: just factor * eye(2)
+        mat = [self.factor, 0, 0, self.factor]
+
+        # TODO: output seems to always be normalized -- no-op?
+        target.sim_reg.mtrx(mat, target.addr)

src/bloqade/squin/op/__init__.py

@@ -11,11 +11,14 @@
 def kron(lhs: types.Op, rhs: types.Op) -> types.Op: ...
 
 
-# FIXME: should we just rewrite the py.binop.mult instead?
 @_wraps(stmts.Mult)
 def mult(lhs: types.Op, rhs: types.Op) -> types.Op: ...
 
 
+@_wraps(stmts.Scale)
+def scale(op: types.Op, factor: complex) -> types.Op: ...
+
+
 @_wraps(stmts.Adjoint)
 def adjoint(op: types.Op) -> types.Op: ...
 

test/pyqrack/test_squin.py

@@ -133,10 +133,28 @@ def main():
     assert result == [1, 1]
 
 
+def test_scale():
+    @squin.kernel
+    def main():
+        q = squin.qubit.new(1)
+        x = squin.op.x()
+
+        # TODO: replace by 2 * x once we have the rewrite
+        s = squin.op.scale(x, 2)
+
+        squin.qubit.apply(s, q)
+        return squin.qubit.measure(q)
+
+    target = PyQrack(1)
+    result = target.run(main)
+    assert result == [1]
+
+
 # TODO: remove
 # test_qubit()
 # test_x()
 # test_basic_ops("x")
 # test_cx()
 # test_mult()
 # test_kron()
+# test_scale()