From bdd201088c1b710a250de0f0db59943be513b387 Mon Sep 17 00:00:00 2001
From: Eric Schweitz <eschweitz@nvidia.com>
Date: Tue, 25 Apr 2023 14:25:33 -0700
Subject: [PATCH] Add new ops to quake. (#109)

* Add new ops to quake.

This adds 3 new ops for #98. This is a preliminary PR and just adds the
ops and a simple test.

* Fix typos.

* Remove future bits from test.
---
 .../cudaq/Optimizer/Dialect/Quake/QuakeOps.td | 62 +++++++++++++++++++
 test/Quake/quake-ops.qke                      | 15 +++++
 2 files changed, 77 insertions(+)

diff --git a/include/cudaq/Optimizer/Dialect/Quake/QuakeOps.td b/include/cudaq/Optimizer/Dialect/Quake/QuakeOps.td
index a56b31d77a..842d8b8eca 100644
--- a/include/cudaq/Optimizer/Dialect/Quake/QuakeOps.td
+++ b/include/cudaq/Optimizer/Dialect/Quake/QuakeOps.td
@@ -353,6 +353,68 @@ def quake_ApplyOp : QuakeOp<"apply",
   }];
 }
 
+//===----------------------------------------------------------------------===//
+// Memory and register conversion instructions: These operations are useful for
+// intermediate conversions between memory-SSA and value-SSA semantics and vice
+// versa of the IR. They mainly exist during the conversion process.
+//===----------------------------------------------------------------------===//
+
+def quake_UnwrapOp : QuakeOp<"unwrap"> {
+  let summary = "Unwrap a reference to a wire and return the wire value.";
+  let description = [{
+    A quantum reference is an SSA-value that is associated with a volatile
+    quantum wire. The unwrap operation allows conversion from the reference
+    value semantics (memory SSA) to the volatile quantum wire value semantics
+    when/as desired. The binding of a reference value corresponds to a
+    particular data flow of volatile quantum wire values.
+
+    Unwrap and wrap operations should (typically) form pairs as in the following
+    example.
+
+    ```mlir
+      %0 = ... : !quake.qref
+      %1 = quake.unwrap %0 : (!quake.qref) -> !quake.wire
+      %2 = quake.rx (%dbl) %1 : (f64, !quake.wire) -> !quake.wire
+      quake.wrap %2 to %0 : !quake.wire, !quake.qref
+    ```
+  }];
+
+  let arguments = (ins Arg<QRefType,"",[MemRead]>:$ref_value);
+  let results = (outs WireType);
+  let assemblyFormat = [{
+    $ref_value `:` functional-type(operands, results) attr-dict
+  }];
+}
+
+def quake_WrapOp : QuakeOp<"wrap"> {
+  let summary = "Wrap a wire value with its reference.";
+  let description = [{
+    Each data flow of a volatile wire in a quantum circuit may be associated
+    and identified with an invariant and unique quantum reference value of type
+    `!qref`. The wrap operation provides a means of returning a quantum value,
+    a wire, back to the reference value domain.
+  }];
+
+  let arguments = (ins
+    WireType:$wire_value,
+    Arg<QRefType,"",[MemWrite]>:$ref_value
+  );
+  let assemblyFormat = [{
+    $wire_value `to` $ref_value `:` qualified(type(operands)) attr-dict
+  }];
+}
+
+def quake_NullWireOp : QuakeOp<"null_wire", [Pure]> {
+  let summary = "Initial state of a wire.";
+  let description = [{
+    |0> - the initial state of a wire when first constructed. A wire is assumed
+    to be defined in the |0> state as its initial state.
+  }];
+
+  let results = (outs WireType);
+  let assemblyFormat = "attr-dict";
+}
+
 //===----------------------------------------------------------------------===//
 // Base quantum instructions
 //===----------------------------------------------------------------------===//
diff --git a/test/Quake/quake-ops.qke b/test/Quake/quake-ops.qke
index 28d8e363af..13612c75d8 100644
--- a/test/Quake/quake-ops.qke
+++ b/test/Quake/quake-ops.qke
@@ -360,3 +360,18 @@ func.func @wire_type() -> !quake.wire {
   %0 = cc.undef !quake.wire
   return %0 : !quake.wire
 }
+
+// CHECK-LABEL:   func.func @unwrap_wrap(
+// CHECK-SAME:      %[[VAL_0:.*]]: !quake.qref) {
+// CHECK:           %[[VAL_1:.*]] = quake.unwrap %[[VAL_0]] : (!quake.qref) -> !quake.wire
+// CHECK:           %[[VAL_2:.*]] = quake.null_wire
+// CHECK:           quake.wrap %[[VAL_2]] to %[[VAL_0]] : !quake.wire, !quake.qref
+// CHECK:           return
+// CHECK:         }
+
+func.func @unwrap_wrap(%0 : !quake.qref) {
+  %1 = quake.unwrap %0 : (!quake.qref) -> !quake.wire
+  %2 = quake.null_wire
+  quake.wrap %2 to %0 : !quake.wire, !quake.qref
+  return
+}