|
|
"""
|
|
|
Threshold Network for 2-Bit Ripple Carry Adder
|
|
|
|
|
|
Adds two 2-bit numbers with carry in, producing 2-bit sum and carry out.
|
|
|
Built from two cascaded full adders.
|
|
|
"""
|
|
|
|
|
|
import torch
|
|
|
from safetensors.torch import load_file
|
|
|
|
|
|
|
|
|
def heaviside(x):
|
|
|
return (x >= 0).float()
|
|
|
|
|
|
|
|
|
class ThresholdRippleCarry2Bit:
|
|
|
"""
|
|
|
2-bit ripple carry adder: (a1,a0) + (b1,b0) + cin = (cout, s1, s0)
|
|
|
"""
|
|
|
|
|
|
def __init__(self, weights_dict):
|
|
|
self.weights = weights_dict
|
|
|
|
|
|
def half_adder(self, prefix, a, b):
|
|
|
inputs = torch.tensor([float(a), float(b)])
|
|
|
|
|
|
or_out = heaviside((inputs * self.weights[f'{prefix}.sum.layer1.or.weight']).sum() +
|
|
|
self.weights[f'{prefix}.sum.layer1.or.bias'])
|
|
|
nand_out = heaviside((inputs * self.weights[f'{prefix}.sum.layer1.nand.weight']).sum() +
|
|
|
self.weights[f'{prefix}.sum.layer1.nand.bias'])
|
|
|
layer1 = torch.tensor([or_out, nand_out])
|
|
|
sum_out = heaviside((layer1 * self.weights[f'{prefix}.sum.layer2.weight']).sum() +
|
|
|
self.weights[f'{prefix}.sum.layer2.bias'])
|
|
|
carry_out = heaviside((inputs * self.weights[f'{prefix}.carry.weight']).sum() +
|
|
|
self.weights[f'{prefix}.carry.bias'])
|
|
|
|
|
|
return int(sum_out.item()), int(carry_out.item())
|
|
|
|
|
|
def full_adder(self, prefix, a, b, cin):
|
|
|
s1, c1 = self.half_adder(f'{prefix}.ha1', a, b)
|
|
|
sum_out, c2 = self.half_adder(f'{prefix}.ha2', s1, cin)
|
|
|
carry_inputs = torch.tensor([float(c1), float(c2)])
|
|
|
cout = heaviside((carry_inputs * self.weights[f'{prefix}.carry_or.weight']).sum() +
|
|
|
self.weights[f'{prefix}.carry_or.bias'])
|
|
|
return int(sum_out), int(cout.item())
|
|
|
|
|
|
def __call__(self, a0, a1, b0, b1, cin=0):
|
|
|
|
|
|
s0, c0 = self.full_adder('fa0', a0, b0, cin)
|
|
|
|
|
|
s1, cout = self.full_adder('fa1', a1, b1, c0)
|
|
|
return s0, s1, cout
|
|
|
|
|
|
@classmethod
|
|
|
def from_safetensors(cls, path="model.safetensors"):
|
|
|
return cls(load_file(path))
|
|
|
|
|
|
|
|
|
if __name__ == "__main__":
|
|
|
model = ThresholdRippleCarry2Bit.from_safetensors("model.safetensors")
|
|
|
|
|
|
print("2-Bit Ripple Carry Adder Tests:")
|
|
|
print("-" * 50)
|
|
|
|
|
|
|
|
|
for a in range(4):
|
|
|
for b in range(4):
|
|
|
a0, a1 = a & 1, (a >> 1) & 1
|
|
|
b0, b1 = b & 1, (b >> 1) & 1
|
|
|
s0, s1, cout = model(a0, a1, b0, b1, 0)
|
|
|
result = s0 + 2*s1 + 4*cout
|
|
|
expected = a + b
|
|
|
status = "OK" if result == expected else "FAIL"
|
|
|
print(f"{a} + {b} = {result} (s1s0={s1}{s0}, cout={cout}) [{status}]")
|
|
|
|