diff --git a/Cargo.toml b/Cargo.toml index 590508ce0a..aa59a4a1ea 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -123,6 +123,17 @@ powdr-autoprecompiles = { path = "./autoprecompiles" } powdr-expression = { path = "./expression" } powdr-constraint-solver = { path = "./constraint-solver" } +# Local fork: 1-line visibility flip (`pub(crate) mod rules` → `pub mod rules` +# in cuda-backend/src/logup_zerocheck/mod.rs) so the GKR Rule/DAG +# infrastructure is callable from powdr for the bus_kernel DAG-eval prototype. +[patch."https://github.com/powdr-labs/stark-backend.git"] +openvm-stark-sdk = { path = "/home/steve/stark-backend-dag-prototype/crates/stark-sdk" } +openvm-stark-backend = { path = "/home/steve/stark-backend-dag-prototype/crates/stark-backend" } +openvm-cpu-backend = { path = "/home/steve/stark-backend-dag-prototype/crates/cpu-backend" } +openvm-cuda-backend = { path = "/home/steve/stark-backend-dag-prototype/crates/cuda-backend" } +openvm-cuda-builder = { path = "/home/steve/stark-backend-dag-prototype/crates/cuda-builder" } +openvm-cuda-common = { path = "/home/steve/stark-backend-dag-prototype/crates/cuda-common" } + [profile.pr-tests] inherits = "dev" opt-level = 3 diff --git a/dag_eval_plan_v2.md b/dag_eval_plan_v2.md new file mode 100644 index 0000000000..bc0dbbb10b --- /dev/null +++ b/dag_eval_plan_v2.md @@ -0,0 +1,123 @@ +# Plan v2: GKR-style DAG eval for `apc_apply_bus_kernel` + +Priorities (strict order): +1. **Runtime speed** of bus_kernel on pairing APC=500. +2. **Reuse** stark-backend's GKR/Rule/SymbolicRulesBuilder infrastructure. +3. **Minimize diff** in powdr and avoid upstream stark-backend changes. + +## Non-negotiable: CSE *must* be the GKR Rule/DAG mechanism + +The entire point of this rewrite is to emulate the GKR DAG approach to CSE. Any path that re-invents a CSE scheme (custom intermediates buffer, ad-hoc opcodes added to the stack VM, host-side substitution before bytecode emit) is **out of scope** even if it benchmarks faster. Concretely this means: + +- Bytecode format: the upstream `Rule { low, high }` 128-bit struct from `codec.cuh`, **not** an extension of the current `OpCode` enum. +- Slot allocation: `SymbolicRulesBuilder` from `cuda-backend/src/logup_zerocheck/rules/mod.rs`, **not** a powdr-side hand-rolled priority queue or use-count heuristic. +- Intermediates buffer: the `Fp inter[K]` register / strided-global `DeviceBuffer` pattern from `gkr_input.cu`, **not** a per-thread stack VM extended with `LOAD_INTER`/`STORE_INTER` opcodes. +- Output identification: the existing `buffer_vars=true` + `constraint_idx` mechanism for forcing buffered slots, **not** a parallel "output table" defined outside the rule scheduler. + +Anything we add to powdr is plumbing (algebraic→symbolic conversion, FFI signature, runtime flag, Fp scalar-eval helper) — the CSE algorithm itself is upstream's, unchanged. + +## Current baseline (HEAD = `7b64604fa`) + +The other agent already shipped two non-DAG bus optimisations: +- `4d5f13c6b` — peephole folds in `emit_expr`: `x*1 → x`, `x*0 → 0`, `x*(-1) → -x`, `x±0 → x`, `Neg(Number c) → Number(-c)`. +- `7b64604fa` — `INTR_FLAG_STATIC_MULT_1`: 81% of pairing-APC interactions have `mult = is_valid * 1`; the kernel now skips that bytecode walk and m==0 short-circuit. + +CSE potential measured **on the raw `AlgebraicExpression` AST** (host-side `analyze_bus_cse`, pairing APC=500, first ≥150-interaction chip): +- 196 interactions, 970 arg expressions, **676 distinct subexpressions**, 3663 raw ops, **944 ops removable by perfect CSE (25.8%)**. +- BUT **81% of those 944 ops are trivial peephole cases** (`is_valid * 1`, `is_valid * -1`, `15360 * timestamp`, `-1`) that the peephole at `4d5f13c6b` already eliminates at bytecode-emit time. After peephole, those become 2-op leaves whose CSE payoff is zero (single-eval cost = `PushIntermediate` cost = 2 ops). +- Residual *structural* CSE potential (e.g. `mem_ptr_limbs__0_X + 65536 * mem_ptr_limbs__1_X` doubletons, ~11 patterns at 6 saved ops each) ≈ **70–100 ops out of ~2700 post-peephole bytecode**, i.e. **~3% on top of the peephole**. + +That's the ceiling for the DAG rewrite on this workload. Expected wall-time win on bus_kernel: ≤10 ms out of ~213 ms (≤0.4% of total prove time). Going ahead per user instruction because every other lever has been pulled. + +## Reused infrastructure (no upstream patch) + +| component | source | role | +|---|---|---| +| `Rule` (16 B), `RuleHeader`, `SourceInfo`, `OperationType`, `decode_rule_header`, `decode_y`, `decode_z_index` | `stark-backend/.../cuda-backend/cuda/include/codec.cuh` | device-side rule decode | +| `SymbolicRulesBuilder`, `SymbolicRulesGpu`, `Rule`, `Source`, `RuleWithFlag` | `stark-backend/.../cuda-backend/src/logup_zerocheck/rules/{mod,codec}.rs` | host: hash-cons, lifetime analysis, slot-priority-queue, 128-bit encoding | +| `SymbolicExpressionDag`, `SymbolicExpressionNode`, `SymbolicVariable::{Main}` | `stark-backend/.../stark-backend/src/air_builders/symbolic/dag.rs` + `symbolic_variable.rs` | input to `SymbolicRulesBuilder::new` | +| Kernel pattern (per-row DAG walk with intermediates buffer, local-vs-global mode) | `stark-backend/.../cuda-backend/cuda/src/logup_zerocheck/gkr_input.cu` | template for our Fp-flavour kernel | + +All public in our existing dep graph — no fork. **`codec.cuh` is vendored** (header-only, ~157 LOC) because `cuda-backend` doesn't expose `DEP_CUDA_BACKEND_CUDA_INCLUDE` (no `links` + no `cargo:include=` emit in its build.rs). + +## Reviewer-flagged fixes folded in + +Independent review of v2 found four correctness gotchas, one reuse improvement, and one bench-fairness gap. Folded into the file plan below; flagged here so the implementer doesn't regress them. + +1. **Slot lookup is not a single `decode_z_index` call.** `SymbolicRulesGpu::dag_idx_to_rule_idx` gives `rule_idx` (`rules/mod.rs:408-412`), and only `accumulate=true` dag_idxs get entries (`rules/mod.rs:400`). `buffer_vars=true` only forces a slot for **shared** Variables — `use_count > 1` (`rules/mod.rs:281`). For each interaction output we must inspect the resulting rule and case-split: + - `Rule::Add/Sub/Mul/Neg` (always `intermediate=true`, slot present): read `inter[decode_z_index]`. + - `Rule::BufferVar` (shared Variable hoisted to a slot): read `inter[decode_z_index]`. + - `Rule::Variable(Source::Constant(c))` (constant outputs — e.g. `Number(1)` mult collapsed by peephole): no slot, **embed the constant in the per-output dispatch table** so the kernel uses it directly. + - `Rule::Variable(Source::Var(v))` with `use_count == 1` (bare-column output that nobody else uses): no slot, **embed the column-base in the per-output dispatch table** so the kernel reads `d_output[base + r]` directly. + - The dispatch table emitted by the host becomes: `enum { Const(Fp), Col(u32 base), Inter(u32 slot) }` per output (flat, 8 bytes/output max). + +2. **`d_main` is an indirection table, not the trace.** `evaluate_dag_entry_gkr` in `gkr_input.cu:33` reads `(Fp*)d_main[src.part]` — i.e. `d_main` is `uint64_t*` pointing at an array of column-set pointers. To reuse upstream's helper signature, the launcher allocates a one-element `uint64_t[1] = {(uint64_t)d_output}` device array and passes its address. ~10 LOC at host launch time. Alternative (slightly smaller): skip the helper's `ENTRY_MAIN` arm entirely and have our `evaluate_dag_entry_bus` take a plain `const Fp* d_main` directly. Choose this — see (3). + +3. **Reuse upstream helper via templatisation, not rewrite.** A `template` over the return type lets us reuse `evaluate_dag_entry_gkr`'s body verbatim — `T(...)` constructors work for both `Fp` and `FpExt`. **But** the helper reads `d_preprocessed`/`d_main[part]`/`d_public_values`/`d_challenges`/`d_intermediates` which we don't all have or need. Cheaper path: a small `evaluate_dag_entry_bus` that takes only `(SourceInfo, row, d_main: const Fp*, d_inter: const Fp*, stride, H)` — 4 EntryType arms instead of 11 (`ENTRY_MAIN`, `SOURCE_CONSTANT`, `SRC_INTERMEDIATE`, default → assert). This is ~25 LOC and avoids the `d_main[part]` indirection entirely. The cost of "not reusing" the helper here is small relative to the cost of building the indirection table. + +4. **`SymbolicRulesBuilder::new` debug-asserts `dag.constraint_idx.iter().is_sorted()`** (`rules/mod.rs:134`). After collecting per-interaction output dag_idxs in `[i0_mult, i0_arg0, …, i1_mult, …]` order, sort the union for `constraint_idx`. Keep the original per-output ordering on the side (for emit) in `Vec>` indexed by interaction. + +5. **Bracket the global-mode intermediates allocation.** When `buffer_size > LOCAL_K`: hoist the `DeviceBuffer` of size `TASK_SIZE × buffer_size` to a per-chip cache (allocated once on the first call, kept alive in `PowdrTraceGeneratorGpu`). The per-launch cost becomes zero. Without this, alloc cost would leak out of `bus_kernel_us` and bias the A/B. + +## Diff (estimated +~600 / -180) + +### New +1. `openvm/cuda/include/codec.cuh` — vendored verbatim from stark-backend SHA `8d36ad26`, header comment cites source. ~157 LOC. +2. `openvm/cuda/src/apc_apply_bus_dag.cu` — new kernel + host launcher. ~250 LOC. + - `evaluate_dag_entry_bus(SourceInfo, row, d_main, d_inter, stride, H) -> Fp` (~30 LOC) — mirror of `evaluate_dag_entry_gkr` minus FpExt and minus the preprocessed/public/challenge/is_first/is_last/is_transition arms (bus exprs only use `ENTRY_MAIN`, `SOURCE_CONSTANT`, `SRC_INTERMEDIATE`). + - `apc_apply_bus_dag_kernel` — one thread per row; walks `d_rules` to fill `inter_ptr[buffer_slot * stride]`; after all rules done, reads each interaction's `mult` + args from `inter_ptr` and runs the existing range/tuple2/bitwise histogram dispatch verbatim. + - Local mode: `Fp inter[K]` with `K=24` compile-time. Global mode if `buffer_size > K`: thread-strided `DeviceBuffer` of size `TASK_SIZE × buffer_size`. +3. `openvm/src/powdr_extension/trace_generator/cuda/expr_dag.rs` — converter. ~120 LOC. + - `algebraic_to_symbolic_dag(bus_interactions, id_to_apc_index, apc_height) -> (SymbolicExpressionDag, Vec)` where `InteractionOutputDagIdxs = { bus_id, mult_dag_idx, arg_dag_idxs: Vec, num_args }`. + - Hash-cons over `AlgebraicExpression` using `(op, left_dag_idx, right_dag_idx)` as key (refs + numbers as leaves). + - Mapping: `Number(c) → Constant(c)`; `Reference(r) → Variable(SymbolicVariable{ entry: Entry::Main{ part_index: 0, offset: 0 }, index: id_to_apc_index[&r.id] })`; `UnaryOp(Minus, e) → Neg{ idx }`; `BinaryOp(Add|Sub|Mul, l, r) → Add|Sub|Mul{ left_idx, right_idx }`. + - Populate `dag.constraint_idx` with the dag_idx of every interaction output (mult + each arg). + - The `degree_multiple` field on `SymbolicExpressionNode` is required (the type carries it) — compute conservatively as `left.deg + right.deg` for Mul, `max(left.deg, right.deg)` for Add/Sub, equal-to-input for Neg, 0 for leaves. (Used only for `max_rotation()` queries on the host; the kernel doesn't read it.) + +### Modified +4. `openvm/build.rs` — no change. `codec.cuh` is already under `cuda/include` (vendored), and the existing `.include("cuda/include")` picks it up. +5. `openvm/src/cuda_abi.rs` — add new extern `_apc_apply_bus_dag(...)` and safe wrapper `apc_apply_bus_dag(...)`. New args: `d_rules: *const Rule` (16 B/elt), `n_rules: usize`, `d_intermediates: *mut Fp` (global-mode buffer; pass null/empty when local mode), `buffer_size: u32`, `d_interactions_dag: *const DevInteractionDag` (length n_interactions), `d_output_dag_idxs: *const u32` (flat, one entry per interaction-output, length = sum_i (1+num_args[i])). Existing histogram args unchanged. **Do not delete `_apc_apply_bus`** — both paths coexist behind a runtime flag. +6. `openvm/src/powdr_extension/trace_generator/cuda/mod.rs` — replace `compile_bus_to_gpu` with a thin selector: + - Read `POWDR_BUS_KERNEL` once at process start (`OnceLock`). Values: `"vm"` (default, current stack-VM) or `"dag"` (new path). + - If `"vm"`: existing path, unchanged. + - If `"dag"`: call `algebraic_to_symbolic_dag` → `dag`. `SymbolicRulesGpu::new(&dag, /*buffer_vars=*/true)` → `(rules: Vec, buffer_size: usize, dag_idx_to_rule_idx)`. For each interaction-output `dag_idx`, decode the corresponding rule's `z_index` (or, for Variable outputs which the builder now buffers thanks to `buffer_vars=true`, the same decode applies) to get its buffer slot. Emit `DevInteractionDag { bus_id, num_args, output_idxs_off }` and a flat `output_dag_idxs: Vec` ordered `[i0_mult, i0_arg0, …, i1_mult, …]`. + - Host-allocate the intermediates `DeviceBuffer` only if `buffer_size > LOCAL_K`; otherwise pass an empty `DeviceBuffer` and the kernel uses its register array. + - Same `bus_compile_h2d` / `bus_kernel` `timed_substage!` envelopes. + +### Deleted +None initially. After the perf bench confirms a winner, a follow-up commit deletes the loser. + +## Why this fix-set addresses the reviewer's flagged risks + +| reviewer flag | resolution in plan v2 | +|---|---| +| **Variable outputs not buffered with `accumulate=true` alone** | Use `SymbolicRulesGpu::new(&dag, true)` — `buffer_vars=true` forces a slot for shared Variable nodes (`mod.rs:178-184`). Confirmed via `mod.rs:281` condition: `intermediate \|\| (buffer_var && use_count > 1)`. | +| **`DEP_CUDA_BACKEND_CUDA_INCLUDE` doesn't exist** | Vendor `codec.cuh` from day 1. Comment cites source SHA so future updates are traceable. | +| **Global-mode buffer bandwidth risk with ~120 slots** | Threshold check at host: `if buffer_size > LOCAL_K { go_global = true; emit warning to tracing; }`. After build runs once, we know the actual `buffer_size` — if global, we measure A/B and revert if regress >5%. | +| **CSE-ratio test doesn't predict speed** | Already measured the ratio (25.8% raw, ~3% post-peephole). Per user's call we proceed anyway. The bench (step 6 below) is the real gate. | +| **No Fp `evaluate_dag_entry`** | We write our own (~30 LOC). Plan accounts. | +| **App-replay non-determinism could break A/B in ncu** | We're not relying on ncu for the bench. Use `POWDR_TRACE_PROFILE=1` + `bus_kernel_us` counter, which is summed across all 704 launches per prove and doesn't require kernel replay. | + +## Sequence of work + +1. **Vendor `codec.cuh`** — copy + provenance comment. 5 min. +2. **Write `algebraic_to_symbolic_dag`** + unit test (round-trip a small expression). 30 min. +3. **One-off host-side measurement** — call `SymbolicRulesGpu::new` on a real pairing APC=500 chip via a `#[ignore]` test or a `POWDR_DUMP_BUS_DAG=1` debug path; log `rules_len`, `buffer_size`, `accumulate_count`. **Gating decision**: if `buffer_size ≤ 24` → proceed local-mode (high confidence speedup); if 24 < `buffer_size` ≤ 64 → proceed cautiously (uncertain); if > 64 → reconsider (global-mode likely regresses given the small structural CSE pool). 15 min. +4. **Write `apc_apply_bus_dag.cu`** + Fp `evaluate_dag_entry_bus` helper. 60–90 min. +5. **Wire FFI + Rust selector** behind `POWDR_BUS_KERNEL=dag`. 30 min. +6. **Bit-equal smoke test** — pairing APC=10 with `POWDR_BUS_KERNEL=dag`. Compare `public_values_commit` against the same prove with default `vm`. 10 min. +7. **Perf bench** — pairing APC=500 with `POWDR_TRACE_PROFILE=1`, both kernels (set the env var to switch between runs). Compare: + - `bus_kernel_us` total (host-side timer, microsecond resolution) + - `bus_compile_h2d_us` total (DAG-build adds host work; expected ~+30%) + - nsys top-kernel slice `apc_apply_bus_dag_kernel` vs `apc_apply_bus_kernel` (% GPU, total ns, avg per launch) + - Emit a `combined_metrics.json` snapshot keyed `{vm_baseline, dag_v1}`. +8. **Ship/kill** — given the ~3% theoretical ceiling, the ship bar drops to **any net win** on the combined `bus_kernel_us + bus_compile_h2d_us`. A regression on either ⇒ revert. The point of running the experiment isn't a guaranteed win; it's to establish empirically whether the GKR DAG approach can beat a heavily-optimised stack VM on a workload with structurally low CSE potential. Either result is publishable. + +## Estimated bound on the prize + +Conservative: peephole already captured ~50% of redundant bytecode. Residual structural CSE is ~70-100 ops out of ~2700 per-row bytecode = ~3% reduction. If that translates linearly to bus_kernel time (it won't — global-mode buffer overhead could eat it), upper bound is ~6 ms saved on 213 ms = **~2.8% bus_kernel speedup**. Below the 10% ship criterion. Plan ships only if reality beats theory. + +## Rollout + +- One commit adds the entire DAG path behind `POWDR_BUS_KERNEL=dag` (default off). Mergeable to the parent branch independent of the bench outcome. +- A follow-up commit either flips the default to `dag` and deletes the VM path, or reverts the DAG-path commit entirely. diff --git a/openvm/cuda/include/codec.cuh b/openvm/cuda/include/codec.cuh new file mode 100644 index 0000000000..35c75ae106 --- /dev/null +++ b/openvm/cuda/include/codec.cuh @@ -0,0 +1,156 @@ +/* + * Source: https://github.com/scroll-tech/plonky3-gpu (private repo) + * Status: BASED ON plonky3-gpu/gpu-backend/src/cuda/kernels/codec.h + * Imported: 2025-01-25 by @gaxiom + */ + +// A very simple custom codec for constraints wroten in the AIR/RAP frontend language. +#pragma once + +#include + +typedef struct { + uint64_t low; + uint64_t high; +} Rule; + +typedef enum { OP_ADD, OP_SUB, OP_MUL, OP_NEG, OP_VAR, OP_INV } OperationType; + +typedef enum { + ENTRY_PREPROCESSED, + ENTRY_MAIN, + ENTRY_PERMUTATION, + ENTRY_PUBLIC, + ENTRY_CHALLENGE, + ENTRY_EXPOSED, + SRC_INTERMEDIATE, + SRC_CONSTANT, + SRC_IS_FIRST, + SRC_IS_LAST, + SRC_IS_TRANSITION +} EntryType; + +typedef struct { + EntryType type; + uint8_t part; + uint8_t offset; + uint32_t index; +} SourceInfo; + +typedef struct { + bool is_constraint; + bool buffer_result; + OperationType op; + SourceInfo x; + SourceInfo y; + uint32_t z_index; +} DecodedRule; + +// Lightweight header: only op, flags, and x (always needed). +// Use this for lazy decoding to reduce register pressure. +typedef struct { + bool is_constraint; + bool buffer_result; + OperationType op; + SourceInfo x; +} RuleHeader; + +__host__ __device__ __forceinline__ SourceInfo decode_source(uint64_t encoded); +__host__ __device__ __forceinline__ DecodedRule decode_rule(Rule encoded); + +// Lazy decoding functions - decode only what's needed +__host__ __device__ __forceinline__ RuleHeader decode_rule_header(Rule encoded); +__host__ __device__ __forceinline__ SourceInfo decode_y(Rule encoded); +__host__ __device__ __forceinline__ uint32_t decode_z_index(Rule encoded); + +static const uint64_t ENTRY_SRC_MASK = 0xF; +static const uint64_t ENTRY_PART_SHIFT = 4; +static const uint64_t ENTRY_PART_MASK = 0xFF; +static const uint64_t ENTRY_OFFSET_SHIFT = 12; +static const uint64_t ENTRY_OFFSET_MASK = 0xF; +static const uint64_t ENTRY_INDEX_SHIFT = 16; +static const uint64_t ENTRY_INDEX_MASK = 0xFFFFFFFF; +static const uint64_t SOURCE_INTERMEDIATE_SHIFT = 4; +static const uint64_t SOURCE_INTERMEDIATE_MASK = 0xFFFFF; +static const uint64_t SOURCE_CONSTANT_SHIFT = 16; +static const uint64_t SOURCE_CONSTANT_MASK = 0xFFFFFFFF; +static const uint64_t LOW_48_BITS_MASK = 0xFFFFFFFFFFFF; +static const int Y_HIGH_SHIFT = 16; +static const uint64_t Y_HIGH_MASK = 0xFFFFFFFF; +static const int Z_LOW_SHIFT = 32; +static const uint64_t Z_LOW_MASK = 0xFFFFFF; +static const uint64_t OP_MASK = 0x3F; +static const int OP_SHIFT = 56; +static const uint64_t BUFFER_RESULT_MASK = 0x4000000000000000; +static const uint64_t IS_CONSTRAINT_MASK = 0x8000000000000000; + +const uint64_t one = 1; +static_assert(LOW_48_BITS_MASK == (one << 48) - 1, "LOW_48_BITS_MASK must be (1 << 48) - 1"); +static_assert(Y_HIGH_MASK == (one << 32) - 1, "Y_HIGH_MASK must be (1 << 32) - 1"); +static_assert(Z_LOW_MASK == (one << 24) - 1, "Z_LOW_MASK must be (1 << 24) - 1"); +static_assert(OP_MASK == (one << 6) - 1, "OP_MASK must be (1 << 6) - 1"); +static_assert(BUFFER_RESULT_MASK == one << 62, "BUFFER_RESULT_MASK must be (1 << 62)"); +static_assert(IS_CONSTRAINT_MASK == one << 63, "IS_CONSTRAINT_MASK must be (1 << 63)"); + +__host__ __device__ __forceinline__ SourceInfo decode_source(uint64_t encoded) { + SourceInfo src; + src.type = (EntryType)(encoded & ENTRY_SRC_MASK); + src.part = (encoded >> ENTRY_PART_SHIFT) & ENTRY_PART_MASK; + src.offset = (encoded >> ENTRY_OFFSET_SHIFT) & ENTRY_OFFSET_MASK; + src.index = (encoded >> ENTRY_INDEX_SHIFT) & ENTRY_INDEX_MASK; + + if (src.type == SRC_INTERMEDIATE) { + src.part = 0; + src.offset = 0; + src.index = (encoded >> SOURCE_INTERMEDIATE_SHIFT) & SOURCE_INTERMEDIATE_MASK; + } else if (src.type == SRC_CONSTANT) { + src.index = (encoded >> SOURCE_CONSTANT_SHIFT) & SOURCE_CONSTANT_MASK; + } + return src; +} + +__host__ __device__ __forceinline__ DecodedRule decode_rule(Rule encoded) { + DecodedRule rule; + + uint64_t x_encoded = (encoded.low & LOW_48_BITS_MASK); + rule.x = decode_source(x_encoded); + + uint64_t y_encoded = ((encoded.low >> 48) | ((encoded.high & Y_HIGH_MASK) << Y_HIGH_SHIFT)); + rule.y = decode_source(y_encoded); + + uint64_t z_encoded = (encoded.high >> Z_LOW_SHIFT) & Z_LOW_MASK; + rule.z_index = (z_encoded >> SOURCE_INTERMEDIATE_SHIFT) & SOURCE_INTERMEDIATE_MASK; + + rule.op = (OperationType)((encoded.high >> OP_SHIFT) & OP_MASK); + + rule.buffer_result = (encoded.high & BUFFER_RESULT_MASK) != 0; + rule.is_constraint = (encoded.high & IS_CONSTRAINT_MASK) != 0; + + return rule; +} + +// Decode only header (op, flags, x) - for lazy decoding pattern +__host__ __device__ __forceinline__ RuleHeader decode_rule_header(Rule encoded) { + RuleHeader header; + + uint64_t x_encoded = (encoded.low & LOW_48_BITS_MASK); + header.x = decode_source(x_encoded); + + header.op = (OperationType)((encoded.high >> OP_SHIFT) & OP_MASK); + header.buffer_result = (encoded.high & BUFFER_RESULT_MASK) != 0; + header.is_constraint = (encoded.high & IS_CONSTRAINT_MASK) != 0; + + return header; +} + +// Decode y operand on demand (only needed for binary ops: ADD, SUB, MUL) +__host__ __device__ __forceinline__ SourceInfo decode_y(Rule encoded) { + uint64_t y_encoded = ((encoded.low >> 48) | ((encoded.high & Y_HIGH_MASK) << Y_HIGH_SHIFT)); + return decode_source(y_encoded); +} + +// Decode z_index on demand (only needed when buffer_result is true) +__host__ __device__ __forceinline__ uint32_t decode_z_index(Rule encoded) { + uint64_t z_encoded = (encoded.high >> Z_LOW_SHIFT) & Z_LOW_MASK; + return (z_encoded >> SOURCE_INTERMEDIATE_SHIFT) & SOURCE_INTERMEDIATE_MASK; +} diff --git a/openvm/cuda/src/apc_apply_bus_dag.cu b/openvm/cuda/src/apc_apply_bus_dag.cu new file mode 100644 index 0000000000..e9deda0ac0 --- /dev/null +++ b/openvm/cuda/src/apc_apply_bus_dag.cu @@ -0,0 +1,277 @@ +#include +#include +#include +#include +#include +#include "primitives/buffer_view.cuh" +#include "primitives/constants.h" +#include "primitives/trace_access.h" +#include "primitives/histogram.cuh" +#include "codec.cuh" // vendored from stark-backend (see openvm/cuda/include/codec.cuh) + +// ============================================================================= +// GKR-style DAG evaluation kernel for bus interactions (GLOBAL-buffer variant). +// +// One thread per row. Each row: +// 1. Walks the per-chip Rule[] (CSE-deduped 128-bit encoded rules from +// stark-backend's SymbolicRulesBuilder), filling a per-thread slice of +// `d_intermediates` with intermediate results. Layout is slot-major +// coalesced: thread `t`'s slot `s` lives at `d_intermediates[s*stride + t]` +// with `stride = total_threads`, so a warp's 32 threads accessing the +// same slot read 32 contiguous u32s → one cache line. +// 2. For each bus interaction, reads `(mult, arg0, ..., argN)` from its +// pre-computed dispatch table — either from `inter[slot]` (buffered +// sub-expression in d_intermediates), from `d_output[col_base + r]` +// (single-use bare column that the scheduler didn't buffer), or as an +// inline `Fp(constant)`. +// 3. Dispatches to range / tuple2 / bitwise histograms — identical logic +// to apc_apply_bus_kernel. +// +// The earlier local-mode variant (`Fp inter[LOCAL_K]`) blew up the per-thread +// stack frame to 8 KB on chips with buffer_size > 1024 (max observed: 1762), +// trashing L1 with per-thread random-access local-memory traffic. This global +// variant trades that for coalesced DRAM reads/writes through a per-launch +// `DeviceBuffer` allocated at `total_threads * buffer_size` Fps. +// ============================================================================= + +static constexpr uint32_t BITWISE_NUM_BITS = 8u; + +extern "C" { + typedef struct { + uint32_t bus_id; // bus id this interaction targets + uint32_t num_args; // number of arg expressions (mult is implicit; total outputs = 1 + num_args) + uint32_t outputs_off; // index into `d_output_descs[]` where this interaction's outputs begin + uint32_t flags; // bitfield (currently unused — reserved) + } DevInteractionDag; + + // Per-output dispatch descriptor. The host emitter inspects each rule + // pointed to by `dag_idx_to_rule_idx` and emits one of three kinds: + // kind == 0 (Inter): value = slot index into `inter[]` + // kind == 1 (Col) : value = column base (apc_col_idx * H), read from d_output[value + r] + // kind == 2 (Const): value = constant as canonical u32 + typedef struct { + uint32_t kind; + uint32_t value; + } OutputDesc; +} + +// Decode source for the bus DAG kernel. Bus expressions only use ENTRY_MAIN, +// SOURCE_CONSTANT, and SRC_INTERMEDIATE (no preprocessed / public / +// challenge / is_first / is_last / is_transition). The fourth arm is +// unreachable for valid host-emitted rules. +__device__ __forceinline__ Fp evaluate_dag_entry_bus( + const SourceInfo &src, + uint32_t row, + const Fp *__restrict__ d_main, // APC trace buffer (column-major) + const Fp *inter, // points at thread's slot-0 entry in d_intermediates + uint32_t inter_stride, // total threads in grid (slot-major stride) + uint32_t H // APC height (for column-major stride) +) { + switch (src.type) { + case ENTRY_MAIN: + return d_main[H * src.index + row]; + case SRC_INTERMEDIATE: + return inter[src.index * inter_stride]; + case SRC_CONSTANT: + return Fp(src.index); // src.index holds the constant value + default: + return Fp(0); + } +} + +__global__ void apc_apply_bus_dag_kernel( + // APC trace + const Fp *__restrict__ d_output, + int num_apc_calls, + uint32_t apc_height, + + // Rule list (CSE-deduped, encoded by stark-backend's codec.rs) + const Rule *__restrict__ d_rules, + uint32_t n_rules, + + // Interaction-output dispatch table + const DevInteractionDag *__restrict__ d_interactions, + uint32_t n_interactions, + const OutputDesc *__restrict__ d_output_descs, + + // Global-mode intermediates buffer. Layout: slot-major coalesced. + // Size = total_threads * buffer_size Fps. + Fp *__restrict__ d_intermediates, + + // Periphery histograms (unchanged from stack-VM kernel) + uint32_t var_range_bus_id, + uint32_t *__restrict__ d_var_hist, + uint32_t var_num_bins, + uint32_t tuple2_bus_id, + uint32_t *__restrict__ d_tuple2_hist, + uint32_t tuple2_sz0, + uint32_t tuple2_sz1, + uint32_t bitwise_bus_id, + uint32_t *__restrict__ d_bitwise_hist +) { + const uint32_t total_threads = gridDim.x * blockDim.x; + const uint32_t thread_id = blockIdx.x * blockDim.x + threadIdx.x; + const uint32_t r = thread_id; + if (r >= (uint32_t)num_apc_calls) return; + + Fp *inter = d_intermediates + thread_id; + const uint32_t inter_stride = total_threads; + + // -------- Phase 1: walk rules, fill `inter[]` -------- + for (uint32_t ri = 0; ri < n_rules; ++ri) { + Rule rule = d_rules[ri]; + RuleHeader hdr = decode_rule_header(rule); + + Fp val; + switch (hdr.op) { + case OP_VAR: + val = evaluate_dag_entry_bus(hdr.x, r, d_output, inter, inter_stride, apc_height); + break; + case OP_ADD: { + SourceInfo y = decode_y(rule); + Fp a = evaluate_dag_entry_bus(hdr.x, r, d_output, inter, inter_stride, apc_height); + Fp b = evaluate_dag_entry_bus(y, r, d_output, inter, inter_stride, apc_height); + val = a + b; + break; + } + case OP_SUB: { + SourceInfo y = decode_y(rule); + Fp a = evaluate_dag_entry_bus(hdr.x, r, d_output, inter, inter_stride, apc_height); + Fp b = evaluate_dag_entry_bus(y, r, d_output, inter, inter_stride, apc_height); + val = a - b; + break; + } + case OP_MUL: { + SourceInfo y = decode_y(rule); + Fp a = evaluate_dag_entry_bus(hdr.x, r, d_output, inter, inter_stride, apc_height); + Fp b = evaluate_dag_entry_bus(y, r, d_output, inter, inter_stride, apc_height); + val = a * b; + break; + } + case OP_NEG: { + Fp a = evaluate_dag_entry_bus(hdr.x, r, d_output, inter, inter_stride, apc_height); + val = -a; + break; + } + default: + val = Fp(0); + } + + if (hdr.buffer_result) { + uint32_t slot = decode_z_index(rule); + inter[slot * inter_stride] = val; + } + } + + // -------- Phase 2: dispatch to histograms -------- + auto read_output = [&](const OutputDesc &d) -> Fp { + switch (d.kind) { + case 0: return inter[d.value * inter_stride]; + case 1: return d_output[d.value + r]; // value = col_base = col_idx * H + case 2: return Fp(d.value); + default: + return Fp(0); + } + }; + + for (uint32_t i = 0; i < n_interactions; ++i) { + DevInteractionDag intr = d_interactions[i]; + const OutputDesc *outs = d_output_descs + intr.outputs_off; + + Fp mult_fp = read_output(outs[0]); + uint32_t m = mult_fp.asUInt32(); + if (m == 0u) continue; + + if (intr.bus_id == var_range_bus_id) { + // expect [mult, value, max_bits] + uint32_t value = read_output(outs[1]).asUInt32(); + uint32_t max_bits = read_output(outs[2]).asUInt32(); + lookup::Histogram hist(d_var_hist, var_num_bins); + uint32_t idx = (1u << max_bits) + value - 1u; + for (uint32_t k = 0; k < m; ++k) hist.add_count(idx); + } else if (intr.bus_id == tuple2_bus_id) { + // expect [mult, v0, v1] + uint32_t v0 = read_output(outs[1]).asUInt32(); + uint32_t v1 = read_output(outs[2]).asUInt32(); + lookup::Histogram hist(d_tuple2_hist, tuple2_sz0 * tuple2_sz1); + uint32_t idx = v0 * tuple2_sz1 + v1; + for (uint32_t k = 0; k < m; ++k) hist.add_count(idx); + } else if (intr.bus_id == bitwise_bus_id) { + // expect [mult, x, y, x_xor_y, selector] — arg[2] (x_xor_y) skipped + uint32_t x = read_output(outs[1]).asUInt32(); + uint32_t y = read_output(outs[2]).asUInt32(); + uint32_t selector = read_output(outs[4]).asUInt32(); + BitwiseOperationLookup bl(d_bitwise_hist, BITWISE_NUM_BITS); + for (uint32_t k = 0; k < m; ++k) { + if (selector == 0u) bl.add_range(x, y); + else if (selector == 1u) bl.add_xor(x, y); + else { assert(false && "Invalid selector"); } + } + } + } +} + +// ============================================================================= +// Host launcher — callable from Rust FFI +// ============================================================================= + +extern "C" int _apc_apply_bus_dag( + const Fp *d_output, + int num_apc_calls, + uint32_t apc_height, + + const Rule *d_rules, + uint32_t n_rules, + + const DevInteractionDag *d_interactions, + uint32_t n_interactions, + const OutputDesc *d_output_descs, + + // d_intermediates: size = total_threads * buffer_size Fps, slot-major coalesced + Fp *d_intermediates, + + uint32_t var_range_bus_id, + uint32_t *d_var_hist, + uint32_t var_num_bins, + uint32_t tuple2_bus_id, + uint32_t *d_tuple2_hist, + uint32_t tuple2_sz0, + uint32_t tuple2_sz1, + uint32_t bitwise_bus_id, + uint32_t *d_bitwise_hist +) { + if (num_apc_calls <= 0) return cudaSuccess; + const int block_x = 128; + const dim3 block(block_x, 1, 1); + const uint32_t g_size = (uint32_t)((num_apc_calls + block_x - 1) / block_x); + const dim3 grid(g_size, 1, 1); + + // POWDR_NCU_BIG=1 wraps "big" launches in cudaProfilerStart/Stop so + // `ncu --profile-from-start off` captures only those, regardless of + // launch count or template specialization. Threshold (`n_rules` count + // is a cheap proxy for chip size) defaults to 1000, override via + // POWDR_NCU_BIG_THRESHOLD. + static const bool ncu_mode = getenv("POWDR_NCU_BIG") != nullptr; + static const uint32_t big_threshold = []() -> uint32_t { + const char *e = getenv("POWDR_NCU_BIG_THRESHOLD"); + return e ? (uint32_t)atoi(e) : 1000u; + }(); + const bool capture = ncu_mode && (n_rules > big_threshold); + if (capture) cudaProfilerStart(); + + apc_apply_bus_dag_kernel<<>>( + d_output, num_apc_calls, apc_height, + d_rules, n_rules, + d_interactions, n_interactions, d_output_descs, + d_intermediates, + var_range_bus_id, d_var_hist, var_num_bins, + tuple2_bus_id, d_tuple2_hist, tuple2_sz0, tuple2_sz1, + bitwise_bus_id, d_bitwise_hist + ); + + if (capture) { + cudaDeviceSynchronize(); // ensure ncu pages on the launch we just enqueued + cudaProfilerStop(); + } + return (int)cudaGetLastError(); +} diff --git a/openvm/src/cuda_abi.rs b/openvm/src/cuda_abi.rs index b268e02608..072e7bf629 100644 --- a/openvm/src/cuda_abi.rs +++ b/openvm/src/cuda_abi.rs @@ -61,8 +61,76 @@ extern "C" { bitwise_bus_id: u32, // bus id for the bitwise lookup d_bitwise_hist: *mut u32, // device histogram for bitwise lookup ) -> i32; + + /// Launches the GKR-DAG variant of the bus kernel. + /// + /// Safety: All pointers must be valid device pointers for the specified lengths. + pub fn _apc_apply_bus_dag( + // APC trace + d_output: *const BabyBear, + num_apc_calls: i32, + apc_height: u32, + + // Rules: encoded 128-bit `Rule { low, high }` array + d_rules: *const DevRule, + n_rules: u32, + + // Interaction-output dispatch table + d_interactions: *const DevInteractionDag, + n_interactions: u32, + d_output_descs: *const OutputDesc, + + // Global intermediates buffer: size = total_threads * buffer_size Fps, + // slot-major coalesced layout. + d_intermediates: *mut BabyBear, + + // Histograms + var_range_bus_id: u32, + d_var_hist: *mut u32, + var_num_bins: u32, + tuple2_bus_id: u32, + d_tuple2_hist: *mut u32, + tuple2_sz0: u32, + tuple2_sz1: u32, + bitwise_bus_id: u32, + d_bitwise_hist: *mut u32, + ) -> i32; +} + +/// Device-side encoded rule. Matches `Rule { uint64_t low; uint64_t high; }` +/// in `codec.cuh` and the `u128` produced by stark-backend's +/// `RuleWithFlag::encode()`. +#[repr(C)] +#[derive(Clone, Copy)] +pub struct DevRule { + pub low: u64, + pub high: u64, } +/// GPU-side bus interaction descriptor for the DAG kernel. +#[repr(C)] +#[derive(Clone, Copy)] +pub struct DevInteractionDag { + pub bus_id: u32, + pub num_args: u32, + pub outputs_off: u32, + pub flags: u32, +} + +/// Per-output dispatch descriptor. The kernel reads `kind` to choose between +/// (Inter: read inter\[value\]), (Col: read d_output\[value + r\]), or +/// (Const: emit Fp(value)). +#[repr(C)] +#[derive(Clone, Copy)] +pub struct OutputDesc { + pub kind: u32, + pub value: u32, +} + +pub const OUTPUT_KIND_INTER: u32 = 0; +pub const OUTPUT_KIND_COL: u32 = 1; +pub const OUTPUT_KIND_CONST: u32 = 2; + #[repr(C)] #[derive(Clone, Copy, Debug)] pub struct OriginalAir { @@ -229,3 +297,45 @@ pub fn apc_apply_bus( )) } } + +/// High-level safe wrapper for `_apc_apply_bus_dag`. +#[allow(clippy::too_many_arguments)] +pub fn apc_apply_bus_dag( + output: &DeviceMatrix, + num_apc_calls: usize, + apc_height: usize, + rules: &DeviceBuffer, + interactions: &DeviceBuffer, + output_descs: &DeviceBuffer, + intermediates: &DeviceBuffer, + var_range_bus_id: u32, + var_range_count: &DeviceBuffer, + tuple2_bus_id: u32, + tuple2_count: &DeviceBuffer, + tuple2_sizes: [u32; 2], + bitwise_bus_id: u32, + bitwise_count: &DeviceBuffer, +) -> Result<(), CudaError> { + unsafe { + CudaError::from_result(_apc_apply_bus_dag( + output.buffer().as_ptr(), + num_apc_calls as i32, + apc_height as u32, + rules.as_ptr(), + rules.len() as u32, + interactions.as_ptr(), + interactions.len() as u32, + output_descs.as_ptr(), + intermediates.as_mut_ptr(), + var_range_bus_id, + var_range_count.as_mut_ptr() as *mut u32, + var_range_count.len() as u32, + tuple2_bus_id, + tuple2_count.as_mut_ptr() as *mut u32, + tuple2_sizes[0], + tuple2_sizes[1], + bitwise_bus_id, + bitwise_count.as_mut_ptr() as *mut u32, + )) + } +} diff --git a/openvm/src/powdr_extension/trace_generator/cuda/expr_dag.rs b/openvm/src/powdr_extension/trace_generator/cuda/expr_dag.rs new file mode 100644 index 0000000000..b2649fb0c5 --- /dev/null +++ b/openvm/src/powdr_extension/trace_generator/cuda/expr_dag.rs @@ -0,0 +1,397 @@ +//! Convert powdr `AlgebraicExpression` bus interactions into stark-backend's +//! `SymbolicExpressionDag` so we can reuse the upstream `SymbolicRulesBuilder` +//! for CSE / slot allocation / 128-bit `Rule` encoding. +//! +//! Spike entry point: `dump_bus_dag_stats` — when env `POWDR_DUMP_BUS_DAG=1` is +//! set, runs once on the first chip with >= 150 bus interactions and prints +//! `rules_len`, `buffer_size`, `accumulate_count`. Gates whether the full DAG +//! kernel rewrite is worth doing. +#![allow(dead_code)] + +use std::{collections::BTreeMap, sync::Arc}; + +use openvm_cuda_backend::logup_zerocheck::rules::{ + codec::Codec, Rule, Source, SymbolicRulesGpu, +}; +use openvm_stark_backend::{ + air_builders::symbolic::{ + symbolic_expression::SymbolicExpression, + symbolic_variable::{Entry, SymbolicVariable}, + SymbolicDagBuilder, SymbolicExpressionDag, SymbolicExpressionNode, + }, + p3_field::PrimeField32, +}; +use openvm_stark_sdk::p3_baby_bear::BabyBear; +use powdr_autoprecompiles::{expression::AlgebraicExpression, symbolic_machine::SymbolicBusInteraction}; +use powdr_expression::{AlgebraicBinaryOperator, AlgebraicUnaryOperator}; + +use crate::cuda_abi::{ + DevInteractionDag, DevRule, OutputDesc, OUTPUT_KIND_COL, OUTPUT_KIND_CONST, OUTPUT_KIND_INTER, +}; + +/// dag_idxs for one interaction's outputs, in fixed order `[mult, arg0, arg1, ...]`. +pub struct InteractionOutputs { + pub bus_id: u32, + pub num_args: u32, + pub mult_dag_idx: usize, + pub arg_dag_idxs: Vec, +} + +/// Convert one expression to an Arc-wrapped `SymbolicExpression`. Each +/// `AlgebraicExpression::Reference` becomes a `SymbolicVariable::Main` whose +/// `index` is the column's index in the APC trace; `offset=0`, `part_index=0` +/// (we have a single main partition). +fn to_symbolic( + expr: &AlgebraicExpression, + id_to_apc_index: &BTreeMap, +) -> Arc> { + match expr { + AlgebraicExpression::Number(c) => Arc::new(SymbolicExpression::Constant(*c)), + AlgebraicExpression::Reference(r) => { + let var = SymbolicVariable::new( + Entry::Main { + part_index: 0, + offset: 0, + }, + id_to_apc_index[&r.id], + ); + Arc::new(SymbolicExpression::Variable(var)) + } + AlgebraicExpression::UnaryOperation(u) => { + let child = to_symbolic(&u.expr, id_to_apc_index); + match u.op { + AlgebraicUnaryOperator::Minus => { + let deg = child.degree_multiple(); + Arc::new(SymbolicExpression::Neg { + x: child, + degree_multiple: deg, + }) + } + } + } + AlgebraicExpression::BinaryOperation(b) => { + let l = to_symbolic(&b.left, id_to_apc_index); + let r = to_symbolic(&b.right, id_to_apc_index); + let l_deg = l.degree_multiple(); + let r_deg = r.degree_multiple(); + let node = match b.op { + AlgebraicBinaryOperator::Add => SymbolicExpression::Add { + x: l, + y: r, + degree_multiple: l_deg.max(r_deg), + }, + AlgebraicBinaryOperator::Sub => SymbolicExpression::Sub { + x: l, + y: r, + degree_multiple: l_deg.max(r_deg), + }, + AlgebraicBinaryOperator::Mul => SymbolicExpression::Mul { + x: l, + y: r, + degree_multiple: l_deg + r_deg, + }, + }; + Arc::new(node) + } + } +} + +/// Build a `SymbolicExpressionDag` over all bus interactions and return it +/// along with the per-interaction output dag_idxs (in `[mult, arg0, ...]` +/// order). `SymbolicDagBuilder` already performs structural dedup + +/// constant-folding peephole (see its docs at `dag.rs:128-141`), so we get the +/// upstream-equivalent of the host-side peephole "for free" inside DAG build. +pub fn algebraic_to_symbolic_dag( + bus_interactions: &[SymbolicBusInteraction], + id_to_apc_index: &BTreeMap, +) -> (SymbolicExpressionDag, Vec) { + let mut builder = SymbolicDagBuilder::::new(); + let mut per_interaction = Vec::with_capacity(bus_interactions.len()); + + // CRITICAL: `SymbolicDagBuilder.expr_to_idx` caches by raw `*const + // SymbolicExpression` Arc pointer (see its docstring — assumes the + // input Arcs stay live for the duration of the build). If we drop an Arc + // between two `add_expr` calls, the next `Arc::new` may reuse the freed + // address, and the pointer cache returns the wrong idx. To keep the + // invariant intact, retain every Arc we hand to `add_expr` in `keepalive` + // until the whole DAG has been built. + let mut keepalive: Vec>> = Vec::new(); + let trace_args = std::env::var("POWDR_DUMP_BUS_ARGS").is_ok(); + for (intr_i, bi) in bus_interactions.iter().enumerate() { + let mult_expr = to_symbolic(&bi.mult, id_to_apc_index); + let mult_dag_idx = builder.add_expr(&mult_expr); + keepalive.push(mult_expr); + + let mut arg_dag_idxs = Vec::with_capacity(bi.args.len()); + for arg in &bi.args { + let a = to_symbolic(arg, id_to_apc_index); + arg_dag_idxs.push(builder.add_expr(&a)); + keepalive.push(a); + } + + // Spot-print: for the failing bitwise interaction (intr_i == 14, 4 + // args), dump the AST of each arg so we can see whether they're + // genuinely structurally equal (legitimate dedup) or whether our + // converter is collapsing distinct exprs. + if trace_args && intr_i == 14 && bi.args.len() == 4 { + eprintln!( + "[bus_args] intr=14 dag_idxs=({}, [{}]) mult={:?}", + mult_dag_idx, + arg_dag_idxs + .iter() + .map(usize::to_string) + .collect::>() + .join(","), + bi.mult + ); + for (ai, (a, d)) in bi.args.iter().zip(arg_dag_idxs.iter()).enumerate() { + eprintln!("[bus_args] arg{} dag_idx={} expr={:?}", ai, d, a); + } + } + + per_interaction.push(InteractionOutputs { + bus_id: bi.id as u32, + num_args: bi.args.len() as u32, + mult_dag_idx, + arg_dag_idxs, + }); + } + + // `SymbolicRulesBuilder::new` debug-asserts `constraint_idx.is_sorted()`. + // Collect all output dag_idxs, sort, dedup. + let mut output_idxs: Vec = per_interaction + .iter() + .flat_map(|i| std::iter::once(i.mult_dag_idx).chain(i.arg_dag_idxs.iter().copied())) + .collect(); + output_idxs.sort_unstable(); + output_idxs.dedup(); + + let dag = SymbolicExpressionDag { + nodes: builder.nodes, + constraint_idx: output_idxs, + }; + (dag, per_interaction) +} + +/// Compile bus interactions for the DAG kernel. Returns three device-ready +/// vectors: +/// - `rules`: 128-bit encoded `DevRule { low, high }` array (CSE-deduped by +/// `SymbolicRulesBuilder` with `buffer_vars=true`). +/// - `interactions`: `DevInteractionDag` per bus interaction (bus_id, num_args, +/// outputs_off). +/// - `output_descs`: flat dispatch table — one `OutputDesc` per (interaction, +/// slot) in order `[i0_mult, i0_arg0, ..., i1_mult, ...]`. Each descriptor +/// tells the kernel where to read the value: `inter[slot]`, `d_output[col + +/// r]`, or inline `Fp(const)`. +pub fn compile_bus_to_gpu_dag( + bus_interactions: &[SymbolicBusInteraction], + id_to_apc_index: &std::collections::BTreeMap, + apc_height: usize, +) -> (Vec, Vec, Vec, usize) { + let (dag, outputs) = algebraic_to_symbolic_dag(bus_interactions, id_to_apc_index); + let rules_gpu: SymbolicRulesGpu = SymbolicRulesGpu::new(&dag, /*buffer_vars=*/ true); + + // Validate that every source we'll send to the kernel is one of the three + // types our kernel handles: Source::Intermediate, Source::Var, Source::Constant. + // Anything else (IsFirst/IsLast/IsTransition or any Entry::{Preprocessed,Permutation,Public,Challenge,Exposed}) + // would trip the device-side assert. + if std::env::var("POWDR_VALIDATE_BUS_DAG").is_ok() { + for (i, r) in rules_gpu.rules.iter().enumerate() { + let sources: Vec<&Source> = match &r.inner { + Rule::Add(x, y, z) | Rule::Sub(x, y, z) | Rule::Mul(x, y, z) => vec![x, y, z], + Rule::Neg(x, z) => vec![x, z], + Rule::Variable(x) => vec![x], + Rule::BufferVar(x, z) => vec![x, z], + }; + for s in sources { + match s { + Source::Intermediate(_) + | Source::TerminalIntermediate + | Source::Constant(_) => {} + Source::Var(v) => match v.entry { + Entry::Main { part_index, offset } => { + if part_index != 0 || offset != 0 { + eprintln!("[bus_dag_check] rule {}: Var with non-default Main entry: part={} offset={}", i, part_index, offset); + } + } + other => eprintln!("[bus_dag_check] rule {}: Var with non-Main entry: {:?}", i, other), + }, + other => eprintln!( + "[bus_dag_check] rule {}: unsupported source {:?}", + i, other + ), + } + } + } + } + + // Encode rules to device format. + let dev_rules: Vec = rules_gpu + .rules + .iter() + .map(|r| { + let encoded: u128 = r.encode(); + DevRule { + low: encoded as u64, + high: (encoded >> 64) as u64, + } + }) + .collect(); + + // For each output (mult + each arg) we need the rule for the corresponding + // dag_idx, then case-split on the rule shape to emit an `OutputDesc`. + // + // dag_idx_to_rule_idx only contains accumulate=true entries — which we + // marked via `constraint_idx`. So every output dag_idx is present. + let mut output_descs: Vec = Vec::new(); + let mut dev_interactions: Vec = Vec::with_capacity(outputs.len()); + + for out in &outputs { + let outputs_off = output_descs.len() as u32; + + // Push mult dispatch desc first, then each arg's desc. + let all_outs = std::iter::once(out.mult_dag_idx).chain(out.arg_dag_idxs.iter().copied()); + for dag_idx in all_outs { + let rule_idx = rules_gpu + .dag_idx_to_rule_idx + .get(&dag_idx) + .copied() + .unwrap_or_else(|| panic!("output dag_idx={} has no rule_idx (not accumulated?)", dag_idx)); + let rule = &rules_gpu.rules[rule_idx].inner; + output_descs.push(rule_to_output_desc(rule, apc_height)); + } + + dev_interactions.push(DevInteractionDag { + bus_id: out.bus_id, + num_args: out.num_args, + outputs_off, + flags: 0, + }); + } + + (dev_rules, dev_interactions, output_descs, rules_gpu.buffer_size) +} + +fn rule_to_output_desc(rule: &Rule, apc_height: usize) -> OutputDesc { + // Helper that converts a `Source` to a `(kind, value)` pair for outputs. + // For Source::Intermediate(slot) → (Inter, slot) + // For Source::Var(v) → (Col, v.index * apc_height) — single-use bare col + // For Source::Constant(c) → (Const, c.as_canonical_u32()) + // Other sources are unreachable for bus outputs. + let source_to_desc = |src: &Source| -> OutputDesc { + match src { + Source::Intermediate(slot) => OutputDesc { + kind: OUTPUT_KIND_INTER, + value: *slot as u32, + }, + Source::Var(v) => OutputDesc { + kind: OUTPUT_KIND_COL, + value: (v.index * apc_height) as u32, + }, + Source::Constant(c) => OutputDesc { + kind: OUTPUT_KIND_CONST, + value: c.as_canonical_u32(), + }, + Source::TerminalIntermediate => { + panic!("output rule has Source::TerminalIntermediate — accumulate=true should have allocated a slot") + } + Source::IsFirst | Source::IsLast | Source::IsTransition => { + panic!("bus outputs should never reference IsFirst/IsLast/IsTransition") + } + } + }; + + // For a rule with a `z` operand (Add/Sub/Mul/Neg/BufferVar), z is the + // result location — that's what we want for downstream lookup. + // For `Rule::Variable(x)`, the rule produces the value of `x` inline (no + // slot), so we case on `x` directly. + match rule { + Rule::Add(_, _, z) + | Rule::Sub(_, _, z) + | Rule::Mul(_, _, z) + | Rule::Neg(_, z) + | Rule::BufferVar(_, z) => source_to_desc(z), + Rule::Variable(x) => source_to_desc(x), + } +} + +/// Run the full host-side DAG build + GKR rule scheduler on every chip whose +/// `bus_interactions.len()` clears the threshold, and accumulate stats. +/// Reports one summary line per chip, then a final "max across chips" line so +/// we can decide local vs global buffer mode. +/// Gated by env `POWDR_DUMP_BUS_DAG=1`. Runs once per chip per process. +pub fn dump_bus_dag_stats( + bus_interactions: &[SymbolicBusInteraction], + id_to_apc_index: &BTreeMap, +) { + if std::env::var("POWDR_DUMP_BUS_DAG").is_err() { + return; + } + let threshold: usize = std::env::var("POWDR_DUMP_BUS_DAG_MIN") + .ok() + .and_then(|s| s.parse().ok()) + .unwrap_or(50); + if bus_interactions.len() < threshold { + return; + } + + use openvm_cuda_backend::logup_zerocheck::rules::{SymbolicRulesBuilder, SymbolicRulesGpu}; + + let (dag, outputs) = algebraic_to_symbolic_dag(bus_interactions, id_to_apc_index); + + let n_nodes = dag.nodes.len(); + let n_outputs: usize = outputs.iter().map(|i| 1 + i.arg_dag_idxs.len()).sum(); + let n_distinct_outputs = dag.constraint_idx.len(); + + let (mut n_const, mut n_var, mut n_add, mut n_sub, mut n_mul, mut n_neg) = + (0usize, 0, 0, 0, 0, 0); + for node in &dag.nodes { + match node { + SymbolicExpressionNode::Constant(_) => n_const += 1, + SymbolicExpressionNode::Variable(_) => n_var += 1, + SymbolicExpressionNode::Add { .. } => n_add += 1, + SymbolicExpressionNode::Sub { .. } => n_sub += 1, + SymbolicExpressionNode::Mul { .. } => n_mul += 1, + SymbolicExpressionNode::Neg { .. } => n_neg += 1, + _ => {} + } + } + + let mut rb = SymbolicRulesBuilder::new(&dag, /*buffer_vars=*/ true); + rb.set_range(0, dag.nodes.len()); + let rules: SymbolicRulesGpu = rb.to_rules(); + + // Maintain a global max across chips so we can spot the worst-case buffer_size. + use std::sync::atomic::{AtomicUsize, Ordering}; + static MAX_BUF: AtomicUsize = AtomicUsize::new(0); + static MAX_RULES: AtomicUsize = AtomicUsize::new(0); + static MAX_NODES: AtomicUsize = AtomicUsize::new(0); + let prev_buf = MAX_BUF.fetch_max(rules.buffer_size, Ordering::Relaxed); + let prev_rules = MAX_RULES.fetch_max(rules.rules.len(), Ordering::Relaxed); + let prev_nodes = MAX_NODES.fetch_max(n_nodes, Ordering::Relaxed); + + eprintln!( + "[bus_dag] chip: n_intr={:>4} outs={:>4} dag_nodes={:>4} (C{:>3} V{:>3} A{:>3} S{:>2} M{:>3} N{:>2}) rules={:>4} buffer_size={:>3}", + bus_interactions.len(), + n_outputs, + n_nodes, + n_const, n_var, n_add, n_sub, n_mul, n_neg, + rules.rules.len(), + rules.buffer_size, + ); + // Note when this chip raised any of the running maxima (cheap signal of the + // worst-case chip for global-vs-local decision). + let new_max = + prev_buf < rules.buffer_size || prev_rules < rules.rules.len() || prev_nodes < n_nodes; + if new_max { + eprintln!( + "[bus_dag] ^ new max: max_buffer_size={} max_rules={} max_nodes={}", + MAX_BUF.load(Ordering::Relaxed), + MAX_RULES.load(Ordering::Relaxed), + MAX_NODES.load(Ordering::Relaxed), + ); + } + // Silence the unused-`n_distinct_outputs` warning while keeping it computed + // (handy for debugging if needed). + let _ = n_distinct_outputs; +} diff --git a/openvm/src/powdr_extension/trace_generator/cuda/mod.rs b/openvm/src/powdr_extension/trace_generator/cuda/mod.rs index 69f0d1e23a..8162baa5fa 100644 --- a/openvm/src/powdr_extension/trace_generator/cuda/mod.rs +++ b/openvm/src/powdr_extension/trace_generator/cuda/mod.rs @@ -31,6 +31,7 @@ use crate::{ BabyBearSC, GpuBackend, }; +mod expr_dag; mod inventory; mod periphery; @@ -370,6 +371,7 @@ pub fn compile_bus_to_gpu( apc_height: usize, ) -> (Vec, Vec, Vec) { analyze_bus_cse(bus_interactions); + expr_dag::dump_bus_dag_stats(bus_interactions, apc_poly_id_to_index); let mut interactions = Vec::with_capacity(bus_interactions.len()); let mut arg_spans = Vec::new(); @@ -638,19 +640,61 @@ impl PowdrTraceGeneratorGpu { cuda_abi::apc_apply_derived_expr(&mut output, d_specs, d_bc, num_apc_calls).unwrap(); }); - // Encode bus interactions for GPU consumption - let (bus_interactions, arg_spans, bytecode) = timed_substage!("bus_compile_h2d", { - let (bus_interactions, arg_spans, bytecode) = compile_bus_to_gpu( - &self.apc.machine.bus_interactions, - &apc_poly_id_to_index, - height, - ); - ( - bus_interactions.to_device().unwrap(), - arg_spans.to_device().unwrap(), - bytecode.to_device().unwrap(), - ) - }); + // Read kernel selector once per process. `POWDR_BUS_KERNEL=dag` uses + // the GKR-DAG kernel; anything else (including unset) uses the stack-VM. + fn use_dag_kernel() -> bool { + static ENABLED: std::sync::OnceLock = std::sync::OnceLock::new(); + *ENABLED.get_or_init(|| { + std::env::var("POWDR_BUS_KERNEL") + .map(|v| v == "dag") + .unwrap_or(false) + }) + } + + // Encode bus interactions. Both paths build their device buffers under + // the same `bus_compile_h2d` substage so the two kernels A/B fairly. + let bus_dag_inputs: Option<(_, _, _, _)> = if use_dag_kernel() { + Some(timed_substage!("bus_compile_h2d", { + let (rules, interactions, output_descs, buffer_size) = + expr_dag::compile_bus_to_gpu_dag( + &self.apc.machine.bus_interactions, + &apc_poly_id_to_index, + height, + ); + // Global-mode intermediates buffer: slot-major coalesced. + // total_threads = grid_size * block_size = ceil(num_apc_calls / 128) * 128. + let block_x = 128usize; + let total_threads = num_apc_calls.div_ceil(block_x) * block_x; + let inter_len = total_threads * buffer_size.max(1); + let intermediates = + openvm_cuda_common::d_buffer::DeviceBuffer::::with_capacity(inter_len); + ( + rules.to_device().unwrap(), + interactions.to_device().unwrap(), + output_descs.to_device().unwrap(), + intermediates, + ) + })) + } else { + None + }; + + let bus_vm_inputs: Option<(_, _, _)> = if !use_dag_kernel() { + Some(timed_substage!("bus_compile_h2d", { + let (bus_interactions, arg_spans, bytecode) = compile_bus_to_gpu( + &self.apc.machine.bus_interactions, + &apc_poly_id_to_index, + height, + ); + ( + bus_interactions.to_device().unwrap(), + arg_spans.to_device().unwrap(), + bytecode.to_device().unwrap(), + ) + })) + } else { + None + }; // Gather GPU inputs for periphery (bus ids, count device buffers) let periphery = &self.periphery.real; @@ -675,26 +719,41 @@ impl PowdrTraceGeneratorGpu { // the next kernel function after the prior (`apc_tracegen`) returns. // This is important because bus evaluation depends on trace results. timed_substage!("bus_kernel", { - cuda_abi::apc_apply_bus( - // APC related - &output, - num_apc_calls, - // Interaction related - bytecode, - bus_interactions, - arg_spans, - // Variable range checker related - var_range_bus_id, - var_range_count, - // Tuple range checker related - tuple2_bus_id, - tuple2_count_u32, - tuple2_sizes, - // Bitwise related - bitwise_bus_id, - bitwise_count_u32, - ) - .unwrap(); + if let Some((rules, interactions, output_descs, intermediates)) = bus_dag_inputs { + cuda_abi::apc_apply_bus_dag( + &output, + num_apc_calls, + height, + &rules, + &interactions, + &output_descs, + &intermediates, + var_range_bus_id, + var_range_count, + tuple2_bus_id, + tuple2_count_u32, + tuple2_sizes, + bitwise_bus_id, + bitwise_count_u32, + ) + .unwrap(); + } else if let Some((bus_interactions, arg_spans, bytecode)) = bus_vm_inputs { + cuda_abi::apc_apply_bus( + &output, + num_apc_calls, + bytecode, + bus_interactions, + arg_spans, + var_range_bus_id, + var_range_count, + tuple2_bus_id, + tuple2_count_u32, + tuple2_sizes, + bitwise_bus_id, + bitwise_count_u32, + ) + .unwrap(); + } }); Some(output)