Skip to content

llm-refactor-plan.md

Latest commit

 

History

History
264 lines (189 loc) · 9.12 KB

llm-refactor-plan.md

File metadata and controls

264 lines (189 loc) · 9.12 KB
Raw

LLM Refactor Plan — petite-ad

Status: In progress on branch llm-refactor. Last updated: 2026-05-11 Commits: ffe4984 (annotations/re-exports), 7e4c7e1 (split compiled.rs), 84dd468 (extract expr/tape)

Generated by the LLM AI Coding Agent audit checklist. Each finding is evaluated against the principles in the skill: predictable structure, flat explicit architecture, linear functions with explicit state, regenerability, and comment hygiene.

High Priority

1. ✅ Split src/multi/compiled.rs into backend modules

File: src/multi/compiled.rs (4,615 lines)

This single file holds: WGPU backend, SIMD x2/x4 backends, scalar backend, device planning types, batch buffer types, CompiledGraph, BackendKind dispatch, OpCode/FlatInstruction, and all execution logic. Five distinct concerns are interleaved in one file.

Proposed split:

src/multi/
  backend/
    mod.rs           → pub mod + re-exports
    types.rs         → OpCode, FlatInstruction, BackendCapabilities, Instruction
    scalar.rs        → ScalarBackend
    simd.rs          → SimdBackend (f64x2, f64x4)
    wgpu.rs          → WgpuBackend, WgpuBuffer, WgpuBufferSet (feature-gated)
    device.rs        → DeviceBatchPlan, DeviceBuffer, DeviceBufferSet,
                       DeviceMemoryLocation, DeviceTransferKind, etc.
    dispatch.rs      → BackendKind, BackendSupportReport, auto-dispatch logic
    mock.rs          → MockDeviceBackend, DeviceExecutionTrace
  compiled.rs        → keep: CompiledGraph, CompiledWorkspace, BatchInputs,
                       BatchValues, BatchGradients, metadata

Benefit: Each backend can be regenerated independently without understanding the others. A future LLM working on WGPU won't need to parse SIMD dispatch logic.

2. ✅ Decompose gradient_batch_auto_into

File: src/multi/compiled.rs, line ~1,852

The single largest function in the codebase. It dispatches to Scalar, SIMD x2, SIMD x4, WGPU paths with inline per-backend logic.

Proposed fix:

fn gradient_batch_auto_into(compiled, batch, buf) -> Result<()> {
    let backend = compiled.recommended_batch_gradient_backend();
    match backend {
        BackendKind::Scalar => gradient_batch_scalar_into(compiled, batch, buf),
        BackendKind::SimdF64x2 => gradient_batch_simd_f64x2_into(compiled, batch, buf),
        // ...
    }
}
// Each arm becomes a standalone 20-50 line function.

Benefit: A future LLM replacing the SIMD backend only needs to read one small function, not parse 1,000 lines of unrelated dispatch code.

3. ✅ Split src/multi/graph.rs into concerns

File: src/multi/graph.rs (6,315 lines)

Contains: Graph (core API, 1,550-line impl), Tape/TapeWorkspace, ExprGraph/ExprNode (expression graph with operator overloading), GradientCheckEntry/GradientCheckReport, GraphNode, NodeId, DomainPolicy, GraphStats, WGPU helpers (cfg-gated), serde support (cfg-gated).

Proposed split:

src/multi/
  graph/
    mod.rs        → pub mod + re-exports
    core.rs       → Graph, GraphNode, NodeId, DomainPolicy, GraphStats
    tape.rs       → Tape, TapeWorkspace
    expr.rs       → ExprGraph, ExprNode, operator overloads
    check.rs      → GradientCheckEntry, GradientCheckReport
  graph.rs        → keep if preferred, or delete

Benefit: The Graph impl shrinks from 1,550 lines to ~400 lines of core logic. ExprGraph changes won't touch Graph internals.

4. ✅ Remove unused import lint suppressions in multi.rs

File: src/multi.rs, lines 51-55

#[allow(unused_imports)]
pub use multi_ad_fr::MultiAD2FR;
#[allow(unused_imports)]
pub use multi_ad_rf::MultiAD2RF;
#[allow(unused_imports)]
pub use multi_ad_rr::MultiAD2RR;

These types are re-exported but never used internally in multi.rs. The #[allow(unused_imports)] is noise. Replace with a module-level #![allow(unused_imports)] or better: re-export them through lib.rs only and remove from multi.rs since they're already in lib.rs's pub use block.

5. ✅ Clean up #[allow(dead_code)] on public API methods

Affected files: src/mono/mono_fn.rs, src/multi/multi_fn.rs, src/mono/types.rs, src/forward.rs, src/multi/builder.rs

Public library methods tagged #[allow(dead_code)] signal that no internal code calls them. This is expected for a library, but the annotations clutter the code. Options:

  • Move to a single #![allow(dead_code)] at crate root
  • Or remove the annotations (dead_code doesn't apply to pub items when building as a library, only as a binary)

Medium Priority

6. Unify duplicate Dual and OpKind types

Files: src/mono/mono_hessian_common.rs and src/multi/multi_hessian_common.rs

Both define:

struct Dual { val: f64, tan: f64 }
enum OpKind { Sin, Cos, Tan, Exp, Neg, Ln, Sqrt, Abs, ... }

The two Dual types are nearly identical (mono adds Dual::constant, multi adds Add/Sub/Mul/Div impls). The OpKind enums have different variants (mono has 8, multi has 16+).

Proposal: Move Dual to src/types.rs with both constant and arithmetic impls. Keep OpKind separate since the variants differ, but name them MonoOpKind and MultiOpKind to clarify which is which.

7. Consolidate re-export blocks

Files: src/lib.rs (30-line pub use block), src/multi.rs (30-line pub use block)

Half the types are re-exported in both files. A future LLM needs to read both to understand what's publicly available.

Proposal: Use pub use multi::*; in lib.rs for most types, with explicit re-exports only for items that need renaming or feature-gating. Keep multi.rs as the single source of truth for what multi exports.

8. Inline #[cfg(test)] example function files

Files: src/mono/mf1.rs through mf4.rs (4 files), src/multi/f1.rs through f3.rs (3 files)

Each file is ~30 lines: one struct + one trait impl, used in exactly one test. Seven extra files for 210 lines of test code.

Proposal: Move the struct definitions into the test module that uses them (src/mono/tests.rs, src/multi/tests.rs) as module-level items, or consolidate into a single src/mono/tests/examples.rs.

9. Split src/forward.rs by variable arity

File: src/forward.rs (840 lines)

Contains both ForwardAD::differentiate (mono) and ForwardAD::directional_derivative (multi), plus graph-based forward mode. Three related but distinct APIs in one file.

Proposal: Keep src/forward.rs for the ForwardAD struct and shared types. Move mono differentiation to src/mono/forward.rs and multi directional derivative to src/multi/forward.rs.

10. Remove architecture notes from test code

File: src/tests_comprehensive.rs, lines ~630-640

// Note: F1, F2, F3 are private modules, tested through multi::tests
// since MF1-4 are private implementation details

These comments document where other types are tested — information that belongs in those test modules, not in the comprehensive test file. Remove or move.

11. Extract per-backend plan constructors from batch_plan

File: src/multi/compiled.rs, fn batch_plan (135 lines)

This function builds a DeviceBatchPlan for each backend variant with inline code for every variant. Extract:

fn plan_scalar(graph, batch_size) -> DeviceBatchPlan { ... }
fn plan_simd_f64x2(graph, batch_size) -> DeviceBatchPlan { ... }
fn plan_wgpu(graph, batch_size) -> DeviceBatchPlan { ... }

Lower Priority

12. Trailing comma in macro invocation

File: src/mono/mf4.rs, line 19

&mono_ops![sin, neg,]

The trailing comma after neg is technically valid in Rust macros but inconsistent with other mono_ops! calls. Remove for consistency.

13. Module-level docs missing in src/macros.rs

File: src/macros.rs (107 lines)

The file defines mono_ops! and multi_ops! macros but has no module-level doc comment. Add a brief explanation of the macro API and syntax.

14. Ambiguous module name types.rs in both mono and multi

Files: src/mono/types.rs, src/multi/types.rs

Both define different types but share the same name. When reading code via use super::types::*, it's ambiguous which module is being referenced.

Proposal: Rename to mono_types.rs and multi_types.rs, or use mono::types / multi::types consistently in imports (already done, but the filename convention would help).

15. Redundant #[cfg(test)] on items inside #[cfg(test)] modules

Files: src/mono/mf1.rs through mf4.rs, src/multi/f1.rs through f3.rs

These files are declared as #[cfg(test)] mod mf1; in mono.rs, but each struct inside also has #[cfg(test)]. The outer annotation already gates the entire module. Remove the inner annotations.

Summary

Priority Items Effort Estimate
High 5 2-3 days (split compiled.rs + graph.rs)
Medium 6 1-2 days
Lower 4 0.5 day

The highest-value change is splitting compiled.rs into backend modules. This alone addresses the core violation of the "predictable structure" and "flat, explicit architecture" principles. A future LLM can work on the WGPU backend without touching SIMD logic, and vice versa.