BDD implementation of some very basic circuit verification proving various 64-bit adder architectures equivalent

bdd_adders.py

import functools
from typing import Optional

class BddContext:
    def __init__(self):
        # 0/1 leaf nodes, it's convenient to just refer to them as 0 and 1
        # a node is a (var id, lo_node, hi_node) tuple
        self.nodes = []
        self.vars = []
        self.node_dict = {}
        self._addNode(None, 0, 0)
        self._addNode(None, 1, 1)

    def _addNode(self, var: int, lo: int, hi: int):
        """Returns the id of the specified node if it already exists, else creates it."""
        key = (var, lo, hi)
        existing = self.node_dict.get(key)
        if existing is not None:
            return existing

        newId = len(self.nodes)
        self.nodes.append(key)
        self.node_dict[key] = newId
        return newId

    def addVar(self, name: str) -> int:
        """Makes a new variable as well as a function representing it, returns the function"""
        varid = len(self.vars)
        self.vars.append(name)
        return self._addNode(varid, 0, 1) # build func that represents that var

    # Memoizing this is load-bearing! It's common to see the same sub-dags many times along a path.
    @functools.cache
    def _muxInternal(self, a: int, b: int, c: int) -> int:
        """The guts of Mux, which run after the initial simplifications."""
        # Recurse into cofactors
        ai, alo, ahi = self.nodes[a]
        bi, blo, bhi = self.nodes[b]
        ci, clo, chi = self.nodes[c]

        # Determine next variable to work on.
        # Conceptually, the 0/1 literals have a variable index of infinity.
        i = ai
        if b >= 2 and bi < i:
            i = bi
        if c >= 2 and ci < i:
            i = ci

        lo = self.Mux((alo if ai == i else a), (blo if bi == i else b), (clo if ci == i else c))
        hi = self.Mux((ahi if ai == i else a), (bhi if bi == i else b), (chi if ci == i else c))

        # If both are the same, that's our result
        if lo == hi:
            return lo

        # Create new node branching to the cofactors
        return self._addNode(i, lo, hi)

    def Mux(self, a: int, b: int, c: int) -> int:
        """If-Then-Else: a ? b : c"""
        # A few basic simplifications
        if b == c: # Pointless mux
            return b

        # Terminal cases
        if a < 2: # a (select) term is constant
            return b if a == 1 else c
        elif b == 1 and c == 0: # a ? 1 : 0 is just a
            return a

        # Interesting case! This is memoized.
        return self._muxInternal(a, b, c)

    # All logic ops are implemented in terms of Mux
    def Not(self, a: int) -> int:
        return self.Mux(a, 0, 1)

    def And(self, a: int, b: int) -> int:
        return self.Mux(a, b, 0)

    def Or(self, a: int, b: int) -> int:
        return self.Mux(a, 1, b)

    def Xor(self, a: int, b: int) -> int:
        return self.Mux(a, self.Not(b), b)

    def AnyDiff(self, a: list[int], b: list[int]) -> Optional[int]:
        assert len(a) == len(b)
        for i, (ai, bi) in enumerate(zip(a, b)):
            if ai != bi:
                return i
        return None

# ---- That's it! Anything below here is, essentially, test code

def half_adder(bdd: BddContext, a: int, b: int) -> tuple[int, int]:
    """Returns sum, carry bits"""
    sum = bdd.Xor(a, b)
    car = bdd.And(a, b)
    return sum, car

def full_adder(bdd: BddContext, a: int, b: int, c: int) -> tuple[int, int]:
    """Returns sum, carry bits"""
    sum0, car0 = half_adder(bdd, a, b)
    sum1, car1 = half_adder(bdd, sum0, c)
    car = bdd.Or(car0, car1)
    return sum1, car

def bit_adder(bdd: BddContext, a: int, b: int, cin: Optional[int] = None) -> tuple[int, int]:
    if cin is None:
        return half_adder(bdd, a, b)
    else:
        return full_adder(bdd, a, b, cin)

def ripple_carry_adder(bdd: BddContext, a: list[int], b: list[int], cin: Optional[int] = None) -> list[int]:
    """Builds a ripple-carry adder"""
    assert len(a) == len(b) and a != []

    carry = cin
    result = []
    for ai, bi in zip(a, b):
        sum, carry = bit_adder(bdd, ai, bi, carry)
        result.append(sum)

    if carry is not None:
        result.append(carry)
    return result

def conditional_sum_adder(bdd: BddContext, a: list[int], b: list[int], cin: Optional[int] = None) -> list[int]:
    """Recursively builds a conditional-sum adder"""
    N = len(a)
    assert N == len(b) and N > 0

    # Base case
    if N == 1:
        return list(bit_adder(bdd, a[0], b[0], cin))

    # Split into halves and recurse
    H = N // 2
    lo_sum = conditional_sum_adder(bdd, a[:H], b[:H], cin)
    hi_sum0 = conditional_sum_adder(bdd, a[H:], b[H:], 0) # sum with carry-in of 0
    hi_sum1 = conditional_sum_adder(bdd, a[H:], b[H:], 1) # sum with carry-in of 1

    # Combine with conditional select based on low half carry
    lo_carry_out = lo_sum[-1]
    result = lo_sum[:-1]

    for hi0, hi1 in zip(hi_sum0, hi_sum1):
        result.append(bdd.Mux(lo_carry_out, hi1, hi0))

    return result

def prefix_adder(bdd: BddContext, a: list[int], b: list[int], cin: Optional[int], pg_prop) -> list[int]:
    """Builds a prefix adder"""
    N = len(A)
    assert N == len(b) and N > 0

    if cin is not None:
        # Reduce case with carry-in to case without with a minimum of fuss
        # we just prepend cin to a and b both, generating a -1'th digit that carries if cin is set, and
        # then strip the extra digit from the result
        return prefix_adder(bdd, [cin] + a, [cin] + b, pg_prop, None)[1:]

    # Initial propagate/generate bits
    pg = [(bdd.Or(ai, bi), bdd.And(ai, bi)) for ai, bi in zip(a, b)]

    # PG propagation - this is the only difference between prefix adder families
    pg_prop(bdd, pg)

    # Result sum bits
    result = [bdd.Xor(bdd.Xor(ai, bi), ci) for (_, ci), (ai, bi) in zip([(0, 0)] + pg[:-1], zip(a, b))]
    # Final carry
    result.append(pg[-1][1])

    return result

def pg_cell(bdd: BddContext, pg: tuple[int, int], pg_prev: tuple[int, int]) -> tuple[int, int]:
    """Propagate/generate cell."""
    pi, gi = pg
    pj, gj = pg_prev
    return bdd.And(pi, pj), bdd.Or(gi, bdd.And(pi, gj))

def kogge_stone_adder(bdd: BddContext, a: list[int], b: list[int], cin: Optional[int] = None) -> list[int]:
    """Builds a Kogge-Stone prefix adder"""
    def pg_prop(bdd: BddContext, pg: list[tuple[int, int]]):
        N = len(pg)
        group_step = 1
        while group_step < N:
            for i in range(N):
                if i >= group_step:
                    pg[i] = pg_cell(bdd, pg[i], pg[i - group_step])

            group_step <<= 1

    return prefix_adder(bdd, a, b, cin, pg_prop)

def ladner_fischer_adder(bdd: BddContext, a: list[int], b: list[int], cin: Optional[int] = None) -> list[int]:
    """Builds a Ladner-Fischer prefix adder"""
    def pg_prop(bdd: BddContext, pg: list[tuple[int, int]]):
        N = len(pg)
        group_step = 1
        while group_step < N:
            for i in range(N):
                if (i & group_step) != 0:
                    pg[i] = pg_cell(bdd, pg[i], pg[(i - group_step) | (group_step - 1)])

            group_step <<= 1

    return prefix_adder(bdd, a, b, cin, pg_prop)

if __name__ == '__main__':
    bdd = BddContext()

    # Create vars for inputs. These need to be interleaved in this order for the
    # BDD size to stay small!
    A = []
    B = []
    N = 64
    for i in range(N):
        A.append(bdd.addVar(f'a{i}'))
        B.append(bdd.addVar(f'b{i}'))

    SumRCA = ripple_carry_adder(bdd, A, B)
    print(f'N={N}, BDD size is {len(bdd.nodes)} after building RCA')

    SumCSA = conditional_sum_adder(bdd, A, B)
    print(f'BDD size is {len(bdd.nodes)} after adding CSA')

    SumLFA = ladner_fischer_adder(bdd, A, B)
    print(f'BDD size is {len(bdd.nodes)} after adding LFA')

    SumKSA = kogge_stone_adder(bdd, A, B)
    print(f'BDD size is {len(bdd.nodes)} after adding KSA')

    Difference = bdd.AnyDiff(SumRCA, SumCSA)
    print(f'First difference between RCA and Conditional-Sum in bit: {Difference}')

    Difference = bdd.AnyDiff(SumRCA, SumLFA)
    print(f'First difference between RCA and LFA in bit: {Difference}')

    Difference = bdd.AnyDiff(SumRCA, SumKSA)
    print(f'First difference between RCA and KSA in bit: {Difference}')