Custom CUDA operation (optional) + bug fix of DPointNet training

[shixnya]

Summary

This PR adds an optional fused CUDA implementation for DPointNet synaptic-current computation and fixes several mixed-precision training issues.

This change speeds up the comuputation ~50% and saves ~20% memory. It requires compiling C++ code.

Changes

Fused CUDA operation

• Adds CSR-based TensorFlow CUDA operations for recurrent and input synaptic currents.
• Implements float16 and float32 forward and gradient kernels.
• Supports uint32 and int64 connectivity metadata.
• Preserves spike counts for input currents and computes gradients only where required.
• Validates GPU compute capability and requires exactly one visible GPU.
• Releases CUDA connectivity resources during model rebuild and cleanup.

The feature is opt-in through use_fused_cuda:

• false: use the TensorFlow implementation.
• true: require the fused CUDA operation.
• "auto": use CUDA when available, otherwise fall back to TensorFlow.

Training fixes

• Corrects Keras 3 loss-scaling detection. A regular optimizer's no-op scale_loss() method is no longer mistaken for active mixed-precision loss scaling.
• Synchronizes recurrent and trainable input compute-weight shadows after optimizer updates.
• Reorders fused CSR weight shadows whenever master weights change.
• Prevents uint32 overflow while calculating lexicographic edge-sort keys.

Build and packaging

The operation can be built from the active environment with:

python -m bmtk.simulator.dpointnet.custom_ops.build

[Copilot]

Pull request overview

Adds an optional fused CUDA implementation for DPointNet synaptic currents, alongside mixed-precision fixes, build tooling, documentation, and tests.

Changes:

• Adds CSR CUDA kernels, Python bindings, resource cleanup, and packaging.
• Integrates fused execution and weight-shadow synchronization.
• Fixes loss scaling and edge-sorting behavior.

Reviewed changes

Copilot reviewed 15 out of 16 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Reviewed scope and final comments
tests/simulator/dpointnet/test_training.py	Training regression tests.
tests/simulator/dpointnet/test_cuda_custom_ops.py	CUDA operation and validation tests.
setup.py	Packages native sources and build entry point.
MANIFEST.in	Includes native build sources.
docs/autodocs/source/dpointnet_guide.rst	Documents CUDA setup and configuration.
bmtk/simulator/dpointnet/rnn_model.py	Manages CUDA resources and optimizer integration.
bmtk/simulator/dpointnet/optimizers.py	Corrects loss-scaling detection.
bmtk/simulator/dpointnet/network_adaptor.py	Moderate (3 votes): The int64 sort-key fallback can overflow for valid uint32 indices; use overflow-safe sorting.
bmtk/simulator/dpointnet/custom_ops/csr_spike_ops.py	Critical (1 vote): Reject metadata outside the int64 range before casting.
bmtk/simulator/dpointnet/custom_ops/csr_spike_ops.cu.cc	Moderate (4 votes): Avoid zero-dimension CUDA launches for empty gradients.
bmtk/simulator/dpointnet/custom_ops/csr_spike_ops.cc	Defines TensorFlow custom operations.
bmtk/simulator/dpointnet/custom_ops/build.sh	Moderate (3 votes): Resolve nvcc via PATH or CUDA_HOME before rejecting the compiler.
bmtk/simulator/dpointnet/custom_ops/build.py	Provides the Python build wrapper.
bmtk/simulator/dpointnet/custom_ops/__init__.py	Exports CUDA operation utilities.
bmtk/simulator/dpointnet/cell_models/glif3_cell.py	Integrates fused currents and shadow synchronization.
.gitignore	Ignores generated CUDA artifacts.

Suppressed comments (2)

bmtk/simulator/dpointnet/cell_models/glif3_cell.py:482

• When use_fused_cuda=True is requested with an unsupported compute or variable dtype, fused_available is false but cuda_op_status() can still return loaded from .... The resulting error says the operator is unavailable while reporting that it loaded, and does not tell the user which dtype requirement failed. Report the dtype incompatibility separately from CUDA availability.

        if use_fused_cuda is True and not fused_available:
            raise RuntimeError(
                'use_fused_cuda=True but the fused DPointNet CUDA operator '
                f'is unavailable: {cuda_op_status()}'

bmtk/simulator/dpointnet/rnn_model.py:610

• RNN.__init__ stores self.dtype as a tf.DType returned by get_precision_policy_and_dtype, not the string 'float16', so this guard is false for mixed-float16 runs. The new wrapping/logging path therefore never executes for a regular optimizer, leaving Keras 3's no-op scale_loss() active and allowing gradients to underflow. Compare against tf.float16 (or use .name) here.

                io.log_info(
                    'Mixed-precision optimizer loss scaling enabled with '
                    'LossScaleOptimizer.'
                )

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[kaeldai]

Is there a reason it has to be a python "str" and not also be a byte, unicode, numpy or other type of string (many of which will fail). For instance sometimes if the value is read in from a npz or hdf5 file you can have _fused == 'auto' and still fail.