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

.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*

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 *

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',
+]

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()

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"

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();
+    });