[Vixel2006]

• porting basic tensor datatype to zig
• porting node and op with vtable to zig
• some fixing
• finish porting the full cpu version
• adding simple jit compilation and scheduling
• finishing the scheduler
• finishing cpu api and training xor in it
• porting full cpu stuff with logical devices and allocators

Summary by CodeRabbit

• New Features

  • Added unified CPU and CUDA device support with device-aware memory allocation.
  • Added arena-based memory management for efficient temporary allocations.
  • Added graph equality checks and device-specific execution scheduling.
  • Added Kaiming uniform tensor initialization.
  • Added configurable Softmax dimension and dimension retention behavior.

• Improvements

  • Updated neural-network layers, losses, optimizers, and tensor operations to support explicit devices and allocators.
  • Improved graph fusion, JIT caching, and memory cleanup.

• Breaking Changes

  • Updated public constructors and operation methods to require device and allocator arguments.

[coderabbitai]

Caution

Review failed

The pull request is closed.

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro Plus

Run ID: e4819940-6011-40b8-a547-1fac5dfb8a4b

📥 Commits

Reviewing files that changed from the base of the PR and between fb173e0 and 4cc8ddc.

📒 Files selected for processing (28)

• src/allocator.zig
• src/device.zig
• src/graph.zig
• src/nn/activation.zig
• src/nn/batchnorm.zig
• src/nn/dropout.zig
• src/nn/flatten.zig
• src/nn/linear.zig
• src/nn/loss.zig
• src/nn/rnn.zig
• src/nn/sequential.zig
• src/nn/softmax.zig
• src/node.zig
• src/ops/element_wise.zig
• src/ops/functions.zig
• src/ops/fused.zig
• src/ops/linear.zig
• src/ops/reduce.zig
• src/optimizer/mod.zig
• src/optimizers/adam.zig
• src/optimizers/adamw.zig
• src/optimizers/rmsprop.zig
• src/optimizers/sgd.zig
• src/root.zig
• src/scheduler/fusion.zig
• src/scheduler/jit.zig
• src/scheduler/scheduler.zig
• src/tensor.zig

---

📝 Walkthrough

Walkthrough

This PR replaces allocator-centric execution with explicit device and allocator plumbing across tensors, graph operations, schedulers, neural-network modules, optimizers, and the training entry point. It adds CPU/CUDA allocation abstractions, device resource caching, graph equality, device dispatch validation, and updated lifecycle APIs.

Changes

Device-aware execution refactor

Layer / File(s)	Summary
Allocator, device, and tensor foundation src/allocator.zig, src/device.zig, src/tensor.zig	Adds CPU/CUDA allocators, arena allocation, device initialization/resource caching, and device-backed tensor construction and cleanup.
Graph and operation construction src/graph.zig, src/node.zig, src/ops/*	Threads gpa and Device through graph nodes and operation builders, adds graph equality, and validates backend device selection.
Scheduling, fusion, and JIT execution src/scheduler/*	Updates scheduling, fusion, graph execution, JIT caching, and forward/backward lifecycle handling for device-aware graphs.
Device-aware neural-network modules src/nn/*	Updates module constructors, calls, tensor creation, parameter handling, and cleanup to use explicit devices and allocators.
Device-aware optimizer state and updates src/optimizer/*, src/optimizers/*	Separates host allocator state from device buffers and constructs optimizer update computations on the selected device.
Network and training integration src/root.zig	Exports device and allocator APIs, updates Net signatures, and switches training execution to device scheduling and JIT.
Estimated code review effort: 5 (Critical)	~120 minutes

Sequence Diagram(s)

sequenceDiagram
  participant Main
  participant Device
  participant Net
  participant Scheduler
  participant Graph
  Main->>Device: initialize CPU device
  Main->>Net: construct device-aware network
  Main->>Net: call with gpa
  Net->>Graph: build allocator-backed operations
  Main->>Device: schedule forward/backward
  Device->>Scheduler: dispatch execution
  Scheduler->>Graph: run graph

Loading

Possibly related PRs

• Vixel2006/plast#235: Refactors the same tensor, node, operation, graph, and scheduler abstractions toward the new device/allocator model.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch port

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}