From ec44271415fdecc0aca94005dff4dc2f38b3493b Mon Sep 17 00:00:00 2001
From: Alessandro Candido <candido.ale@gmail.com>
Date: Thu, 13 Feb 2025 18:20:59 +0100
Subject: [PATCH] fix: Handle negative numbers in registers

Caused by sweepers initialization and update
---
 .../_core/instruments/qblox/sequence/loops.py |  9 ++-
 .../instruments/qblox/sequence/transpile.py   | 76 ++++++++++++++++---
 2 files changed, 71 insertions(+), 14 deletions(-)

diff --git a/src/qibolab/_core/instruments/qblox/sequence/loops.py b/src/qibolab/_core/instruments/qblox/sequence/loops.py
index 01f3a0fef..629d6fa32 100644
--- a/src/qibolab/_core/instruments/qblox/sequence/loops.py
+++ b/src/qibolab/_core/instruments/qblox/sequence/loops.py
@@ -19,6 +19,7 @@
     Reference,
     Register,
     ResetPh,
+    Sub,
     Wait,
 )
 from .asm import Registers
@@ -131,8 +132,12 @@ def _sweep_update(p: Param, channel: ChannelId) -> Block:
     """
     return (
         Line(
-            instruction=Add(a=p.reg, b=p.step, destination=p.reg),
-            comment=f"increment {p.description}",
+            instruction=(
+                Add(a=p.reg, b=p.step, destination=p.reg)
+                if p.step >= 0
+                else Sub(a=p.reg, b=-p.step, destination=p.reg)
+            ),
+            comment=f"shift {p.description}",
         ),
         *(
             update_instructions(p.kind, p.reg)
diff --git a/src/qibolab/_core/instruments/qblox/sequence/transpile.py b/src/qibolab/_core/instruments/qblox/sequence/transpile.py
index eb47ffaed..5df520e9a 100644
--- a/src/qibolab/_core/instruments/qblox/sequence/transpile.py
+++ b/src/qibolab/_core/instruments/qblox/sequence/transpile.py
@@ -1,3 +1,5 @@
+from typing import Optional, cast
+
 from ..q1asm.ast_ import (
     Block,
     Instruction,
@@ -6,6 +8,8 @@
     Move,
     Program,
     Reference,
+    Register,
+    Sub,
     Wait,
 )
 from .asm import Registers
@@ -15,7 +19,7 @@
 MAX_WAIT = 2**16 - 1
 
 
-def long_wait(duration: int) -> Block:
+def _long_wait(duration: int) -> Block:
     iterations = duration // MAX_WAIT
     remainder = duration % MAX_WAIT
     register = Registers.wait.value
@@ -27,18 +31,66 @@ def long_wait(duration: int) -> Block:
     ]
 
 
-def decompose(line: Line) -> Block:
-    if not isinstance(line.instruction, Wait):
-        return [line]
-    duration = line.instruction.duration
+def _decompose_wait(instr: Wait) -> Optional[Block]:
+    duration = instr.duration
     if not isinstance(duration, int) or duration <= MAX_WAIT:
+        return None
+    return _long_wait(duration)
+
+
+def _negative_move(instr: Move):
+    """Compile negative value sets.
+
+    Apparently, the only place where negative numbers are not allowed
+    are registers, otherwise they are handled by the internal compiler.
+
+    https://docs.qblox.com/en/main/tutorials/q1asm_tutorials/intermediate/nco_control_adv.html#:~:text=Internally,%20the%20processor%20stores
+
+    Thus, we compile instructions setting negative values as suggested:
+    first setting them to 0, than subtracting the desired amount. This
+    is more reliable than manually complementing the number, since it
+    makes no assumption about the registers size.
+
+    https://docs.qblox.com/en/main/cluster/troubleshooting.html#:~:text=How%20do%20I%20set%20negative%20numbers
+    """
+    src = cast(int, instr.source)
+    dest = instr.destination
+    return [Move(source=0, destination=dest), Sub(a=dest, b=abs(src), destination=dest)]
+
+
+def _decompose_move(instr: Move) -> Optional[Block]:
+    src = instr.source
+    if isinstance(src, Register):
+        return None
+    assert isinstance(src, int)
+    if src >= 0:
+        return None
+    return _negative_move(instr)
+
+
+def _decompose(line: Line) -> Block:
+    instr = line.instruction
+    block = (
+        _decompose_wait(instr)
+        if isinstance(instr, Wait)
+        else _decompose_move(instr)
+        if isinstance(instr, Move)
+        else None
+    )
+
+    # default
+    if block is None:
         return [line]
-    wait = long_wait(duration)
-    assert isinstance(wait[0], Instruction)
-    return [
-        Line(instruction=wait[0], label=line.label, comment=line.comment),
-        *(el for el in wait[1:]),
-    ]
+
+    assert isinstance(block[0], Instruction)
+    return (
+        [
+            Line(instruction=block[0], label=line.label, comment=line.comment),
+            *(el for el in block[1:]),
+        ]
+        if block is not None
+        else [line]
+    )
 
 
 def transpile(prog: Block) -> Program:
@@ -46,6 +98,6 @@ def transpile(prog: Block) -> Program:
         elements=[
             el if isinstance(el, Line) else Line.instr(el)
             for oel in prog
-            for el in (decompose(oel) if isinstance(oel, Line) else [oel])
+            for el in (_decompose(oel) if isinstance(oel, Line) else [oel])
         ]
     )