diff options
| author | Tanner Robison <[email protected]> | 2026-06-26 19:13:44 -0700 |
|---|---|---|
| committer | Tanner Robison <[email protected]> | 2026-06-29 19:44:13 -0700 |
| commit | 3a17db637fd41015df26f1cceb0d1c098e85d436 (patch) | |
| tree | 34b6913348444b6ff2b47218eafa7ad9748e574f /spires_memristor_sim | |
| parent | 3c9a3d7d92d155545ce4892524f00dcd092b52d3 (diff) | |
spires python memTorch interface
Implemented an interface for memTorch to talk to spires. memTorch is
used for calculating the weights and are passed to spires for
determining the spikes
readout layer has still not been implemented, no way of actually reading
the output
Diffstat (limited to 'spires_memristor_sim')
| -rw-r--r-- | spires_memristor_sim/README.md | 12 | ||||
| -rw-r--r-- | spires_memristor_sim/memTorch_SNN.py | 28 | ||||
| -rw-r--r-- | spires_memristor_sim/requirements.txt | 71 | ||||
| -rw-r--r-- | spires_memristor_sim/spires_interface.py | 119 | ||||
| -rw-r--r-- | spires_memristor_sim/torch_reservoir.py | 35 |
5 files changed, 213 insertions, 52 deletions
diff --git a/spires_memristor_sim/README.md b/spires_memristor_sim/README.md index bc89498..e5717b6 100644 --- a/spires_memristor_sim/README.md +++ b/spires_memristor_sim/README.md @@ -1,5 +1,17 @@ # Spires reservoir simulated on memristor crossbar using memTorch +## overview +### simulation +Here we are using both memTorch and the Spiresrc libraries working together to simulate +a spires reservoir on a memristor crossbar. memTorch is acting as the weights, simulating +at a physics level including device to device differences, conductance drag, etc.... +Spires is acting as the actual reservoir neurons and determining the spikes. + +### benchmark + +## setup + + ## Goal The goal is to able to simulate a spires reservoir on a memristor crossbar using a custom memristor spice model(HIL), thus you can run multiple simulations with diff --git a/spires_memristor_sim/memTorch_SNN.py b/spires_memristor_sim/memTorch_SNN.py deleted file mode 100644 index 67e9cb6..0000000 --- a/spires_memristor_sim/memTorch_SNN.py +++ /dev/null @@ -1,28 +0,0 @@ -import torch -import memtorch - -# 1. Standard memTorch setup (Static crossbar weights) -ann_layer = torch.nn.Linear(100, 10) -patched_layer = memtorch.mn.Module.patch_model(ann_layer, memristor_model) - -# 2. Your custom SNN wrapper loop -def forward_snn(input_spikes_over_time): - # input_spikes_over_time shape: (time_steps, batch_size, input_dim) - time_steps = input_spikes_over_time.shape[0] - v_mem = torch.zeros(batch_size, 10) # Hidden neuron membrane potentials - output_spikes = [] - - for t in range(time_steps): - # Pass binary spikes through memTorch's physical crossbar simulation - current_in = patched_layer(input_spikes_over_time[t]) - - # Leaky Integrate-and-Fire (LIF) logic (Written by you!) - v_mem = 0.9 * v_mem + current_in # Leak & Integrate - - # Fire threshold - spike = (v_mem >= 1.0).float() - v_mem[v_mem >= 1.0] = 0.0 # Reset - - output_spikes.append(spike) - - return torch.stack(output_spikes) diff --git a/spires_memristor_sim/requirements.txt b/spires_memristor_sim/requirements.txt new file mode 100644 index 0000000..c9b7376 --- /dev/null +++ b/spires_memristor_sim/requirements.txt @@ -0,0 +1,71 @@ +asteval==1.0.9 +asttokens==3.0.1 +contourpy==1.3.3 +cuda-bindings==13.3.1 +cuda-pathfinder==1.5.5 +cuda-toolkit==13.0.2 +cycler==0.12.1 +decorator==5.3.1 +dill==0.4.1 +executing==2.2.1 +filelock==3.29.4 +fonttools==4.63.0 +fsspec==2026.6.0 +ipython==9.15.0 +ipython_pygments_lexers==1.1.1 +jedi==0.20.0 +Jinja2==3.1.6 +joblib==1.5.3 +kiwisolver==1.5.0 +lmfit==1.3.4 +MarkupSafe==3.0.3 +matplotlib==3.11.0 +matplotlib-inline==0.2.2 +memtorch-cpu==1.1.6 +mpmath==1.3.0 +narwhals==2.22.1 +networkx==3.6.1 +numpy==2.5.0 +nvidia-cublas==13.1.1.3 +nvidia-cuda-cupti==13.0.85 +nvidia-cuda-nvrtc==13.0.88 +nvidia-cuda-runtime==13.0.96 +nvidia-cudnn-cu13==9.20.0.48 +nvidia-cufft==12.0.0.61 +nvidia-cufile==1.15.1.6 +nvidia-curand==10.4.0.35 +nvidia-cusolver==12.0.4.66 +nvidia-cusparse==12.6.3.3 +nvidia-cusparselt-cu13==0.8.1 +nvidia-nccl-cu13==2.29.7 +nvidia-nvjitlink==13.0.88 +nvidia-nvshmem-cu13==3.4.5 +nvidia-nvtx==13.0.85 +packaging==26.2 +pandas==3.0.4 +parso==0.8.7 +pexpect==4.9.0 +pillow==12.2.0 +prompt_toolkit==3.0.52 +psutil==7.2.2 +ptyprocess==0.7.0 +pure_eval==0.2.3 +Pygments==2.20.0 +pyparsing==3.3.2 +python-dateutil==2.9.0.post0 +scikit-learn==1.9.0 +scipy==1.18.0 +seaborn==0.13.2 +setuptools==81.0.0 +six==1.17.0 +sklearn==0.0.post12 +stack-data==0.6.3 +sympy==1.14.0 +threadpoolctl==3.6.0 +torch==2.12.1 +torchvision==0.27.1 +traitlets==5.15.1 +triton==3.7.1 +typing_extensions==4.15.0 +uncertainties==3.2.3 +wcwidth==0.8.1 diff --git a/spires_memristor_sim/spires_interface.py b/spires_memristor_sim/spires_interface.py index 3f1e302..b7e5fff 100644 --- a/spires_memristor_sim/spires_interface.py +++ b/spires_memristor_sim/spires_interface.py @@ -1,35 +1,116 @@ import ctypes import numpy import torch +import random -spires_lib = ctypes.CDLL("../spires/build/.libspires.so") +LIF_DISCRETE = 0 -spires_lib.spires_reservoir_step.argtypes = [ctypes.POINTER(ctypes.c_float)] -spires_lib.spires_reservoir_step.restype = None +#load the spires library +spires_lib = ctypes.CDLL("../spires/lib/libspires.so") + +# C signautes for creating reservoir +spires_lib.create_reservoir.argtypes = [ + ctypes.c_size_t, + ctypes.c_size_t, + ctypes.c_size_t, + ctypes.c_double, + ctypes.c_double, + ctypes.c_double, + ctypes.c_double, + ctypes.c_double, + ctypes.c_int, + ctypes.c_int, + ctypes.POINTER(ctypes.c_double) +] +spires_lib.create_reservoir.restype = ctypes.c_void_p + +# C signatures for reservoir destruction +spires_lib.free_reservoir.argtypes = [ctypes.c_void_p] +spires_lib.free_reservoir.restype = None + +# C signatures for step reservoir function +spires_lib.update_neuron.argtypes = [ + ctypes.c_void_p, #Pointer to spires reservoir struct + ctypes.c_int, + ctypes.c_double, + ctypes.c_double +] +spires_lib.update_neuron.restype = None + +# C signatures for get neuron spike function +spires_lib.read_reservoir_spikes.argtypes = [ + ctypes.c_void_p, + ctypes.POINTER(ctypes.c_float) +] +spires_lib.read_reservoir_spikes.restype = None + +def init_spires_reservoir(reservoir_size): + print("Creating spires reservoir") + + neuron_parameters = (ctypes.c_double * 4)(0.0, 1.0, 0.2, 0.5) + c_neuron_parameters = ctypes.cast(neuron_parameters, ctypes.POINTER(ctypes.c_double)) + + reservoir_ptr = spires_lib.create_reservoir( + ctypes.c_size_t(int(reservoir_size)), # num neurons + ctypes.c_size_t(int(reservoir_size)), # num neurons + ctypes.c_size_t(2), # num_outputs + ctypes.c_double(0.9), # spectral radius + ctypes.c_double(0.8), # ei_ratio + ctypes.c_double(1.0), # input_strength + ctypes.c_double(0.1), # connectivity + ctypes.c_double(1.0), # dt + ctypes.c_int(1), # connectivity type ( 1 = sparse) + ctypes.c_int(LIF_DISCRETE), # neuron type + c_neuron_parameters # neuron params + ) + + #allocate empty void pointer, (reservoir will go here) + if not reservoir_ptr: + raise RuntimeError(f"Spires faile to initialize the reservoir: {status_code}") + else: + print("Spires reservoir initialized") + + # reservoir_ptr._keep_alive = neuron_parameters + + return reservoir_ptr + +def free_spires_reservoir(reservoir_ptr): + print("Freeing the spires reservoir") + spires_lib.free_reservoir(reservoir_ptr) + return 0 + #change currents from tensor to a flat C pointer array for spires to read -def send_currents_to_spires(currents): +def send_currents_to_spires(reservoir_ptr, currents_tensor): + print("Sending currents to spires") # ----- extract from pytorhc graph ----- # .detach() removes it from auto gradient tracking # .cpu() make sure data is in RAM, not VRAM # .numpy() maps it to numpy array - numpy_array = torch_tensor.detach().cpu().numpy() - - # ----- makes suren layout matches 32 bit float C-array ----- - # .astype(npfloat32) forces standard CC float precistion + # ----- makes suren layout matches 64 bit double C-array ----- + # .astype(npfloat64) forces standard CC float precistion # .flatten() makes sure the memory is a 1D block - contiguous_array = numpy_array.astype(np.float32).flatten() + numpy_array = currents_tensor.detach().cpu().numpy().astype(numpy.float64).flatten() - # ----- Ensure raw memory pointer ----- - c_float_ptr = contiguous_array.ctypes.data_as(ctypes.POINTER(ctypes.c_float)) - array_size = contiguous_array.size + some_ptr = ctypes.cast(reservoir_ptr, ctypes.POINTER(ctypes.POINTER(ctypes.c_void_p))) - # ----- Call the C library ----- - spires_lib.spires_reservoir_step(c_float_ptr, array_size) - return contiguous_array + neurons_array = some_ptr[0] - # NEED TO MODIFY SPIRES FOR THIS?? :(( + print("updating the neurons") + for i in range(len(numpy_array)): + neuron_ptr = neurons_array[i] + input_current = numpy_array[i] -def read_spikes_from_spires(): - placeholder = 0 - return 0 + spires_lib.update_neuron(neuron_ptr, LIF_DISCRETE, input_current, 1.0) + + return numpy_array + +def read_spikes_from_spires(reservoir_ptr, size=0): + print("Recieved spikes from spires") + #didnt do any safety checking womp womp + returned_spikes = numpy.zeros(size, dtype=numpy.float64) + c_spike_ptr = returned_spikes.ctypes.data_as(ctypes.POINTER(ctypes.c_float)) + + spires_lib.read_reservoir_spikes(reservoir_ptr, c_spike_ptr) + + return returned_spikes diff --git a/spires_memristor_sim/torch_reservoir.py b/spires_memristor_sim/torch_reservoir.py index c3017c5..9d58d43 100644 --- a/spires_memristor_sim/torch_reservoir.py +++ b/spires_memristor_sim/torch_reservoir.py @@ -1,16 +1,22 @@ import torch import memtorch from spires_interface import ( + free_spires_reservoir, + init_spires_reservoir, send_currents_to_spires, read_spikes_from_spires, ) NUM_INPUTS = 4 +NUM_OUTPUTS = 4 NUM_NEURONS = 800 input_layer = torch.nn.Linear(NUM_INPUTS, NUM_NEURONS, bias=False) #this is a recurrent layer sort of??? reservoir_layer = torch.nn.Linear(NUM_NEURONS, NUM_NEURONS, bias=False) +readout_layer = torch.nn.Linear(NUM_NEURONS, NUM_OUTPUTS, bias=False) + +spires_reservoir = init_spires_reservoir(NUM_NEURONS) #make the reservoir sparse and random with torch.no_grad(): @@ -22,22 +28,41 @@ with torch.no_grad(): # patch layers into memristor crossbars with memTorch memristor_model = memtorch.bh.memristor.VTEAM -mem_input_layer = memtorch.mn.Module.patch_model(input_layer, memristor_model) -mem_reservoir_layer = memtorch.mn.Module.patch_model(reservoir_layer, memristor_model) +memristor_model_params = { + 'time_series_resolution': 1e-3, + 'r_on': 50, + 'r_off': 1000, +} + +mem_input_layer = memtorch.mn.Module.patch_model( + input_layer, + memristor_model, + memristor_model_params, +) + +mem_reservoir_layer = memtorch.mn.Module.patch_model( + reservoir_layer, + memristor_model, + memristor_model_params, +) #run spiking loop with spires previous_spikes = torch.zeros(1, NUM_NEURONS) for step in range(500): + cartpole_state = [0, 1, 2, 3] #this isn't final input_tensor = torch.tensor(cartpole_state).float().unsqueeze(0) currents_in = mem_input_layer(input_tensor) currents_recv = mem_reservoir_layer(previous_spikes) + #may need to scale currents? total_currents = currents_in + currents_recv - send_currents_to_spires(total_currents) - - current_spikes_np = read_spikes_from_spires() + send_currents_to_spires(spires_reservoir, total_currents) + + current_spikes_np = read_spikes_from_spires(spires_reservoir, NUM_NEURONS) previous_spikes = torch.from_numpy(current_spikes_np).float().unsqueeze(0) +free_spires_reservoir(spires_reservoir) #pass current spikes np to readout layer + |
