diff --git a/.gitignore b/.gitignore index a688bc7d..196db64d 100644 --- a/.gitignore +++ b/.gitignore @@ -44,8 +44,11 @@ benchmarks/.asv **/*.code-workspace **/.coverage **/.DS_Store -examples/dpointnet_v1/GLIF_network/network/ +**/_csr_spike_ops.so +**/_csr_spike_ops.archs +examples/dpointnet_v1/GLIF_network/network examples/dpointnet_v1/GLIF_network/cached_networks/ +examples/dpointnet_v1/lgn_cache/ examples/dpointnet_v1/Neuropixels_data/ examples/dpointnet_v1/Synchronization_data/ examples/dpointnet_*/*callbacks* \ No newline at end of file diff --git a/MANIFEST.in b/MANIFEST.in index 4cc6ecf1..fa330cda 100644 --- a/MANIFEST.in +++ b/MANIFEST.in @@ -1,2 +1,2 @@ -recursive-include bmtk *.py *.md *.txt *.cfg **/*.json **/*.hoc *.csv *.swc *.mod *.yaml +recursive-include bmtk *.py *.md *.txt *.cfg **/*.json **/*.hoc *.csv *.swc *.mod *.yaml *.cc *.h *.sh recursive-exclude bmtk/tests * diff --git a/bmtk/simulator/dpointnet/cell_models/glif3_cell.py b/bmtk/simulator/dpointnet/cell_models/glif3_cell.py index 91e3b91c..2232ec6e 100644 --- a/bmtk/simulator/dpointnet/cell_models/glif3_cell.py +++ b/bmtk/simulator/dpointnet/cell_models/glif3_cell.py @@ -4,6 +4,13 @@ import pickle as pkl from pathlib import Path from bmtk.simulator.dpointnet.io_tools import io +from bmtk.simulator.dpointnet.custom_ops import ( + build_csr_connectivity, + cuda_op_status, + fused_cuda_available, + fused_spike_currents, + reorder_csr_values, +) try: @@ -410,6 +417,35 @@ def straight_through_dampen(x, dampening): return x * (tf_one - dampening) + tf.stop_gradient(x * dampening) +def _validate_fused_cuda_option(value): + if isinstance(value, np.ndarray) and value.ndim == 0: + value = value.item() + if value is True or value is False: + return value + if isinstance(value, (bytes, np.bytes_)): + try: + value = value.decode('utf-8') + except UnicodeDecodeError: + pass + if isinstance(value, (str, np.str_)) and value == 'auto': + return 'auto' + raise ValueError('use_fused_cuda must be true, false, or "auto".') + + +def _fused_cuda_dtype_error(compute_dtype, variable_dtype): + compute_dtype = tf.as_dtype(compute_dtype) + variable_dtype = tf.as_dtype(variable_dtype) + if (compute_dtype in (tf.float16, tf.float32) + and variable_dtype == tf.float32): + return None + return ( + 'the fused operator requires float16 or float32 computation and ' + 'float32 variables; got ' + f'compute_dtype={compute_dtype.name}, ' + f'variable_dtype={variable_dtype.name}' + ) + + class GLIF3Cell(tf.keras.layers.Layer): @@ -451,11 +487,34 @@ def __init__( hard_reset=False, tau_basis=None, synaptic_basis_weights=None, + use_fused_cuda=False, # current_input=False, ): super().__init__() self.__seq_idx = 0 + use_fused_cuda = _validate_fused_cuda_option(use_fused_cuda) + fused_dtype_error = _fused_cuda_dtype_error( + self.compute_dtype, self.variable_dtype + ) + fused_available = fused_cuda_available() and fused_dtype_error is None + if use_fused_cuda is True and not fused_available: + unavailable_reason = fused_dtype_error or cuda_op_status() + raise RuntimeError( + 'use_fused_cuda=True but the fused DPointNet CUDA operator ' + f'is unavailable: {unavailable_reason}' + ) + self._use_fused_cuda = fused_available and ( + use_fused_cuda is True or use_fused_cuda == 'auto' + ) + if use_fused_cuda == 'auto' and not self._use_fused_cuda: + unavailable_reason = fused_dtype_error or cuda_op_status() + io.log_warning( + 'DPointNet fused CUDA currents are unavailable; using the ' + f'TensorFlow fallback. Status: {unavailable_reason}' + ) + elif self._use_fused_cuda: + io.log_info(f'DPointNet fused CUDA currents enabled ({cuda_op_status()}).') _node_params = dict(glif_network['node_params']) @@ -606,7 +665,20 @@ def __init__( # Define the Tensorflow variables self.recurrent_indices = tf.Variable(indices, dtype=tf.int64, trainable=False) #dtype necessary for sparse dense matmul - self.pre_ind_table = make_pre_ind_table(indices, n_source_neurons=self.recurrent_dense_shape[1]) # dtype int32 + if self._use_fused_cuda: + self.recurrent_fused_connectivity = build_csr_connectivity( + indices, + syn_ids, + self.recurrent_dense_shape[1], + self._n_neurons, + self.synaptic_basis_weights.shape[0], + ) + self.pre_ind_table = None + else: + self.recurrent_fused_connectivity = None + self.pre_ind_table = make_pre_ind_table( + indices, n_source_neurons=self.recurrent_dense_shape[1] + ) recurrent_weight_positive = tf.constant(weights >= 0, dtype=tf.bool) @@ -647,6 +719,21 @@ def __init__( self.recurrent_weight_values_compute = self._untracked_variable(recurrent_weight_values_compute) else: self.recurrent_weight_values_compute = self.recurrent_weight_values + if self._use_fused_cuda: + recurrent_csr_weights = tf.Variable( + reorder_csr_values( + self.recurrent_weight_values_compute, + self.recurrent_fused_connectivity, + ), + name='sparse_recurrent_weights_csr_compute', + trainable=False, + dtype=self.compute_dtype, + ) + self.recurrent_csr_weight_values_compute = self._untracked_variable( + recurrent_csr_weights + ) + else: + self.recurrent_csr_weight_values_compute = None self.syn_ids = tf.constant(syn_ids, dtype=tf.int64) # this needs to be int64 for efficiency # self.recurrent_weights_factors = tf.gather(self.synaptic_basis_weights, self.syn_ids, axis=0) # TensorShape([23525415, 5]) @@ -705,30 +792,46 @@ def __init__( else: input_props['input_weight_values_compute'] = input_props['input_weight_values'] input_props['input_syn_ids'] = tf.constant(input_syn_ids, dtype=tf.int64) # for efficiency this needs to be in int64 - # if not self._current_input: - # input_props['pre_input_ind_table'] = make_pre_ind_table( - # input_indices, - # n_source_neurons=input_dense_shape[1] - # ) input_props['input_type'] = input_type if input_type == 'spikes': - input_props['pre_input_ind_table'] = make_pre_ind_table( - input_indices, - n_source_neurons=input_dense_shape[1] - ) end_indx = self.inputs_idx[idx] + n_input_nodes elif input_type in ('poisson_spikes_internal', 'noisy_current'): firing_rate = input_options.get('firing_rate', 250.0) input_props['spike_prob'] = tf.constant(firing_rate * dt / 1000.0, dtype=self.compute_dtype) - input_props['pre_input_ind_table'] = make_pre_ind_table( - input_indices, - n_source_neurons=input_dense_shape[1] - ) end_indx = self.inputs_idx[idx] elif input_type == 'current': end_indx = self.inputs_idx[idx] + n_input_nodes else: raise ValueError(f'Unknown input type {input_type}') + if input_type in ('spikes', 'poisson_spikes_internal', 'noisy_current'): + if self._use_fused_cuda: + input_props['fused_connectivity'] = build_csr_connectivity( + input_indices, + input_syn_ids, + input_dense_shape[1], + self._n_neurons, + self.synaptic_basis_weights.shape[0], + ) + input_props['pre_input_ind_table'] = None + input_csr_weights = tf.Variable( + reorder_csr_values( + input_props['input_weight_values_compute'], + input_props['fused_connectivity'], + ), + name=f'{input_name}_input_weights_csr_compute', + trainable=False, + dtype=self.compute_dtype, + ) + input_props['csr_weight_values_compute'] = self._untracked_variable( + input_csr_weights + ) + else: + input_props['fused_connectivity'] = None + input_props['csr_weight_values_compute'] = None + input_props['pre_input_ind_table'] = make_pre_ind_table( + input_indices, + n_source_neurons=input_dense_shape[1], + ) io.log_debug(f' > Added "{input_name}" input synapses: indices = {len(input_indices)}, trainble = {input_trainable}') self.inputs_idx[idx+1] = end_indx @@ -816,6 +919,17 @@ def refresh_recurrent_weight_shadow(self): self.recurrent_weight_values_compute.assign( tf.cast(self.recurrent_weight_values, self.compute_dtype) ) + recurrent_csr_shadow = getattr( + self, 'recurrent_csr_weight_values_compute', None + ) + if (recurrent_csr_shadow is not None + and self.recurrent_weight_values.trainable): + recurrent_csr_shadow.assign( + reorder_csr_values( + tf.cast(self.recurrent_weight_values, self.compute_dtype), + self.recurrent_fused_connectivity, + ) + ) # Also sync any trainable input-weight shadows (e.g. trainable background weights), # which the input-current custom gradient reads in its forward pass. @@ -825,8 +939,35 @@ def refresh_recurrent_weight_shadow(self): if shadow is None or shadow is master or not master.trainable: continue shadow.assign(tf.cast(master, self.compute_dtype)) + csr_shadow = input_net.get('csr_weight_values_compute') + if csr_shadow is not None: + csr_shadow.assign( + reorder_csr_values( + tf.cast(master, self.compute_dtype), + input_net['fused_connectivity'], + ) + ) + + def close_fused_cuda(self): + connectivity = getattr(self, 'recurrent_fused_connectivity', None) + if connectivity is not None: + connectivity.close() + for input_net in self.inputs.values(): + connectivity = input_net.get('fused_connectivity') + if connectivity is not None: + connectivity.close() def calculate_i_rec_with_custom_grad(self, rec_z_buf): + if self._use_fused_cuda: + return fused_spike_currents( + rec_z_buf, + self.recurrent_weight_values, + self.recurrent_csr_weight_values_compute, + self.recurrent_fused_connectivity, + self.synaptic_basis_weights, + self._n_neurons, + compute_spike_gradient=True, + ) return calculate_synaptic_currents( rec_z_buf, self.recurrent_indices, @@ -863,6 +1004,16 @@ def calculate_input_current_from_firing_probabilities(self, x_t, input_net): return i_in_flat def calculate_input_current_from_spikes(self, x_t, input_net): + if self._use_fused_cuda: + return fused_spike_currents( + tf.cast(x_t, self.compute_dtype), + input_net['input_weight_values'], + input_net['csr_weight_values_compute'], + input_net['fused_connectivity'], + self.synaptic_basis_weights, + self._n_neurons, + compute_spike_gradient=False, + ) # Memory-efficient input current via the @tf.custom_gradient module function: the forward # reads the compute-dtype shadow and the backward recomputes the cheap basis gather (instead # of retaining per-timestep activations), while the gradient targets the trainable master. @@ -934,9 +1085,6 @@ def _dense_update_impl(self, batch_size, prev_z, v, r, asc, psc_rise, psc, rec_i return new_v, new_r, new_asc, new_psc_rise, new_psc def calculate_noise_current(self, batch_size, input_net): - # n_post_neurons = self.bkg_input_dense_shape[0] - n_post_neurons = input_net['input_dense_shape'][0] - # Use persistent self.noise_step instead of the passed-in noise_step (which is now unused) step_seed = tf.cast(self.noise_step, tf.int32) base_seed = tf.cast(self.noise_seed, tf.int32) replica_context = tf.distribute.get_replica_context() @@ -959,64 +1107,7 @@ def calculate_noise_current(self, batch_size, input_net): lam=input_net['spike_prob'], dtype=tf.int32, ) - - # Keep noise indexing in int64 - non_zero_indices = tf.where(rest_of_brain > 0) - - batch_indices = non_zero_indices[:, 0] #int64 - pre_neuron_indices = non_zero_indices[:, 1] #int64 - # Get the indices into self.recurrent_indices for each pre_neuron_index - # self.pre_ind_table is a RaggedTensor or a list of lists mapping pre_neuron_index to indices in recurrent_indices - new_indices, new_weights, new_syn_ids, post_in_degree, _ = get_new_inds_table( - input_net['input_indices'], - input_net['input_weight_values'], - input_net['input_syn_ids'], - pre_neuron_indices, - input_net['pre_input_ind_table'], - # self.bkg_input_indices, - # self.bkg_input_weights, - # self.bkg_input_syn_ids, - # pre_neuron_indices, - # self.pre_bkg_ind_table - ) - # Expand batch_indices to match the length of inds_flat - batch_indices_per_connection = tf.repeat(batch_indices, post_in_degree) - # Get post-synaptic neuron indices - post_neuron_indices = new_indices[:, 0] - # Compute segment IDs - segment_ids = batch_indices_per_connection * n_post_neurons + post_neuron_indices - segment_ids = tf.cast(segment_ids, dtype=tf.int32) - num_segments = tf.cast(batch_size * n_post_neurons, dtype=tf.int32) - - # # Alternative (slower): Gather spike counts once per active input and broadcast by connection degree. - # # Keeping it here only as reference. - # active_spike_counts = tf.gather_nd(rest_of_brain, non_zero_indices) - # n_pre_spikes = tf.cast(tf.repeat(active_spike_counts, post_in_degree), dtype=self.variable_dtype) - - presynaptic_indices = tf.stack( - [batch_indices_per_connection, new_indices[:, 1]], - axis=1 - ) # gather works better with int64 - n_pre_spikes = tf.cast(tf.gather_nd(rest_of_brain, presynaptic_indices), dtype=self.compute_dtype) - # n_pre_spikes = tf.cast(tf.gather_nd(rest_of_brain, presynaptic_indices), dtype=self.compute_dtype) - - # Compute weighted basis factors - basis_factors = tf.gather(self.synaptic_basis_weights, new_syn_ids, axis=0) - # Accumulate currents in variable_dtype for better numerical fidelity, then cast once. - new_weights = tf.cast(new_weights, dtype=self.compute_dtype) # shape (n_active_connections,) - # new_weights = tf.cast(new_weights * n_pre_spikes, self.compute_dtype) - # new_weights_final = new_weights[:, tf.newaxis] * basis_factors - # new_weights_final = tf.cast(new_weights * n_pre_spikes, dtype=self.compute_dtype) - new_weights_final = (new_weights * n_pre_spikes)[:, tf.newaxis] * basis_factors - - # Calculate input currents - i_in_flat = tf.math.unsorted_segment_sum(new_weights_final, segment_ids, num_segments) - - # # Cast back to compute_dtype to keep downstream state buffers compact. - # if i_in_flat.dtype != self.compute_dtype: - # i_in_flat = tf.cast(i_in_flat, dtype=self.compute_dtype) - - return i_in_flat + return self.calculate_input_current_from_spikes(rest_of_brain, input_net) def call(self, inputs, states): # lgn_inputs = inputs[:, :self.input_dim] @@ -1099,4 +1190,3 @@ def zero_state(self, batch_size, dtype, with_names=False): return (z0_buf, v0, r0, asc, psc_rise0, psc0), ('z0_buf', 'v0', 'r0', 'asc', 'psc_rise0', 'psc0') else: return z0_buf, v0, r0, asc, psc_rise0, psc0 - diff --git a/bmtk/simulator/dpointnet/custom_ops/__init__.py b/bmtk/simulator/dpointnet/custom_ops/__init__.py new file mode 100644 index 00000000..12f98f85 --- /dev/null +++ b/bmtk/simulator/dpointnet/custom_ops/__init__.py @@ -0,0 +1,16 @@ +from .csr_spike_ops import ( + build_csr_connectivity, + cuda_op_status, + fused_cuda_available, + fused_spike_currents, + reorder_csr_values, +) + + +__all__ = [ + 'build_csr_connectivity', + 'cuda_op_status', + 'fused_cuda_available', + 'fused_spike_currents', + 'reorder_csr_values', +] diff --git a/bmtk/simulator/dpointnet/custom_ops/build.py b/bmtk/simulator/dpointnet/custom_ops/build.py new file mode 100644 index 00000000..fbf707c5 --- /dev/null +++ b/bmtk/simulator/dpointnet/custom_ops/build.py @@ -0,0 +1,15 @@ +import os +import subprocess +import sys +from pathlib import Path + + +def main(): + script = Path(__file__).with_name('build.sh') + environment = os.environ.copy() + environment['PYTHON'] = sys.executable + subprocess.run(['bash', str(script)], check=True, env=environment) + + +if __name__ == '__main__': + main() diff --git a/bmtk/simulator/dpointnet/custom_ops/build.sh b/bmtk/simulator/dpointnet/custom_ops/build.sh new file mode 100755 index 00000000..b3b9b1de --- /dev/null +++ b/bmtk/simulator/dpointnet/custom_ops/build.sh @@ -0,0 +1,81 @@ +#!/usr/bin/env bash +set -euo pipefail + +custom_ops_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +python="${PYTHON:-python}" +prefix="$("$python" -c 'import sys; print(sys.prefix)')" +cxx="${CXX:-$prefix/bin/x86_64-conda-linux-gnu-g++}" +build_dir="${DPOINTNET_CUSTOM_OP_BUILD_DIR:-$custom_ops_dir/build}" +output="${DPOINTNET_CUSTOM_OP_OUTPUT:-$custom_ops_dir/_csr_spike_ops.so}" + +if [[ -n "${NVCC:-}" ]]; then + nvcc="$NVCC" +elif [[ -x "$prefix/bin/nvcc" ]]; then + nvcc="$prefix/bin/nvcc" +elif [[ -n "${CUDA_HOME:-}" && -x "$CUDA_HOME/bin/nvcc" ]]; then + nvcc="$CUDA_HOME/bin/nvcc" +elif [[ -n "${CUDA_PATH:-}" && -x "$CUDA_PATH/bin/nvcc" ]]; then + nvcc="$CUDA_PATH/bin/nvcc" +elif nvcc_path="$(command -v nvcc 2>/dev/null)"; then + nvcc="$nvcc_path" +else + echo "nvcc was not found in the Python environment, CUDA_HOME, CUDA_PATH, or PATH" >&2 + exit 1 +fi +if [[ ! -x "$nvcc" ]]; then + echo "nvcc is not executable at $nvcc" >&2 + exit 1 +fi +if [[ ! -x "$cxx" ]]; then + cxx="${CXX:-c++}" +fi + +mapfile -t tf_compile_flags < <( + "$python" -c 'import tensorflow as tf; print(*tf.sysconfig.get_compile_flags(), sep="\n")' +) +mapfile -t tf_link_flags < <( + "$python" -c 'import tensorflow as tf; print(*tf.sysconfig.get_link_flags(), sep="\n")' +) + +read -r -a cuda_archs <<<"${DPOINTNET_CUDA_ARCHS:-70 75 80 86 89 90}" + +gencode_flags=() +for arch in "${cuda_archs[@]}"; do + gencode_flags+=("-gencode=arch=compute_${arch},code=sm_${arch}") +done +highest_arch="${cuda_archs[${#cuda_archs[@]}-1]}" +gencode_flags+=( + "-gencode=arch=compute_${highest_arch},code=compute_${highest_arch}" +) + +mkdir -p "$build_dir" +"$cxx" -std=c++17 -fPIC -O3 \ + -I"$prefix/include" \ + "${tf_compile_flags[@]}" \ + -c "$custom_ops_dir/csr_spike_ops.cc" \ + -o "$build_dir/csr_spike_ops.o" + +"$nvcc" -ccbin "$cxx" -std=c++17 -x cu -Xcompiler=-fPIC -O3 \ + --expt-relaxed-constexpr \ + -DGOOGLE_CUDA=1 \ + -I"$prefix/include" \ + "${tf_compile_flags[@]}" \ + "${gencode_flags[@]}" \ + -c "$custom_ops_dir/csr_spike_ops.cu.cc" \ + -o "$build_dir/csr_spike_ops.cu.o" + +"$cxx" -shared \ + "$build_dir/csr_spike_ops.o" \ + "$build_dir/csr_spike_ops.cu.o" \ + "${tf_link_flags[@]}" \ + -L"$prefix/lib" -lcudart \ + -Wl,-rpath,"$prefix/lib" \ + -o "$output" + +arch_file="${output%.so}.archs" +{ + printf 'sm=%s\n' "${cuda_archs[*]}" + printf 'ptx=%s\n' "$highest_arch" +} > "$arch_file" + +echo "$output" diff --git a/bmtk/simulator/dpointnet/custom_ops/csr_spike_ops.cc b/bmtk/simulator/dpointnet/custom_ops/csr_spike_ops.cc new file mode 100644 index 00000000..c9b67c5e --- /dev/null +++ b/bmtk/simulator/dpointnet/custom_ops/csr_spike_ops.cc @@ -0,0 +1,88 @@ +#include "tensorflow/core/framework/op.h" +#include "tensorflow/core/framework/shape_inference.h" + +using tensorflow::shape_inference::DimensionHandle; +using tensorflow::shape_inference::InferenceContext; +using tensorflow::shape_inference::ShapeHandle; + +REGISTER_OP("DpointnetCsrReorder") + .Input("values: T") + .Input("metadata: resource") + .Attr("T: {half, float}") + .Attr("Tindex: {uint32, int64}") + .Attr("n_edges: int >= 0") + .Output("reordered: T") + .SetShapeFn([](InferenceContext* context) -> absl::Status { + ShapeHandle values; + TF_RETURN_IF_ERROR(context->WithRank(context->input(0), 1, &values)); + context->set_output(0, values); + return absl::OkStatus(); + }); + +REGISTER_OP("DpointnetCsrSpikeForward") + .Input("spikes: T") + .Input("master_weights: Tmaster") + .Input("metadata: resource") + .Input("weights: T") + .Input("basis: T") + .Attr("T: {half, float}") + .Attr("Tmaster: {half, float}") + .Attr("Tindex: {uint32, int64}") + .Attr("n_post: int >= 1") + .Attr("n_edges: int >= 0") + .Attr("compute_spike_gradient: bool") + .Output("currents: T") + .SetShapeFn([](InferenceContext* context) -> absl::Status { + ShapeHandle spikes; + ShapeHandle basis; + TF_RETURN_IF_ERROR(context->WithRank(context->input(0), 2, &spikes)); + TF_RETURN_IF_ERROR(context->WithRank(context->input(4), 2, &basis)); + int n_post; + TF_RETURN_IF_ERROR(context->GetAttr("n_post", &n_post)); + DimensionHandle flattened_batch; + TF_RETURN_IF_ERROR(context->Multiply( + context->Dim(spikes, 0), n_post, &flattened_batch)); + context->set_output( + 0, context->Matrix(flattened_batch, context->Dim(basis, 1))); + return absl::OkStatus(); + }); + +REGISTER_OP("DpointnetCsrSpikeGrad") + .Input("spikes: T") + .Input("current_grad: T") + .Input("metadata: resource") + .Input("weights: T") + .Input("basis: T") + .Attr("T: {half, float}") + .Attr("Tindex: {uint32, int64}") + .Attr("n_post: int >= 1") + .Attr("n_edges: int >= 0") + .Output("spike_grad: T") + .Output("weight_grad: float") + .SetShapeFn([](InferenceContext* context) -> absl::Status { + ShapeHandle spikes; + ShapeHandle weights; + TF_RETURN_IF_ERROR(context->WithRank(context->input(0), 2, &spikes)); + TF_RETURN_IF_ERROR(context->WithRank(context->input(3), 1, &weights)); + context->set_output(0, spikes); + context->set_output(1, weights); + return absl::OkStatus(); + }); + +REGISTER_OP("DpointnetCsrWeightGrad") + .Input("spikes: T") + .Input("current_grad: T") + .Input("metadata: resource") + .Input("basis: T") + .Attr("T: {half, float}") + .Attr("Tindex: {uint32, int64}") + .Attr("n_post: int >= 1") + .Attr("n_edges: int >= 0") + .Output("weight_grad: float") + .SetShapeFn([](InferenceContext* context) -> absl::Status { + ShapeHandle edge_ids; + int64_t n_edges; + TF_RETURN_IF_ERROR(context->GetAttr("n_edges", &n_edges)); + context->set_output(0, context->Vector(n_edges)); + return absl::OkStatus(); + }); diff --git a/bmtk/simulator/dpointnet/custom_ops/csr_spike_ops.cu.cc b/bmtk/simulator/dpointnet/custom_ops/csr_spike_ops.cu.cc new file mode 100644 index 00000000..b3f0c2fa --- /dev/null +++ b/bmtk/simulator/dpointnet/custom_ops/csr_spike_ops.cu.cc @@ -0,0 +1,589 @@ +#if GOOGLE_CUDA + +#define EIGEN_USE_GPU + +#include +#include + +#include "tensorflow/core/framework/op_kernel.h" +#include "tensorflow/core/framework/register_types.h" +#include "tensorflow/core/framework/resource_mgr.h" +#include "tensorflow/core/framework/resource_var.h" +#include "tensorflow/core/platform/errors.h" +#include "tensorflow/core/platform/types.h" +#include "tensorflow/core/util/gpu_kernel_helper.h" + +namespace tensorflow { + +using GPUDevice = Eigen::GpuDevice; + +template +__device__ inline float ToFloat(T value) { + return static_cast(value); +} + +template +__device__ inline T FromFloat(float value) { + return static_cast(value); +} + +template +__global__ void SetZeroKernel(int64_t count, T* values) { + for (int64_t index : GpuGridRangeX(count)) { + values[index] = FromFloat(0.0f); + } +} + +template +__global__ void CsrReorderKernel( + int64_t count, const T* values, const Index* edge_ids, T* reordered) { + for (int64_t index : GpuGridRangeX(count)) { + reordered[index] = values[edge_ids[index]]; + } +} + +inline int BlockCountFor(int64_t count, int threads, const GPUDevice& device) { + const int64_t requested = (count + threads - 1) / threads; + const int maximum = device.getNumGpuMultiProcessors() * 8; + return static_cast( + std::max(1, std::min(requested, maximum))); +} + +template +__device__ inline void FastAtomicAdd(T* address, T value) { + GpuAtomicAdd(address, value); +} + +template <> +__device__ inline void FastAtomicAdd( + Eigen::half* address, Eigen::half value) { + atomicAdd( + reinterpret_cast<__half*>(address), + __float2half(ToFloat(value))); +} + +template +__global__ void CsrSpikeForwardKernel( + int count, int n_pre, int n_post, int n_basis, const T* spikes, + const Index* post_ids, const T* weights, + const Index* synapse_types, const T* basis, + const Index* row_splits, T* currents) { + const int index = blockIdx.x; + if (index >= count) { + return; + } + const float spike = ToFloat(spikes[index]); + if (spike <= 0.0f) { + return; + } + const int batch = index / n_pre; + const int pre = index - batch * n_pre; + const Index start = row_splits[pre]; + const int64_t work_items = + static_cast(row_splits[pre + 1] - start) * n_basis; + for (int64_t item = threadIdx.x; item < work_items; item += blockDim.x) { + const Index edge = start + static_cast(item / n_basis); + const int receptor = item % n_basis; + const int post = static_cast(post_ids[edge]); + const int synapse_type = static_cast(synapse_types[edge]); + const T weighted_spike = FromFloat(spike) * weights[edge]; + const T value = + weighted_spike * basis[synapse_type * n_basis + receptor]; + const int64_t output_index = + (static_cast(batch) * n_post + post) * n_basis + receptor; + FastAtomicAdd(currents + output_index, value); + } +} + +template +__global__ void CsrSpikeGradKernel( + int count, int n_pre, int n_post, int n_basis, const T* spikes, + const T* current_grad, const Index* post_ids, const T* weights, + const Index* synapse_types, const T* basis, + const Index* row_splits, const Index* edge_ids, T* spike_grad, + float* weight_grad) { + __shared__ float partial_gradients[128]; + const int index = blockIdx.x; + if (index >= count) { + return; + } + const int batch = index / n_pre; + const int pre = index - batch * n_pre; + const float spike = ToFloat(spikes[index]); + float pre_gradient = 0.0f; + for (int64_t edge = + static_cast(row_splits[pre]) + threadIdx.x; + edge < static_cast(row_splits[pre + 1]); + edge += blockDim.x) { + const int post = static_cast(post_ids[edge]); + const int synapse_type = static_cast(synapse_types[edge]); + const int64_t gradient_base = + (static_cast(batch) * n_post + post) * n_basis; + const int basis_base = synapse_type * n_basis; + float edge_gradient = 0.0f; + for (int receptor = 0; receptor < n_basis; ++receptor) { + edge_gradient += + ToFloat(current_grad[gradient_base + receptor]) * + ToFloat(basis[basis_base + receptor]); + } + pre_gradient += edge_gradient * ToFloat(weights[edge]); + if (spike > 0.0f) { + GpuAtomicAdd( + weight_grad + static_cast(edge_ids[edge]), + edge_gradient * spike); + } + } + partial_gradients[threadIdx.x] = pre_gradient; + __syncthreads(); + for (int offset = blockDim.x / 2; offset > 0; offset /= 2) { + if (threadIdx.x < offset) { + partial_gradients[threadIdx.x] += + partial_gradients[threadIdx.x + offset]; + } + __syncthreads(); + } + if (threadIdx.x == 0) { + spike_grad[index] = FromFloat(partial_gradients[0]); + } +} + +template +__global__ void CsrWeightGradKernel( + int count, int n_pre, int n_post, int n_basis, const T* spikes, + const T* current_grad, const Index* post_ids, + const Index* synapse_types, const T* basis, + const Index* row_splits, const Index* edge_ids, + float* weight_grad) { + const int index = blockIdx.x; + if (index >= count) { + return; + } + const float spike = ToFloat(spikes[index]); + if (spike <= 0.0f) { + return; + } + const int batch = index / n_pre; + const int pre = index - batch * n_pre; + for (int64_t edge = + static_cast(row_splits[pre]) + threadIdx.x; + edge < static_cast(row_splits[pre + 1]); + edge += blockDim.x) { + const int post = static_cast(post_ids[edge]); + const int synapse_type = static_cast(synapse_types[edge]); + const int64_t gradient_base = + (static_cast(batch) * n_post + post) * n_basis; + const int basis_base = synapse_type * n_basis; + float edge_gradient = 0.0f; + for (int receptor = 0; receptor < n_basis; ++receptor) { + edge_gradient += + ToFloat(current_grad[gradient_base + receptor]) * + ToFloat(basis[basis_base + receptor]); + } + GpuAtomicAdd( + weight_grad + static_cast(edge_ids[edge]), + edge_gradient * spike); + } +} + +inline void RequireVector( + OpKernelContext* context, const Tensor& tensor, const char* name) { + OP_REQUIRES( + context, TensorShapeUtils::IsVector(tensor.shape()), + errors::InvalidArgument(name, " must be rank 1, got ", tensor.shape())); +} + +inline void RequireMatrix( + OpKernelContext* context, const Tensor& tensor, const char* name) { + OP_REQUIRES( + context, TensorShapeUtils::IsMatrix(tensor.shape()), + errors::InvalidArgument(name, " must be rank 2, got ", tensor.shape())); +} + +template +bool LookupVariable( + OpKernelContext* context, int input, const char* name, + core::RefCountPtr* variable) { + const absl::Status status = + LookupResource(context, HandleFromInput(context, input), variable); + if (!status.ok()) { + context->SetStatus(status); + return false; + } + if (!(*variable)->is_initialized) { + context->SetStatus( + errors::FailedPrecondition(name, " is uninitialized.")); + return false; + } + if ((*variable)->tensor()->dtype() != DataTypeToEnum::v()) { + context->SetStatus(errors::InvalidArgument( + name, " has dtype ", + DataTypeString((*variable)->tensor()->dtype()), ", expected ", + DataTypeString(DataTypeToEnum::v()), ".")); + return false; + } + return true; +} + +template +class DpointnetCsrReorderOp : public OpKernel { + public: + explicit DpointnetCsrReorderOp(OpKernelConstruction* context) + : OpKernel(context) { + OP_REQUIRES_OK(context, context->GetAttr("n_edges", &n_edges_)); + } + + void Compute(OpKernelContext* context) override { + const Tensor& values = context->input(0); + core::RefCountPtr metadata_variable; + if (!LookupVariable( + context, 1, "metadata", &metadata_variable)) { + return; + } + const Tensor& metadata = *metadata_variable->tensor(); + RequireVector(context, values, "values"); + RequireVector(context, metadata, "metadata"); + if (!context->status().ok()) { + return; + } + OP_REQUIRES( + context, values.NumElements() == n_edges_, + errors::InvalidArgument("values length must equal n_edges.")); + OP_REQUIRES( + context, metadata.NumElements() >= 3 * n_edges_ + 1, + errors::InvalidArgument("metadata is too short for n_edges.")); + const int64_t row_splits_size = + metadata.NumElements() - 3 * n_edges_; + const Index* edge_ids = + metadata.flat().data() + 2 * n_edges_ + row_splits_size; + + Tensor* reordered = nullptr; + OP_REQUIRES_OK( + context, + context->allocate_output(0, values.shape(), &reordered)); + const int64_t count = values.NumElements(); + if (count == 0) { + return; + } + const GPUDevice& device = context->eigen_device(); + constexpr int threads = 256; + OP_REQUIRES_OK( + context, + GpuLaunchKernel( + CsrReorderKernel, + BlockCountFor(count, threads, device), threads, 0, + device.stream(), count, values.flat().data(), + edge_ids, reordered->flat().data())); + } + + private: + int64_t n_edges_; +}; + +template +class DpointnetCsrSpikeForwardOp : public OpKernel { + public: + explicit DpointnetCsrSpikeForwardOp(OpKernelConstruction* context) + : OpKernel(context) { + OP_REQUIRES_OK(context, context->GetAttr("n_post", &n_post_)); + OP_REQUIRES_OK(context, context->GetAttr("n_edges", &n_edges_)); + } + + void Compute(OpKernelContext* context) override { + const Tensor& spikes = context->input(0); + const Tensor& master_weights = context->input(1); + const Tensor& weights = context->input(3); + const Tensor& basis = context->input(4); + core::RefCountPtr metadata_variable; + if (!LookupVariable( + context, 2, "metadata", &metadata_variable)) { + return; + } + const Tensor& metadata = *metadata_variable->tensor(); + RequireMatrix(context, spikes, "spikes"); + RequireVector(context, master_weights, "master_weights"); + RequireVector(context, weights, "weights"); + RequireMatrix(context, basis, "basis"); + RequireVector(context, metadata, "metadata"); + if (!context->status().ok()) { + return; + } + + const int64_t batch = spikes.dim_size(0); + const int64_t n_pre = spikes.dim_size(1); + const int64_t n_basis = basis.dim_size(1); + OP_REQUIRES( + context, weights.NumElements() == n_edges_, + errors::InvalidArgument("weights length must equal n_edges.")); + OP_REQUIRES( + context, master_weights.NumElements() == n_edges_, + errors::InvalidArgument( + "master_weights length must equal n_edges.")); + OP_REQUIRES( + context, + metadata.NumElements() == 3 * n_edges_ + n_pre + 1, + errors::InvalidArgument( + "metadata length does not match n_edges and n_pre.")); + OP_REQUIRES( + context, basis.dim_size(0) > 0 && n_basis > 0, + errors::InvalidArgument("basis must have non-zero dimensions.")); + OP_REQUIRES( + context, + batch * n_pre <= std::numeric_limits::max(), + errors::InvalidArgument("Tensor size exceeds CUDA kernel index range.")); + + const Index* metadata_values = metadata.flat().data(); + const Index* post_ids = metadata_values; + const Index* synapse_types = metadata_values + n_edges_; + const Index* row_splits = metadata_values + 2 * n_edges_; + Tensor* currents = nullptr; + OP_REQUIRES_OK( + context, + context->allocate_output( + 0, TensorShape({batch * n_post_, n_basis}), ¤ts)); + const GPUDevice& device = context->eigen_device(); + const int64_t output_count = currents->NumElements(); + constexpr int zero_threads = 256; + OP_REQUIRES_OK( + context, + GpuLaunchKernel( + SetZeroKernel, + BlockCountFor(output_count, zero_threads, device), + zero_threads, 0, device.stream(), output_count, + currents->flat().data())); + + const int work_count = static_cast(batch * n_pre); + OP_REQUIRES_OK( + context, + GpuLaunchKernel( + CsrSpikeForwardKernel, work_count, 128, 0, + device.stream(), + work_count, + static_cast(n_pre), n_post_, static_cast(n_basis), + spikes.flat().data(), post_ids, + weights.flat().data(), synapse_types, + basis.flat().data(), row_splits, + currents->flat().data())); + } + + private: + int n_post_; + int64_t n_edges_; +}; + +template +class DpointnetCsrSpikeGradOp : public OpKernel { + public: + explicit DpointnetCsrSpikeGradOp(OpKernelConstruction* context) + : OpKernel(context) { + OP_REQUIRES_OK(context, context->GetAttr("n_post", &n_post_)); + OP_REQUIRES_OK(context, context->GetAttr("n_edges", &n_edges_)); + } + + void Compute(OpKernelContext* context) override { + const Tensor& spikes = context->input(0); + const Tensor& current_grad = context->input(1); + const Tensor& weights = context->input(3); + const Tensor& basis = context->input(4); + core::RefCountPtr metadata_variable; + if (!LookupVariable( + context, 2, "metadata", &metadata_variable)) { + return; + } + const Tensor& metadata = *metadata_variable->tensor(); + RequireMatrix(context, spikes, "spikes"); + RequireMatrix(context, current_grad, "current_grad"); + RequireVector(context, weights, "weights"); + RequireMatrix(context, basis, "basis"); + RequireVector(context, metadata, "metadata"); + if (!context->status().ok()) { + return; + } + + const int64_t batch = spikes.dim_size(0); + const int64_t n_pre = spikes.dim_size(1); + const int64_t n_basis = basis.dim_size(1); + OP_REQUIRES( + context, weights.NumElements() == n_edges_, + errors::InvalidArgument("weights length must equal n_edges.")); + OP_REQUIRES( + context, + metadata.NumElements() == 3 * n_edges_ + n_pre + 1, + errors::InvalidArgument( + "metadata length does not match n_edges and n_pre.")); + OP_REQUIRES( + context, + current_grad.dim_size(0) == batch * n_post_ && + current_grad.dim_size(1) == n_basis, + errors::InvalidArgument( + "current_grad shape does not match the forward output.")); + OP_REQUIRES( + context, batch * n_pre <= std::numeric_limits::max(), + errors::InvalidArgument("Tensor size exceeds CUDA kernel index range.")); + + const Index* metadata_values = metadata.flat().data(); + const Index* post_ids = metadata_values; + const Index* synapse_types = metadata_values + n_edges_; + const Index* row_splits = metadata_values + 2 * n_edges_; + const Index* edge_ids = row_splits + n_pre + 1; + Tensor* spike_grad = nullptr; + Tensor* weight_grad = nullptr; + OP_REQUIRES_OK( + context, + context->allocate_output(0, spikes.shape(), &spike_grad)); + OP_REQUIRES_OK( + context, + context->allocate_output(1, weights.shape(), &weight_grad)); + const GPUDevice& device = context->eigen_device(); + constexpr int zero_threads = 256; + OP_REQUIRES_OK( + context, + GpuLaunchKernel( + SetZeroKernel, + BlockCountFor(n_edges_, zero_threads, device), + zero_threads, 0, device.stream(), n_edges_, + weight_grad->flat().data())); + + const int work_count = static_cast(batch * n_pre); + OP_REQUIRES_OK( + context, + GpuLaunchKernel( + CsrSpikeGradKernel, work_count, 128, 0, + device.stream(), + work_count, + static_cast(n_pre), n_post_, static_cast(n_basis), + spikes.flat().data(), current_grad.flat().data(), + post_ids, weights.flat().data(), + synapse_types, basis.flat().data(), + row_splits, edge_ids, + spike_grad->flat().data(), weight_grad->flat().data())); + } + + private: + int n_post_; + int64_t n_edges_; +}; + +template +class DpointnetCsrWeightGradOp : public OpKernel { + public: + explicit DpointnetCsrWeightGradOp(OpKernelConstruction* context) + : OpKernel(context) { + OP_REQUIRES_OK(context, context->GetAttr("n_post", &n_post_)); + OP_REQUIRES_OK(context, context->GetAttr("n_edges", &n_edges_)); + } + + void Compute(OpKernelContext* context) override { + const Tensor& spikes = context->input(0); + const Tensor& current_grad = context->input(1); + const Tensor& basis = context->input(3); + core::RefCountPtr metadata_variable; + if (!LookupVariable( + context, 2, "metadata", &metadata_variable)) { + return; + } + const Tensor& metadata = *metadata_variable->tensor(); + RequireMatrix(context, spikes, "spikes"); + RequireMatrix(context, current_grad, "current_grad"); + RequireMatrix(context, basis, "basis"); + RequireVector(context, metadata, "metadata"); + if (!context->status().ok()) { + return; + } + + const int64_t batch = spikes.dim_size(0); + const int64_t n_pre = spikes.dim_size(1); + const int64_t n_basis = basis.dim_size(1); + OP_REQUIRES( + context, + metadata.NumElements() == 3 * n_edges_ + n_pre + 1, + errors::InvalidArgument( + "metadata length does not match n_edges and n_pre.")); + OP_REQUIRES( + context, + current_grad.dim_size(0) == batch * n_post_ && + current_grad.dim_size(1) == n_basis, + errors::InvalidArgument( + "current_grad shape does not match the forward output.")); + OP_REQUIRES( + context, + batch * n_pre <= std::numeric_limits::max(), + errors::InvalidArgument("Tensor size exceeds CUDA kernel index range.")); + + const Index* metadata_values = metadata.flat().data(); + const Index* post_ids = metadata_values; + const Index* synapse_types = metadata_values + n_edges_; + const Index* row_splits = metadata_values + 2 * n_edges_; + const Index* edge_ids = row_splits + n_pre + 1; + Tensor* weight_grad = nullptr; + OP_REQUIRES_OK( + context, + context->allocate_output( + 0, TensorShape({n_edges_}), &weight_grad)); + const GPUDevice& device = context->eigen_device(); + constexpr int zero_threads = 256; + OP_REQUIRES_OK( + context, + GpuLaunchKernel( + SetZeroKernel, + BlockCountFor(n_edges_, zero_threads, device), + zero_threads, 0, device.stream(), n_edges_, + weight_grad->flat().data())); + + const int work_count = static_cast(batch * n_pre); + OP_REQUIRES_OK( + context, + GpuLaunchKernel( + CsrWeightGradKernel, work_count, 128, 0, + device.stream(), + work_count, + static_cast(n_pre), n_post_, static_cast(n_basis), + spikes.flat().data(), current_grad.flat().data(), + post_ids, synapse_types, basis.flat().data(), + row_splits, edge_ids, + weight_grad->flat().data())); + } + + private: + int n_post_; + int64_t n_edges_; +}; + +#define REGISTER_GPU_KERNELS(T, Index) \ + REGISTER_KERNEL_BUILDER( \ + Name("DpointnetCsrReorder") \ + .Device(DEVICE_GPU) \ + .TypeConstraint("T") \ + .TypeConstraint("Tindex"), \ + DpointnetCsrReorderOp); \ + REGISTER_KERNEL_BUILDER( \ + Name("DpointnetCsrSpikeForward") \ + .Device(DEVICE_GPU) \ + .TypeConstraint("T") \ + .TypeConstraint("Tindex"), \ + DpointnetCsrSpikeForwardOp); \ + REGISTER_KERNEL_BUILDER( \ + Name("DpointnetCsrSpikeGrad") \ + .Device(DEVICE_GPU) \ + .TypeConstraint("T") \ + .TypeConstraint("Tindex"), \ + DpointnetCsrSpikeGradOp); \ + REGISTER_KERNEL_BUILDER( \ + Name("DpointnetCsrWeightGrad") \ + .Device(DEVICE_GPU) \ + .TypeConstraint("T") \ + .TypeConstraint("Tindex"), \ + DpointnetCsrWeightGradOp); + +#define REGISTER_GPU_KERNELS_FOR_TYPE(T) \ + REGISTER_GPU_KERNELS(T, uint32); \ + REGISTER_GPU_KERNELS(T, int64_t); + +TF_CALL_half(REGISTER_GPU_KERNELS_FOR_TYPE); +TF_CALL_float(REGISTER_GPU_KERNELS_FOR_TYPE); + +#undef REGISTER_GPU_KERNELS_FOR_TYPE +#undef REGISTER_GPU_KERNELS + +} // namespace tensorflow + +#endif // GOOGLE_CUDA diff --git a/bmtk/simulator/dpointnet/custom_ops/csr_spike_ops.py b/bmtk/simulator/dpointnet/custom_ops/csr_spike_ops.py new file mode 100644 index 00000000..43860cba --- /dev/null +++ b/bmtk/simulator/dpointnet/custom_ops/csr_spike_ops.py @@ -0,0 +1,330 @@ +import os +import itertools +import weakref +from collections.abc import Mapping +from pathlib import Path + +import numpy as np +import tensorflow as tf +from tensorflow.python.framework import ops + + +_LIBRARY_PATH = Path(__file__).with_name('_csr_spike_ops.so') +_ARCHITECTURE_PATH = Path(__file__).with_name('_csr_spike_ops.archs') +_OPS = None +_LOAD_ERROR = None +_RESOURCE_COUNTER = itertools.count() +_RESOURCE_CONTAINER = 'bmtk_dpointnet_csr' + + +def _read_built_architectures(): + if not _ARCHITECTURE_PATH.exists(): + return (), None + values = {} + for line in _ARCHITECTURE_PATH.read_text().splitlines(): + key, separator, value = line.partition('=') + if separator: + values[key] = value + sm_architectures = tuple( + int(value) for value in values.get('sm', '').split() + ) + ptx_architecture = values.get('ptx') + return sm_architectures, ( + int(ptx_architecture) if ptx_architecture else None + ) + + +_SM_ARCHITECTURES, _PTX_ARCHITECTURE = _read_built_architectures() + + +def _gpu_compatibility_error(): + visible_gpus = tf.config.get_visible_devices('GPU') + if len(visible_gpus) != 1: + return f'fused CUDA requires exactly one visible GPU; found {len(visible_gpus)}' + if not _SM_ARCHITECTURES and _PTX_ARCHITECTURE is None: + return f'build architecture metadata is missing at {_ARCHITECTURE_PATH}' + details = tf.config.experimental.get_device_details(visible_gpus[0]) + capability = details.get('compute_capability') + if capability is None: + return f'compute capability is unavailable for {visible_gpus[0].name}' + architecture = int(capability[0]) * 10 + int(capability[1]) + if architecture in _SM_ARCHITECTURES: + return None + if _PTX_ARCHITECTURE is not None and architecture >= _PTX_ARCHITECTURE: + return None + return ( + f'GPU compute capability sm_{architecture} is incompatible with ' + f'sm targets {_SM_ARCHITECTURES} and compute_{_PTX_ARCHITECTURE} PTX' + ) + + +def _destroy_metadata_resource(handle): + try: + tf.raw_ops.DestroyResourceOp( + resource=handle, ignore_lookup_error=True + ) + except (tf.errors.OpError, RuntimeError): + pass + + +class CsrConnectivity(Mapping): + def __init__(self, values): + self._values = values + self._finalizer = weakref.finalize( + self, _destroy_metadata_resource, values['metadata_handle'] + ) + + def __getitem__(self, key): + return self._values[key] + + def __iter__(self): + return iter(self._values) + + def __len__(self): + return len(self._values) + + def close(self): + self._finalizer() + + +def _environment_flag(name): + value = os.environ.get(name, '').strip().lower() + return value in ('1', 'true', 'yes', 'on') + + +if not _environment_flag('BMTK_DPOINTNET_DISABLE_FUSED_CUDA'): + if _LIBRARY_PATH.exists(): + try: + _OPS = tf.load_op_library(str(_LIBRARY_PATH)) + except (tf.errors.NotFoundError, OSError) as exc: + _LOAD_ERROR = exc + else: + _LOAD_ERROR = FileNotFoundError( + f'Fused DPointNet CUDA library does not exist at {_LIBRARY_PATH}.' + ) + + +def cuda_op_status(): + if _environment_flag('BMTK_DPOINTNET_DISABLE_FUSED_CUDA'): + return 'disabled by BMTK_DPOINTNET_DISABLE_FUSED_CUDA' + if _OPS is not None: + compatibility_error = _gpu_compatibility_error() + if compatibility_error is not None: + return f'loaded, but {compatibility_error}' + return f'loaded from {_LIBRARY_PATH}' + return str(_LOAD_ERROR) + + +def fused_cuda_available(): + return _OPS is not None and _gpu_compatibility_error() is None + + +def _csr_index_dtype( + n_edges, n_source_neurons, n_target_neurons, n_synapse_types): + maximum_value = max( + int(n_edges), + int(n_source_neurons) - 1, + int(n_target_neurons) - 1, + int(n_synapse_types) - 1, + ) + return tf.uint32 if maximum_value <= np.iinfo(np.uint32).max else tf.int64 + + +def _create_metadata_resource(metadata, index_dtype): + resource_name = f'csr_{os.getpid()}_{next(_RESOURCE_COUNTER)}' + metadata_handle = tf.raw_ops.VarHandleOp( + dtype=index_dtype, + shape=metadata.shape, + container=_RESOURCE_CONTAINER, + shared_name=resource_name, + ) + tf.raw_ops.AssignVariableOp( + resource=metadata_handle, + value=tf.constant(metadata, dtype=index_dtype), + ) + return metadata_handle + + +def build_csr_connectivity( + indices, + synapse_types, + n_source_neurons, + n_target_neurons, + n_synapse_types): + indices = np.asarray(indices) + synapse_types = np.asarray(synapse_types) + for dimension, name in ( + (n_source_neurons, 'n_source_neurons'), + (n_target_neurons, 'n_target_neurons'), + (n_synapse_types, 'n_synapse_types')): + if not isinstance(dimension, (int, np.integer)) or dimension <= 0: + raise ValueError(f'{name} must be a positive integer.') + if indices.ndim != 2 or indices.shape[1] != 2: + raise ValueError(f'indices must have shape [n_edges, 2], got {indices.shape}.') + if synapse_types.shape != (indices.shape[0],): + raise ValueError( + 'synapse_types must contain one value per edge, got ' + f'{synapse_types.shape} for {indices.shape[0]} edges.' + ) + for values, name in ( + (indices, 'indices'), + (synapse_types, 'synapse_types')): + if not np.issubdtype(values.dtype, np.number): + raise TypeError(f'{name} must contain numeric integer values.') + if not np.all(np.isfinite(values)) or not np.all(values == np.floor(values)): + raise ValueError(f'{name} must contain finite integer values.') + int64_info = np.iinfo(np.int64) + if np.any(values < int64_info.min) or np.any(values > int64_info.max): + raise ValueError(f'{name} values must be within the int64 range.') + + indices = indices.astype(np.int64, copy=False) + synapse_types = synapse_types.astype(np.int64, copy=False) + pre_ids = indices[:, 1] + if np.any(pre_ids < 0) or np.any(pre_ids >= n_source_neurons): + raise ValueError('Presynaptic indices are outside the declared source dimension.') + post_ids = indices[:, 0] + if np.any(post_ids < 0) or np.any(post_ids >= n_target_neurons): + raise ValueError('Postsynaptic indices are outside the declared target dimension.') + if np.any(synapse_types < 0) or np.any(synapse_types >= n_synapse_types): + raise ValueError('Synapse type indices are outside the basis table.') + index_dtype = _csr_index_dtype( + indices.shape[0], + n_source_neurons, + n_target_neurons, + n_synapse_types, + ) + numpy_index_dtype = index_dtype.as_numpy_dtype + edge_ids = np.argsort(pre_ids, kind='stable').astype( + numpy_index_dtype, copy=False + ) + sorted_pre_ids = pre_ids[edge_ids] + counts = np.bincount(sorted_pre_ids, minlength=n_source_neurons) + row_splits = np.empty( + n_source_neurons + 1, dtype=numpy_index_dtype + ) + row_splits[0] = 0 + np.cumsum(counts, dtype=np.int64, out=row_splits[1:]) + + device = '/GPU:0' if tf.config.get_visible_devices('GPU') else '/CPU:0' + with tf.device(device): + post_ids = indices[edge_ids, 0].astype( + numpy_index_dtype, copy=False + ) + sorted_synapse_types = synapse_types[edge_ids].astype( + numpy_index_dtype, copy=False + ) + metadata = np.concatenate( + (post_ids, sorted_synapse_types, row_splits, edge_ids) + ) + metadata_handle = _create_metadata_resource( + metadata, index_dtype + ) + return CsrConnectivity({ + 'metadata_handle': metadata_handle, + 'index_dtype': index_dtype.name, + 'n_edges': int(indices.shape[0]), + 'n_sources': int(n_source_neurons), + 'n_post': int(n_target_neurons), + 'n_synapse_types': int(n_synapse_types), + }) + + +def reorder_csr_values(values, connectivity): + if _OPS is None: + raise RuntimeError(f'Fused DPointNet CUDA operator is unavailable: {cuda_op_status()}') + return _OPS.dpointnet_csr_reorder( + values, + connectivity['metadata_handle'], + Tindex=tf.dtypes.as_dtype(connectivity['index_dtype']), + n_edges=connectivity['n_edges'], + ) + + +@ops.RegisterGradient('DpointnetCsrSpikeForward') +def _fused_spike_currents_gradient(op, current_grad): + compute_spike_gradient = op.get_attr('compute_spike_gradient') + n_post = op.get_attr('n_post') + index_dtype = op.get_attr('Tindex') + if compute_spike_gradient: + spike_grad, weight_grad = _OPS.dpointnet_csr_spike_grad( + op.inputs[0], + current_grad, + op.inputs[2], + op.inputs[3], + op.inputs[4], + Tindex=index_dtype, + n_post=n_post, + n_edges=op.get_attr('n_edges'), + ) + else: + spike_grad = None + weight_grad = _OPS.dpointnet_csr_weight_grad( + op.inputs[0], + current_grad, + op.inputs[2], + op.inputs[4], + Tindex=index_dtype, + n_post=n_post, + n_edges=op.get_attr('n_edges'), + ) + return ( + spike_grad, + tf.cast(weight_grad, op.inputs[1].dtype), + None, + None, + None, + ) + + +def fused_spike_currents( + spikes, + master_weights, + csr_weights, + connectivity, + basis, + n_post, + compute_spike_gradient): + if _OPS is None: + raise RuntimeError(f'Fused DPointNet CUDA operator is unavailable: {cuda_op_status()}') + if spikes.dtype not in (tf.float16, tf.float32): + raise TypeError( + f'Fused DPointNet CUDA operator requires float16 or float32 spikes, got {spikes.dtype}.' + ) + if csr_weights.dtype != spikes.dtype or basis.dtype != spikes.dtype: + raise TypeError('spikes, csr_weights, and basis must have the same dtype.') + if master_weights.shape.rank != 1 or csr_weights.shape.rank != 1: + raise ValueError('master_weights and csr_weights must be rank 1.') + n_edges = connectivity['n_edges'] + if (master_weights.shape[0] is not None + and master_weights.shape[0] != n_edges): + raise ValueError( + f'master_weights has {master_weights.shape[0]} values; ' + f'connectivity has {n_edges} edges.' + ) + if csr_weights.shape[0] is not None and csr_weights.shape[0] != n_edges: + raise ValueError( + f'csr_weights has {csr_weights.shape[0]} values; ' + f'connectivity has {n_edges} edges.' + ) + + if connectivity['n_post'] != n_post: + raise ValueError( + f'Connectivity targets {connectivity["n_post"]} neurons, got n_post={n_post}.' + ) + if connectivity['n_synapse_types'] != basis.shape[0]: + raise ValueError( + 'Connectivity synapse types do not match the basis table: ' + f'{connectivity["n_synapse_types"]} != {basis.shape[0]}.' + ) + + return _OPS.dpointnet_csr_spike_forward( + spikes, + master_weights, + connectivity['metadata_handle'], + csr_weights, + basis, + Tindex=tf.dtypes.as_dtype(connectivity['index_dtype']), + n_post=n_post, + n_edges=connectivity['n_edges'], + compute_spike_gradient=compute_spike_gradient, + ) diff --git a/bmtk/simulator/dpointnet/network_adaptor.py b/bmtk/simulator/dpointnet/network_adaptor.py index fc9e8e6c..c40aa1e7 100644 --- a/bmtk/simulator/dpointnet/network_adaptor.py +++ b/bmtk/simulator/dpointnet/network_adaptor.py @@ -14,11 +14,7 @@ def lex_sort_order_np(indices): - max_ind = np.max(indices) + 1 - if np.iinfo(indices.dtype).max < max_ind * (max_ind + 1) : - indices = indices.astype(np.int64) - q = indices[:, 0] * max_ind + indices[:, 1] - return np.argsort(q) + return np.lexsort((indices[:, 1], indices[:, 0])) def lex_sort_indices_np(indices, *arrays): diff --git a/bmtk/simulator/dpointnet/optimizers.py b/bmtk/simulator/dpointnet/optimizers.py index 559e7648..ffd45ddf 100644 --- a/bmtk/simulator/dpointnet/optimizers.py +++ b/bmtk/simulator/dpointnet/optimizers.py @@ -23,10 +23,16 @@ def build_learning_rate(lr_schedule, **lr_params): def optimizer_supports_loss_scaling(optimizer): - return ( - hasattr(optimizer, 'scale_loss') or - (hasattr(optimizer, 'get_scaled_loss') and hasattr(optimizer, 'get_unscaled_gradients')) + if hasattr(optimizer, 'get_scaled_loss') and hasattr(optimizer, 'get_unscaled_gradients'): + return True + # Keras 3 exposes scale_loss() on every optimizer, even when + # loss_scale_factor is None and the method is a no-op. + loss_scale_optimizer = getattr( + tf.keras.mixed_precision, 'LossScaleOptimizer', () ) + if isinstance(optimizer, loss_scale_optimizer): + return True + return getattr(optimizer, 'loss_scale_factor', None) is not None def scale_loss_for_optimizer(optimizer, loss): diff --git a/bmtk/simulator/dpointnet/rnn_model.py b/bmtk/simulator/dpointnet/rnn_model.py index d04784c0..0871bd00 100644 --- a/bmtk/simulator/dpointnet/rnn_model.py +++ b/bmtk/simulator/dpointnet/rnn_model.py @@ -328,6 +328,16 @@ def build(self, rebuild=False, seq_len=None, dtype=tf.float32, use_dummy_state_i if self._model_built and not rebuild: io.log_debug('Model already built. Skipping.') return + if rebuild and self._cell is not None: + close_fused_cuda = getattr(self._cell, 'close_fused_cuda', None) + if close_fused_cuda is not None: + close_fused_cuda() + self.model = None + self.extractor_model = None + self._state_only_model = None + self._rsnn_layer = None + self.zero_state = None + self._model_built = False _batch_size = batch_size or self.adjusted_batch_size _seq_len = seq_len or self.adjusted_seq_len @@ -571,6 +581,10 @@ def run_inference(self, spikes=None, initial_state=None, inference=None, **kwarg def cleanup(self): for inference in self._inferences: inference.close() + if self._cell is not None: + close_fused_cuda = getattr(self._cell, 'close_fused_cuda', None) + if close_fused_cuda is not None: + close_fused_cuda() def train(self, training_engine=None): if self.extractor_model is None: @@ -584,12 +598,16 @@ def train(self, training_engine=None): ## Build the optimizer (in strategy scope so its slot variables are created correctly) with self.strategy.scope(): optimizer = training_engine.optimizer - if self.dtype == 'float16' and not optimizers.optimizer_supports_loss_scaling(optimizer): + if self.dtype == tf.float16 and not optimizers.optimizer_supports_loss_scaling(optimizer): # Prevent gradient underflow in mixed-float16 training. The wrapped optimizer # must be built and applied as the active optimizer, especially under Keras 3. from tensorflow.keras import mixed_precision as mixed_precision_module optimizer = mixed_precision_module.LossScaleOptimizer(optimizer) training_engine.set_optimizer(optimizer) + io.log_info( + 'Mixed-precision optimizer loss scaling enabled with ' + 'LossScaleOptimizer.' + ) optimizer.build(self.model.trainable_variables) training_engine.train() diff --git a/docs/autodocs/source/dpointnet_guide.rst b/docs/autodocs/source/dpointnet_guide.rst index d0ba7fd8..be72bb94 100644 --- a/docs/autodocs/source/dpointnet_guide.rst +++ b/docs/autodocs/source/dpointnet_guide.rst @@ -31,6 +31,32 @@ or without a gpu: $ pip install tensorflow +Optional fused CUDA operator +---------------------------- + +DPointNet can use a fused CUDA operator for recurrent and input synaptic currents. This optional operator +requires an NVIDIA CUDA toolkit with ``nvcc``, a C++17 compiler, and a GPU-enabled TensorFlow installation. +Build it from the same environment in which BMTK and TensorFlow are installed: + +:: + + $ python -m bmtk.simulator.dpointnet.custom_ops.build + +This module form ensures that the build uses the active Python environment and does not require a console script on +``PATH``. Installing this version of BMTK also creates ``bmtk-build-dpointnet-cuda`` in the environment's executable +directory; the shortcut is available when that environment is activated. + +The build targets compute capabilities 7.0, 7.5, 8.0, 8.6, 8.9, and 9.0 by default, with PTX for the +highest target. To build for a different set of architectures, provide space-separated architecture numbers: + +:: + + $ DPOINTNET_CUDA_ARCHS="80 86 90" python -m bmtk.simulator.dpointnet.custom_ops.build + +The operator requires exactly one visible GPU. Set ``use_fused_cuda`` to ``true`` in ``rnn_cell_params`` to +require the operator, or to ``"auto"`` to use it when available and otherwise fall back to TensorFlow. The +default is ``false``. Rebuild the operator after changing TensorFlow or CUDA installations. + Overview ======== @@ -207,6 +233,9 @@ the `GLIF point-neuron models + * - use_fused_cuda + - Use the optional fused CUDA synaptic-current operator. ``True`` requires it; ``"auto"`` falls back to TensorFlow when unavailable. + - False Setting the Network Model diff --git a/setup.py b/setup.py index 65e47acc..f265fb3b 100644 --- a/setup.py +++ b/setup.py @@ -70,6 +70,15 @@ def read(*filenames, **kwargs): packages=find_packages(exclude=['bmtk.tests', 'bmtk.tests.*', '*tests*']), # package_data={'': ['*.md', '*.txt', '*.cfg', '**/*.json', '**/*.hoc']}, include_package_data=True, + package_data={ + 'bmtk.simulator.dpointnet.custom_ops': ['*.cc', '*.cu.cc', '*.sh'], + }, + entry_points={ + 'console_scripts': [ + 'bmtk-build-dpointnet-cuda=' + 'bmtk.simulator.dpointnet.custom_ops.build:main', + ], + }, platforms='any' ) diff --git a/tests/simulator/dpointnet/test_cuda_custom_ops.py b/tests/simulator/dpointnet/test_cuda_custom_ops.py new file mode 100644 index 00000000..e19ca3fe --- /dev/null +++ b/tests/simulator/dpointnet/test_cuda_custom_ops.py @@ -0,0 +1,424 @@ +import numpy as np +import pytest + +tf = pytest.importorskip('tensorflow') + +from bmtk.simulator.dpointnet.custom_ops import ( + build_csr_connectivity, + fused_cuda_available, + fused_spike_currents, + reorder_csr_values, +) +from bmtk.simulator.dpointnet.custom_ops import csr_spike_ops +from bmtk.simulator.dpointnet.custom_ops.csr_spike_ops import ( + _csr_index_dtype, +) +from bmtk.simulator.dpointnet.cell_models.glif3_cell import ( + _fused_cuda_dtype_error, + _validate_fused_cuda_option, +) + + +INDICES = np.array([ + [0, 0], + [1, 0], + [1, 2], + [0, 1], +], dtype=np.int64) +SYNAPSE_TYPES = np.array([0, 1, 0, 1], dtype=np.int64) + + +def _reference_currents(spikes, master_weights, basis): + post_ids = tf.constant(INDICES[:, 0], tf.int32) + pre_ids = tf.constant(INDICES[:, 1], tf.int32) + edge_basis = tf.gather(basis, SYNAPSE_TYPES) + compute_weights = tf.cast(master_weights, spikes.dtype) + batch_outputs = [] + for batch in range(spikes.shape[0]): + edge_values = ( + tf.gather(spikes[batch], pre_ids)[:, tf.newaxis] + * compute_weights[:, tf.newaxis] + * edge_basis + ) + batch_outputs.append( + tf.math.unsorted_segment_sum(edge_values, post_ids, 2) + ) + return tf.reshape(tf.stack(batch_outputs), [-1, basis.shape[1]]) + + +def _metadata_values(connectivity): + return tf.raw_ops.ReadVariableOp( + resource=connectivity['metadata_handle'], + dtype=tf.dtypes.as_dtype(connectivity['index_dtype']), + ) + + +def test_build_csr_connectivity_groups_edges_by_source(): + connectivity = build_csr_connectivity(INDICES, SYNAPSE_TYPES, 3, 2, 2) + metadata = _metadata_values(connectivity).numpy() + + np.testing.assert_array_equal(metadata[:4], [0, 1, 0, 1]) + np.testing.assert_array_equal(metadata[4:8], [0, 1, 1, 0]) + np.testing.assert_array_equal(metadata[8:12], [0, 2, 3, 4]) + np.testing.assert_array_equal(metadata[12:], [0, 1, 3, 2]) + assert connectivity['index_dtype'] == 'uint32' + assert _metadata_values(connectivity).dtype == tf.uint32 + + +def test_csr_index_dtype_retains_int64_beyond_uint32(): + uint32_max = np.iinfo(np.uint32).max + + assert _csr_index_dtype(uint32_max, 1, 1, 1) == tf.uint32 + assert _csr_index_dtype(uint32_max + 1, 1, 1, 1) == tf.int64 + + +@pytest.mark.parametrize( + ('indices', 'synapse_types', 'message'), + [ + (np.array([[2, 0]]), np.array([0]), 'Postsynaptic indices'), + (np.array([[0, 0]]), np.array([2]), 'Synapse type indices'), + ], +) +def test_build_csr_connectivity_rejects_unsafe_indices( + indices, synapse_types, message): + with pytest.raises(ValueError, match=message): + build_csr_connectivity(indices, synapse_types, 1, 2, 2) + + +@pytest.mark.parametrize( + ('indices', 'synapse_types'), + [ + (np.array([[0.0, 0.5]]), np.array([0])), + (np.array([[0, 0]]), np.array([0.5])), + (np.array([[0, 0]]), np.array([np.nan])), + ], +) +def test_build_csr_connectivity_rejects_non_integer_metadata( + indices, synapse_types): + with pytest.raises(ValueError, match='finite integer values'): + build_csr_connectivity(indices, synapse_types, 1, 1, 1) + + +@pytest.mark.parametrize( + ('indices', 'synapse_types', 'n_source_neurons', 'n_synapse_types'), + [ + ( + np.array([[0, np.iinfo(np.uint64).max]], dtype=np.uint64), + np.array([0]), + int(np.iinfo(np.uint64).max) + 1, + 1, + ), + ( + np.array([[0, 0]]), + np.array([np.iinfo(np.uint64).max], dtype=np.uint64), + 1, + int(np.iinfo(np.uint64).max) + 1, + ), + ], +) +def test_build_csr_connectivity_rejects_metadata_outside_int64( + indices, synapse_types, n_source_neurons, n_synapse_types): + with pytest.raises(ValueError, match='int64 range'): + build_csr_connectivity( + indices, + synapse_types, + n_source_neurons, + 1, + n_synapse_types, + ) + + +@pytest.mark.parametrize( + 'value', + [ + 'auto', + np.str_('auto'), + b'auto', + np.bytes_('auto'), + np.array('auto'), + np.array(b'auto'), + ], +) +def test_fused_cuda_option_accepts_string_scalars(value): + assert _validate_fused_cuda_option(value) == 'auto' + + +@pytest.mark.parametrize('value', [0, 1, np.bool_(False), np.bool_(True)]) +def test_fused_cuda_option_rejects_non_boolean_lookalikes(value): + with pytest.raises(ValueError, match='use_fused_cuda'): + _validate_fused_cuda_option(value) + + +def test_fused_cuda_dtype_error_describes_unsupported_policy(): + assert _fused_cuda_dtype_error(tf.float16, tf.float32) is None + message = _fused_cuda_dtype_error(tf.bfloat16, tf.float32) + assert 'compute_dtype=bfloat16' in message + assert 'variable_dtype=float32' in message + + +def test_fused_availability_respects_visible_devices(monkeypatch): + monkeypatch.setattr(tf.config, 'get_visible_devices', lambda device_type: []) + + assert not csr_spike_ops.fused_cuda_available() + + +def test_fused_availability_rejects_multiple_visible_gpus(monkeypatch): + monkeypatch.setattr( + tf.config, 'get_visible_devices', lambda device_type: [object(), object()] + ) + + assert not csr_spike_ops.fused_cuda_available() + + +def test_fused_availability_rejects_incompatible_gpu(monkeypatch): + gpu = object() + monkeypatch.setattr( + tf.config, 'get_visible_devices', lambda device_type: [gpu] + ) + monkeypatch.setattr( + tf.config.experimental, + 'get_device_details', + lambda device: {'compute_capability': (6, 0)}, + ) + + assert not csr_spike_ops.fused_cuda_available() + + +def test_csr_connectivity_close_releases_resource(): + connectivity = build_csr_connectivity(INDICES, SYNAPSE_TYPES, 3, 2, 2) + handle = connectivity['metadata_handle'] + + connectivity.close() + + with pytest.raises((tf.errors.NotFoundError, tf.errors.FailedPreconditionError)): + tf.raw_ops.ReadVariableOp(resource=handle, dtype=tf.uint32) + + +@pytest.mark.skipif(not fused_cuda_available(), reason='Fused CUDA op is unavailable.') +@pytest.mark.parametrize('dtype', [tf.float16, tf.float32]) +@pytest.mark.parametrize( + 'spike_values', + [ + [[1.0, 0.0, 2.0], [0.0, 3.0, 0.0]], + [[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]], + ], +) +def test_fused_recurrent_currents_match_forward_and_gradients(dtype, spike_values): + connectivity = build_csr_connectivity(INDICES, SYNAPSE_TYPES, 3, 2, 2) + master_weights = tf.Variable([1.0, 2.0, 3.0, 4.0], dtype=tf.float32) + csr_weights = reorder_csr_values( + tf.cast(master_weights, dtype), connectivity + ) + basis = tf.constant([[1.0, 0.5], [0.25, 2.0]], dtype=dtype) + spikes = tf.Variable(spike_values, dtype=dtype) + + with tf.GradientTape() as fused_tape: + fused = fused_spike_currents( + spikes, + master_weights, + csr_weights, + connectivity, + basis, + n_post=2, + compute_spike_gradient=True, + ) + fused_loss = tf.reduce_sum(fused) + fused_gradients = fused_tape.gradient( + fused_loss, [spikes, master_weights] + ) + + with tf.GradientTape() as reference_tape: + reference = _reference_currents(spikes, master_weights, basis) + reference_loss = tf.reduce_sum(reference) + reference_gradients = reference_tape.gradient( + reference_loss, [spikes, master_weights] + ) + + tolerance = 2e-3 if dtype == tf.float16 else 1e-6 + np.testing.assert_allclose( + fused.numpy(), reference.numpy(), rtol=tolerance, atol=tolerance + ) + for fused_gradient, reference_gradient in zip( + fused_gradients, reference_gradients): + np.testing.assert_allclose( + fused_gradient.numpy(), + reference_gradient.numpy(), + rtol=tolerance, + atol=tolerance, + ) + + +@pytest.mark.skipif(not fused_cuda_available(), reason='Fused CUDA op is unavailable.') +def test_fused_input_currents_preserve_counts_and_only_differentiate_weights(): + connectivity = build_csr_connectivity(INDICES, SYNAPSE_TYPES, 3, 2, 2) + master_weights = tf.Variable([1.0, 2.0, 3.0, 4.0], dtype=tf.float32) + csr_weights = reorder_csr_values(master_weights, connectivity) + basis = tf.constant([[1.0, 0.5], [0.25, 2.0]], dtype=tf.float32) + spike_counts = tf.Variable( + [[2.0, 0.0, 3.0], [0.0, 4.0, 0.0]], dtype=tf.float32 + ) + + with tf.GradientTape() as tape: + currents = fused_spike_currents( + spike_counts, + master_weights, + csr_weights, + connectivity, + basis, + n_post=2, + compute_spike_gradient=False, + ) + loss = tf.reduce_sum(currents) + spike_gradient, weight_gradient = tape.gradient( + loss, [spike_counts, master_weights] + ) + + with tf.GradientTape() as reference_tape: + reference = _reference_currents(spike_counts, master_weights, basis) + reference_loss = tf.reduce_sum(reference) + reference_weight_gradient = reference_tape.gradient( + reference_loss, master_weights + ) + + assert spike_gradient is None + np.testing.assert_allclose(currents.numpy(), reference.numpy()) + np.testing.assert_allclose( + weight_gradient.numpy(), reference_weight_gradient.numpy() + ) + + +@pytest.mark.skipif(not fused_cuda_available(), reason='Fused CUDA op is unavailable.') +def test_fused_currents_execute_inside_tf_function(): + connectivity = build_csr_connectivity(INDICES, SYNAPSE_TYPES, 3, 2, 2) + master_weights = tf.Variable([1.0, 2.0, 3.0, 4.0], dtype=tf.float32) + csr_weights = reorder_csr_values(master_weights, connectivity) + basis = tf.constant([[1.0, 0.5], [0.25, 2.0]], dtype=tf.float32) + + @tf.function + def run(spikes): + return fused_spike_currents( + spikes, + master_weights, + csr_weights, + connectivity, + basis, + n_post=2, + compute_spike_gradient=True, + ) + + result = run(tf.constant([[1.0, 0.0, 2.0]], tf.float32)) + reference = _reference_currents( + tf.constant([[1.0, 0.0, 2.0]], tf.float32), + master_weights, + basis, + ) + np.testing.assert_allclose(result.numpy(), reference.numpy()) + + +@pytest.mark.skipif(not fused_cuda_available(), reason='Fused CUDA op is unavailable.') +def test_fused_currents_support_int64_connectivity(): + connectivity = build_csr_connectivity(INDICES, SYNAPSE_TYPES, 3, 2, 2) + metadata = _metadata_values(connectivity).numpy().astype(np.int64) + with tf.device('/GPU:0'): + metadata_handle = csr_spike_ops._create_metadata_resource( + metadata, tf.int64 + ) + connectivity = csr_spike_ops.CsrConnectivity({ + **connectivity, + 'metadata_handle': metadata_handle, + 'index_dtype': 'int64', + }) + master_weights = tf.Variable([1.0, 2.0, 3.0, 4.0], dtype=tf.float32) + csr_weights = reorder_csr_values(master_weights, connectivity) + basis = tf.constant([[1.0, 0.5], [0.25, 2.0]], dtype=tf.float32) + spikes = tf.constant([[1.0, 0.0, 2.0]], dtype=tf.float32) + + result = fused_spike_currents( + spikes, + master_weights, + csr_weights, + connectivity, + basis, + n_post=2, + compute_spike_gradient=True, + ) + + np.testing.assert_allclose( + result.numpy(), + _reference_currents(spikes, master_weights, basis).numpy(), + ) + + +@pytest.mark.skipif(not fused_cuda_available(), reason='Fused CUDA op is unavailable.') +def test_fused_currents_support_empty_connectivity_and_gradients(): + connectivity = build_csr_connectivity( + np.empty((0, 2), dtype=np.int64), + np.empty((0,), dtype=np.int64), + 3, + 2, + 2, + ) + master_weights = tf.Variable([], dtype=tf.float32) + csr_weights = reorder_csr_values(master_weights, connectivity) + basis = tf.constant([[1.0, 0.5], [0.25, 2.0]], dtype=tf.float32) + spikes = tf.Variable([[1.0, 0.0, 2.0]], dtype=tf.float32) + + with tf.GradientTape() as tape: + currents = fused_spike_currents( + spikes, + master_weights, + csr_weights, + connectivity, + basis, + n_post=2, + compute_spike_gradient=True, + ) + loss = tf.reduce_sum(currents) + spike_gradient, weight_gradient = tape.gradient( + loss, [spikes, master_weights] + ) + + np.testing.assert_array_equal(currents.numpy(), np.zeros((2, 2))) + np.testing.assert_array_equal(spike_gradient.numpy(), np.zeros((1, 3))) + assert weight_gradient.shape == (0,) + + +@pytest.mark.skipif(not fused_cuda_available(), reason='Fused CUDA op is unavailable.') +def test_fused_currents_reject_mismatched_master_shape(): + connectivity = build_csr_connectivity(INDICES, SYNAPSE_TYPES, 3, 2, 2) + master_weights = tf.Variable([1.0], dtype=tf.float32) + csr_weights = tf.ones([4], dtype=tf.float32) + basis = tf.constant([[1.0, 0.5], [0.25, 2.0]], dtype=tf.float32) + + with pytest.raises(ValueError, match='master_weights'): + fused_spike_currents( + tf.ones([1, 3]), + master_weights, + csr_weights, + connectivity, + basis, + n_post=2, + compute_spike_gradient=True, + ) + + +@pytest.mark.skipif(not fused_cuda_available(), reason='Fused CUDA op is unavailable.') +def test_fused_currents_reject_dynamic_mismatched_master_shape(): + connectivity = build_csr_connectivity(INDICES, SYNAPSE_TYPES, 3, 2, 2) + csr_weights = tf.ones([4], dtype=tf.float32) + basis = tf.constant([[1.0, 0.5], [0.25, 2.0]], dtype=tf.float32) + + @tf.function(input_signature=[tf.TensorSpec([None], tf.float32)]) + def run(master_weights): + return fused_spike_currents( + tf.ones([1, 3]), + master_weights, + csr_weights, + connectivity, + basis, + n_post=2, + compute_spike_gradient=True, + ) + + with pytest.raises(tf.errors.InvalidArgumentError, match='master_weights'): + run(tf.ones([1], dtype=tf.float32)) diff --git a/tests/simulator/dpointnet/test_training.py b/tests/simulator/dpointnet/test_training.py new file mode 100644 index 00000000..8712f3c4 --- /dev/null +++ b/tests/simulator/dpointnet/test_training.py @@ -0,0 +1,175 @@ +from types import SimpleNamespace + +import numpy as np +import pytest + +tf = pytest.importorskip("tensorflow") + +from bmtk.simulator.dpointnet import training +from bmtk.simulator.dpointnet.cell_models.glif3_cell import GLIF3Cell +from bmtk.simulator.dpointnet.network_adaptor import lex_sort_order_np +from bmtk.simulator.dpointnet.rnn_model import RNN +from bmtk.simulator.dpointnet.optimizers import ( + ExponentiatedAdam, + optimizer_supports_loss_scaling, + scale_loss_for_optimizer, +) + + +def test_refresh_weight_shadows_after_multiple_optimizer_steps(): + recurrent_master = tf.Variable([1.0], dtype=tf.float32, trainable=True) + recurrent_shadow = tf.Variable([1.0], dtype=tf.float16, trainable=False) + input_master = tf.Variable([2.0], dtype=tf.float32, trainable=True) + input_shadow = tf.Variable([2.0], dtype=tf.float16, trainable=False) + cell = SimpleNamespace( + recurrent_weight_values=recurrent_master, + recurrent_weight_values_compute=recurrent_shadow, + compute_dtype=tf.float16, + inputs={ + "bkg": { + "input_weight_values": input_master, + "input_weight_values_compute": input_shadow, + } + }, + ) + optimizer = tf.keras.optimizers.SGD(learning_rate=1.0) + + for expected_recurrent, expected_input in ((0.75, 1.5), (0.5, 1.0)): + optimizer.apply_gradients([ + (tf.constant([0.25]), recurrent_master), + (tf.constant([0.5]), input_master), + ]) + GLIF3Cell.refresh_recurrent_weight_shadow(cell) + + np.testing.assert_allclose(recurrent_shadow.numpy(), [expected_recurrent]) + np.testing.assert_allclose(input_shadow.numpy(), [expected_input]) + + +def test_training_refreshes_weight_shadows_after_each_step(monkeypatch): + class FakeDataIterator: + def __init__(self, *args, **kwargs): + self.closed = False + + def close(self): + self.closed = True + + class FakeCell: + def __init__(self): + self.refresh_count = 0 + + def refresh_recurrent_weight_shadow(self): + self.refresh_count += 1 + + callbacks = SimpleNamespace( + on_train_begin=lambda: None, + on_epoch_start=lambda: None, + on_step_start=lambda: None, + on_step_end=lambda loss: None, + on_epoch_end=lambda loss: False, + on_train_end=lambda **kwargs: None, + ) + cell = FakeCell() + engine = object.__new__(training.TrainingEngine) + engine._parameters = [ + SimpleNamespace(input_generators=[], batch_size=1, seq_len=2) + ] + engine.rnn = SimpleNamespace(ordered_inputs_populations=[], _cell=cell) + engine.regenerate_initial_state_each_epoch = False + engine._init_state_mod = SimpleNamespace(get_state=lambda: None) + engine._prepare_normalizers = lambda: None + engine._normalizers = None + engine.gradient_checkpointing = False + engine._extractor_forward = None + engine._callbacks = callbacks + engine.n_epochs = 1 + engine.steps_per_epoch = 3 + engine._training_approach = "single" + engine._next_spikes_with_retry = lambda input_itr: ("spikes", "targets") + engine._distributed_train_step = ( + lambda spikes, targets, init_state: {"loss": tf.constant(0.0)} + ) + engine._distributed_validation_step = lambda *args, **kwargs: tf.constant(0.0) + monkeypatch.setattr(training, "DataIterator", FakeDataIterator) + + engine.train() + + assert cell.refresh_count == engine.steps_per_epoch + + +def test_lex_sort_order_avoids_integer_overflow(): + indices = np.array( + [ + [np.iinfo(np.uint32).max, np.iinfo(np.uint32).max], + [np.iinfo(np.uint32).max - 1, np.iinfo(np.uint32).max], + [0, 0], + ], + dtype=np.uint32, + ) + + order = lex_sort_order_np(indices) + + np.testing.assert_array_equal(order, [2, 1, 0]) + + +def test_rnn_wraps_float16_optimizer_with_loss_scaling(): + class FakeTrainingEngine: + def __init__(self): + self.optimizer = tf.keras.optimizers.SGD() + self.trained = False + + def set_optimizer(self, optimizer): + self.optimizer = optimizer + + def train(self): + self.trained = True + + rnn = object.__new__(RNN) + rnn.extractor_model = object() + rnn.strategy = tf.distribute.get_strategy() + rnn.dtype = tf.float16 + rnn.model = SimpleNamespace(trainable_variables=[]) + engine = FakeTrainingEngine() + + rnn.train(engine) + + assert isinstance( + engine.optimizer, tf.keras.mixed_precision.LossScaleOptimizer + ) + assert engine.trained + + +def test_base_optimizer_scale_loss_method_is_not_active_loss_scaling(): + optimizer = ExponentiatedAdam(learning_rate=0.005) + + assert not optimizer_supports_loss_scaling(optimizer) + np.testing.assert_allclose( + scale_loss_for_optimizer(optimizer, tf.constant(1.0)).numpy(), + 1.0, + ) + + +def test_loss_scale_optimizer_is_active_loss_scaling(): + optimizer = tf.keras.mixed_precision.LossScaleOptimizer( + ExponentiatedAdam(learning_rate=0.005) + ) + + assert optimizer_supports_loss_scaling(optimizer) + assert scale_loss_for_optimizer(optimizer, tf.constant(1.0)).numpy() > 1.0 + + +def test_loss_scaling_preserves_small_gradient_through_float16_activation(): + activation = tf.Variable(1.0, dtype=tf.float16) + with tf.GradientTape() as tape: + loss = tf.cast(activation, tf.float32) * tf.constant(1.0e-8) + unscaled_gradient = tape.gradient(loss, activation) + + optimizer = tf.keras.mixed_precision.LossScaleOptimizer( + ExponentiatedAdam(learning_rate=0.005) + ) + with tf.GradientTape() as tape: + loss = tf.cast(activation, tf.float32) * tf.constant(1.0e-8) + scaled_loss = scale_loss_for_optimizer(optimizer, loss) + scaled_gradient = tape.gradient(scaled_loss, activation) + + assert unscaled_gradient.numpy() == 0.0 + assert scaled_gradient.numpy() > 0.0