Graph capture fix

.gitattributes

@@ -0,0 +1,39 @@
+* text=auto eol=lf
+
+# though text=auto should already catch these
+*.cpp   text eol=lf
+*.cc    text eol=lf
+*.h     text eol=lf
+*.hpp   text eol=lf
+*.cu    text eol=lf
+*.cuh   text eol=lf
+*.py    text eol=lf
+*.cmake text eol=lf
+CMakeLists.txt text eol=lf
+*.md    text eol=lf
+*.txt   text eol=lf
+*.json  text eol=lf
+*.yml   text eol=lf
+*.yaml  text eol=lf
+*.toml  text eol=lf
+*.sh    text eol=lf
+.clang-format text eol=lf
+
+*.png    binary
+*.jpg    binary
+*.jpeg   binary
+*.pdf    binary
+*.gz     binary
+*.bin    binary
+*.so     binary
+*.pyd    binary
+*.dll    binary
+*.pyc    binary
+
+# GitHub/Linguist
+data/**                         linguist-generated=true
+resources/**                    linguist-generated=true
+**/__pycache__/**                linguist-generated=true
+experiments/step-tuning/**       linguist-generated=true
+docs/html/**                     linguist-generated=true
+docs/latex/**                    linguist-generated=true

bindings/LayerBindings.cpp

@@ -270,7 +270,10 @@ void bind_layers(nb::module_& m)
                             nb::arg("activation_deriv") = "drelu");
   gsLayerCls.def("update_state", &Deep::GaussSeidelPCLayer::UpdateState)
       .def("compute_prediction", &Deep::GaussSeidelPCLayer::ComputePrediction)
-      .def("compute_error", &Deep::GaussSeidelPCLayer::ComputeError)
+      .def("compute_error",
+           static_cast<float (Deep::GaussSeidelPCLayer::*)(bool) noexcept>(
+               &Deep::GaussSeidelPCLayer::ComputeError),
+           nb::arg("need_energy") = true)
       .def_prop_ro("mu",
                    [](Deep::GaussSeidelPCLayer& self)
                    {
@@ -672,7 +675,8 @@ void bind_layers(nb::module_& m)
           nb::arg("lmbda") = 1e-2f,
           nb::arg("activation") = "relu",
           nb::arg("activation_deriv") = "drelu")
-      .def("calculate_state", &Deep::ConvPCLayer::CalculateState)
+      .def("calculate_state",
+           static_cast<float (Deep::ConvPCLayer::*)() noexcept>(&Deep::ConvPCLayer::CalculateState))
       .def("update_state", &Deep::ConvPCLayer::UpdateState)
       .def("update_weights", &Deep::ConvPCLayer::UpdateWeights)
       .def("update_precision", &Deep::ConvPCLayer::UpdatePrecision)

bindings/NetworkBindings.cpp

@@ -104,7 +104,18 @@ void bind_networks(nb::module_& m)
   auto discNetCls = nb::class_<Deep::DiscriminativePCNetwork>(
       m, "DiscriminativePCNetwork", "Predictive Coding Network.");
   BindCommonPCNetwork<Deep::DiscriminativePCNetwork>(discNetCls, "DiscriminativePCNetwork");
-  discNetCls.def(nb::init<>(), "Construct a network with automatic batch-size detection.")
+  discNetCls
+      .def(
+          "__init__",
+          [](Deep::DiscriminativePCNetwork* self, const std::string& device)
+          {
+            Deep::DeviceType dt = (device == "cuda" || device == "gpu")
+                                      ? Deep::DeviceType::DEVICE_GPU
+                                      : Deep::DeviceType::DEVICE_CPU;
+            new (self) Deep::DiscriminativePCNetwork(dt);
+          },
+          nb::arg("device") = "cpu",
+          "Construct a network with automatic batch-size detection, on the given device.")
       .def(
           "add_layer",
           [](Deep::DiscriminativePCNetwork& self,
@@ -298,7 +309,17 @@ void bind_networks(nb::module_& m)
 
   nb::class_<Deep::GaussSeidelPCNetwork>(
       m, "GaussSeidelPCNetwork", "Predictive Coding Network with Gauss-Seidel settling dynamics.")
-      .def(nb::init<int>(), nb::arg("batch_size"))
+      .def(
+          "__init__",
+          [](Deep::GaussSeidelPCNetwork* self, int batch_size, const std::string& device)
+          {
+            Deep::DeviceType dt = (device == "cuda" || device == "gpu")
+                                      ? Deep::DeviceType::DEVICE_GPU
+                                      : Deep::DeviceType::DEVICE_CPU;
+            new (self) Deep::GaussSeidelPCNetwork(batch_size, dt);
+          },
+          nb::arg("batch_size"),
+          nb::arg("device") = "cpu")
       .def(
           "add_layer",
           [](Deep::GaussSeidelPCNetwork& self,
@@ -947,7 +968,18 @@ void bind_networks(nb::module_& m)
            });
 
   nb::class_<Deep::ConvPCNetwork>(m, "ConvPCNetwork", "Convolutional Predictive Coding Network.")
-      .def(nb::init<int>(), nb::arg("batch_size"), "Construct a network with a fixed batch size.")
+      .def(
+          "__init__",
+          [](Deep::ConvPCNetwork* self, int batch_size, const std::string& device)
+          {
+            Deep::DeviceType dt = (device == "cuda" || device == "gpu")
+                                      ? Deep::DeviceType::DEVICE_GPU
+                                      : Deep::DeviceType::DEVICE_CPU;
+            new (self) Deep::ConvPCNetwork(batch_size, dt);
+          },
+          nb::arg("batch_size"),
+          nb::arg("device") = "cpu",
+          "Construct a network with a fixed batch size, on the given device.")
       .def(
           "add_layer",
           [](Deep::ConvPCNetwork& self,

imagenet.py

@@ -2,7 +2,7 @@
 import os
 import sys
 from time import perf_counter
-from pydeepity import dy
+from pydeepity import dy  # pyright: ignore[reportAttributeAccessIssue] -- dy is the compiled native extension; see pydeepity/_backend.py and pyrightconfig.json's pydeepity/bindings suppressions for the same reason
 from PIL import Image
 
 IMG_SIZE = 64

include/deepity/backend/CPUBackend.h

@@ -23,6 +23,8 @@ class CPUBackend : public IComputeBackend
   } // nothing to fail on CPU
   /// @copydoc Deep::IComputeBackend::ReplayGraph
   void ReplayGraph() noexcept override {}
+  /// @copydoc Deep::IComputeBackend::Synchronize
+  void Synchronize() noexcept override {}
 
   /// @copydoc Deep::IComputeBackend::Allocate
   float* Allocate(size_t numFloats) override;
@@ -54,6 +56,8 @@ class CPUBackend : public IComputeBackend
               const float* B, int ldb, float beta, float* C, int ldc) noexcept override;
   /// @copydoc Deep::IComputeBackend::SumRows
   void SumRows(float* dst, const float* src, size_t batchSize, size_t width) noexcept override;
+  /// @copydoc Deep::IComputeBackend::Sum
+  float Sum(const float* buf, size_t n) noexcept override;
 
   /// @copydoc Deep::IComputeBackend::Scale
   void Scale(float* buf, size_t n, float alpha) noexcept override;
@@ -153,6 +157,9 @@ class CPUBackend : public IComputeBackend
   float ComputePrecisionWeightedErrorAndEnergy(float* e, const float* z, const float* mu,
                                                const float* p, size_t batchSize,
                                                size_t width) noexcept override;
+  /// @copydoc Deep::IComputeBackend::ComputePrecisionWeightedError
+  void ComputePrecisionWeightedError(float* e, const float* z, const float* mu, const float* p,
+                                     size_t batchSize, size_t width) noexcept override;
   /// @copydoc Deep::IComputeBackend::AxpyBroadcastInto
   void AxpyBroadcastInto(float* y, const float* x, const float* factor, size_t batchSize,
                          size_t width, float alpha) noexcept override;

include/deepity/backend/CUDABackend.h

@@ -23,6 +23,8 @@ class CUDABackend : public IComputeBackend
   bool EndGraphCapture() noexcept override;
   /// @copydoc Deep::IComputeBackend::ReplayGraph
   void ReplayGraph() noexcept override;
+  /// @copydoc Deep::IComputeBackend::Synchronize
+  void Synchronize() noexcept override;
 
   /// @copydoc Deep::IComputeBackend::Allocate
   float* Allocate(size_t numFloats) override;
@@ -45,6 +47,8 @@ class CUDABackend : public IComputeBackend
 
   /// @copydoc Deep::IComputeBackend::SumRows
   void SumRows(float* dst, const float* src, size_t batchSize, size_t width) noexcept override;
+  /// @copydoc Deep::IComputeBackend::Sum
+  float Sum(const float* buf, size_t n) noexcept override;
 
   /// @brief Must be called once, before Compile()'s first
   /// BeginGraphCapture(), for any backend that needs to prepare
@@ -146,6 +150,9 @@ class CUDABackend : public IComputeBackend
   float ComputePrecisionWeightedErrorAndEnergy(float* e, const float* z, const float* mu,
                                                const float* p, size_t batchSize,
                                                size_t width) noexcept override;
+  /// @copydoc Deep::IComputeBackend::ComputePrecisionWeightedError
+  void ComputePrecisionWeightedError(float* e, const float* z, const float* mu, const float* p,
+                                     size_t batchSize, size_t width) noexcept override;
   /// @copydoc Deep::IComputeBackend::AxpyBroadcastInto
   void AxpyBroadcastInto(float* y, const float* x, const float* factor, size_t batchSize,
                          size_t width, float alpha) noexcept override;
@@ -167,5 +174,14 @@ class CUDABackend : public IComputeBackend
   float* onesVector = nullptr;
   size_t onesCapacity = 0;
   float* workspace = nullptr;
+
+  /// @brief Single-float device scratch buffer for scalar reductions
+  /// (ComputePrecisionWeightedErrorAndEnergy(), Sum()). Allocated once in
+  /// the constructor instead of per-call: cudaMalloc/cudaFree are
+  /// synchronizing driver calls that stall the whole pipeline, and both
+  /// callers already block on cudaStreamSynchronize() before returning,
+  /// so reusing one persistent buffer across calls is safe -- no two
+  /// reductions on this backend instance can ever be in flight at once.
+  float* scalarScratch = nullptr;
 };
 } // namespace Deep

include/deepity/backend/IComputeBackend.h

@@ -31,8 +31,19 @@ class IComputeBackend
   virtual bool EndGraphCapture() noexcept = 0;
   /// @brief Launches a graph previously captured and instantiated by
   /// BeginGraphCapture()/EndGraphCapture(). No-op if EndGraphCapture()
-  /// never returned true.
+  /// never returned true. cudaGraphLaunch() is asynchronous: this
+  /// returns as soon as the launch is ENQUEUED, not once it completes.
+  /// Ordinarily that's fine (CUDA guarantees same-stream ordering, so a
+  /// later call that needs the result just needs to be queued on the
+  /// same stream too) -- but call Synchronize() afterward if you need
+  /// the replayed work to have definitely finished, e.g. before reading
+  /// results back on the host through any path that ISN'T itself a
+  /// synchronous backend call.
   virtual void ReplayGraph() noexcept = 0;
+  /// @brief Blocks until all previously enqueued work on this backend's
+  /// stream has completed. A no-op on CPUBackend (every CPU call is
+  /// already synchronous).
+  virtual void Synchronize() noexcept = 0;
 
   // Memory
 
@@ -103,6 +114,14 @@ class IComputeBackend
   /// [width,batchSize] with no data movement (verified numerically).
   virtual void SumRows(float* dst, const float* src, size_t batchSize, size_t width) noexcept = 0;
 
+  /// @brief Returns the plain (signed) sum of every element in buf.
+  /// Synchronous on CUDABackend (reads the reduction back to a host
+  /// float before returning) -- same constraint as ComputeErrorAndEnergy
+  /// and friends: never call this from inside a captured CUDA graph
+  /// region, only from code that runs outside any
+  /// BeginGraphCapture()/EndGraphCapture() pair.
+  virtual float Sum(const float* buf, size_t n) noexcept = 0;
+
   // Elementwise scalar ops
 
   /// @brief buf[i] *= alpha for all i in [0, n).
@@ -198,15 +217,26 @@ class IComputeBackend
                                                size_t batchSize, size_t nextSize) noexcept = 0;
 
   /// @brief Attempts a fused forward pass (GEMM + bias + activation in
-  /// one kernel, via CUTLASS on GPU) for RELU or LINEAR activation types
-  /// only, see CUDABackend's implementation for why other activation
-  /// types aren't attempted yet. Returns true if the fused path was used
-  /// (mu is fully computed, including bias and activation); false if the
-  /// caller should fall back to the existing MatMul+AddBiasBroadcast+
-  /// Activation sequence (always false on CPUBackend, CPU has no fused
-  /// path, this is a GPU-only optimization).
-  /// @param actType Activation to fuse in; only RELU and LINEAR take the
-  /// fused path, everything else returns false immediately.
+  /// one kernel, via CUTLASS on GPU) for RELU, LINEAR, GELU, TANH, or
+  /// SIGMOID activation types; anything else returns false immediately.
+  /// Returns true if the fused path was used (mu is fully computed,
+  /// including bias and activation); false if the caller should fall
+  /// back to the existing MatMul+AddBiasBroadcast+Activation sequence
+  /// (always false on CPUBackend, CPU has no fused path, this is a
+  /// GPU-only optimization).
+  /// @warning As of this writing, every existing call site
+  /// (FullPCLayer::ComputeMuOnly(), DirectKPPCLayer::ComputeMuOnly())
+  /// passes actType=LINEAR unconditionally -- their real activation is
+  /// applied via a separate ActivationInto() pass on the INPUT before
+  /// this GEMM runs, not fused into its epilogue. The RELU/GELU/TANH/
+  /// SIGMOID branches below are real and GPU-tested-correct in
+  /// isolation, but exercising them for an actual layer requires moving
+  /// that activation into this call instead (see CUDABackendGemm.cu's
+  /// file comment before doing so: it's a data-flow change to a settling
+  /// loop, not just a new branch).
+  /// @param actType Activation to fuse in; RELU, LINEAR, GELU, TANH, and
+  /// SIGMOID take the fused path, everything else returns false
+  /// immediately.
   /// @param zF Activated input, shape [batchSize, size], row-major.
   /// @param W Weight matrix, shape [nextSize, size], row-major.
   /// @param bias Bias vector, shape [nextSize].
@@ -337,6 +367,17 @@ class IComputeBackend
   virtual float ComputePrecisionWeightedErrorAndEnergy(float* e, const float* z, const float* mu,
                                                        const float* p, size_t batchSize,
                                                        size_t width) noexcept = 0;
+  /// @brief Same as ComputePrecisionWeightedErrorAndEnergy(), without
+  /// computing (or returning) the energy, for callers that only need e.
+  /// On CUDABackend this additionally means: no blocking
+  /// cudaStreamSynchronize() -- the ...AndEnergy() variant must
+  /// synchronously read the energy reduction back to a host float, which
+  /// is illegal while a CUDA graph is being captured, exactly why this
+  /// counterpart exists (mirrors ComputeError() alongside
+  /// ComputeErrorAndEnergy() above).
+  virtual void ComputePrecisionWeightedError(float* e, const float* z, const float* mu,
+                                             const float* p, size_t batchSize,
+                                             size_t width) noexcept = 0;
 
   /// @brief y[b*width+i] += alpha * x[b*width+i] * factor[i], for every b
   /// in [0,batchSize), i in [0,width), AxpyInto with an added

include/deepity/layers/ConvPCLayer.h

@@ -70,7 +70,14 @@ class ConvPCLayer : public Layer {
 
   /// @brief Calculate energy/prediction errors for this layer.
   /// @return This layer's energy contribution at the current state.
-  float CalculateState() noexcept override;
+  float CalculateState() noexcept override { return CalculateState(true); }
+  /// @brief Same as CalculateState(), but can skip computing (and
+  /// returning) the energy. On CUDABackend, needEnergy=false also means
+  /// no blocking host sync, required for any settling-loop caller that
+  /// runs this inside a captured CUDA graph region (see
+  /// ConvPCNetwork::TrainStep()'s GPU branch).
+  /// @param needEnergy Whether to compute and return the energy.
+  float CalculateState(bool needEnergy) noexcept;
   /// @brief Update latent beliefs (z/r) via inference gradient.
   void UpdateState() noexcept override;
   /// @brief Hebbian/gradient weight update.
@@ -165,6 +172,22 @@ class ConvPCLayer : public Layer {
   /// @brief Whether ClampState() is currently active on this layer.
   bool IsClamped() const noexcept { return isClamped; }
 
+  /// @brief Forces the next ComputeMuOnly() call to recompute instead of
+  /// reusing cachedMu, without otherwise touching clamp state.
+  ///
+  /// ClampState() normally does this (muCacheValid = false), and runs
+  /// fresh at the top of every TrainStep() call. But a CUDA-graph (re)capture
+  /// replays ClampState() only implicitly, by re-executing the whole
+  /// settling loop's C++ as part of recording the graph: that loop's last
+  /// ComputeMuOnly() call leaves muCacheValid = true as a side effect, and
+  /// nothing clears it again before the graph is replayed. The following
+  /// uncaptured, fresh CalculateState() read (used for the returned/printed
+  /// energy) then wrongly takes the cached branch, returning mu from
+  /// BEFORE this call's weight update instead of after it. Callers that
+  /// just finished a capture should invalidate the cache here before that
+  /// read, exactly like a fresh ClampState() call would have.
+  void InvalidateMuCache() noexcept { muCacheValid = false; }
+
   /// @brief Sets the layer immediately above this one in the network.
   /// @param above Pointer to the layer above; may be nullptr for a
   /// terminal layer.

include/deepity/layers/DirectKPPCLayer.h

@@ -173,6 +173,14 @@ class DirectKPPCLayer : public Layer {
   /// first time by moving ProjectForward() inside graph capture.
   bool IsClamped() const noexcept { return isClamped; }
 
+  /// @brief Forces the next ComputeMuOnly() call to recompute mu instead
+  /// of reusing a cached value, without otherwise touching clamp state.
+  /// Needed after a CUDA-graph (re)capture: the recording pass runs the
+  /// settling loop as real C++, leaving muCacheValid=true as a side
+  /// effect, which would otherwise make the following fresh, uncaptured
+  /// energy read reuse mu from before this call's own weight update.
+  void InvalidateMuCache() noexcept { muCacheValid = false; }
+
   /// @brief Sets the optimizer for W. Call before Compile().
   void SetOptimizer(OptimizerType o) noexcept { opt = o; }
   /// @brief Sets the optimizer for Psi. Call before Compile().

include/deepity/layers/DiscriminativePCLayer.h

@@ -83,7 +83,14 @@ namespace Deep
         /// E = \sum_l 1/2 p^{(l)} ||z^{(l)} - \mu^{(l)}||^2 - 1/2 \log p^{(l)}
         /// \f]
         /// @return This layer's energy contribution at the current state.
-        float CalculateState() noexcept override;
+        float CalculateState() noexcept override { return CalculateState(true); }
+        /// @brief Same as CalculateState(), but can skip computing (and
+        /// returning) the energy. On CUDABackend, needEnergy=false also
+        /// means no blocking host sync, required for any settling-loop
+        /// caller that runs this inside a captured CUDA graph region
+        /// (see DiscriminativePCNetwork::TrainStep()'s GPU branch).
+        /// @param needEnergy Whether to compute and return the energy.
+        float CalculateState(bool needEnergy) noexcept;
 
         /// @brief Computes the state derivatives for inference.
         ///
@@ -119,6 +126,15 @@ namespace Deep
         /// @brief Whether ClampState() is currently active on this layer.
         bool IsClamped() const noexcept { return isClamped; }
 
+        /// @brief Forces the next ComputeMuOnly() call to recompute mu
+        /// instead of reusing a cached value, without otherwise touching
+        /// clamp state. Needed after a CUDA-graph (re)capture: the
+        /// recording pass runs the settling loop as real C++, leaving
+        /// muCacheValid=true as a side effect, which would otherwise make
+        /// the following fresh, uncaptured energy read reuse mu from
+        /// before this call's own weight update.
+        void InvalidateMuCache() noexcept { muCacheValid = false; }
+
         /// @brief Returns beliefs.
         /// @return float *z
         float *GetBeliefs() noexcept override { return z; }