Architecture hardening: autograd core, storage soundness, security, and module layout - #30
Architecture hardening: autograd core, storage soundness, security, and module layout#30neuralsorcerer wants to merge 13 commits into
Conversation
Architectural refactor based on a full review of the engine, bindings and Python layers (assessment and remaining migration plan documented in docs/architecture_review.md): Autograd: - Plan backward passes over the subgraph reachable from the loss tensor instead of topologically sorting every node recorded on the thread (ComputationGraph::plan_backward + execute_backward_plan). - Stop cloning the full gradient map on every backward call; backward now returns Result<()> and gradients are read via get_gradient. Tests and diagnostics that want the full map use the new backward_collect. - Make gradient-recording policy explicit with a thread-local grad mode (NoGradGuard / is_grad_enabled). Previously, node registration during gradient-kernel execution was suppressed only because add_to_graph used try_borrow_mut while the graph happened to be borrowed - a silent-failure invariant. add_to_graph is now gated on the grad mode and borrows loudly. - Release the tensors saved for backward as soon as a non-retaining backward() completes (release_saved_subgraph), instead of holding all activations until the next optimizer step; this also makes the "computation graph has been freed" error truthful. - DivBackward: compute both gradients with 4 kernels and no scratch ones tensor (was 6 kernels plus a ones allocation). Tensor/storage invariants: - Tensor::view now rejects non-contiguous tensors; Tensor::reshape and operations::shape_ops::reshape materialise a contiguous copy instead. Previously expand(...).reshape(...) silently produced a tensor whose shape claimed more elements than its storage held (masked at the Python boundary by a blanket copy in PyTensor::from_tensor, but corrupt for any engine-level consumer). squeeze/unsqueeze/flatten route through reshape. - Remove TensorData's vestigial manual reference count (inc_ref/dec_ref were never used outside their own tests) and make Drop unconditionally return raw device buffers to the allocator; sharing is Arc's job. - Tensor::data_mut takes the sound Arc::get_mut path whenever storage is uniquely owned; the intentional shared-parameter in-place update case is now documented at the call site. Kernels: - Replace the raw-pointer per-element parallel loops in add_inplace (gradient accumulation hot path) with safe chunked par_chunks_mut/zip loops, removing 10 unsafe blocks without changing semantics. Validation: full Rust workspace suite and Python pytest suite (778 passed) pass; interleaved A/B benchmarks on release wheels show no regression and a consistent ~15% improvement on gradient-accumulation-heavy backward passes (numbers in docs/architecture_review.md). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01SuYLZdiACFPJWs3v5vtboz
Second tranche of the architecture plan in docs/architecture_review.md: Gradient-recording mode: - Add user-facing no_grad()/enable_grad() context managers plus is_grad_enabled()/set_grad_enabled() in Python, backed by the engine's thread-local grad mode. Recording is gated centrally (Tensor::new, Tensor::set_grad_fn, add_to_graph): results inside no_grad are detached leaves, no operands are saved for backward, and inference no longer grows the autograd graph without bound. requires_grad_(True) remains an explicit, ungated opt-in. - Tensor::zero_grad now clears only that tensor's gradient instead of wiping every gradient on the thread (autograd::clear_gradient). - PyTensor.requires_grad_ returns self so calls chain like PyTorch. Correctness/security hardening: - Shape::numel and Strides::from_shape use checked multiplication in all build profiles. A wrapped product could under-allocate storage while stride-based kernels still trusted the dimensions, turning absurd shapes into out-of-bounds access; now they fail loudly. - nn/init.rs no longer creates &mut slices over uninitialized memory (instant UB for bool, unsound for all dtypes): all twelve initializer kernels write through Vec::extend. clone_data uses Vec::clone. - Op output buffers stay uninitialized behind one documented helper (owned_output_buffer) with an explicit write-before-read contract; full zero-initialization was implemented and measured at 30-35% slower on large element-wise ops (memset doubles output write traffic on reused allocations), so it was rejected in favor of the tracked MaybeUninit migration. Bool buffers are always zeroed. Lint debt: - Remove #![allow(clippy::all)] from the engine; fix or justify every warning (workspace is clean under -D warnings). Three style lints are allowed crate-wide with documented reasons (needless_range_loop, too_many_arguments, items_after_test_module). Device now implements std::str::FromStr. - Add a clippy job (-D warnings) to the lints workflow. Validation: 640 Rust tests and 789 Python tests pass (11 new grad-mode tests); interleaved A/B benchmarks on release wheels show performance parity with the previous commit. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01SuYLZdiACFPJWs3v5vtboz
Third tranche of the architecture plan in docs/architecture_review.md: Storage soundness: - TensorData's buffer now lives in an UnsafeCell, and Tensor::data_mut returns a DataMut access token instead of fabricating &mut TensorData from a shared Arc (undefined behavior whenever any other handle to the parameter existed - i.e. always, since the caller itself holds one). DataMut::Unique is ordinary exclusive access after copy-on-write; DataMut::Shared is the single documented exception - in-place parameter updates that must stay visible through every Arc handle (PyTorch semantics) - and goes through UnsafeCell-backed accessors with an explicit aliasing contract. Call sites are unchanged in shape: the token consumes itself and returns a slice borrowing from the tensor. - The ten hand-copied typed slice accessors (plus five new shared-mutation variants) are generated by one typed_slice_accessors! macro - the first application of the dtype-dispatch dedup planned for the ops layer. Autograd efficiency/correctness: - AddBackward/SubBackward/MulBackward/DivBackward now carry per-input requires_grad flags (like min/max/where/matmul already did) and skip the whole gradient chain for frozen inputs: no multiply, no broadcast reduction, no gradient map entry. This fixes get_gradient returning gradients for tensors that never required them (root cause, for arithmetic ops) and measurably speeds up graphs with frozen operands - scalar-broadcast backward work disappears, and the wide fan-out benchmark improved 15-30% in interleaved A/B runs. Validation: 641 Rust tests (new frozen-input regression test) and 789 Python tests pass; workspace clippy-clean under -D warnings; interleaved A/B release-wheel benchmarks show parity everywhere else. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01SuYLZdiACFPJWs3v5vtboz
Fourth tranche of the architecture plan in docs/architecture_review.md: - The eighteen hand-copied *_direct binary kernels (add/sub/mul/div across dtypes) are now generated by two macros (binary_kernel! / binary_kernel_simd!, the latter keeping the f32/f64 SIMD fast path), and neg's five dtype arms by one macro - ~440 lines of copy-paste removed. Interleaved A/B release-wheel benchmarks show exact performance parity, as expected for a purely mechanical expansion. - Remove operations/fusion.rs (623 lines): compiled and re-exported but never called by any execution path, the bindings, or the Python API. BREAKING (engine crate only; the Python package surface is unchanged): code depending on engine::operations::fusion should pin engine 0.2.1 or vendor the module from git history. Rationale and migration note in docs/architecture_review.md. Validation: 637 Rust tests pass (the 4 removed with fusion were its own unit tests) and 789 Python tests pass; workspace clippy-clean under -D warnings. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01SuYLZdiACFPJWs3v5vtboz
Fifth tranche of the architecture plan in docs/architecture_review.md: All seven loss backward functions (MSE, MAE, Huber, CrossEntropy, BCE, KLDiv, Focal) now carry per-input requires_grad flags, matching the arithmetic/minmax/matmul gradient functions. Previously MSE/MAE/Huber allocated and negated a full-size target gradient on every training step, and KLDiv ran a log/sub/add/mul chain for the target distribution - all discarded work whenever targets are constants, which is the overwhelmingly common case. Prediction gradients are likewise skipped when predictions are frozen (e.g. evaluating a loss against trainable targets only). Also fixes get_gradient reporting gradients for constant target tensors. New regression test: test_loss_targets_receive_no_gradient (frozen targets get no gradient; trainable targets still do). Validation: 638 Rust tests and 789 Python tests pass; workspace clippy-clean under -D warnings; interleaved A/B release-wheel benchmarks show parity-to-slightly-better on the MLP training step (the saving scales with target size). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01SuYLZdiACFPJWs3v5vtboz
Sixth tranche of the architecture plan in docs/architecture_review.md: - LogAddExpBackward and ConcatBackward - the last multi-input gradient functions without per-input requires_grad flags - now skip work for frozen inputs. logaddexp previously ran an exp/sub/mul/broadcast-reduce chain per side unconditionally; concatenate extracted a gradient slice for every input including constants (e.g. constant features concatenated with trainable embeddings). PowBackward and LayerNormBackward already had flags; single-input gradient functions need none since they are only reached when their input requires grad. This completes the audit item from the plan, with a new regression test (test_concat_frozen_inputs_receive_no_gradient). - The five hand-copied cmp_* comparison kernels collapsed into one cmp_kernel! macro (511 -> 424 lines), continuing the dedup pattern established for storage accessors and arithmetic kernels. Validation: 639 Rust tests and 789 Python tests pass; workspace clippy-clean under -D warnings; benchmarks within established ranges. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01SuYLZdiACFPJWs3v5vtboz
Seventh tranche of the architecture plan in docs/architecture_review.md: The autograd subsystem's seven include!-merged files are now proper submodules (core, arithmetic, linalg, shape, reduction, activation, tests) with their own imports, explicit pub(crate) boundaries for shared helpers, and glob re-exports in autograd/mod.rs keeping every crate::autograd::X path unchanged - no callers were touched. The merged namespace had been hiding real layout problems that the conversion surfaced and fixed: - the broadcast-reduction helper used by four gradient-function families lived in the tests include file (moved to core); - shared helpers had no visibility story at all (now explicit pub(crate): accumulate_grad, create_scalar_tensor, the softmax backward kernels, expand_reduction_grad, repeat_interleave_backward_impl); - per-file imports were unstated, resolving through whichever sibling happened to import them (each module now declares exactly what it uses, enforced by -D warnings). Also documents that PowBackward's trait impl lives in a different file than its struct - visible now that files are real modules. This is the template for converting the remaining include! clusters (tensor/mod, operations, bindings pytensor). Validation: 639 Rust tests and 789 Python tests pass; workspace clippy-clean under -D warnings. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01SuYLZdiACFPJWs3v5vtboz
Eighth tranche of the architecture plan in docs/architecture_review.md: The tensor subsystem's five include!-merged files are now proper modules. The method-impl files (ops, indexing, autograd, utils) become children of the module that declares the Tensor struct (children-of-core layout), so they retain access to its private fields - encapsulation is preserved exactly; no field had to widen to pub(crate). Public paths are unchanged via re-exports (crate::tensor::X still works; no callers touched). As with the autograd conversion, real module boundaries made the implicit explicit: the strided-copy helper shared between indexing and core now has a declared pub(super) visibility, and every file states its own imports (enforced by -D warnings) instead of resolving through whichever sibling happened to import a name. Remaining include! clusters (operations/*, bindings pytensor/*) follow the same two templates. Validation: 639 Rust tests and 789 Python tests pass; workspace clippy-clean under -D warnings. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01SuYLZdiACFPJWs3v5vtboz
Ninth tranche of the architecture plan in docs/architecture_review.md: All seven operations include! clusters (arithmetic, simd, shape_ops, loss, linalg, activation, reduction - 27 files) are now proper submodules with declared imports, following the autograd/tensor templates. Public paths are unchanged via re-exports; kernel-only submodules re-export at pub(crate) so internal helpers no longer leak into the crate's public API surface. As before, real module boundaries forced the implicit into the open: - ~190 cross-file helpers, layout structs, and macro-generated kernels now carry explicit pub(crate) visibility instead of resolving through the merged namespace by accident; - loss's private log/sign helpers shadowed the public activation ops with the same names - renamed (log_tensor/sign_tensor) after the ambiguity surfaced as an ambiguous-glob-reexport error; - every file declares its own imports, enforced by -D warnings. Engine-side include! conversion is complete. Remaining: the feature-gated backends/opencl pair (not compiled by default, so a conversion cannot be validated in this environment) and the bindings clusters. Validation: 639 Rust tests and 789 Python tests pass; workspace clippy-clean under -D warnings. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01SuYLZdiACFPJWs3v5vtboz
Tenth tranche of the architecture plan in docs/architecture_review.md. The bindings' two include! clusters are now real modules: - tensor.rs -> 19 files across pytensor/, python/, and creation/ become children of the `preamble` module (which declares PyTensor and the shared conversion helpers), so their `impl PyTensor` blocks keep access to the private `inner` field and the shared imports; - nn.rs -> layers.rs (pyclass wrappers + module registration) becomes a child of module.rs (pyclass structs + #[pyfunction]s), so `impl PyReLU` and the `wrap_pyfunction!` calls resolve against the definitions. The children-of-core layout (also used for tensor/ in the previous change) was necessary here because these files are one cohesive module split only for file size: siblings could not see each other's private fields, methods, or pyo3 macro-generated companion modules. Only genuinely cross-file items were widened to pub(crate) (helpers, a handful of PyTensor methods like is_leaf/eq_from_py, the FanInitKind enum, and the functional #[pyfunction]s). Public paths are unchanged via re-exports; no Python-facing behavior changed. With this, every include! in the default build is gone (only the feature-gated backends/opencl pair remains, untouched because it cannot be compiled/validated here). The crate-wide items_after_test_module allow, an artifact of the old layout, is removed - clippy passes without it. Validation: 639 Rust tests and 789 Python tests pass; workspace clippy-clean under -D warnings. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01SuYLZdiACFPJWs3v5vtboz
Eleventh tranche of the architecture plan in docs/architecture_review.md. The 34 uniform float-unary elementwise kernels in activation/hyperbolic.rs (exp/log/sinh/cosh/tanh/log1p/... - each a fetch-input, fetch-output, unary_apply(closure) wrapper) collapse into one float_unary_kernel! macro parameterized by the per-op mapping closure. ~300 lines removed, continuing the dedup pattern from the storage accessors, arithmetic, and comparison kernels. Each closure (including multi-line custom ones like log1p's -1/<-1 branch handling) was extracted verbatim by the conversion, so behavior is byte-identical. Verified by the full Rust and Python suites plus a numerical spot-check of exp/log/sqrt/sinh/cosh/tanh against numpy and log1p's special-value branches (-inf at -1, nan below -1). The non-uniform activation kernels (softplus/gelu/elu, which take extra scalar parameters) and the reduction kernels keep their per-dtype copies. Validation: 612 engine tests + full workspace suite and 789 Python tests pass; workspace clippy-clean under -D warnings. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01SuYLZdiACFPJWs3v5vtboz
Twelfth tranche of the architecture plan in docs/architecture_review.md.
The four prod_along_dim_{f32,f64,i32,i64} kernels are byte-identical apart
from the element type and multiplicative identity; they collapse into one
prod_along_dim_kernel! macro. The float/int sum_along_dim kernels are
deliberately NOT touched - their float versions carry distinct SIMD-style
paths that the int versions don't, so they are genuinely per-dtype, not
duplication.
Verified with the full Rust suite, workspace clippy under -D warnings, and
a numerical spot-check of prod along every dim for f32/f64/i32/i64 against
numpy; 789 Python tests pass.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01SuYLZdiACFPJWs3v5vtboz
Thirteenth tranche of the architecture plan in docs/architecture_review.md. Addresses the two remaining validatable items. Soundness (finding 16, now fully fixed): Operation output buffers were allocated uninitialized and handed out as &mut [f32]/&mut [i64]/... by every kernel. Forming a &mut T reference to uninitialized memory is undefined behavior - the referent must be a valid T, and a float/int read from uninitialized memory is not a valid value even though every bit pattern is representable. Output buffers are now zero-initialized (owned_buffer_for_dtype always zeroes; the unsound owned_output_buffer helper is removed). Per the project's stated priority order (correctness before performance), soundness wins. The cost is measured honestly via interleaved A/B on release wheels: ~25-35% on the pure element-wise microbenchmark (add+sum over 2000x2000, where the extra memset raises output write traffic from ~3N to ~4N) and within noise on the matmul-dominated training step. The zero-cost alternative (MaybeUninit-typed kernel writes across the ~71 output sites) is documented as opt-in future work. Module layout (migration now complete): Converted the feature-gated backends/opencl include! pair to real modules (ops_impl as a child of context, matching the impl/struct split pattern), validated under --features opencl. This also surfaced and fixed pre-existing latent bugs in opencl code that CI never compiled (default features only): a private device_type field access, a missing CL_MEM_* import, an unused allocator parameter, and several clippy lints (div_ceil, collapsible if, non-Send/Sync Arc documented as intentional for the single-threaded backend). Every include! in the codebase is now gone. Validation: 639 Rust tests and 789 Python tests pass; clippy clean under -D warnings for both the default build and --features opencl. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01SuYLZdiACFPJWs3v5vtboz
| runs-on: ubuntu-latest | ||
| steps: | ||
| - uses: actions/checkout@v7 | ||
| - uses: actions/setup-python@v6 | ||
| with: | ||
| python-version: "3.12" | ||
| - uses: dtolnay/rust-toolchain@stable | ||
| with: | ||
| components: clippy | ||
| - uses: Swatinem/rust-cache@v2 | ||
| - name: Run clippy (deny warnings) | ||
| run: cargo clippy --workspace --all-targets -- -D warnings |
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #30 +/- ##
=======================================
Coverage 99.54% 99.55%
=======================================
Files 22 23 +1
Lines 5534 5623 +89
Branches 129 129
=======================================
+ Hits 5509 5598 +89
Misses 19 19
Partials 6 6 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: e00bb170a1
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| if grad_fn.is_some() && !autograd::is_grad_enabled() { | ||
| return; |
There was a problem hiding this comment.
Clear requires_grad when suppressing grad_fn
When grad mode is disabled this branch only skips attaching the grad function; it does not clear an already-true requires_grad flag. Any op that builds its result by cloning/reflagging an input instead of going through Tensor::new (for example with mt.no_grad(): y = x.reshape(x.shape) for a grad-requiring x, or the MAE/Huber paths that call requires_grad_(true)) can return requires_grad=True while add_to_graph also no-ops, leaving a tensor that looks differentiable but has no recorded graph and violating the new no_grad contract.
Useful? React with 👍 / 👎.
Summary
A staged, fully-validated refactor of the engine, bindings, and Python layers based on a complete architectural assessment. Every change is documented in
docs/architecture_review.md(findings, target architecture, migration plan, measurements). Work was done in reviewable tranches, each independently green.Priority order followed: correctness → security → maintainability → performance. Every performance claim is backed by an interleaved A/B measurement on release wheels; one optimization was rejected because it was measured slower, and one throughput optimization was reverted in favor of soundness (see below).
Correctness
view/reshapedata corruption fixed.viewre-strided unconditionally, soexpand(...).reshape(...)produced a tensor whose shape claimed more elements than its storage held (reproduced).viewnow rejects non-contiguous tensors;reshapematerializes a contiguous copy.logaddexp, andconcatskip gradient work for inputs that don't require it. Fixesget_gradientleaking gradients for constant tensors (e.g. loss targets) and speeds up graphs with frozen operands.zero_grad(was wiping every gradient on the thread), chainablerequires_grad_, overflow-checkedShape::numel/Strides.Soundness (all identified UB removed)
Tensor::data_mutfabricated&mut TensorDatafrom a sharedArcvia raw-pointer cast (UB whenever storage was shared). Storage moved behindUnsafeCell;data_mutreturns aDataMuttoken whose shared-parameter-update path has a documented aliasing contract.&mut [bool]/&mut [T]over uninitialized memory (instant UB for bool). Rewritten to write throughVec::extend.&mut [f32]/… — forming a reference to invalid values, i.e. UB. Now zero-initialized. Measured cost (A/B): ~25–35% on the allocation-bound element-wise microbenchmark, noise on training; the zero-costMaybeUninitalternative is documented as opt-in future work. Also removed a vestigial manual refcount whoseDropleaked device buffers.Security
Checked multiplication in
Shape::numel/Strides::from_shapein all build profiles — a wrapped element count could under-allocate storage while stride-based kernels still trusted the dimensions (out-of-bounds access from an absurd shape).API / features
no_grad()/enable_grad()/is_grad_enabled()/set_grad_enabled()(inference mode: no graph growth, results are detached leaves).operations::fusionmodule (623 lines, never called by any path). Migration note in the review doc.Maintainability
include!→real-modules migration complete — 50 files across autograd, tensor, alloperationsclusters, the bindings, and the feature-gatedbackends/openclpair. Two layout patterns (siblings +pub(crate)re-exports; children-of-core where impls touch private fields), preserving all field privacy.-D warnings); removed the crate-wide#[allow(clippy::all)]. Fixed real issues it surfaced, including latent bugs in never-compiled opencl test code.Validation
-D warningsfor both the default build and--features opencl.Deliberately left to the maintainer
Removing the
hardware/pooled-allocator subsystems (semver-major, still have in-repo consumers) and the optionalMaybeUninitthroughput recovery — both documented with rationale rather than forced.🤖 Generated with Claude Code
Generated by Claude Code