diff options
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 + |