:compute Profile Roadmap

:compute Profile Roadmap (Janus Haiku Way)

Executive Summary

Janus already has the first stones for :compute: device and memory-space enums, a lightweight tensor view, depth sorting for splat rendering, polar embedding primitives, the std.compute.vector scalar/reference kernel layer, positional TurboIndex v0, stable-ID IdMapIndex v0, fixed .jvi / .jvim header encoding, and std.compute.coordination for deterministic multi-agent planning with Gaussian influence fields. The missing center is clear: no kernel emission, no tensor op kernel library, no autodiff, no optimizers, no tokenizer stack, no training runtime.

The recommended path is:

1. CPU-first deterministic core (std.compute.tensor + std.compute.autograd + std.compute.optim),
2. Kenya-budget discipline from day zero (explicit memory budgets and bounded tensor plans),
3. Optional accelerator backends (CUDA/WebGPU) behind stable Janus-native APIs,
4. No mandatory foreign ML runtime dependencies (no PyTorch/JAX/cuDNN lock-in).

This yields a sovereign minimal ML framework that is small, auditable, and shippable on constrained hardware.

---

1) Current State Audit

1.1 State from STATE_OF_JANUS-2604.md

Documented status:

• :compute profile status: Types ready.
• Explicitly listed: tensor, device targeting, memory spaces.
• Explicitly pending: kernel emission.

1.2 Audited stdlib reality (janus/std)

Present and compilable

• std.compute.device

  • Device { cpu, gpu, npu, auto }
  • MemorySpace { host, device, shared, sram, dram, vram }
  • resolve_device(.auto) currently resolves to .cpu.

• std.compute.tensor

  • Runtime metadata tensor view:
    • shape: *usize, stride: *usize, ndim: usize, data: *f32
  • Pure queries: rank, numel
  • Marked gap: compile-time shape params deferred.

• std.compute.sort

  • sort_by_depth(... ) on device(auto) for Gaussian splat ordering.
  • Current implementation uses insertion sort fallback.
  • Notes internal spec gap about function-level on device(...) placement.

• std.compute.polar

  • PolarEmbedding + cartesian/polar conversion primitives.
  • Numeric basis for future ML embedding flows.

• std.compute.vector

  • Phase 1 scalar/reference kernel layer.
  • Ships dot, norm_f32, l2, cosine, normalize_f32.
  • Ships seeded rotation/unrotation, b2/b4 bit packing, fixed Lloyd-Max b2/b4 quantization, and caller-owned top-k heap helpers.
  • Ships positional TurboIndex v0: add/search/remove over retained packed-code storage allocated through an explicit std.mem.page_allocator parent.
  • Ships stable-ID IdMapIndex v0: external u64 IDs, delete repair, search, and allowlist-before-heap search.
  • Ships .jvi / .jvim fixed-header encode/decode.
  • Payload save/load, CID verification, and Lloyd-Max tables are live; SIMD backends remain open SPEC-237 work.

• std.compute.coordination

  • Deterministic CPU-first multi-agent planner.
  • Ships Gaussian coverage/avoidance fields, hard obstacle rejection, stale-state memory windows, greedy one-zone-per-agent assignment, and compact DecisionFrame encode/decode for STL inline effects.
  • Ships AgentReadiness, agent_stale_ticks, agent_memory_remaining_ticks, agent_readiness, and agent_can_travel for memory-age, freshness-budget, and readiness preflight before planner scoring.
  • Ships CandidateFacts and candidate_status helpers for distance, travel time, zone idle age, basic travel preflight, one-pair status classification, and the public COORD_UNREACHABLE_TRAVEL_TIME sentinel for stopped agents.
  • Ships ZoneReadiness, zone_readiness, and zone_blocked for hard-obstacle zone preflight using the same rejection predicate as the planner.
  • Emits unsafe_assignment when a ready agent has no safe unassigned zone because hard geometry blocks the candidate set.
  • Ships assignment query/stats/health helpers for row rejection/conflict checks, status buckets, normalized planner-output rates, compact CoordHealth classification, and agent/zone lookup over caller-owned planner rows.
  • Ships DecisionLedger replay reducers plus total-row, rows-per-frame, and assignment/rejected/conflict/stale-agent rate helpers for decoded decision-frame streams, including replay tick-span reporting.
  • Ships std.compute.coordination_stl as the optional STL Event bridge.
  • Proves LSM-backed STL append, lookup, decode, replay merge, flush, and post-flush readback via test-coordination-stl-lsm.
  • Ships CoordCell, canonical JCK1 key bytes, and deterministic 3x3 neighborhood helpers for read-heavy field-map and zone-catalog indexes.
  • Storage is intentionally external: LSM/STL for replay and audit, while LMX callers own the actual B+ tree storage behind those catalog keys.

• std.compute.splat_pipeline

  • Present in janus/std/compute/splat_pipeline.jan.
  • Pipeline surface for splat-oriented compute flows.

• std.math.linalg

  • Currently includes l2_norm utility and vector/matrix/quaternion structs.

• std.math.sh

  • Minimal spherical harmonics type surface (SH3f) and count utility.

GPU/NPU code observed outside stdlib API

• Compiler-side NPU backend scaffolding exists:
  • compiler/npu_backend.zig exports simulator + MLIR emitter.
  • QTJIR examples lower tensor ops to backend hooks like @npu_tensor_matmul.
• This is backend infrastructure, not yet a stable std.compute runtime API.

Key current-state conclusion

std.compute is now a small but real kernel surface, not yet a full ML substrate. The right next step is to extend the CPU reference execution model around concrete primitives like std.compute.vector, then build AD/core training pieces before backend acceleration.

---

2) Reference Stack Mapping (Rust ML ecosystems)

2.1 Target stack features requested (mni-ml style)

Target capabilities:

• Transformer architecture
• BPE tokenizer
• Flash Attention (custom kernels)
• AdamW
• Training loop
• WebGPU fallback

Note: direct repository identification for mni-ml was not conclusively resolved in this pass; design mapping below uses requested feature set plus validated patterns from Burn, dfdx, and tch-rs.

2.2 What Rust references demonstrate

• Burn: backend abstraction + autodiff as composable backend decorator + kernel fusion + WebGPU/CUDA pathways.
• dfdx: strong compile-time shape checking, reverse-mode AD tape ownership discipline, training ergonomics with SGD/Adam/AdamW.
• tch-rs: practical training loop and optimizer ergonomics, but depends on libtorch (non-sovereign dependency model).

2.3 Mapping to Janus needs

Reference capability	Janus requirement
Transformer blocks	std.compute.nn primitives over tensor ops (matmul, attention, norm, activation)
BPE tokenizer	std.tokenizer.bpe with train/encode/decode and compact vocab storage
Flash Attention	fused attention kernel path (CPU tiled baseline, CUDA/WebGPU optimized variants)
AdamW/SGD	std.compute.optim stateful optimizers with explicit parameter/grad buffers
Training loop	std.compute.train (forward, loss, backward, step, checkpoint hooks)
WebGPU fallback	backend adapter under std.compute.device + shader runtime

---

3) Proposed :compute Profile Design

3.1 std.compute.tensor (expand from view-only to engine)

Mandatory data model

• Tensor[T] handle with:
  • shape, stride, rank
  • dtype (f32, f16, bf16, i32, u8 minimum set)
  • device + memory-space metadata
  • ownership policy (owned buffer vs view)

Mandatory ops (Phase 0-2 core)

• Elementwise: add/sub/mul/div, exp/log, max/min, clamp
• Structural: reshape/view, slice, concat, transpose, broadcast
• Reduction: sum/mean/max over axes
• Dense kernels: matmul (baseline + tiled), batched matmul
• NN essentials: softmax, layer_norm, GELU
• Convolution: minimal 1D/2D baseline (CPU reference first)

3.2 std.compute.device

Unify backend dispatch API:

• CPU backend (reference, always available)
• CUDA backend (optional capability)
• WebGPU backend (optional capability)
• NPU backend (future, compiler/backend dependent)

Add explicit allocator hooks:

• alloc_tensor(device, dtype, shape, memory_space)
• free_tensor(...)
• transfer ops: to_device, to_host, copy_async (where supported)

3.3 NEW std.compute.autograd

Recommended MVP:

• Reverse-mode tape AD
• Node records: op, inputs, output metadata, backward closure id
• Gradient map keyed by tensor id
• APIs:
  • require_grad(t)
  • backward(loss)
  • grad(param)
  • zero_grad(params)

3.4 NEW std.compute.optim

Required:

• SGD (with optional momentum)
• AdamW

State layout should be explicit and allocator-owned:

• SGD: velocity buffers optional
• AdamW: m, v, timestep, decoupled weight decay

3.5 NEW std.tokenizer.bpe

Required modules:

• Corpus stats and pair counting
• Merge rule training
• Deterministic encode/decode
• Vocabulary serialization format for on-device loading

---

4) Technical Design Decisions and Trade-offs

4.1 Autodiff architecture

Decision

Start with tape-based reverse-mode dynamic graph.

Why

• Fastest path to training support.
• Matches Janus ownership philosophy: explicit intermediate lifetimes.
• Lower implementation risk than source-transformed gradients.

Future path

• Add optional static graph lowering for inference and kernel fusion.
• Later: source-transform AD for hot kernels if profiling justifies it.

4.2 Static graph vs dynamic graph

• Dynamic (MVP): define-by-run; easier debugging and language integration.
• Static (later): capture and optimize stable subgraphs for deployment.

Recommendation: hybrid. Author in dynamic mode; compiler/runtime can freeze stable segments.

4.3 Memory model (Kenya-first)

Decision

Adopt bounded arenas + explicit tensor views.

Rules:

• Views never own base storage.
• Base storage has a single owner allocator context.
• Slices/views carry offset + stride and borrow lifetime from owner.
• In-place ops require unique ownership or explicit mutable borrow.

Kenya constraint policy (64MB)

• Runtime memory budget object required for training contexts.
• Tensor planner rejects graphs that exceed budget.
• Default dtypes on Kenya: prefer f16/bf16 activations, quantized optimizer states where possible.

4.4 GPU memory as separate allocator context

Decision

Yes. Device memory must be isolated allocator domains.

Rationale:

• Distinct lifetime and synchronization semantics.
• Cleaner transfer accounting and debugging.
• Enables deterministic fallback when device alloc fails.

4.5 CUDA integration options

Option A: graft C to CUDA runtime (recommended Phase 3)

• Pros: fastest delivery, low compiler complexity.
• Cons: external toolchain dependency.

Option B: Janus->PTX native emission

• Pros: maximal sovereignty long-term.
• Cons: high compiler/backend cost.

Option C: LLVM GPU path

• Pros: reuses LLVM stack.
• Cons: backend complexity and weaker direct control initially.

Recommendation: A first, then B/C as strategic R&D.

4.6 WebGPU integration

Decision

Implement as optional backend with shader IR generation + runtime adapter.

• WASM-targetable design is desirable for edge/browser deployment.
• Kenya fit is possible only with strict kernel set and memory caps.

4.7 Numerical safety policy

• NaN handling: deterministic policy per op (propagate by default; optional checked mode).
• Overflow behavior: explicit for integer ops; floating behavior documented and test-covered.
• Dtype ladder:
  • Phase 0: f32, i32
  • Phase 1: f16 (storage + compute where supported)
  • Phase 2+: bf16

---

5) Phased Implementation Roadmap

Phase 0: Foundation (CPU reference)

Deliverables:

• Complete tensor core API and dtype system.
• Elementwise ops + reductions.
• Memory-safe tensor views and ownership rules.
• CPU kernels for baseline correctness.
• Conformance tests and numerical test vectors.

Exit criteria:

• Deterministic tests for all primitive ops.
• matmul baseline correctness.
• Kenya memory budget smoke tests.

Phase 1: Core ML

Deliverables:

• Cache-tiled CPU matmul.
• std.compute.autograd tape engine.
• SGD optimizer.
• Minimal MLP and tiny transformer block forward/backward examples.

Exit criteria:

• End-to-end gradient test (finite-diff checks on selected ops).
• Stable loss decrease on toy dataset.

Phase 2: Training Stack

Deliverables:

• AdamW optimizer.
• std.tokenizer.bpe train/encode/decode.
• Training loop utilities (fit, checkpoint save/load hooks, metrics).

Exit criteria:

• Tiny language model toy training loop from text corpus to perplexity trend.

Phase 3: GPU Acceleration

Deliverables:

• CUDA backend via graft C runtime.
• WebGPU fallback backend.
• Fused attention path (flash-style kernel family; baseline then optimized).

Exit criteria:

• Backend parity tests (CPU vs CUDA/WebGPU numerical agreement within tolerance).
• Runtime backend switch with unchanged model code.

Phase 4: Optimization + Kenya Validation

Deliverables:

• Operator fusion pass for common patterns.
• Memory planner (buffer reuse, liveness-based release).
• Kenya profile benchmark suite (64MB budget).

Exit criteria:

• Demonstrable training/inference workloads under Kenya limits.
• Published perf and memory envelopes.

---

6) Dependency and Reuse Analysis

6.1 Reuse candidates in current stdlib

Existing module	Reuse role
std.math.linalg	core numeric helpers, vector/matrix utilities
std.math.trig / std.math.sh	attention variants, embeddings, signal transforms
std.mem + std.mem.allocator	allocator traits, raw/managed memory ops
std.encoding (cbor/msgpack/json)	checkpoint and model metadata serialization
std.crypto	optional tokenizer seeding, reproducibility hashing, artifact integrity

6.2 New code required

• std.compute.autograd (new)
• std.compute.optim (new)
• Expanded tensor kernel library in std.compute.tensor
• std.tokenizer.bpe (new)
• GPU shader/kernel runtime bridges (CUDA/WebGPU)
• Training utility layer (std.compute.train or equivalent)

---

7) Sovereignty Considerations

Non-negotiables:

• No mandatory dependency on PyTorch/JAX/cuDNN/cuBLAS.
• Optional grafts allowed, but Janus-native path must exist and remain first-class.
• Fully self-hostable training/inference workflow.
• Kenya-profile viability is a design constraint, not a later optimization pass.
• Compile-time shape/type checks should eliminate class-level runtime type confusion.

Recommended governance rules:

• Every non-sovereign backend dependency is optional and feature-gated.
• Every core op has a CPU-native reference implementation.
• Every optimization must preserve deterministic fallback behavior.

---

Risks and Mitigations

Risk	Impact	Mitigation
Missing splat_pipeline module despite docs/export references	Doc/code drift, build confusion	Reconcile module index immediately; add CI check for exported module existence
Kernel emission gap remains unresolved	Blocks real acceleration	Keep CPU-first roadmap independent; define backend ABI early
Kenya memory overrun	Violates core mission	Memory budget enforcement in API, not optional tooling
Backend divergence (CPU vs GPU)	Silent correctness bugs	Golden tests with tolerance bounds across backends
Autodiff complexity creep	Delays delivery	Tape AD MVP first; postpone source-transform AD

---

Open Questions

1. Should tensor<T, Dims> compile-time shape typing be required for Phase 0, or can runtime-shape API ship first with migration path?
2. Is :compute allowed to depend on :service memory modules directly, or should a strict :compute memory façade be introduced?
3. For CUDA: is policy to permit optional cuBLAS fast paths, or require pure custom kernels only?
4. For WebGPU: which minimum shader feature set is mandatory for Kenya-class deployments?
5. Should tokenizer training be in-core (std.tokenizer.bpe.train) or tooling-side with runtime-only decode on constrained devices?
6. Where should training orchestration live: std.compute.train vs external package layered on stdlib primitives?

---

Phase 0 Implementation Starter (Concrete)

Minimum implementation package for immediate execution:

1. Tensor ABI freeze (v0)

   • Finalize tensor header layout (shape/stride/dtype/device/ownership flags).

2. CPU kernels v0

   • add, mul, sum(axis), matmul (naive), softmax (stable), layer_norm.

3. Memory budget object

   • ComputeBudget { max_bytes, used_bytes } enforced at allocation points.

4. View safety rules

   • Introduce explicit API for slicing/view creation with aliasing checks.

5. Test harness

   • Property tests for shape/stride correctness.
   • Numerical parity tests against small known vectors/matrices.
   • OOM tests under synthetic Kenya budget.

This Phase 0 scope is enough to begin implementation immediately while keeping the architecture aligned with long-term sovereignty and acceleration goals.