Debugging Transversal Moves. (#57)

* fixing bugs + updating demo

* adding constants to visualization

Author: weinbe58

demo/ghz_moves_demo.py

@@ -1,10 +1,12 @@
 from bloqade.native.upstream import SquinToNative
-from kirin import ir
+from kirin import ir, rewrite
 from kirin.dialects import ilist
 
 from bloqade import qubit, squin
 from bloqade.lanes import visualize
 from bloqade.lanes.arch.gemini import logical
+from bloqade.lanes.arch.gemini.impls import generate_arch
+from bloqade.lanes.arch.gemini.logical.simulation import rewrite as sim_rewrite
 from bloqade.lanes.heuristics import fixed
 from bloqade.lanes.upstream import NativeToPlace, PlaceToMove
 
@@ -40,7 +42,7 @@ def ghz_optimal():
     squin.broadcast.cx(qs[:5], qs[5:])
 
 
-def compile_and_visualize(mt: ir.Method, interactive=True):
+def compile(mt: ir.Method, transversal: bool = False):
     # Compile to move dialect
 
     mt = SquinToNative().emit(mt)
@@ -51,10 +53,26 @@ def compile_and_visualize(mt: ir.Method, interactive=True):
         fixed.LogicalMoveScheduler(),
     ).emit(mt)
 
-    arch_spec = logical.get_arch_spec()
+    if transversal:
+        rewrite.Walk(
+            rewrite.Chain(sim_rewrite.RewriteLocations(), sim_rewrite.RewriteMoves())
+        ).rewrite(mt.code)
+    return mt
 
-    visualize.debugger(mt, arch_spec, interactive=interactive, atom_marker="s")
+
+def compile_and_visualize(mt: ir.Method, interactive=True, transversal: bool = False):
+    # Compile to move dialect
+    mt = compile(mt, transversal=transversal)
+    if transversal:
+        arch_spec = generate_arch(4)
+        marker = "o"
+    else:
+        arch_spec = logical.get_arch_spec()
+        marker = "s"
+
+    visualize.debugger(mt, arch_spec, interactive=interactive, atom_marker=marker)
 
 
-compile_and_visualize(log_depth_ghz)
-# compile_and_visualize(ghz_optimal)
+compile_and_visualize(log_depth_ghz, transversal=False)
+compile_and_visualize(ghz_optimal, transversal=False)
+compile_and_visualize(ghz_optimal, transversal=True)

src/bloqade/lanes/arch/gemini/impls.py

@@ -32,7 +32,7 @@ def site_buses(site_addresses: np.ndarray):
 def hypercube_busses(hypercube_dims: int):
     word_buses: list[Bus] = []
     for shift in range(hypercube_dims):
-        m = 1 << shift
+        m = 1 << (hypercube_dims - shift - 1)
 
         srcs = []
         dsts = []
@@ -45,6 +45,7 @@ def hypercube_busses(hypercube_dims: int):
             dsts.append(dst)
 
         word_buses.append(Bus(tuple(srcs), tuple(dsts)))
+
     return tuple(word_buses)
 
 

src/bloqade/lanes/arch/gemini/logical/simulation/rewrite.py

@@ -16,7 +16,7 @@ def physical_word_id(address: AddressType) -> Iterator[AddressType]:
     if address.word_id == 0:
         yield from (replace(address, word_id=word_id) for word_id in range(7))
     elif address.word_id == 1:
-        yield from (replace(address, word_id=word_id) for word_id in range(9, 16, 1))
+        yield from (replace(address, word_id=word_id) for word_id in range(8, 15, 1))
     else:
         yield address
 

src/bloqade/lanes/visualize.py

@@ -1,6 +1,7 @@
 from itertools import chain
 
 from kirin import ir
+from kirin.dialects import py
 from matplotlib import figure, pyplot as plt
 from matplotlib.axes import Axes
 from matplotlib.widgets import Button
@@ -126,11 +127,22 @@ def get_drawer(mt: ir.Method, arch_spec: ArchSpec, ax: Axes, atom_marker: str =
     y_max += 0.1 * y_width
 
     steps: list[tuple[ir.Statement, AtomState]] = []
-
+    constants = {}
     for stmt in mt.callable_region.walk():
         curr_state = frame.atom_state_map.get(stmt)
         if isinstance(curr_state, AtomState):
             steps.append((stmt, curr_state))
+        elif isinstance(stmt, py.Constant):
+            constants[stmt.result] = stmt.value.unwrap()
+
+    def stmt_text(stmt: ir.Statement) -> str:
+        if len(stmt.args) == 0:
+            return f"{type(stmt).__name__}"
+        return (
+            f"{type(stmt).__name__}("
+            + ", ".join(f"{constants.get(arg,'missing')}" for arg in stmt.args)
+            + ")"
+        )
 
     def draw(step_index: int):
         if len(steps) == 0:
@@ -145,7 +157,7 @@ def draw(step_index: int):
         )
         curr_state.draw_moves(arch_spec, ax=ax, color="orange")
 
-        ax.set_title(f"Step {step_index+1} / {len(steps)}: {type(stmt).__name__}")
+        ax.set_title(f"Step {step_index+1} / {len(steps)}: {stmt_text(stmt)}")
 
         ax.set_xlim(x_min, x_max)
         ax.set_ylim(y_min, y_max)